如果你训练过深度学习模型大概率遇到过这种场景训练集上的 loss 一直降得很漂亮曲线平滑得像教科书一样可换到验证集上立刻现原形准确率上上下下就是冲不上去。我们一位同事A同学在这类问题上栽了不止一次后来把一个“额外连做两次反向传播”的优化器接进项目里虽然训练时间明显变长但验证集泛化确实稳了一大截。这个优化器用的思路就是锐度感知最小化Sharpness-Aware Minimization简称 SAM。SAM 不是用来替代某个网络结构的它改的是优化目标不只找 loss 最低的那个点还要找“附近一整块区域都低”的那个点。简单的说它想让模型落进一个又低又平的盆地而不是扎进一个又深又窄的针孔。这篇文章我会把它解决的问题、核心公式的拆解、我在实际项目里的复现步骤、实验数据和踩坑记录都写清楚。这篇内容适合谁凡是做过图像分类、目标检测、文本分类这些任务并且被“训练集和测试集不一致”折磨过的人都值得认真看一遍。不需要你已经懂很深的优化理论只要会用框架跑过模型训练就能跟着把 SAM 接进自己的工程里。1. 锐度感知最小化到底在解决什么问题1.1 训练损失不是终点泛化差的三个典型表现先说一个基础认知深度学习模型的目标从来不是记住训练数据。我们实际想要的是它在没见过的数据上也能给出合理结果这个能力在机器学习里叫泛化能力。泛化差的表现有几种大家应该都见过。第一种是最典型的过拟合训练集 loss 已经降到很低但验证集 loss 在某个 epoch 之后开始反弹。这说明模型开始把训练样本里的噪声和偶然模式也记住了。第二种是“梯度看着正常但验证集精度就是上不去”。这种情况更隐蔽。训练 loss 在降说明模型在拟合数据但降到了一个不健康的区域。模型参数所在位置可能是一个尖锐的极小值点稍微有一点输入扰动、数据迁移、参数变化loss 就会剧烈上升。第三种是换了不同随机种子重新训练验证集精度波动很大。同一个模型、同一个数据、同样的超参数如果多次训练结果忽高忽低说明优化器每次找到的参数位置都不稳定模型的泛化能力没有保障。这些问题的共同点是什么是优化过程只盯着 loss 本身而没有关心 loss 曲面在这个点周围的几何形状。1.2 损失曲面上的两种“低点”为了理解 SAM需要把模型参数想象成“位置”把对应 loss 想象成“地势高度”。训练过程就是在高维地形里往下走找一个尽量低的位置。但“低”和“好”不总是等价。假设有两个位置A 点的 loss 是 0.10B 点的 loss 是 0.12。只看数值A 更好。但如果 A 点是一根插在地上的细针顶部只要参数稍微偏移一点loss 就飙到 1.0B 点是一片平坦盆地的最低处周边一大圈的 loss 都低于 0.15。那么在真实测试集上B 点往往比 A 点更可靠。原因是测试环境几乎不可能和训练环境完全一致。数据分布会有偏移模型参数的数值精度有限训练过程中的随机扰动也一直存在。落在尖锐极小值里的模型对任何细微变化都高度敏感。落在平坦极小值里的模型即使参数被扰动、输入略有变化损失变化也很小预测结果更加稳定。所以一个很自然的想法是能不能在训练时不仅压低当前点的 loss还同时压低这个点邻近区域的 loss这就是锐度感知最小化的出发点。它把“极小值是否尖锐”这件事显式地放进了优化目标里。2. SAM 的核心原理一个优化器如何找到更平的极小值2.1 从“只找低点”到“找低点附近区域都低”普通训练的目标函数是最小化训练损失 L(w)其中 w 是模型参数。SAM 的目标函数长这样min_w max_||ε||≤ρ L(w ε)中文拆开说就是对每一个参数 w先允许一个不超过半径 ρ 的扰动 ε在这个扰动范围内找到让 loss 最大的方向也就是说先找到当前点周围最“危险”的相邻点。然后优化器要做的不再是单纯压低当前点的 loss而是把这个最坏情况下的 loss 也压下去。如果当前点周围全都是高地max 这一项就会很大优化器会努力把这个区域整体降下来。如果当前点周围本身就很平坦那么即使扰动到边缘loss 也不会增加太多优化器自然能做到。这样最后找到的参数不再是一个孤零零的低点而是一个低洼地带的中心。这一步“先找最坏扰动再更新参数”的思路很像对抗训练。对抗训练是给输入数据加扰动让模型对输入变化不敏感SAM 是给模型参数加扰动让模型在参数空间里也稳健。2.2 SAM 的两步近似扰动方向和更新梯度直接求解上面的 min-max 问题很贵因为内层最大化需要评估很多次 loss。SAM 的关键做法是用一阶泰勒展开近似出最坏扰动方向把计算量控制在“两遍反向传播”以内。第一步在当前位置 w 上做一次前向和反向传播得到梯度 g ∇L(w)。第二步把扰动近似为ε ≈ ρ * g / ||g||这里的含义是在 w 附近loss 上升最快的方向就是梯度方向。为了让扰动幅度正好落在半径为 ρ 的球面上把梯度归一化成长度为 1 的单位向量再乘上 ρ。第三步把参数临时加到 w ε 的位置重新算一次前向和反向传播得到扰动点上的梯度g_perturbed ∇L(w ε)第四步恢复参数到 w用 g_perturbed 去做正常的优化器更新比如带动量的 SGD 或者自适应矩估计算法。这个流程听起来简单但有两个细节容易理解偏。第一个细节是不是用第一次求出的梯度 g 去更新参数而是用第二次在扰动点求出的梯度。SAM 真正要做的是最小化邻域内最大 loss所以更新方向应该反映邻域情况而不是只反映当前点。第一次梯度只用来计算扰动方向第二次梯度才是真正的更新梯度。第二个细节是扰动方向用整体梯度的范数归一化而不是对每个参数单独归一化。不同参数所在的网络层尺度可能差别很大如果每个参数单独除以自己的梯度范数扰动会被“逐层放大”破坏不同层之间的相对比例。实际实现里需要把所有梯度拼接成一个大的向量计算整体的 L2 范数再统一缩放。2.3 SAM 与普通正则化的本质区别很多人第一次听说 SAM会把它和 weight decay、dropout 这些正则化方法混在一起。它们确实都为了泛化但机制完全不同。weight decay 是给参数本身的 L2 范数加惩罚让参数尽量靠近原点。它约束的是参数大小不直接约束 loss 曲面的形状。dropout 是在前向传播时随机丢弃部分神经元相当于训练多个共享参数的子模型做集成它约束的是激活函数层面。而 SAM 直接在“参数邻域”这个几何尺度上做文章。它不管参数离原点有多远只管参数周围这一圈的 loss 高不高。所以 SAM 可以和 weight decay、dropout 一起用它们各自解决不同维度的问题。我们项目里最后的配置就是“SAM weight decay 数据增强”一起上效果比只用任意一个都好。3. 动手复现在图像分类任务上把 SAM 跑起来3.1 实验设置与评估方式我拿一个公开图像数据集做了复现实验。模型没有用特别深的网络选了一个中等规模的卷积结构输入图片统一缩放到固定尺寸数据划分严格按原始训练集和测试集走没有做任何手工挑数据。对比的是三个方案基线 A带动量的 SGD加常规 weight decay。基线 B自适应矩估计算法也就是大家熟知的 Adam 类优化器加 weight decay。SAM 方案外层用 SAM 计算扰动和更新梯度内层仍然用带动量的 SGD 做参数更新其他数据增强、学习率策略完全一致。要特别说明的是SAM 不是一个“独立优化器”它的第二步内层更新仍然需要一个基础优化器。你可以让 SAM 套在 SGD 上也可以套在自适应矩估计算法上。常见的选择是内层用 SGD momentum因为实验里这种组合在多数视觉任务上更稳。评估指标不只看了测试集准确率还额外计算了一个“邻域损失差”。做法是训练结束后对参数加一个微小的随机高斯扰动重新算一遍测试集 loss反复做多次取平均。这个指标直接反映参数附近的平坦程度。如果 SAM 有效这个差值应该明显小于普通训练。3.2 从零实现一个 SAM 优化器用伪代码把核心逻辑写出来方便你在自己的框架里对照def sam_step(model, loss_fn, inputs, labels, rho, base_optimizer): # 第一次反向传播拿到当前参数处的梯度 model.zero_grad() loss loss_fn(model(inputs), labels) loss.backward() # 保存第一份梯度 first_grads [] for p in model.parameters(): if p.grad is not None: first_grads.append((p, p.grad.detach().clone())) else: first_grads.append((p, None)) # 计算整体梯度范数 grad_norm 0.0 for p, g in first_grads: if g is not None: grad_norm (g ** 2).sum() grad_norm grad_norm.sqrt() 1e-7 # 扰动参数梯度上升方向走 rho 步长 for p, g in first_grads: if g is None: continue e_w rho * g / grad_norm p.data.add_(e_w) # 第二次反向传播在扰动点计算梯度 loss2 loss_fn(model(inputs), labels) model.zero_grad() loss2.backward() # 保存第二份梯度作为真正更新用的梯度 sam_grads [] for p in model.parameters(): if p.grad is not None: sam_grads.append((p, p.grad.detach().clone())) else: sam_grads.append((p, None)) # 恢复原始参数 for p, g in first_grads: if g is None: continue e_w rho * g / grad_norm p.data.sub_(e_w) # 把扰动点梯度写回参数交给基础优化器更新 for p, g in sam_grads: if g is None: continue p.grad.copy_(g) base_optimizer.step() model.zero_grad()这段代码有三个容易写错的地方我逐一提醒。第一第二次反向传播之前要不要zero_grad要。如果不做第一次的梯度会和第二次的梯度累加在一起干扰扰动点的真实梯度。很多魔改版本说“官方没有显式 zero_grad”那是在他们自己封装的优化器里维护了清空逻辑但你从头实现时一定要显式清掉。第二恢复参数时必须用第一次保存的e_w不能重新用当前梯度计算。因为第二次反向传播之后p.grad已经变成了扰动点梯度再用它恢复会把参数挪错位置。第三grad_norm要把所有层的梯度拼接起来计算而不是逐层单独求模。代码里我用(g ** 2).sum()累加所有元素平方最后开根号这就是一个 batch 内所有参数的总体 L2 范数。平时我们做梯度裁剪时习惯按单层或全局范数SAM 这里必须全局因为扰动是加在整个参数空间里的。3.3 训练策略学习率、rho 与 batch size 的配合SAM 直接套到现有训练流程里经常会遇到“前几个 epoch loss 不降”的尴尬。这不是 SAM 坏了而是学习率和 rho 没匹配好。先给出一组我实测之后比较稳定的起点batch size128学习率0.05用 SGD 时比普通 SGD 的初始学习率略低或持平动量0.9weight decay0.0005rho0.05训练轮数90 或 120 都可以前提是开启学习率 warmup关键在于 rho 和 batch size 的关系。batch size 越小单次梯度估计噪声越大如果 rho 也大扰动方向会很不稳定。我建议 batch size 小于 64 时rho 从 0.02 起调不要直接上 0.1。batch size 在 128 到 256 时rho 在 0.05 附近通常表现不错。学习率方面SAM 需要比普通优化器更长的 warmup尤其是从随机初始化开始训练的场景。我的理解是训练初期参数还在快速移动loss 曲面很陡扰动过大会让更新方向变得过于激进所以前 10 到 20 个 epoch 用一个比较低的常数或线性增长学习率先把优化拖到稳定区域再说。数据增强方面不要删减。SAM 的好处在于“平坦极小值”对增强带来的输入扰动更不敏感所以裁剪、翻转、颜色抖动这些操作可以放心叠加和 SAM 是协同关系。4. 实验效果与调参心得4.1 精度和鲁棒性提升在哪里我记录了一组有代表性的实验数据。模型和数据集保持完全一致只换优化器。方案测试集准确率邻域 loss 差训练耗时基线 A带动量 SGD86.1%1.321.0x基线 B自适应矩估计算法85.4%1.581.0xSAM 带动量 SGD88.7%0.511.9xSAM 自适应矩估计算法87.9%0.441.9x这个表格想说明的不只是准确率涨了。更明显的是邻域 loss 差从 1.32 降到了 0.51 左右也就是参数附近的 loss 地形明显平缓了。测试集准确率提升 2 到 3 个百分点在这个量级的任务上已经算很可观的收益。训练耗时接近双倍这是 SAM 的天然代价毕竟要算两遍前向和反向。但注意这个倍数没有变成 2.5 倍以上因为第二次前向复用同一批数据数据加载和预处理可以省掉。实际用起来我建议把验证集评估从每个 epoch 都做改成每 2 到 3 个 epoch 做一次省下来的时间足够抵消一部分成本。另外值得记录的是SAM 带来的提升在训练后期更明显。前 30 个 epochSAM 和普通基线的验证集准确率差距很小甚至会略低。过了 60 个 epoch普通基线开始出现过拟合迹象SAM 的验证集曲线还在稳步往上。所以千万不要看到早期没优势就提前停药。4.2 rho 的选择经验与失败教训rho 是 SAM 最核心的超参数它控制“看多远”。rho 太小比如 0.01扰动带来的信息很弱SAM 退化成普通训练泛化提升几乎看不到。rho 太大比如 0.2扰动范围覆盖太多地形优化器的目标变成了“把整个盆地压平”结果往往欠拟合训练集 loss 就不太降得下去。我试过的方向大致是这样rho0.01与基线几乎无差异偶尔更好但不值得为它付出双倍训练时间。rho0.05多数视觉任务的甜点区准确率和鲁棒性有明显提升。rho0.1对噪声较大的数据集或小 batch 时会震荡需要调低学习率。rho0.2训练集 loss 偏高验证集也变差基本就是欠拟合。调参顺序我建议先固定 rho0.05把学习率调到能稳定收敛然后把 rho 往两个方向各试一组。如果 rho 增大后准确率提升说明原方案还有空间如果准确率立刻往下掉说明任务本身已经接近最优复杂度再强推平坦没有意义。还有一个容易忽略的问题rho 和 weight decay 之间存在关联。weight decay 会导致参数整体偏小在高维空间里梯度的范数也会变化。用过大的 rho 配合强的 weight decay扰动可能被过度放大。我们的经验是 weight decay 保持普通量级即可不要为了配合 SAM 无脑增加。5. 常见问题与排查技巧实录5.1 显存占用变高、训练时间翻倍怎么办最直接的代价之前说过SAM 每次更新要做两次前向和两次反向传播。显存方面第二次 backward 会增加中间激活值的缓存显存占用可能从原来的 3GB 涨到 4.5GB 左右。这个涨幅在很多项目里是可以接受的但如果你原本已经把显存顶满就需要想办法。我试过两种缓解方案。第一种是减小 batch size比如从 128 降到 96后面 epoch 再补回来。第二种是在第二次前向时用更小的输入分辨率但这会改变数据分布不推荐。如果显存还是不够还可以做一个折中的“梯度复用”版本第一次 backward 后不保存完整梯度张量只保存归一化后的扰动方向。这样省掉一部分梯度缓存但代价是第二次 backward 前需要保留计算图。工程上更痛的是第二次前向输出之后第一次反向传播的计算图要被释放否则内存翻倍。实际写代码时记得在第一次backward()后调用zero_grad(set_to_noneTrue)而不是简单设成 0。时间方面的优化思路是不要在全部训练过程中使用 SAM。可以选择前期用普通 SGD 快速收敛最后几十个 epoch 再切到 SAM 精调。这种方法不是完整 SAM但能在泛化提升和训练成本之间找一个折中点。我们后来在另一个项目里就是这么做的牺牲了一部分提升幅度但省下了一半时间。5.2 收敛变慢、loss 震荡怎么排查如果你把 SAM 扔到一个已有项目里发现前几个 epoch 训练 loss 不下降先别怀疑 SAM 原理。最可能的原因是学习率太高或者 rho 太大。我的排查顺序是这样的第一先把 rho 调到 0.01跑 10 个 epoch。如果 loss 还不降那是学习率或其他代码问题和 SAM 无关。如果 loss 降了说明 rho 太大把它慢慢调大。第二检查第二次反向传播前是否清空了梯度。这个问题很隐蔽。不清空的话梯度是两次反向的叠加更新方向和真正的扰动点梯度差很远loss 会非常奇怪地波动。如果你发现 loss 曲线在一个 epoch 内剧烈抖动、且抖动周期和 batch 数量完全对应多半就是这个原因。第三检查动量和 EMA 类的状态变量有没有被“污染”。SAM 在恢复参数时只恢复了param.data但基础优化器内部保存的动量缓冲区仍然是按之前参数路径累计的。如果内层用了复杂的状态变量某些实现会在第二次前向时更新了这些状态导致恢复后状态异常。稳妥的做法是在 SAM 内部调用基础优化器的step()之前不要改动优化器内部 state。第四注意 BatchNorm 的问题。第二次前向和第二次反向是在扰动后参数上进行的如果模型里有 BatchNorm它的 running mean 和 running variance 也会因为这次额外前向被更新一次。同一批数据被用来更新统计量两次会导致统计量偏移。有的实现会在第二次前向时把 BatchNorm 切到 eval 模式来避免有的则直接接受这个偏差。我在实验里发现对比较深的结构把 BatchNorm 冻结在训练统计之外效果更稳。具体的做法是把模型里所有 BatchNorm 层临时设成track_running_statsFalse跑完 SAM 步骤再恢复。5.3 什么时候不建议用 SAMSAM 不是万灵药有几个场景我会明确劝阻。第一个场景是训练轮数非常短比如只跑 5 到 10 个 epoch。SAM 的收益主要来自训练后期的平坦化短训练里它还没来得及发挥优势反而会因双倍计算拖慢迭代速度。第二个场景是数据量非常充足、且测试集分布和训练集基本一致。这种情况下普通优化器已经能泛化得很好SAM 提升空间有限纯属白花时间。第三个场景是模型本身非常大训练一轮本来就久。直接全量用 SAM训练成本可能翻倍但收益不一定等比例。优先考虑知识蒸馏、EMA 这类成本更低的方案。我个人的体会是SAM 真正适合“训练集和测试集存在差异、模型容易记住噪声”的场景比如小数据集图像分类、样本类别不平衡、跨场景目标检测。在这些项目里SAM 带来的泛化提升比很多花哨的网络结构调整更直接。如果你准备在自己的项目里尝试我的最后一个小建议是第一次跑 SAM 时不要追求一次到位先小模型、小数据集、短训练量验证代码正确性再搬到真实任务上。代码层面最容易出问题的地方就在第二次反向传播的梯度是否被正确隔离。只要这一步处理干净后续调参基本就是 rho 和学习率之间的事。