1. 为什么平坦极小值比尖锐极小值更值得追求1.1 从一个反直觉的实验现象说起如果你训练过深度神经网络大概率遇到过这种让人挠头的情况训练集上的损失已经降到接近零准确率几乎百分之百但把模型拿到测试集上一跑效果却差了一大截。更气人的是换一个随机种子重新训练同样的网络结构、同样的数据测试精度可能相差好几个百分点。这种“训练满分、考试不及格”的现象就是泛化差距。传统做法通常把矛头指向过拟合然后搬出正则化、Dropout、数据增强、早停这一套组合拳。这些方法确实有用但它们更多是在“限制模型容量”或“增加数据多样性”的层面做文章。而有一类方法换了个视角不去改模型结构也不去加数据而是去改变损失曲面的几何形状让优化过程主动寻找那些“平坦”的极小值点。这个思路的核心洞察非常朴素如果一个参数点附近的损失曲面很平坦那么当测试数据的分布和训练数据有轻微偏移时损失不会突然飙升反过来如果极小值非常尖锐参数稍微一动损失就爆炸测试表现自然脆弱。你可以把它想象成在山谷里放一个球——平坦的谷底球滚来滚去损失变化不大尖锐的V型谷底球稍微偏一点就冲到陡坡上去了。1.2 尖锐度与泛化的数学直觉把“尖锐”和“平坦”量化最直接的方式是看损失函数在参数点附近的曲率。对于参数向量 w损失 L(w)如果在 w 附近做一个半径为 ρ 的扰动 ε那么最坏情况下的损失增量可以写成max_{||ε||≤ρ} L(wε) - L(w)这个量越大说明该点越尖锐。Sharpness-Aware Minimization简称 SAM要做的就是把这个最坏情况下的损失也压下去而不是只压中心点的损失。换句话说SAM 的优化目标从“让 L(w) 尽量小”变成了“让 w 附近一个球内的最大损失尽量小”。这个转变看似微小但影响深远。它意味着优化器不再只盯着一个点而是带着“鲁棒性”的意识去搜索参数空间。从数学上看SAM 的目标函数可以写成min_w max_{||ε||≤ρ} L(wε)这是一个典型的极小极大问题。内层是找当前点附近最坏扰动方向外层是让这个最坏情况下的损失最小。这种形式天然地鼓励模型落到平坦区域因为只有平坦区域才能让“最坏情况”和“平均情况”差距不大。1.3 为什么这个思路值得单独拿出来讲市面上提升泛化的方法很多SAM 之所以值得深入拆解有三个原因。第一它几乎不挑模型结构和任务类型图像分类、目标检测、语义分割、甚至一些自然语言处理任务都能用属于通用型工具。第二它的额外计算开销可控通过近似和并行技巧训练时间通常只增加百分之几十而不是翻倍。第三它的实现门槛低核心代码往往只有十几行不需要改动网络定义也不需要重新设计损失函数。我最初接触 SAM 是在一个图像分类项目里当时模型在验证集上的波动特别大同一个配置跑三次最好和最差能差三个点。试过调学习率、加权重衰减、换优化器效果都不稳定。后来把优化器换成 SAM验证集方差明显收窄最差的那次也比原来最好的那次高。这个经历让我意识到有些泛化问题不是靠“调参”能解决的而是优化目标本身需要换一个写法。2. SAM 的核心机制与实现拆解2.1 内层最大化一步梯度上升找“最坏扰动”SAM 的每一步迭代分两个阶段。第一阶段是找扰动 ε。给定当前参数 w计算损失 L(w) 对 w 的梯度 g ∇L(w)。然后沿着梯度方向走一小步得到扰动后的参数ε ρ * g / ||g||这里 ρ 是一个超参数控制扰动半径。注意这个扰动不是随便加的而是沿着梯度方向——因为梯度方向是损失上升最快的方向所以这个扰动近似于“在半径为 ρ 的球内找到让损失最大的方向”。这一步不需要精确求解内层最大化问题用一阶近似就够了这也是 SAM 高效的关键。实际操作中这个扰动通常施加在权重上而不是所有参数上。对于 BatchNorm 层里的缩放和偏移参数一般不做扰动因为它们的尺度不同扰动会破坏归一化效果。这个细节在论文里可能一笔带过但在代码实现时如果搞错效果会打折扣。2.2 外层最小化在扰动点上更新原始参数找到扰动 ε 之后第二阶段是在扰动后的参数 wε 上重新计算梯度然后用这个梯度去更新原始参数 w。注意更新的是 w不是 wε。也就是说SAM 用“最坏情况下的梯度”来指导原始参数的更新方向。这个设计非常巧妙。如果只用 w 处的梯度更新那就是普通 SGD如果用 wε 处的梯度更新 w相当于让 w 朝着“能同时降低 w 和 wε 损失”的方向移动。由于 wε 是当前点附近最坏的情况所以这个更新方向天然地倾向于把 w 拉向平坦区域。用伪代码表示就是# 第一次前向传播和反向传播 loss1 criterion(model(x), y) loss1.backward() # 保存原始梯度 g [p.grad.clone() for p in model.parameters()] # 计算扰动并施加 for p, gi in zip(model.parameters(), g): p.data.add_(gi, alpharho / (norm(g) 1e-12)) # 清零梯度准备第二次传播 optimizer.zero_grad() # 第二次前向传播和反向传播 loss2 criterion(model(x), y) loss2.backward() # 恢复原始参数 for p, gi in zip(model.parameters(), g): p.data.sub_(gi, alpharho / (norm(g) 1e-12)) # 用第二次的梯度更新 optimizer.step()这段代码里有两个容易踩坑的地方。第一扰动施加后必须恢复原始参数否则参数会被永久改变。第二第二次反向传播得到的梯度是累积在 p.grad 上的所以第一次的梯度必须先保存再清零。很多开源实现用了一个小技巧把第一次的梯度存下来第二次反向传播后直接覆盖这样就不需要手动清零。2.3 与 SGD、Adam 的关系和差异SAM 不是一个独立的优化器它更像是一个“优化器包装器”。你可以把 SAM 套在 SGD 外面也可以套在 Adam 外面。区别在于内层扰动和外层更新用的基础优化器可以不同。实践中常见两种组合SAMSGD 和 SAMAdam。前者在图像分类任务上表现稳定后者在训练初期收敛更快适合大模型或难训练的任务。和普通 SGD 相比SAM 每步多了一次前向传播和一次反向传播计算量大约增加一倍。但通过一些工程优化比如复用第一次前向传播的中间结果、用混合精度训练、或者只在部分层上施加扰动可以把额外开销压到百分之三十到五十。这个代价换来的泛化提升在很多任务上是划算的。和 Adam 相比SAM 的优势在于它不依赖自适应学习率的精细调节。Adam 对学习率、beta 参数比较敏感调不好容易震荡SAM 的 ρ 参数相对鲁棒在一个合理范围内比如 0.01 到 0.1都能工作只是效果略有差异。3. 从零实现一个可用的 SAM 优化器3.1 参数选择ρ 到底设多少合适ρ 是 SAM 最重要的超参数它决定了扰动半径。ρ 太小扰动可以忽略SAM 退化成普通 SGDρ 太大扰动会破坏训练稳定性损失震荡甚至发散。根据我的实测经验ρ 的取值和任务、模型规模、batch size 都有关系。对于图像分类任务ResNet 系列模型batch size 在 128 到 256 之间时ρ 取 0.05 左右比较稳妥。如果模型更大比如 ViT 或 Swin Transformerρ 可以适当减小到 0.02 到 0.03因为大模型对扰动更敏感。如果 batch size 很小比如 32 或 64ρ 也要相应减小否则扰动方向噪声太大。有一个经验公式可以参考ρ 的合理范围和 1/sqrt(batch_size) 成正比。也就是说batch size 翻倍ρ 可以乘以 0.7 左右。这个关系不是严格的数学定理而是从梯度噪声尺度推导出来的近似。我在几个不同 batch size 的设置下验证过基本符合这个趋势。另外ρ 可以随训练过程动态调整。训练初期损失曲面比较陡峭ρ 可以小一点训练后期接近收敛时适当增大 ρ 有助于跳出尖锐区域。不过这个策略需要额外调参如果不想增加复杂度固定 ρ 也能工作。3.2 代码实现一个最小可用的 SAM 类下面是一个基于 PyTorch 的最小实现可以直接嵌入现有训练流程import torch class SAM(torch.optim.Optimizer): def __init__(self, params, base_optimizer, rho0.05, adaptiveFalse, **kwargs): assert rho 0.0, fInvalid rho, should be non-negative: {rho} defaults dict(rhorho, adaptiveadaptive, **kwargs) super(SAM, self).__init__(params, defaults) self.base_optimizer base_optimizer(self.param_groups, **kwargs) self.param_groups self.base_optimizer.param_groups self.defaults.update(self.base_optimizer.defaults) torch.no_grad() def first_step(self, zero_gradFalse): grad_norm self._grad_norm() for group in self.param_groups: scale group[rho] / (grad_norm 1e-12) for p in group[params]: if p.grad is None: continue self.state[p][old_p] p.data.clone() e_w (torch.pow(p, 2) if group[adaptive] else 1.0) * p.grad * scale.to(p) p.add_(e_w) if zero_grad: self.zero_grad() torch.no_grad() def second_step(self, zero_gradFalse): for group in self.param_groups: for p in group[params]: if p.grad is None: continue p.data self.state[p][old_p] self.base_optimizer.step() if zero_grad: self.zero_grad() def _grad_norm(self): shared_device self.param_groups[0][params][0].device norm torch.norm( torch.stack([ ((torch.abs(p) if group[adaptive] else 1.0) * p.grad).norm(p2).to(shared_device) for group in self.param_groups for p in group[params] if p.grad is not None ]), p2 ) return norm def load_state_dict(self, state_dict): super().load_state_dict(state_dict) self.base_optimizer.param_groups self.param_groups这个实现里有一个adaptive选项对应的是 ASAMAdaptive SAM。当 adaptiveTrue 时扰动会根据参数本身的尺度进行缩放适合那些参数尺度差异很大的模型比如包含 BatchNorm 或 LayerNorm 的网络。对于普通 CNNadaptiveFalse 就够了。使用方式也很简单base_optimizer torch.optim.SGD optimizer SAM(model.parameters(), base_optimizer, rho0.05, lr0.1, momentum0.9) for x, y in dataloader: # 第一次前向传播 loss criterion(model(x), y) loss.backward() optimizer.first_step(zero_gradTrue) # 第二次前向传播 criterion(model(x), y).backward() optimizer.second_step(zero_gradTrue)注意first_step和second_step之间的顺序不能颠倒而且第二次前向传播必须重新计算不能复用第一次的结果因为参数已经被扰动了。3.3 与混合精度训练的配合如果你用 AMP自动混合精度训练SAM 的实现需要稍微调整。因为第一次反向传播得到的梯度是 FP16 格式直接用来计算扰动可能会溢出。解决办法是在第一次反向传播后把梯度转成 FP32 再计算扰动施加扰动后再转回 FP16 做第二次前向传播。另外GradScaler 的使用也要注意。第一次反向传播后不要立即调用 scaler.step()而是先做扰动第二次反向传播后再调用 scaler.step() 和 scaler.update()。这个流程如果搞错会出现梯度缩放不一致的问题导致训练不稳定。我在实际项目里踩过这个坑一开始直接把 SAM 套在 AMP 外面结果训练损失忽高忽低后来发现是 GradScaler 在两次反向传播之间做了不该做的缩放。修正之后训练曲线就平滑了。4. 实战中的调参经验与避坑指南4.1 学习率需要重新调整SAM 的外层更新用的是扰动后的梯度这个梯度的尺度和普通梯度不一样。实测下来SAM 的最佳学习率通常比普通 SGD 小一些大约是原来的 0.5 到 0.8 倍。如果你直接从 SGD 切换到 SAM 而不调学习率可能会发现训练初期损失下降很慢甚至震荡。举个例子在一个图像分类任务上普通 SGD 用 lr0.1 效果很好换成 SAM 后 lr0.1 也能跑但验证集精度比 lr0.05 低了将近一个点。后来把学习率降到 0.05再配合余弦退火效果就超过了普通 SGD 的基线。另外权重衰减也需要微调。SAM 本身有正则化效果所以权重衰减可以适当减小否则双重正则化会让模型欠拟合。我通常会把权重衰减减半比如从 5e-4 降到 2.5e-4。4.2 BatchNorm 层的处理细节前面提到过BatchNorm 的缩放和偏移参数通常不做扰动。但在代码实现时怎么识别这些参数是个问题。一种做法是按参数名称过滤包含 bn 或 norm 的层跳过另一种做法是按参数维度过滤一维参数偏置和缩放跳过二维参数卷积核和全连接权重才施加扰动。这两种做法我都试过效果差别不大但按名称过滤更稳妥因为有些自定义层可能不遵循常规命名。如果你不确定可以先把所有参数都加上扰动跑一遍再对比只扰动权重的情况通常后者会好一点。还有一个细节BatchNorm 在训练和推理时的行为不同。SAM 的两次前向传播都在训练模式下进行所以 BatchNorm 的统计量会被更新两次。这本身不是问题但如果你用了同步 BatchNorm 或者分布式训练需要确保两次前向传播的统计量同步逻辑正确。4.3 训练轮数和早停策略SAM 收敛通常比普通 SGD 慢一些因为每步的有效更新幅度被“最坏情况”约束了。所以训练轮数可能需要增加百分之二十到三十。如果你原来训练 100 个 epoch换成 SAM 后可能需要 120 到 130 个 epoch 才能达到同样的训练损失。但好消息是SAM 的验证集曲线通常更平滑过拟合出现得更晚。这意味着你可以用更宽松的早停策略比如验证集损失连续 20 个 epoch 不下降才停止而不是原来的 10 个。这样能充分利用 SAM 的泛化优势。我在一个语义分割任务上做过对比普通 SGD 在第 80 个 epoch 左右验证集精度达到峰值然后开始下降SAM 在第 110 个 epoch 才达到峰值而且峰值比 SGD 高了 1.5 个点。虽然训练时间长了但最终模型质量提升明显。4.4 常见问题速查表问题现象可能原因解决方法训练损失震荡严重ρ 太大或学习率太高减小 ρ 到 0.02-0.03学习率乘以 0.5验证集精度不如普通 SGD学习率未调整或权重衰减过大降低学习率减小权重衰减训练速度太慢两次前向传播开销大使用混合精度或只在部分层施加扰动训练初期损失下降慢SAM 收敛特性增加 warmup 轮数或初期用较小 ρ显存不足两次前向传播需要保存中间结果减小 batch size或用梯度检查点分布式训练不同步扰动计算未跨卡同步确保 grad_norm 在所有卡上做 all_reduce这个表里的问题我都实际遇到过尤其是显存不足和分布式同步这两个。显存问题可以通过梯度检查点缓解但会增加计算时间分布式同步问题需要在计算 grad_norm 时调用 all_reduce否则每张卡上的扰动方向不一致训练会发散。5. SAM 的变体与扩展方向5.1 ASAM自适应尖锐度感知ASAM 的核心改进是让扰动幅度与参数尺度成正比。具体来说扰动从 ε ρ * g / ||g|| 变成 ε ρ * |w| * g / |||w| * g||。这样做的理由是不同参数的尺度差异很大统一扰动半径可能对某些参数过大对另一些过小。ASAM 在包含大量归一化层的模型上表现更好比如 Transformer 系列。实现上ASAM 只需要在计算扰动时乘以 |w|代码改动很小。但效果提升在有些任务上很明显尤其是那些参数分布跨度大的模型。我在一个目标检测任务上对比过ASAM 比原始 SAM 的 mAP 高了 0.8 个点训练稳定性也更好。5.2 与数据增强的协同SAM 和数据增强是互补的。数据增强在输入空间增加多样性SAM 在参数空间寻找平坦区域。两者结合通常能获得比单独使用更好的效果。但要注意强数据增强比如 RandAugment、MixUp、CutMix本身就会让损失曲面变得更复杂此时 ρ 需要适当减小否则扰动会落在不合理的区域。我的经验是如果用了 MixUp 或 CutMixρ 取 0.02 到 0.03 比较合适如果只用随机裁剪和翻转ρ 可以取 0.05。另外数据增强的强度可以和 ρ 一起调通常先固定数据增强策略再调 ρ这样更容易找到最优组合。5.3 在自监督学习和微调中的应用SAM 在自监督学习里的应用也很有意思。自监督任务通常没有标签损失函数是对比损失或重建损失这些损失的曲面和分类损失不同。有实验表明在 SimCLR 或 MoCo 这类框架里加入 SAM可以提升下游任务的微调效果。原因可能是自监督学习更容易陷入尖锐的局部极小值而 SAM 帮助找到了更平坦的区域。在微调场景下SAM 也有用武之地。当你从一个预训练模型出发在少量数据上微调时过拟合风险很高。SAM 可以作为一种正则化手段让微调后的模型保持更好的泛化能力。不过微调时 ρ 要设得小一些因为预训练模型的参数已经在一个比较好的区域扰动太大会破坏已有知识。6. 我个人在实际操作中的几点体会第一次把 SAM 用到生产项目里时我犯了一个低级错误忘了在第二次前向传播前恢复原始参数。结果训练了十几个 epoch损失一直不降排查了半天才发现参数被永久扰动了。这个教训让我养成了一个习惯每次实现新的优化器先用一个极小的数据集跑通确认损失能正常下降再上大规模训练。另一个体会是SAM 不是万能的。如果你的模型本身容量不足或者数据质量很差SAM 也救不了。它解决的是“优化过程导致的泛化差距”而不是“模型表达能力不足”或“数据分布偏移”的问题。所以在上 SAM 之前先确认基线模型没有明显的欠拟合或数据问题。最后分享一个小技巧如果你不想改动训练代码可以先用一个简单的“梯度惩罚”替代 SAM 做快速验证。具体来说在损失里加一项 λ * ||∇L(w)||^2这个惩罚项也会鼓励平坦极小值虽然效果不如 SAM但实现更简单。如果梯度惩罚能带来提升再上 SAM 通常会有更好的效果。这个思路帮我在几个项目里快速判断了是否值得引入 SAM节省了不少实验时间。