Adam优化器源码解析：公式、PyTorch实现与AdamW避坑指南

第一次把Adam的源码从头到尾读一遍是在一个训练到第三轮loss突然变成NaN的深夜。当时我盯着torch.optim.Adam的黑盒心里其实没底——学习率、betas、eps这些参数到底在公式里扮演什么角色为什么换个优化器loss曲线就完全不一样。后来我把Adam的原始论文、PyTorch实现、以及各种复现博客对着看又用NumPy手撸了一版能跑通的优化器才算真正把这块知识补齐。这篇内容就是那次折腾的完整沉淀从Adam优化器背后的理论直觉讲起把每一步公式拆开揉碎再给出可以直接运行、对照验证的代码实现最后聊几个实战里最容易踩的坑。不管你是刚学深度学习、只会调optimizer Adam(model.parameters())的新手还是想搞清楚AdamW和自己手写版差别的老手应该都能从里面找到对自己有用的东西。1. 为什么Adam能挤掉SGD成为默认优化器1.1 手调学习率的痛苦是Adam出现的直接动机用最朴素的随机梯度下降SGD训练模型最折磨人的地方不是公式难而是学习率这个单一的超参数太难伺候。设大了loss在最优解附近来回震荡甚至发散设小了收敛慢得让人怀疑人生。更麻烦的是同一个学习率放在不同的参数上效果完全不一样embedding层可能需要小一点的学习率而某些稀疏特征对应的参数又需要大一点的步长才能动起来。传统的做法是手动设计学习率衰减策略比如每过若干个epoch就乘以0.1或者用步进衰减、余弦退火。这些方法本质上是全局统一调度它假设所有参数共享同一个学习节奏。可现实里每个参数的梯度尺度千差万别有的参数梯度常年在1e-2量级有的可能只有1e-5。用一个统一学习率去更新它们就像用同一把尺子量所有人的鞋码怎么调都不对劲。Adam的思路很直接既然每个参数的梯度尺度不一样那就给每个参数配一个自适应的学习率。梯度大的参数步子迈小一点梯度小但一直朝同一方向走的参数步子迈大一点。这个想法并不新鲜AdaGrad和RMSProp都试过Adam的贡献在于把动量和自适应缩放这两件事漂亮地揉到了一起并且加了偏差校正让它从训练第一步开始数值就稳。1.2 自适应学习率解决的到底是哪一类问题要理解Adam的价值得先分清梯度噪声和梯度尺度这两个概念。梯度噪声指的是mini-batch采样带来的随机波动动量法Momentum通过累积历史梯度来平滑这种噪声让更新方向更稳定。梯度尺度则是不同参数梯度绝对大小差异的问题自适应方法通过除以梯度平方的某种平均来归一化。这两类问题需要两套机制。SGDMomentum只解决了噪声问题没解决尺度问题AdaGrad只解决了尺度问题没解决噪声问题而且它的学习率单调递减训练久了会衰减到几乎为零。Adam同时用了一阶矩梯度的指数移动平均来处理噪声用二阶矩梯度平方的指数移动平均来处理尺度相当于一个动量版RMSProp这也是为什么它在大多数任务上开箱即用的效果好。我个人的经验是Adam最大的优势是降低了对学习率初始值的敏感度。用SGD你可能要在1e-1到1e-4之间试好几轮才能找到能训起来的点而Adam通常1e-3就能给你一个像样的结果。这在快速验证想法、做原型实验的时候省下来的时间非常可观。当然这不代表Adam一定比调好的SGDMomentum泛化更好后面聊到AdamW和泛化差距时我会细说。下面这张表是我自己总结的几种优化器在解决什么问题上的对照方便你建立整体印象优化器处理梯度噪声处理梯度尺度学习率是否自适应主要短板SGD否否否对学习率极其敏感SGDMomentum是否否仍需手工调学习率AdaGrad否是是学习率单调衰减至消失RMSProp否是是无动量方向抖动Adam是是是权重衰减处理有缺陷AdamW是是是需额外关注weight decay取值2. 把Adam公式拆开一阶矩、二阶矩与偏差校正2.1 指数移动平均是什么为什么用它Adam的核心构件是指数移动平均Exponential Moving AverageEMA。它做的事情很简单给最近的梯度更大的权重给久远的梯度越来越小的权重权重按指数衰减。写成递推式就是m_t β * m_{t-1} (1 - β) * g_t这里的β控制记忆长度β越接近1历史信息保留得越多曲线越平滑但反应越迟钝。为什么不用简单的算术平均因为训练是个非平稳过程梯度的分布会随参数更新而漂移早期那些梯度对当前方向已经没有参考价值了。EMA用一个衰减因子把它们自然忘记既保留了平滑效果又不会让陈旧信息拖后腿。这个设计在信号处理领域用了几十年搬到优化里非常自然。一阶矩m_t就是梯度的EMA它承担的是动量的角色让更新方向沿着历史一致的方向走抑制来回摆动。二阶矩v_t是梯度平方的EMA它衡量的是这个参数最近的梯度有多剧烈后续用它来归一化步长。两个量各司其职这就是Adam名字里Adaptive Moment estimation的由来——自适应矩估计。2.2 偏差校正不写它训练开局会出问题很多人看Adam公式时最容易忽略的就是**偏差校正bias correction**这一步但它恰恰是Adam区别于朴素RMSProp的关键细节之一。问题出在初始化m_0和v_0都初始化为0而EMA的递推是从0开始的。在训练最开始几步m_t和v_t都会被这个零初始值拉低导致它们明显小于真实的梯度均值和平方均值。数学上可以推导出第t步的m_t期望是(1 - β1^t) * E[g]v_t期望是(1 - β2^t) * E[g²]。也就是说EMA天然带了一个(1 - β^t)的缩放系数。当t很小比如t1时β1^1 0.9系数是0.1意味着m_1只有真实梯度的十分之一。如果不校正第一步的更新量会被严重低估相当于开局学习率被人为压小了十倍。偏差校正的做法就是把这个系数除掉m_hat m_t / (1 - β1^t)v_hat v_t / (1 - β2^t)。当t增大时β^t趋近0校正项趋近1自动失效。所以它只在前几十步起作用之后几乎无感。我实测过去掉偏差校正后训练前100步的loss下降会明显变慢尤其是在β20.999这种大记忆长度下更严重因为1 - 0.999^1 0.001v_hat被放大了整整1000倍不校正根本没法用。2.3 完整更新公式的逐项解读把前面几块拼起来Adam在第t步对每个参数θ执行的操作完整链条是这样的输入学习率 α衰减率 β1、β2数值稳定项 ε 初始化m_0 0, v_0 0, t 0 每一步 t t 1 g_t 当前mini-batch上对θ的梯度 m_t β1 * m_{t-1} (1 - β1) * g_t # 一阶矩 v_t β2 * v_{t-1} (1 - β2) * g_t² # 二阶矩 m_hat m_t / (1 - β1^t) # 一阶矩偏差校正 v_hat v_t / (1 - β2^t) # 二阶矩偏差校正 θ θ - α * m_hat / (sqrt(v_hat) ε) # 参数更新逐项看α是全局学习率它决定了整体步子的大小是你要调的最重要的参数。β1默认0.9控制动量记忆越大越平滑。β2默认0.999控制二阶矩记忆因为它要估计梯度平方的长期水平所以需要更长的记忆窗口设置得比β1更接近1。ε默认1e-8加在分母上防止除以零同时也限制了当v_hat极小时更新量的上限。我在实际调试中发现一个有意思的细节ε放在根号外还是根号内不同框架不一样。原始论文和PyTorch都是sqrt(v_hat) ε而有些实现写成sqrt(v_hat ε)。当ε很小时两者几乎没区别但在梯度极小的稀疏场景下ε放的位置会影响更新量的上界进而影响收敛行为。如果你在复现别人的实验这一步的差异有时能解释为什么结果对不上。提示ε虽然默认1e-8但在混合精度训练fp16下1e-8会直接下溢成0PyTorch建议改成1e-4到1e-6之间这一点我在用AMP训练时踩过坑。3. 手撸一版Adam再和PyTorch对照着看3.1 用NumPy写一个最小可运行实现看再多公式不如自己动手写一遍。下面这版NumPy实现没有任何依赖把参数当成一个列表来处理虽然简陋但逻辑完整适合用来验证你对公式的理解import numpy as np class AdamNumpy: def __init__(self, params, lr1e-3, betas(0.9, 0.999), eps1e-8): self.params params # list of np.ndarray self.lr lr self.beta1, self.beta2 betas self.eps eps self.m [np.zeros_like(p) for p in params] self.v [np.zeros_like(p) for p in params] self.t 0 def step(self, grads): self.t 1 for i, (p, g) in enumerate(zip(self.params, grads)): self.m[i] self.beta1 * self.m[i] (1 - self.beta1) * g self.v[i] self.beta2 * self.v[i] (1 - self.beta2) * g ** 2 m_hat self.m[i] / (1 - self.beta1 ** self.t) v_hat self.v[i] / (1 - self.beta2 ** self.t) p - self.lr * m_hat / (np.sqrt(v_hat) self.eps)拿一个简单的二次函数测试一下先随便造几个参数和解析梯度跑个一百步看它能不能收敛到最优点w np.array([5.0, -3.0]) opt AdamNumpy([w], lr0.1) for _ in range(200): grad np.array([2 * w[0], 2 * w[1]]) # f w0^2 w1^2 opt.step([grad]) print(w) # 应该非常接近 [0, 0]这段代码跑下来w会稳稳收敛到接近零的位置。我第一次跑通时最大的感受是Adam对学习率的容忍度确实高lr0.1和lr0.01都能收敛只是速度不同这在SGD上是很难想象的。3.2 和torch.optim.Adam逐行对照手写版跑通之后再看PyTorch的实现就清晰多了。PyTorch把状态存进state字典每步的关键代码逻辑其实和上面一模一样只是工程上做了大量优化。我把它简化后的核心逻辑贴出来对照# PyTorch 内部逻辑简化伪代码 exp_avg state[exp_avg] # 对应 m_t exp_avg_sq state[exp_avg_sq] # 对应 v_t step state[step] 1 exp_avg.mul_(beta1).add_(grad, alpha1 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value1 - beta2) bias_correction1 1 - beta1 ** step bias_correction2 1 - beta2 ** step step_size lr / bias_correction1 denom (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps) param.addcdiv_(exp_avg, denom, value-step_size)注意这里有个数学上的等价变形很有意思PyTorch没有分别计算m_hat和v_hat而是把1/bias_correction1折进step_size把1/sqrt(bias_correction2)折进分母。经过这一步重排公式在数值上等价但少算几次数值较大的除法浮点误差更小。这种看起来一样但更稳的工程细节是自己写代码时才容易注意到的。3.3 state_dict里到底存了什么如果你要保存和恢复训练状态理解优化器的state_dict很关键。Adam每个参数会存两个和参数同形状的张量exp_avg和exp_avg_sq再加上一个全局的step计数。记得排除startup的首次假设模型有1亿参数用fp32保存光优化器状态就要额外占约800MB显存这就是为什么大模型训练时优化器状态常常比模型本身还占地方。我在做断点续训时踩过一次坑只保存了模型权重没保存优化器状态结果续训后loss曲线出现明显的抖动回弹。原因是exp_avg和exp_avg_sq被重置成零动量信息和二阶矩估计全部丢失优化器需要重新热身几十到几百步才能找回原来的状态。所以只要你的训练需要中断恢复务必同时保存model.state_dict()和optimizer.state_dict()并且在加载时确认step计数也一起恢复。注意跨设备加载时如果用的是map_location要确保优化器状态里的张量和模型在同一设备上否则会报device mismatch。我见过有人在CPU上加载GPU保存的checkpoint时忘了处理优化器状态训练直接崩在第一步。4. Adam和AdamW的分歧权重衰减该放在哪儿4.1 L2正则和权重衰减根本不是一回事这是理解AdamW最关键的一点也是最容易被搞混的地方。在很多教程里L2正则和权重衰减被当成同义词因为在朴素SGD下它们确实等价——在损失里加一项λ/2 * ||θ||²对θ求导得到λθ加到梯度里更新式就变成了θ θ - α*(g λθ) (1 - αλ)θ - αg这正好等价于每步把参数乘一个(1 - αλ)的衰减因子。但换到Adam身上这个等价关系就断了。因为Adam会用二阶矩v_hat去除梯度当你把λθ加进梯度g后这个正则项的梯度也会被sqrt(v_hat)归一化。结果是梯度大的参数其权重衰减被缩小梯度小的参数权重衰减被放大。这就完全背离了所有参数统一按比例衰减的初衷自适应学习率和L2正则相互干扰导致正则效果变得不可控。4.2 AdamW的改动其实只有一行AdamW的解决方式非常干脆把权重衰减从梯度里拿出来直接在参数更新时单独做衰减不参与任何矩估计和归一化。更新式从Adam的θ θ - α * m_hat / (sqrt(v_hat) ε)变成θ θ - α * m_hat / (sqrt(v_hat) ε) - α * λ * θ就这么一个改动让权重衰减恢复了按比例收缩的语义。虽然数学上只是一项的位置变化但实测效果差别很大尤其是在带weight decay训练Transformer、CNN时AdamW通常能给出更好的泛化性能。这也是为什么现在几乎所有大模型预训练默认都用AdamW而不是Adam。对比项Adam L2正则AdamW衰减项位置混在梯度里单独作用于参数是否被二阶矩归一化是否与自适应学习率的关系相互干扰解耦稀疏参数表现衰减不均统一衰减大模型预训练默认选择否是我在做图像分类微调时对比过两者同样的数据集、同样的学习率AdamW在验证集上比AdamL2大约高0.5到1个百分点而且对weight decay取值的敏感度明显更低。当然这个结论不是绝对的具体任务还得自己试。5. 超参数怎么调lr、betas、eps、amsgrad的取舍5.1 学习率是唯一你必须认真调的参数Adam虽然对学习率不敏感但不代表可以随便设。默认1e-3在大多数任务上是个安全起点但在不同场景下应该这样调整从头训练大模型时通常用1e-4甚至更小因为大模型参数多、梯度累积效应强微调预训练模型时常用2e-5到5e-5这个区间是BERT类模型微调的经验值太大容易破坏预训练学到的特征训练小的全连接网络或做原型实验时1e-3到3e-3都能用。我的习惯是先用1e-3快速验证代码能不能跑通、loss会不会下降确认没问题后再根据任务类型把学习率调到目标区间配合warmup做正式训练。这个先粗后细的流程能省下大量无谓的调参时间。另外一个小技巧如果你发现学习率调小后loss下降很慢但很稳说明模型容量没问题可以适当加大如果loss在初期就剧烈震荡优先检查是不是学习率过高导致梯度爆炸。5.2 betas、eps这类参数什么时候才需要动β1和β2的默认值0.9和0.999在95%的情况下不需要改。只有两种场景值得调整一是训练非常稀疏的模型比如推荐系统里大量embedding只被很少的样本激活可以适当降低β2比如到0.99或0.98让二阶矩更快响应近期梯度变化二是强化学习等非平稳性极强的任务有时把β1降到0.9以下能加快对新策略的响应减少旧梯度的拖累。ε默认1e-8在fp32下基本不用动。但如果你用fp16混合精度训练一定要把它调到1e-6到1e-4量级否则梯度平方在fp16下会下溢。这个坑我在用AMP训视觉模型时踩得很惨——loss直接变NaN查了半天才发现是eps太小。amsgrad选项值得单独提一句。标准Adam的二阶矩是EMA如果某段时间梯度很大v_t会被抬高之后即使梯度回归正常v_t也需要一段时间才能降下来间接导致这段时间的有效学习率偏小。AMSGrad的思路是维护v的历史最大值只在梯度确实变大时才提升分母避免学习率被历史大梯度长期压制。实测在部分任务上AMSGrad能带来一点提升但它多了一份状态、多了一点计算收益并不总是明显我一般只在训练不稳定时才打开试试。6. 实战踩坑NaN、loss震荡、收敛慢的排查链路6.1 训练loss变NaN我一般这样一步一步查NaN是训练里最让人抓狂的问题因为它可能来自很多地方。我总结的排查顺序是这样从最常见到最罕见依次排除。第一步先降学习率再看。把学习率临时降到原来的十分之一如果是学习率过高导致的梯度爆炸这一步通常能立刻见效。很多所谓的NaN bug其实只是学习率没调好。第二步检查数据里有没有脏样本。个别样本的特征值异常大或者标签越界会在一次前向里放大成极端梯度直接把某一步的更新推飞。做法是遍历一遍数据统计每个特征的均值和极值把明显超出合理范围的值过滤或裁剪掉。第三步给梯度加裁剪。这是最省事的保底手段torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)几乎成了我所有训练脚本的标配。即使前面两步没找到根因加了梯度裁剪也能让训练稳住代价只是轻微影响收敛速度。第四步查混合精度和eps设置。如果用了fp16先确认ε不是1e-8再检查loss scaler是否正常工作。fp16的数值范围窄溢出太正常了。第五步检查是否有除零或log(0)。比如自定义损失里的torch.log(x)当x为0时就是负无穷几个batch累积下来就能污染整个训练。给log加一个极小的ε或者用torch.log_softmax这种数值稳定的实现。6.2 loss震荡和收敛慢很多时候不是优化器的锅遇到loss上下剧烈波动或者收敛比预期慢很多人第一反应是换优化器、调学习率但真正的原因往往在数据或模型结构上。我见过太多例子batch size太小导致梯度噪声过大loss曲线自然抖得厉害把batch size从16加到64后立刻平稳数据没有shuffle或者shuffle不充分模型学到的batch顺序里藏着强相关性更新方向来回打架学习率没配合batch size调整batch变大后等效学习率应该相应提高。收敛慢还有一个高频原因是缺少学习率预热warmup。训练最开始模型参数还是随机初始化的梯度方向很不稳定此时如果用完整学习率很可能一步就把参数带偏到很差的区域。Warmup的做法是让学习率从很小的值比如1e-6线性爬升到目标值通常用总步数的5%到10%来爬坡。对于Transformer类模型warmup几乎是必需品不加的话前几百步的loss会疯狂震荡。我把常见的排查方向整理成一张表方便你按图索骥现象优先排查项典型修复手段loss变NaN学习率、脏数据、fp16 eps降lr、清洗数据、eps调大loss剧烈震荡batch size、shuffle、lr加大batch、检查shuffle、降lr收敛过慢warmup、lr过小、梯度消失加warmup、提高lr、换激活函数前期下降后期停滞lr衰减策略、探索不足加余弦退火、定期重置lr训练集好验证集差过拟合、正则不足加weight decay、early stop6.3 一个我常用的Adam训练模板最后把我自己在小项目里反复复用的训练配置整理出来包括warmup和梯度裁剪可以直接抄去改import torch from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR optimizer AdamW(model.parameters(), lr2e-5, weight_decay0.01, betas(0.9, 0.999), eps1e-8) warmup_steps int(0.1 * total_steps) def lr_lambda(current_step): if current_step warmup_steps: return current_step / max(1, warmup_steps) progress (current_step - warmup_steps) / max(1, total_steps - warmup_steps) return max(0.0, 0.5 * (1.0 math.cos(math.pi * progress))) scheduler LambdaLR(optimizer, lr_lambda) for batch in dataloader: optimizer.zero_grad() loss model(batch) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step()这套配置里学习率用余弦退火从峰值平滑降到0warmup占前10%梯度裁剪阈值设1.0。我拿它在几个中等规模的文本分类和图像分类任务上跑过基本不用大改就能稳定收敛。唯一需要注意是total_steps要对如果你的dataloader是流式的没长度得自己估算另外warmup比例也别设太大10%对大部分任务够用设到30%反而会拖慢整体进度。回头看Adam真正值得反复琢磨的地方不在一行optimizer.step()能跑通而在于你清楚它在每一步对每个参数做了什么、哪些状态需要保存、什么时候该切到AdamW、以及遇到异常时该从哪个方向查。把这几个点想通了你会发现不仅是Adam连带着AdaGrad、RMSProp这些优化器的逻辑也变得一目了然。我个人在项目里最深的体会是优化器选对能让你省一半调参时间。