1. 为什么一个找平坦山谷的想法能提升泛化能力如果你训练过深度神经网络大概率遇到过这种让人抓狂的情况训练集上的损失已经降到很低了准确率也刷到了99%以上但一换到验证集或者真实场景的数据上效果就明显掉一截。你调了学习率、加了正则化、换了数据增强策略甚至把模型结构改小了一圈泛化差距依然顽固地存在。这个问题困扰了很多人而有一类方法专门从优化过程本身入手去缓解它其中比较有代表性的思路就是锐度感知最小化Sharpness-Aware Minimization简称SAM。我第一次接触这个思路的时候直觉上觉得它有点反常识。传统的训练目标很直接找到一组参数让训练损失尽可能小。但SAM说光让损失小还不够你还得让这组参数周围的损失曲面尽可能平坦。换句话说它不满足于找到一个谷底的点而是希望找到一个宽阔平坦的谷底区域。这个区别听起来很微妙但实际效果在不少任务上相当明显。这篇文章我会从从业者的视角把SAM这套方法拆开讲清楚它到底在优化什么、为什么平坦解通常泛化更好、具体怎么实现、实操中有哪些坑、以及它和对抗训练、随机权重平均这些方法的区别在哪里。不管你是刚入门想搞懂原理还是已经准备在自己的项目里试一把我都尽量把该说的细节说透。先给一个最直观的类比。想象你在山里找营地一种策略是找到某个局部最低的坑就扎营但这个坑可能很窄稍微挪一步就爬坡了另一种策略是找一个又宽又平的大草甸即使你搭帐篷的位置不是绝对最低点周围一大片区域都很平缓你随便走动都不会掉下去。SAM追求的就是后者。在机器学习里随便走动都不会掉下去意味着参数即使因为数据分布偏移、噪声扰动而轻微变化损失也不会剧烈上升这正好对应了更好的泛化能力。2. 平坦最小值与泛化之间的那层关系2.1 损失曲面的几何直觉要理解SAM得先接受一个前提损失函数的曲面形状和模型的泛化能力是有关联的。这个观点不是SAM首创的早些年就有研究者在探索平坦最小值和尖锐最小值的区别。所谓尖锐最小值指的是参数在某个点附近损失很低但只要参数稍微扰动一下损失就迅速飙升平坦最小值则相反参数在一个较大邻域内损失都保持在较低水平。为什么平坦解更有利于泛化一个被广泛接受的解释是训练数据和测试数据之间存在分布差异你训练时找到的最优点放到测试分布上未必还是最优点。如果你的解处于一个尖锐的谷底那么训练分布和测试分布之间的微小偏移就可能让测试损失大幅上升而如果你的解处于一个平坦区域同样的分布偏移带来的损失变化就小得多。这就像在窄坑里和宽草甸上分别放一个球地面轻微倾斜时窄坑里的球容易滚出去宽草甸上的球则稳得多。这里要强调一点平坦和尖锐并不是绝对的二元对立而是一个连续谱。实际中我们很难精确度量平坦程度因为损失曲面的维度极高你没法可视化。所以SAM的做法不是去直接测量平坦度而是通过一个巧妙的优化目标间接地把模型往平坦区域推。2.2 从数学目标看SAM在优化什么传统经验风险最小化ERM的目标很简单就是最小化训练集上的平均损失。SAM把它改成了一个带扰动的形式核心思想是我不要求当前参数点的损失低我要求当前参数点邻域内最坏情况下的损失也低。用公式表达SAM的优化目标是让参数w在半径为ρ的邻域内损失的最大值尽可能小。这个邻域内最大损失就是所谓的锐度sharpness。如果某个点的邻域内最大损失和该点损失差不多说明这个区域很平如果邻域内最大损失远高于该点损失说明这个点很尖。SAM要做的就是同时压低损失值和锐度。这个目标在数学上可以写成一个min-max问题外层是min找让最坏损失最小的参数内层是max在参数邻域内找让损失最大的扰动方向。求解的时候内层的max通常用一阶近似来估计也就是沿着梯度方向走一小步因为梯度方向正是损失上升最快的方向。这就引出了SAM最核心的两步操作我在下一节会详细拆。2.3 为什么这个目标和对抗训练不是一回事很多人第一次看到SAM会联想到对抗训练Adversarial Training因为两者都涉及扰动和最坏情况。但它们的出发点和操作对象完全不同这点必须分清楚否则容易用错场景。对抗训练扰动的是输入数据它要找的是让模型预测出错最大的输入扰动然后把扰动后的样本喂给模型训练目的是提升模型对恶意输入的鲁棒性。而SAM扰动的是模型参数它要找的是让损失上升最快的参数扰动方向目的是让优化过程倾向于平坦解从而提升泛化。一个在输入空间做文章一个在参数空间做文章虽然数学形式上有相似之处但解决的问题不一样。实操中我见过有人把两者混用结果调参调得一头雾水。我的建议是如果你的目标是提升模型在干净测试集上的泛化表现优先考虑SAM如果你的目标是防御对抗样本攻击那对抗训练更对口。当然两者也可以结合但那是另一个话题了。3. SAM的两步走实现机制拆解3.1 第一步沿梯度方向做一次试探性扰动SAM的一次参数更新实际上包含两次前向和两次反向传播这也是它比普通训练慢一些的原因。第一步是计算当前参数下的梯度然后沿着这个梯度的方向把参数临时推出去一小步推的距离由超参数ρ控制。具体来说假设当前参数是w当前batch的损失梯度是g那么扰动后的参数是w ρ * g / ||g||这里的归一化是为了让扰动幅度只由ρ决定不受梯度大小影响。这一步的意图很明确梯度方向是损失上升最快的方向我沿着这个方向走一小步就到达了当前邻域内损失较高的那个点。这个点就是内层max要找的最坏情况的近似。这里有个细节值得注意扰动是在当前batch上计算的不是在整个训练集上。这意味着每次迭代的扰动方向都带有随机性因为batch不同梯度就不同。这种随机性其实是有好处的它让模型不会针对某个固定的扰动方向过拟合相当于一种隐式的正则。ρ这个超参数的取值很关键。太小了扰动几乎没效果SAM退化成普通训练太大了扰动可能把参数推到完全无关的区域反而破坏训练稳定性。经验上ρ通常取0.01到0.2之间具体要看任务和模型规模。我在图像分类任务上试过ρ取0.05左右比较稳在一些小规模任务上ρ取0.01到0.02就够了。3.2 第二步在扰动点上重新计算梯度并更新第一步只是试探并没有真正更新参数。真正的更新发生在第二步在扰动后的参数w ρ * g / ||g||上重新做一次前向和反向传播得到新的梯度g然后用这个g去更新原始参数w。为什么要用扰动点的梯度而不是原始点的梯度因为扰动点代表了邻域内损失较高的位置用这个位置的梯度来更新相当于在告诉优化器别只盯着当前这个低点你得考虑周围的高点把周围也压下去。这样更新出来的参数自然会更倾向于落在平坦区域。整个流程可以概括成下面这几步我用伪代码的形式写出来方便你对照实现# 假设 model 是模型loss_fn 是损失函数optimizer 是优化器 # rho 是邻域半径超参数 # 第一步计算原始梯度 loss1 loss_fn(model(inputs), targets) loss1.backward() # 保存原始梯度 grad_norm compute_grad_norm(model) # 计算梯度的范数 # 沿梯度方向扰动参数 perturb_parameters(model, rho, grad_norm) # 第二步在扰动后的参数上重新计算梯度 optimizer.zero_grad() loss2 loss_fn(model(inputs), targets) loss2.backward() # 恢复原始参数用扰动点的梯度更新 restore_parameters(model) optimizer.step()这段伪代码里有两个容易出错的点。第一扰动参数后一定要记得恢复原始参数否则你更新的就不是原始参数了。第二两次前向传播用的是同一个batch不要中途换数据否则扰动点和原始点对应的数据不一致逻辑就乱了。3.3 计算开销与显存占用的真实情况SAM最被人诟病的一点就是慢。普通训练一次迭代一次前向一次反向SAM要两次前向两次反向理论上时间开销接近翻倍。实际测下来在GPU上大概比普通训练慢60%到80%因为第二次前向可以和第一次的部分计算复用但反向传播基本是实打实多了一遍。显存方面SAM需要额外保存一份原始参数的副本用于扰动后恢复。对于大模型来说这意味着显存占用会增加一份参数量的开销。如果你本来就在显存边缘徘徊加SAM可能直接OOM。我的建议是如果显存紧张可以配合梯度累积或者混合精度训练来缓解或者把batch size适当调小。不过话说回来SAM带来的泛化提升在很多任务上是实打实的多花的时间往往能换来更好的最终效果甚至有时候能省掉一些调参和正则化的功夫。所以这笔账划不划算得看你的具体场景。如果训练成本不是瓶颈追求极致泛化SAM值得一试如果训练资源极其紧张那可能得权衡一下。4. 把SAM接进现有训练流程的实操细节4.1 优化器层面的改动最小化原则SAM本质上是一个优化策略它不改变模型结构也不改变损失函数只是在参数更新那一步做了文章。所以把它接进现有流程改动量其实很小。你不需要重写整个训练循环只需要把原来的loss.backward(); optimizer.step()替换成SAM的两步逻辑。我一般会封装一个SAM优化器类把原始优化器作为参数传进去对外暴露和普通优化器一样的step()接口。这样训练代码几乎不用改只是把优化器换一下。下面是一个简化的实现思路class SAMOptimizer: def __init__(self, base_optimizer, rho0.05): self.base_optimizer base_optimizer self.rho rho def step(self): # 第一步的梯度已经在外部计算好了 grad_norm self._compute_grad_norm() # 保存当前参数 self._save_parameters() # 扰动参数 self._perturb(rhoself.rho, grad_normgrad_norm) def second_step(self): # 在扰动后的参数上外部会重新计算梯度 # 恢复原始参数 self._restore_parameters() # 用扰动点的梯度更新 self.base_optimizer.step()这种封装方式的好处是你原来的学习率调度、权重衰减、动量等设置都能继续用因为它们都在base_optimizer里。SAM只是在外层套了一个扰动逻辑不干扰原有的优化器行为。4.2 学习率和ρ的联合调节经验SAM对学习率比普通训练更敏感一些。因为它的更新方向来自扰动点梯度的大小和方向都可能和原始点不同如果学习率设得太大容易震荡。我的经验是用SAM时学习率可以比普通训练稍微调小一点或者至少不要一上来就用很大的学习率。ρ和学习率之间也有耦合。ρ大意味着扰动范围大如果学习率也大两次更新的叠加效果可能过猛。反过来ρ小的时候学习率可以相对宽松一些。我通常的做法是先把ρ固定在一个中等值比如0.05然后调学习率等学习率稳定了再微调ρ。如果发现训练损失下降很慢可能是ρ太大导致扰动过强如果发现SAM和普通训练效果差不多可能是ρ太小扰动没起作用。还有一个技巧是对ρ做warmup。训练初期参数离最优点很远这时候大扰动可能有害训练后期参数接近收敛适当增大扰动有助于找到更平坦的解。有些实现会在训练前期用较小的ρ后期逐渐增大实测下来对稳定性有帮助。4.3 batch size对扰动估计的影响前面提到SAM的扰动方向是在当前batch上估计的。batch size越小梯度的随机性越大扰动方向就越噪声化。这未必是坏事因为噪声本身有正则效果但如果batch太小扰动方向可能完全偏离真实的最坏方向SAM的效果就会打折扣。我的建议是用SAM时batch size不要设得太小。如果显存允许尽量用和普通训练相当或稍大的batch size。如果只能用很小的batch可以考虑用梯度累积来模拟大batch的效果或者适当减小ρ让扰动更保守一些。另外有些实现会用多个batch的梯度平均来估计扰动方向这样估计更准但计算开销也更大。对于大多数场景单batch估计已经够用了不必过度追求精确。5. 哪些场景下SAM真正能打哪些场景要谨慎5.1 图像分类与分布偏移场景的实测表现SAM最早引起关注就是在图像分类任务上。在CIFAR、ImageNet这类基准上配合合适的训练策略SAM通常能带来一到两个百分点的准确率提升尤其是在分布偏移distribution shift的测试集上提升更明显。所谓分布偏移就是测试数据的分布和训练数据不完全一致比如训练用的是晴天照片测试用的是阴天照片。这种场景下平坦解的鲁棒性优势就体现出来了。我在一个图像分类项目里试过SAM训练集准确率基本没变但验证集准确率提升了大概1.5个百分点而且验证损失的曲线更平滑过拟合的迹象明显减轻。这个提升幅度不算惊天动地但考虑到改动量很小性价比还是不错的。5.2 标签噪声与长尾分布下的意外收获除了常规的泛化提升SAM在一些困难场景下表现更亮眼。比如标签噪声场景训练数据里有相当比例的标签是错的。普通训练容易去拟合这些错误标签导致泛化变差SAM因为追求平坦解对个别噪声样本的敏感度更低所以抗噪能力更强。有研究显示在标签噪声比例较高时SAM相对普通训练的优势会扩大。长尾分布也是类似。长尾分布指的是少数类别样本很多多数类别样本很少。普通训练容易在头部类别上过拟合尾部类别学不好。SAM的平坦性偏好有助于缓解这种不平衡带来的过拟合让模型在尾部类别上也有不错的表现。如果你手头的任务正好有这些特点SAM值得优先考虑。5.3 大模型微调与计算预算的权衡到了大模型时代SAM的定位有点微妙。一方面大模型微调时泛化问题同样存在SAM理论上也能帮上忙另一方面大模型本身训练成本就高SAM再翻倍开销就很可观了。所以在大模型场景下是否用SAM要仔细算账。一种折中方案是只在微调的后期用SAM前期用普通训练快速收敛后期用SAM精调平坦度。这样既控制了总开销又能享受到SAM带来的泛化收益。另一种方案是只在部分层上应用SAM比如只对分类头或者最后几层做扰动减少计算量。这些变体在实操中都有应用效果因任务而异需要自己试。6. 和随机权重平均、对抗训练这些方法的对比与组合6.1 SAM与随机权重平均的异同随机权重平均SWA是另一个提升泛化的常用技巧它和SAM有相似的目标但手段不同。SWA的做法是在训练后期把多个不同时刻的模型参数做平均得到一个平均模型。这个平均模型往往落在损失曲面的平坦区域因为平均操作本身就有平滑效果。SAM和SWA的区别在于SAM是在训练过程中主动把参数往平坦区域推SWA是在训练结束后对参数做平均。两者可以结合使用先SAM训练再SWA平均效果有时比单用其中一个更好。不过结合使用会增加实现复杂度如果你的项目对简洁性有要求可以先单独试SAM看效果再决定要不要叠加SWA。6.2 与对抗训练的取舍逻辑前面说过对抗训练扰动输入SAM扰动参数。如果你的任务同时面临对抗攻击和泛化不足两个问题理论上可以两者结合但计算开销会进一步增加。我的建议是分清楚主次如果对抗鲁棒性是刚需那对抗训练是基础SAM可以作为补充如果只是想要更好的泛化SAM更直接没必要引入对抗训练的复杂度。还有一个实际考虑是对抗训练对超参数非常敏感调起来很费劲SAM相对好调一些主要就是ρ和学习率两个旋钮。从工程落地的角度SAM的上手成本更低。6.3 组合使用的注意事项如果你决定把SAM和其他正则化手段权重衰减、Dropout、数据增强一起用要注意它们之间可能有叠加效应。比如权重衰减本身就在限制参数范数SAM的扰动也是在参数空间做文章两者叠加可能让有效扰动幅度变小。这时候可以适当增大ρ或者减小权重衰减系数让两者平衡。数据增强和SAM的关系则比较友好两者都在提升泛化叠加使用通常不会冲突反而可能互相增强。Dropout和SAM的叠加需要留意因为Dropout本身引入了随机性SAM的扰动估计会受影响可能需要在扰动计算时把Dropout关掉或者接受一定的噪声。7. 我在实操中踩过的几个坑7.1 扰动后忘记恢复参数导致的诡异现象这个坑我踩过一次印象很深。当时实现SAM的时候扰动参数后直接做了第二次前向和反向然后调用了优化器的step但忘了把参数恢复到扰动前的状态。结果就是模型实际上是在扰动后的参数基础上更新的相当于每次迭代都额外加了一个随机扰动训练损失一直下不去我还以为是ρ设太大了调了半天才发现是恢复逻辑漏了。这个错误的隐蔽性在于训练不会报错只是效果差。所以实现SAM时一定要确保扰动和恢复是成对出现的最好在代码里加个断言或者日志确认参数在第二次前向之前已经恢复。7.2 ρ设得过大导致训练不收敛ρ这个参数我一开始按论文里的推荐值设结果在自己的任务上训练直接不收敛损失震荡得厉害。后来把ρ调小到原来的五分之一才稳定下来。这件事让我意识到论文里的超参数是在特定任务和模型规模下调出来的直接照搬未必合适。我的经验是ρ的取值和模型规模、任务难度都有关系。模型越大、任务越复杂ρ可以适当大一些模型小、任务简单ρ就要小。一个实用的做法是从很小的ρ开始比如0.01逐步增大观察验证集表现找到那个既能提升泛化又不破坏训练稳定性的平衡点。7.3 和BatchNorm层的交互问题SAM和BatchNormBN层一起用时有个细节要注意。BN层在训练时会统计当前batch的均值和方差如果你在扰动参数后做第二次前向BN层会重新统计扰动后参数的统计量这可能导致两次前向的BN行为不一致。有些实现会选择在第二次前向时复用第一次的BN统计量避免这种不一致。这个问题在实际中影响有多大取决于你的模型和任务。我试过两种做法差异不是特别大但在一些敏感任务上复用BN统计量确实更稳定。如果你用SAM时发现训练不稳定可以检查一下BN的处理方式。7.4 学习率调度器的适配SAM和某些学习率调度器比如余弦退火配合时需要注意调度器的步进时机。因为SAM一次参数更新涉及两次前向如果你在错误的位置调用了调度器的step可能导致学习率更新节奏错乱。我的做法是把调度器的step放在真正的参数更新之后也就是第二次反向传播和优化器step都完成之后确保学习率和参数更新同步。这个细节看起来小但影响不小。我曾经因为调度器步进时机不对导致学习率提前衰减模型欠拟合排查了好一阵才找到原因。8. 一个可以直接抄的最小实现框架8.1 核心代码结构下面给一个相对完整的最小实现基于PyTorch风格你可以直接拿去改。这个实现把SAM封装成一个优化器包装类训练循环的改动降到最低。import torch class SAM(torch.optim.Optimizer): def __init__(self, params, base_optimizer, rho0.05, **kwargs): defaults dict(rhorho, **kwargs) super(SAM, self).__init__(params, defaults) self.base_optimizer base_optimizer(self.param_groups, **kwargs) self.param_groups self.base_optimizer.param_groups 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 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([ 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对应的训练循环大概是这样optimizer SAM(model.parameters(), torch.optim.SGD, rho0.05, lr0.1, momentum0.9) for inputs, targets in dataloader: # 第一次前向反向 loss loss_fn(model(inputs), targets) loss.backward() optimizer.first_step(zero_gradTrue) # 第二次前向反向 loss_fn(model(inputs), targets).backward() optimizer.second_step(zero_gradTrue)这个实现里first_step负责计算扰动并保存原始参数second_step负责恢复参数并用扰动点的梯度更新。两次前向用的是同一个batch逻辑清晰。8.2 训练循环的改动清单把普通训练改成SAM训练需要改的地方其实就几处我列个清单方便你对照优化器从普通的torch.optim.SGD或Adam换成SAM包装类base_optimizer传原来的优化器。每次迭代从一次前向反向变成两次第一次后调用first_step第二次后调用second_step。学习率调度器的step放在second_step之后。如果用了梯度裁剪注意裁剪的时机一般放在第二次反向之后、更新之前。日志记录时损失可以记录第二次前向的损失因为它对应扰动点的损失更能反映锐度。这些改动都不复杂但每一处都要确认到位尤其是调度器和梯度裁剪的时机容易出错。8.3 验证SAM是否真的生效怎么判断SAM是不是真的在起作用我一般看两个信号。第一训练损失和验证损失的差距是否缩小。如果SAM生效验证损失通常会比普通训练更低过拟合迹象减轻。第二验证准确率是否稳定提升。如果两者都没变化可能是ρ太小或者实现有问题。还有一个更直接的检查方法在训练结束后对最终参数做一个小扰动看损失上升多少。如果上升很小说明解比较平坦SAM起作用了如果上升很大说明解还是尖锐的SAM可能没生效。这个方法虽然粗糙但能快速给你一个直观判断。9. 一些延伸方向和我的个人体会SAM这套思路其实打开了一扇门优化过程本身可以承载正则化的功能不一定非要靠额外的正则项。沿着这个方向后续有不少变体比如对扰动方式做改进、对ρ做自适应调整、把SAM用到特定结构上等等。如果你对优化和泛化的交叉领域感兴趣SAM是一个很好的切入点。我在实际项目里用SAM的体会是它不是万能药但在泛化差距明显、又不想大动模型结构的时候它是一个性价比很高的选择。改动小、原理清晰、效果可预期这三点在工程落地时很重要。当然它也不是没有代价训练时间增加是实打实的所以要不要用得结合你的资源预算和效果要求来定。最后分享一个小技巧如果你不确定ρ该设多少可以先在一个小规模子集上快速跑几组不同ρ的实验看验证集表现的趋势找到大致范围后再在全量数据上正式训练。这样能省不少时间也避免在大规模训练上浪费算力。踩过几次坑之后我现在做任何超参数调整都会先在小规模上探路这个习惯帮我省了很多冤枉时间。