AI/ML 详解：优化的几何与前沿

训练神经网络 = 在一个上亿维的「地形」里找谷底。但谷底有两种：一种是宽阔平坦的盆地，一种是窄而深的尖缝。类比你的系统配置——平坦极小像容差大的配置（某个参数漂移一点照常工作），尖锐极小像踩在钢丝上的配置（任意参数偏一点点系统就崩）。一个模型泛化好不好，很大程度上取决于它停在哪种谷底。

损失函数 L(w) 是参数 w（动辄上亿维）的函数，训练就是用梯度下降找 L 的极小点。但这里有个谜：两个训练 loss 都接近 0 的模型，为什么一个测试集表现好、一个差？

关键洞察（「平坦极小」概念由 Hochreiter & Schmidhuber 1990 年代提出）：测试集的 loss 曲面相对训练集会有微小平移（因为训练/测试数据分布略有差异）。如果你停在的极小点很「平」，平移后你仍处在低 loss 区——泛化好；如果很「尖」，平移一点 loss 就飙升——泛化差。平坦 = 对数据分布的扰动鲁棒。

Li et al. 2018 用一种叫 filter normalization 的技术，把上亿维曲面投影到 2D 画出来，做出一个漂亮的发现：ResNet 的残差连接（skip connection）会显著把原本尖锐、混乱、布满沟壑的地形「抚平」成光滑盆地——这从几何上解释了为什么深网加了 skip connection 之后突然变得好训练。

import torch, copy

# 量化一个训练好的模型所在极小点的"锐度"：
# 在权重邻域里随机扰动，看 loss 最多涨多少
def sharpness(model, loss_fn, batch, rho=0.05, trials=10):
    base = loss_fn(model, batch).item()
    worst = base
    for _ in range(trials):
        m2 = copy.deepcopy(model)
        with torch.no_grad():
            for p in m2.parameters():
                # 半径正比于参数自身尺度的随机扰动
                p.add_(torch.randn_like(p) * rho * p.norm())
        worst = max(worst, loss_fn(m2, batch).item())
    return worst - base   # 邻域最坏 loss 比中心高多少 = 锐度
# 锐度越小 → 极小点越"平" → 通常泛化越好

如果平坦极小更好，能不能把「找平坦」直接写进训练目标？SAM 就是这么做的。它像混沌工程 / 故障注入——不满足于「当前配置能跑」，而是主动问「在我周围最坏的扰动下还能跑吗？」，专挑邻域里最糟的点来优化，逼自己进入宽盆地。

普通 SGD 只最小化当前点的 loss L(w)。但 Keskar et al. 2016 发现一个著名现象——大 batch 泛化鸿沟：用大 batch 训练容易掉进尖锐极小（因为大 batch 的梯度噪声小，少了把模型「抖出尖缝」的随机扰动），导致泛化变差。

Foret et al. 2020 提出的 SAM 把优化目标从「最小化当前 loss」改成 min-max：

逐符号拆解：w 是参数，ε 是加在权重上的扰动向量，ρ（rho）是邻域半径（一个超参，比如 0.05），‖·‖ 是向量范数。内层 max 的含义是「在半径 ρ 的小球里，loss 最高能到多少」；外层 min 是「最小化那个最坏值」。直觉：不再问脚下 loss 多低，而问邻域内最高的 loss 多低——只有宽平的盆地才能让「邻域最坏值」也很低。

实现上每步要做两次前向：第一次算梯度，找到「最锐方向」 ε̂ = ρ·g/‖g‖（往坡最陡处先爬一步）；第二次在 w+ε̂ 处重新算梯度，这个梯度才用来真正更新。代价是训练慢约一倍，换更好的泛化。

# SAM 的一步：先爬到邻域最差点，再从那里下降（两次前向）
def sam_step(model, loss_fn, batch, optimizer, rho=0.05):
    loss_fn(model, batch).backward()
    # 1) 计算扰动 ê = ρ·g/‖g‖，把权重推到"最锐"方向
    gn = torch.norm(torch.stack(
        [p.grad.norm() for p in model.parameters()]))
    eps = {}
    with torch.no_grad():
        for p in model.parameters():
            e = p.grad * rho / (gn + 1e-12); eps[p] = e; p.add_(e)
    optimizer.zero_grad()
    # 2) 在扰动点重新算梯度——这才是真正用于更新的方向
    loss_fn(model, batch).backward()
    with torch.no_grad():
        for p in model.parameters(): p.sub_(eps[p])  # 退回原点
    optimizer.step()      # 用"邻域最差点"的梯度更新 → 逼向平整区

一阶方法（SGD / Adam）像只看脚下坡度走路——在又窄又长的「山谷」里会左右横跳、走 Z 字，半天下不到底。二阶方法多看一样东西：曲率（地形怎么弯的），于是能直接朝谷底斜切过去。类比：一阶像只看瞬时梯度的拥塞控制，二阶像还掌握了「这条链路的曲率特性」从而一步切到位。

梯度 g 是一阶导（坡度）。理论上最优的 Newton 法用 Hessian 矩阵 H（二阶导，描述曲率）做预条件（preconditioner）：

它能纠正「病态」山谷——即各方向尺度差异巨大、Adam 会横跳的那种地形。但 pain point 致命：H 是 n×n，n = 上亿参数 → 既存不下、也求不了逆（复杂度 O(n²)~O(n³)）。这就像你绝不会去存一个上亿节点的全连接邻接矩阵。

于是有两条「用结构假设把 O(n²) 砍成可承受」的近似路线：

• K-FAC（Martens & Grosse 2015）：把每一层的 Fisher 矩阵（曲率的一种）近似成两个小矩阵的 Kronecker 积 A⊗B——把一个大块分解成两个小因子，既存得下、又可逆。
• Shampoo（Gupta et al. 2018）：为参数张量的每个维度各维护一个小预条件矩阵，而不是一个跨所有参数的巨阵。

两者本质都是「用结构化分解逼近全曲率」。2024–2025 年这类方法（及其变体，如 SOAP）在大模型预训练里重新走红，因为它们能用更少的步数收敛——在超大规模下，步数省下来就是真金白银的 wall-clock。

# Adam 只用对角"自适应"（把每个参数当独立的）；
# 二阶方法用曲率做预条件，能感知参数间的耦合。
from torch_optimizer import Shampoo  # pip install torch-optimizer

opt = Shampoo(
    model.parameters(), lr=1e-3,
    # 为张量每个维度各维护一个小预条件矩阵，
    # 而非一个 n×n 巨阵——把"全曲率"分解成小因子
    update_freq=20)   # 每 20 步才更新一次预条件子（摊薄开销）

for batch in loader:
    opt.zero_grad()
    loss_fn(model, batch).backward()
    opt.step()

Adam 统治了近十年后，2023 年冒出两个挑战者。Lion 像「极简主义重构」——砍掉 Adam 一半的状态（只留动量），更新只取方向的符号；Sophia 像「轻量二阶」——偷偷估一点曲率，但只花极小代价。

Adam 给每个参数存两个状态：一阶矩 m（动量）和二阶矩 v（梯度平方的滑动平均）。显存 ≈ 2× 参数量。pain point：到百亿参数时，优化器状态本身就吃掉巨量显存。两个挑战者从不同方向找便宜的更新规则：

• Lion（Chen et al. 2023，「符号动量」）：它本身是用程序搜索（把「找优化算法」当成在程序空间里搜索）自动发现的，不是人手设计。它只存动量 m（显存减半），更新量取符号：update = sign(β₁·m + (1−β₁)·g)。sign 意味着每个参数的步长幅度都一样（只有方向不同），这本身是一种隐式正则。因为符号更新的等效步长偏大，论文指出 Lion 通常要用比 Adam 小 3–10 倍的学习率。
• Sophia（Liu et al. 2023）：专为 LLM 预训练设计，每隔几步才轻量估一次对角 Hessian 当预条件子，再对更新做逐元素裁剪（clipping）控制最坏步长——用很小的二阶信息换更快收敛。

共同主题：在「Adam 够用、但不够省 / 不够快」的缝隙里，从显存（Lion）和曲率（Sophia）两个方向各找了一条出路。

from lion_pytorch import Lion  # pip install lion-pytorch

# Lion 只存动量（不存 Adam 的二阶矩 v）→ 优化器显存减半
opt = Lion(model.parameters(),
           lr=1e-4,            # 通常取 Adam 的 1/3~1/10
           weight_decay=1e-2)

# 核心更新规则（概念示意）：
#   update = sign(β1·m + (1-β1)·g)   ← 只取符号，步长幅度统一
#   m      = β2·m + (1-β2)·g          ← 动量用另一组系数更新
for batch in loader:
    opt.zero_grad()
    loss_fn(model, batch).backward()
    opt.step()

• Li et al. 2018 — Visualizing the Loss Landscape of Neural Nets — 损失景观可视化的奠基作，解释 skip connection 如何抚平地形
• Keskar et al. 2016 — On Large-Batch Training: Sharp Minima — 「大 batch 泛化鸿沟」与尖锐极小的经典实证
• Foret et al. 2020 — Sharpness-Aware Minimization (SAM) — 把「找平坦」写进优化目标的 min-max 方法
• Gupta et al. 2018 — Shampoo: Preconditioned Tensor Optimization — 按张量维度做预条件的二阶方法
• Chen et al. 2023 — Symbolic Discovery of Optimization Algorithms (Lion) — 用程序搜索自动发现优化器
• Liu et al. 2023 — Sophia: Scalable Second-order Optimizer — 面向 LLM 预训练的轻量二阶方法