企业官网建设流程全解析

1. 为什么平坦极小值能让模型泛化更稳

第一次读到这篇工作的标题时，我正在为一个图像分类项目头疼。训练集准确率冲到99%以上，验证集却始终在88%附近晃荡，换了几种数据增强、调了几轮学习率，提升都非常有限。后来一位做优化的朋友跟我说，你试试从损失曲面的几何形状入手，别老盯着学习率调度。这句话把我引向了Sharpness-Aware Minimization，也就是常说的SAM。

SAM要解决的问题很具体：深度网络在训练集上能轻松降到很低的损失，但换到没见过的数据上表现就掉下来。传统做法是让模型参数朝着让训练损失下降最快的方向走，也就是梯度下降。但梯度下降只关心当前点的损失值，不关心这个点周围是不是一个陡峭的坑。如果当前点落在一个又窄又深的极小值里，训练损失确实很低，可只要输入数据稍微扰动一下，损失就会急剧上升，泛化自然差。SAM的核心洞察就是：与其找一个损失值最低的点，不如找一个周围一整片区域损失都低的点，也就是平坦极小值。

这个思路听起来简单，但落地时有个关键难点：怎么高效地找到平坦区域？如果对每个参数都去搜索邻域内的最坏情况，计算量会爆炸。SAM给出的方案是分两步走：先沿着梯度方向做一次扰动，找到邻域内损失较高的那个点，再在这个扰动点上计算梯度来更新原参数。这样一次迭代需要两次前向和两次反向传播，计算开销大约是普通训练的两倍，但换来的是明显更稳的泛化能力。

适合读这篇内容的人，我大致分三类。第一类是正在做图像分类、目标检测、语义分割的工程师，模型在验证集上卡住了，想从优化角度找突破口。第二类是研究机器学习理论的同学，想理解平坦极小值和泛化之间的关联。第三类是做模型部署的从业者，关心模型在分布偏移场景下的鲁棒性。不管你是哪一类，只要你的模型存在过拟合或者验证集波动大的问题，SAM都值得花时间研究。

我自己的体会是，SAM不是那种“调个参数就能涨点”的trick，它改变的是训练的目标函数本身。你原来的损失函数是L(θ)，SAM把它变成了max over neighborhood L(θ+ε)，也就是在参数邻域内找最坏情况的损失。这个改动让模型不再只追求单点最优，而是追求区域最优。理解这一点，后面所有的实现细节就都顺了。

2. SAM的核心机制与数学推导拆解

2.1 从单点损失到邻域最坏损失

普通训练的目标很直接：最小化训练集上的平均损失。写成公式就是min L(θ)，其中θ是模型参数，L是损失函数。梯度下降沿着∇L(θ)的反方向更新，每一步都让当前点的损失下降。问题在于，这个过程中模型可能钻进一个非常窄的极小值，就像山谷里的一条细缝，训练数据刚好落在缝底，但测试数据稍微偏一点就撞到缝壁上了。

SAM把目标改成了min max L(θ+ε)，其中ε是参数空间里的扰动向量，max表示在某个邻域内找让损失最大的那个扰动。换句话说，SAM要求模型不仅在当前参数点表现好，还要在它周围一小片区域内都表现好。这个邻域的大小由超参数ρ控制，ρ越大，要求的平坦区域越宽，模型越鲁棒，但训练难度也越大。

这个min-max问题的求解并不需要真的去搜索整个邻域。SAM做了一个一阶近似：在θ附近，L(θ+ε)约等于L(θ) + ε^T ∇L(θ)。要让这个近似值最大，ε应该取∇L(θ)的方向，因为梯度方向是损失上升最快的方向。所以最优扰动近似为ε = ρ * ∇L(θ) / ||∇L(θ)||，也就是把梯度归一化后乘以ρ。这个近似虽然粗糙，但实际效果出奇地好，而且计算成本可控。

2.2 两步更新的具体流程

有了扰动方向，SAM的一次完整更新分两步。第一步叫“上升步”：计算当前参数θ下的梯度g = ∇L(θ)，然后构造扰动ε = ρ * g / ||g||，把参数临时改成θ + ε。第二步叫“下降步”：在扰动后的参数θ + ε上重新计算梯度g' = ∇L(θ + ε)，然后用这个g'来更新原始参数θ，更新规则和普通SGD一样，θ ← θ - η * g'，其中η是学习率。

这里有个细节容易搞混：最终更新的是原始参数θ，而不是扰动后的参数。扰动只是用来“探测”邻域内哪个方向损失更高，真正走的方向是扰动点上的梯度。你可以把这个过程想象成：先往最陡的方向迈一小步看看情况，然后根据看到的情况决定真正往哪走。这样走出来的路径，会倾向于避开那些陡峭的极小值，因为一旦靠近陡峭区域，扰动后的损失会急剧上升，梯度g'会给出一个把参数拉离陡峭区域的方向。

计算开销方面，一次SAM更新需要两次前向传播和两次反向传播。第一次前向和反向用来算g，第二次前向和反向用来算g'。相比普通训练，计算量翻倍，显存占用也会增加，因为需要保存两次计算的中间结果。这是SAM最大的代价，也是很多人在实际项目中犹豫要不要用的原因。

2.3 为什么ρ的取值如此关键

ρ是SAM最重要的超参数，它控制邻域的半径。ρ太小，扰动几乎不起作用，SAM退化成普通SGD；ρ太大，扰动后的参数可能跑到损失曲面完全不同的区域，梯度g'失去参考意义，训练反而会不稳定。根据原论文和后续大量复现实验，ρ的典型取值在0.01到0.2之间，具体取决于模型大小、数据集难度和训练轮数。

我自己的经验是，ρ和模型参数量大致成反比。小模型参数少，损失曲面相对简单，ρ可以取大一点，比如0.1到0.2；大模型参数多，曲面维度高，ρ要取小一点，比如0.01到0.05。另外，ρ和学习率也有耦合关系。学习率大时，参数更新步长大，ρ可以适当大一点来匹配；学习率小时，ρ也要相应减小，否则扰动会盖过更新本身。

还有一个容易被忽略的点：ρ的取值和Batch Normalization的配合。如果模型里有BN层，扰动后的参数会改变BN的统计量，导致g'的计算出现偏差。原论文建议在使用SAM时，对BN层的参数不做扰动，或者用单独的BN统计量来处理。这个细节在复现时经常被漏掉，导致效果打折扣。

3. 动手实现：从零搭建SAM训练流程

3.1 环境准备与依赖选择

实现SAM不需要特殊的库，PyTorch或TensorFlow都能做。我用PyTorch比较多，所以下面的示例以PyTorch为主。基础环境就是常规的深度学习环境：Python 3.8以上，PyTorch 1.10以上，CUDA可用的话更好。不需要额外安装包，SAM的核心逻辑就是一个自定义优化器，几十行代码就能写完。

如果你不想自己写，也有一些开源实现可以直接用，比如某些模型库里的SAM优化器封装。但我建议至少自己实现一遍，因为SAM的细节比较多，自己写一遍能理解每一步在做什么，后面调参和排查问题会容易很多。自己实现还有一个好处：可以根据项目需求灵活调整，比如对某些层不做扰动，或者动态调整ρ。

3.2 自定义优化器的代码实现

下面是一个最小可用的SAM优化器实现，基于PyTorch的Optimizer基类。核心逻辑就是前面说的两步：先算梯度、构造扰动、更新参数到扰动点；再算梯度、恢复参数、用新梯度更新。

import torch from torch.optim import Optimizer class SAM(Optimizer): def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs): defaults = dict(rho=rho, adaptive=adaptive, **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_grad=False): 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_grad=False): 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(p=2).to(shared_device) for group in self.param_groups for p in group["params"] if p.grad is not None ]), p=2 ) return norm def load_state_dict(self, state_dict): super().load_state_dict(state_dict) self.base_optimizer.param_groups = self.param_groups

这段代码的关键点有几个。first_step里保存了原始参数old_p，然后计算扰动并加到参数上。second_step里先把参数恢复到原始值，再用扰动点上的梯度做真正的更新。_grad_norm计算的是所有参数梯度的全局范数，用来归一化扰动。adaptive参数控制是否使用自适应版本的SAM，这个后面会讲。

3.3 训练循环的改造要点

有了优化器，训练循环需要相应改造。普通训练是loss.backward()之后直接optimizer.step()，SAM需要拆成两步。下面是一个典型的训练循环片段：

model = MyModel() base_optimizer = torch.optim.SGD optimizer = SAM(model.parameters(), base_optimizer, lr=0.1, momentum=0.9, rho=0.05) criterion = torch.nn.CrossEntropyLoss() for epoch in range(num_epochs): for inputs, targets in train_loader: inputs, targets = inputs.cuda(), targets.cuda() # 第一次前向和反向 outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.first_step(zero_grad=True) # 第二次前向和反向 outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.second_step(zero_grad=True)

这里有个坑要注意：第一次反向传播后不能清零梯度，因为first_step需要用到梯度来构造扰动。first_step内部会在构造完扰动后清零梯度，为第二次反向做准备。第二次反向传播后，second_step会用新的梯度更新参数并清零梯度。整个流程走下来，一次迭代做了两次完整的前向和反向。

另外，如果你用的是混合精度训练，SAM的两步流程需要特别处理。第一次前向和反向可以用autocast，但first_step里的参数扰动操作要在float32下做，否则精度损失会影响扰动方向。第二次前向和反向同样用autocast，但second_step里的参数恢复和更新也要在float32下做。这个细节在混合精度场景下很容易出错，建议先用纯float32跑通再上混合精度。

4. 调参实战：ρ、学习率与训练轮数的配合

4.1 ρ的搜索策略与经验取值

ρ是SAM最需要调的超参数。我的做法是先固定其他参数，在一个合理的范围内做网格搜索。对于CIFAR-10上的ResNet-18，我试过ρ从0.01到0.3，步长0.01。结果发现ρ在0.05到0.1之间效果最好，验证集准确率比普通SGD高1.5到2个百分点。ρ小于0.02时提升不明显，ρ大于0.2时训练后期会出现震荡，验证集准确率反而下降。

对于更大的模型，比如ImageNet上的ResNet-50，ρ的合适范围会小一些，大概在0.01到0.05之间。这是因为大模型的损失曲面维度更高，同样大小的ρ在高维空间里对应的邻域体积更大，扰动更容易把参数推到不相关的区域。一个经验公式是ρ和模型参数量的平方根成反比，但实际调参时还是以验证集表现为准。

还有一个技巧是ρ可以随训练轮数动态调整。训练初期损失曲面比较平滑，ρ可以大一点帮助探索；训练后期接近收敛，ρ要小一点避免破坏已经找到的平坦区域。我试过线性衰减ρ，从0.1降到0.02，效果比固定ρ略好，但提升幅度不大，大概0.2个百分点。如果不想增加调参复杂度，固定ρ也够用。

4.2 学习率与SAM的耦合关系

SAM对学习率比普通SGD更敏感。因为SAM的更新方向是扰动点上的梯度，这个梯度的尺度和原始梯度不完全一样。如果学习率太大，SAM的更新步长会过大，参数可能在平坦区域边缘来回跳；如果学习率太小，SAM的优势体现不出来，因为扰动的影响被小步长稀释了。

我的经验是，SAM的学习率可以比普通SGD略大一点，但不要大太多。比如普通SGD用0.1，SAM可以用0.1到0.15。如果用了warmup，warmup阶段可以短一点，因为SAM本身就有一定的稳定训练的作用。另外，SAM和余弦退火学习率调度配合得很好，因为余弦退火在训练后期学习率很小，正好匹配SAM后期需要小ρ的需求。

还有一个细节：SAM的两次梯度计算会引入额外的噪声，如果学习率调度里有基于梯度的自适应方法（比如Adam），需要小心。SAM和Adam的配合不如和SGD那么直接，因为Adam的自适应学习率会放大扰动点梯度的尺度差异。如果非要用Adam，建议把ρ调小一点，或者用AdamW并降低权重衰减。

4.3 训练轮数与早停策略

SAM的训练轮数通常比普通训练多10%到20%。因为SAM每轮的有效更新次数和普通训练一样，但每次更新考虑的信息更多，收敛速度会稍慢一些。在CIFAR-10上，普通SGD大概200轮收敛，SAM可能需要220到250轮。如果训练轮数不够，SAM可能还没走到平坦区域就停了，效果反而不如普通训练。

早停策略也要调整。普通训练看验证集损失，SAM的验证集损失曲线通常更平滑，波动更小，所以早停的耐心值可以设大一点。我一般设20到30轮，如果验证集损失连续这么多轮不下降才停。另外，SAM的验证集准确率往往比损失更早达到峰值，所以如果只看准确率，早停可以更早触发。

还有一个值得注意的现象：SAM在训练后期的验证集表现非常稳定，不像普通训练那样会出现明显的过拟合波动。这意味着如果你有多个随机种子，SAM的方差通常比普通训练小。我在一个项目里跑了5个种子，普通训练的验证集准确率标准差是0.3%，SAM只有0.15%。这个稳定性对于需要复现的实验很重要。

5. 常见问题排查与避坑指南

5.1 训练不收敛或损失震荡

第一次用SAM时，最常见的问题是训练不收敛，损失曲线上下震荡。我踩过的坑主要有三个。第一个是ρ设得太大，扰动后的参数跑到了损失曲面很远的区域，梯度g'的方向和原始梯度g差异太大，更新方向混乱。解决办法是把ρ降到0.05以下，或者从0.01开始慢慢往上加。

第二个坑是学习率没调。SAM的更新步长和ρ、学习率都有关，如果学习率沿用了普通训练的值，可能偏大。我建议先用普通训练的学习率的一半跑几轮，看看损失是否稳定下降，再逐步调大。第三个坑是Batch Normalization的处理。如果模型有BN层，扰动后的参数会改变BN的running statistics，导致第二次前向的均值和方差和第一次不一致。解决办法是在first_step之后、第二次前向之前，不要更新BN的running statistics，或者对BN层单独处理。

5.2 显存不足与计算开销优化

SAM的计算开销是普通训练的两倍左右，显存占用也会增加，因为需要保存两次前向的中间激活值。如果显存不够，有几个优化方向。第一是减小batch size，但要注意batch size太小会影响BN的效果，可以配合梯度累积。第二是用混合精度训练，把激活值和梯度存成float16，能省不少显存。第三是只对部分层做SAM扰动，比如只对最后的分类层或者某些关键层做扰动，其他层用普通SGD。这个做法在有些任务上效果不错，但需要实验验证。

计算开销方面，如果训练时间紧张，可以考虑每隔几步做一次SAM更新，其他步用普通SGD。比如每4步里做1次SAM、3次SGD，这样计算开销只增加25%，但泛化提升能保留大部分。这个策略在大规模训练里比较实用，因为大规模训练的瓶颈往往是总计算量而不是单步时间。

5.3 与数据增强、正则化的配合

SAM和数据增强是互补的，不是替代关系。数据增强在输入空间做扰动，SAM在参数空间做扰动，两者结合通常比单独用任何一个都好。我在CIFAR-10上试过，单独用Cutout能涨1个点，单独用SAM能涨1.5个点，两者一起用能涨2.5个点。但要注意，如果数据增强太强，SAM的扰动可能被淹没，因为模型已经对输入扰动很鲁棒了，参数空间的平坦性提升空间就小了。

和权重衰减的配合也有讲究。权重衰减本身就是在鼓励参数靠近原点，间接起到平滑损失曲面的作用。SAM和权重衰减一起用时，权重衰减的系数可以适当减小，因为SAM已经在做类似的事情了。我试过把权重衰减从5e-4降到1e-4，配合SAM，效果比原来更好。和Dropout的配合则比较微妙，Dropout引入的随机性和SAM的扰动有叠加效应，如果Dropout率太高，训练会变得很不稳定。建议用SAM时把Dropout率降低一半，或者只在全连接层保留Dropout。

5.4 常见问题速查表

6. 进阶变体与适用场景扩展

6.1 自适应SAM与参数尺度处理

原版SAM对所有参数用同一个ρ，但不同参数的尺度差异很大。比如卷积层的权重通常在0.01到0.1量级，全连接层的权重可能在0.1到1量级。用同一个ρ，对小尺度参数扰动过大，对大尺度参数扰动过小。自适应SAM（ASAM）的思路是让ρ和参数尺度成正比，具体做法是在计算扰动时乘以参数的绝对值。这样大参数得到大扰动，小参数得到小扰动，扰动的相对幅度更一致。

ASAM的实现只需要在first_step里把扰动项从p.grad * scale改成torch.abs(p) * p.grad * scale。这个改动很小，但在一些任务上提升明显。我在一个类别不平衡的数据集上试过，ASAM比SAM的少数类召回率高2个百分点。不过ASAM对ρ更敏感，因为参数尺度会放大ρ的影响，调参时要更小心。

6.2 SAM在目标检测与分割中的应用

SAM最初是在图像分类上验证的，但它的思路可以推广到任何用梯度下降训练的任务。目标检测和语义分割的损失函数比分类复杂，通常有多个子损失，比如分类损失、回归损失、分割损失。这些子损失的尺度不一样，如果直接对总损失做SAM，扰动方向可能被某个大尺度子损失主导。

我的做法是对每个子损失分别做SAM扰动，然后合并。具体来说，先算每个子损失的梯度，分别构造扰动，把扰动加权平均后再更新参数。权重可以根据子损失的重要性来定，比如分类损失权重0.5，回归损失权重0.3，分割损失权重0.2。这个做法比直接对总损失做SAM更稳定，因为每个子损失的扰动方向都被考虑到了。

在分割任务上，SAM对边界区域的提升比较明显。因为边界区域的损失曲面通常更陡峭，普通训练容易在边界上过拟合，SAM的平坦性约束能让边界预测更平滑。我在一个遥感图像分割项目里用SAM，边界区域的IoU提升了3个百分点，内部区域的提升只有0.5个百分点。

6.3 什么时候不适合用SAM

SAM不是万能的。如果你的模型本身已经很小，参数量在百万级别以下，损失曲面比较简单，SAM的提升可能很有限，但计算开销翻倍，性价比不高。如果你的训练数据量非常大，比如上亿张图片，训练本身已经很充分，SAM的边际收益也会下降。另外，如果你的任务对训练速度要求极高，比如在线学习或者强化学习的实时训练，SAM的两倍开销可能无法接受。

还有一个场景要注意：如果验证集和测试集的分布与训练集差异极大，SAM的平坦性约束可能帮助有限。因为SAM保证的是训练分布邻域内的平坦性，如果测试分布完全在训练分布之外，平坦性约束覆盖不到。这种情况下，领域自适应或者分布鲁棒优化可能更合适。SAM更适合的是训练集和测试集同分布但存在过拟合的场景。

7. 我个人的实操体会与后续扩展方向

SAM是我近几年用过的最有性价比的优化技巧之一。它的实现不复杂，调参也不难，但带来的泛化提升在多个项目里都稳定复现了。我印象最深的一次是在一个医学图像分类项目里，数据量只有几千张，普通训练过拟合严重，验证集准确率卡在82%。用了SAM之后，验证集准确率到了86%，而且不同种子之间的方差从1.2%降到了0.4%。这个稳定性对于医学应用来说比单纯的准确率提升更有价值。

踩过的坑也不少。最开始没注意BN层的处理，训练损失一直震荡，排查了两天才发现是扰动改变了BN统计量。后来在混合精度训练里又遇到精度问题，扰动在float16下计算导致方向偏差，改成float32计算扰动才解决。这些细节在原论文里没有展开，但实际落地时每一个都可能卡住你。

后续扩展方向有几个我比较感兴趣。一是把SAM和知识蒸馏结合，用SAM训练教师模型，再用教师模型指导小模型，看看平坦性能不能传递到小模型上。二是把SAM用到自监督学习里，自监督的损失函数和分类不一样，平坦性的定义可能需要调整。三是研究SAM和模型量化的关系，平坦的损失曲面是否意味着量化后精度损失更小，这个方向目前实验还不多。

最后分享一个小技巧：如果你不确定ρ设多少，可以先跑一个短训练，比如10个epoch，用几个不同的ρ值，看验证集准确率的变化趋势。ρ太小和太大都会让验证集准确率下降，中间会有一个峰值。找到峰值附近的ρ，再用完整训练轮数跑一遍。这个方法比直接网格搜索省时间，而且能快速排除明显不合适的ρ值。