前面的文章已经把损失函数沿计算图反向求导,PyTorch 用一行 loss.backward() 将结果累积到各个参数的 grad 中。训练循环紧接着还有一行 optimizer.step()。模型参数真正发生变化,就在这一步。
梯度只描述当前位置附近,损失增长最快的方向。优化器要把它变成有限长度的更新,还要处理小批量采样带来的波动、不同参数的尺度差异,以及此前若干步留下的信息。SGD 直接使用当前梯度;Momentum 记住近期方向;RMSProp 为每个坐标估计梯度尺度;Adam 同时维护方向和尺度;AdamW 又把权重衰减从自适应更新中分离出来。
本文从同一条参数更新公式出发,依次拆开这些状态。目标是读懂每个超参数在更新中处于什么位置,以及训练异常时应该先检查什么。示意实验只解释算法的更新几何,不比较模型精度,也不据此判断哪种优化器普遍更好。
梯度估计与参数更新
设训练集包含 个样本,模型参数为 ,平均经验损失写成
完整梯度需要遍历全部样本。神经网络训练通常从数据中取一个大小为 的小批量 ,在第 步计算
当样本按相同概率随机抽取时, 的期望对应完整梯度;某一个批次仍会偏离它。批量越小,这种波动通常越明显。数据分布、采样方式、数据增强和损失缩放也会改变实际看到的梯度,不能只用批量大小解释所有噪声。
优化器根据当前梯度与内部状态产生更新向量 ,再用学习率 控制步长:
一次这样的更新称为一个 step 或 iteration。一个 epoch 表示训练数据大致被遍历一遍,通常包含许多 step。学习率调度器究竟按 step 还是按 epoch 更新,必须与具体实现保持一致;混淆两者会让衰减速度相差一个 epoch 所含的批次数。
SGD
随机梯度下降(Stochastic Gradient Descent,SGD)的基本形式没有额外状态:
负梯度是当前点在欧氏几何下最陡的下降方向,这个结论只描述足够小的局部移动。学习率把方向变成实际步长。 太大时,更新可能越过低损失区域并持续振荡;太小时,方向即使正确,有限训练预算内也走不了多远。
不同方向的曲率差异会放大学习率的矛盾。考虑一个只有两个参数的二次目标:
方向比 方向陡 12 倍。固定学习率下,SGD 的两个坐标分别乘上 和 。为了让 方向稳定,需要满足 ;这时 方向可能下降得很慢。图中的 SGD 使用 ,所以 会交替越过谷底, 则缓慢接近零。
真实训练中的小批量还会让轨迹围绕平均下降方向波动。波动可能帮助参数离开狭窄区域,也可能使训练在低损失附近迟迟不能稳定。随着训练接近收敛点,常见做法是降低学习率,让更新幅度随之缩小。SGD 本身不会决定何时衰减,这部分由训练配置或学习率调度器负责。
Momentum
Momentum 为每个参数保存一个速度 。一种常见写法是
其中 , 为动量系数。展开递推式可以看到,当前速度是历史梯度的指数加权和:
若某个方向的梯度连续同号,这些项会累积,参数沿该方向走得更快。若梯度在狭长谷底两侧反复变号,旧项与新项会部分抵消,横跨谷底的振荡逐渐减弱。上图中的橙色轨迹仍有过冲,但振幅随更新缩小,并且比同图的 SGD 更快穿过平缓的 方向。
与学习率共同决定速度的尺度。增大动量后沿用原学习率,实际更新可能明显变大,因此两者需要一起调节。不同教材和框架还会把 放在梯度项前,或者用速度的相反数记号;公式外观会变化,历史方向被指数加权这一点不变。
Nesterov Momentum 在计算当前梯度时考虑速度将要到达的位置,使更新能较早感知前方曲率。PyTorch 通过 nesterov=True 提供这一变体,但它仍然依赖合适的学习率、动量与初始化。是否启用应以目标架构的训练配方和验证结果为准。
AdaGrad 与 RMSProp
SGD 和 Momentum 对所有坐标使用同一个学习率。若某些参数的梯度长期很大,另一些参数只偶尔收到稀疏梯度,一个全局尺度很难同时照顾它们。AdaGrad 为每个坐标累计梯度平方:
平方与除法都按元素进行。经常出现大梯度的坐标会积累较大的 ,后续步长随之变小;很少更新的坐标保留相对较大的步长。这一性质适合高维稀疏特征。代价是 只增不减,训练足够久后,所有有效步长都可能变得很小。
RMSProp 将无限累积改成梯度平方的指数移动平均:
较早的梯度会按 的幂次逐渐淡出,分母因而反映近期尺度。 用于避免除零并改善数值稳定性,它在低精度训练和极小梯度下可能实际影响更新,不能只把它看作形式上的常数。
AdaGrad 与 RMSProp 解决的是坐标尺度问题,分子仍以当前梯度为主。若还希望平滑梯度方向,需要再加入动量,或者使用同时维护一阶与二阶状态的 Adam。
Adam
Adam 为每个参数维护两个与参数形状相同的状态: 是梯度的一阶矩估计, 是梯度平方的二阶原始矩估计。
和 都初始化为零,这会让训练早期的移动平均偏向零。假设梯度暂时保持为常量 ,递推后有
当 接近 1 且 很小时, 明显小于 1。Adam 用下面两项消除这部分初始化偏差:
最终更新为
分子汇总近期方向,分母按坐标缩放更新。若一个坐标的梯度稳定同号, 保留较强信号;若符号频繁变化,一阶平均趋近零。梯度整体乘上常数时,分子与分母的尺度大体相消,但 、历史状态和有限精度使这种尺度不敏感性并不严格成立。
Adam 论文给出的常用起点是 、、 和 ,PyTorch 也沿用这组默认值。它们适合用作基线,不代表任何数据和模型上的最佳配置。较大的 会让二阶状态记住更长历史;训练总步数很少、梯度分布快速改变或梯度十分稀疏时,这段记忆可能需要重新评估。
Adam 的自适应缩放也没有消除学习率问题。学习率过大仍会越过低损失区域,过小仍会进展缓慢。原始 Adam 的收敛论证后来还出现了反例,AMSGrad 通过保留二阶状态的历史最大值,为特定在线凸优化条件提供了修正。深度网络属于非凸问题,实践中仍应依靠训练曲线、验证集和成熟配方判断。
AdamW
训练通常还会约束参数规模。给损失增加 L2 正则项后,目标变为
对应梯度为 。在最基本的 SGD 中,把这项放进更新可以整理为
参数每一步都乘上 ,这就是乘法形式的权重衰减。在 SGD 里,L2 正则与权重衰减经过系数换算可以得到相同更新。
Adam 会先用二阶状态对每个坐标缩放梯度。若把 直接加到 ,正则项会进入 和 ,随后再被各坐标不同的分母缩放。参数收缩因而依赖梯度历史,已经不等同于统一乘上 。
AdamW 将两条计算分开。 与 只根据数据损失的梯度更新,参数衰减单独执行:
“解耦”描述的是权重衰减与梯度预条件器的关系,它没有让 失去超参数含义。学习率变化时,每步收缩因子也会变化。许多训练配方还会把偏置和归一化层的缩放参数放进 weight_decay=0 的参数组;这是一项架构相关的经验选择,应遵循所用模型的已验证实现,而不能仅凭参数名称机械划分。
PyTorch 优化器
PyTorch 的训练循环把求导与更新分开。下面的 AdamW 配置使用论文和框架常见的初始取值,具体任务仍需验证学习率与权重衰减:
optimizer = torch.optim.AdamW(
model.parameters(),
lr=3e-4,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=1e-2,
)
for inputs, targets in train_loader:
optimizer.zero_grad(set_to_none=True)
logits = model(inputs)
loss = criterion(logits, targets)
loss.backward()
optimizer.step()
loss.backward() 只计算并累积梯度,optimizer.step() 才读取梯度、更新状态并改写参数。PyTorch 默认累积多次反向传播的结果,所以每个常规训练 step 前要清理旧梯度。set_to_none=True 不为所有参数写入零张量,通常能减少一次内存操作;没有参与当前计算的参数会保持 grad=None,优化器也可能跳过它们。
若使用梯度裁剪,位置应在 loss.backward() 之后、optimizer.step() 之前。学习率调度器的 step() 位置则取决于它按批次还是按 epoch 工作。混合精度训练还要让梯度缩放器先完成反缩放,再执行依赖真实梯度范数的裁剪。把这些操作按固定顺序写清楚,比将它们藏进一个大型训练辅助函数更容易排查问题。
各方法为每个参数保存的主要状态如下。表格忽略了 step 计数、可选的额外动量和框架实现细节。
| 方法 | 主要状态 | 坐标缩放方式 | 更新特点 |
|---|---|---|---|
| SGD | 无 | 全局学习率 | 直接使用当前小批量梯度 |
| SGD Momentum | 速度 | 全局学习率 | 指数累积历史方向 |
| AdaGrad | 平方和 | 分母持续增长,适合稀疏梯度 | |
| RMSProp | 平方移动平均 | 根据近期梯度尺度调整各坐标 | |
| Adam | 与 | 一阶矩除以二阶矩平方根 | 同时平滑方向并调整坐标尺度 |
| AdamW | 与 | 与 Adam 相同 | 自适应更新外单独收缩参数 |
这也解释了内存差异。纯 SGD 不保存与参数等大的优化器状态;Momentum 通常多一个状态张量;Adam 和 AdamW 通常多两个。大模型训练时,参数、梯度、优化器状态和低精度训练使用的主权重需要一起计算,优化器选择因此也会影响可容纳的模型规模。
优化器选择与诊断
成熟模型通常已有经过验证的训练配方。经典卷积网络常见 SGD、Momentum 与学习率衰减的组合,许多 Transformer 配方使用 AdamW、预热和后续衰减。沿用对应实现的基线,再改变一个因素观察结果,通常比从优化器名称开始猜测更可靠。
没有现成配方时,可以先用 AdamW 建立容易运行的基线,同时记录训练损失、验证指标、学习率和梯度范数。若任务对最终泛化、显存或既有部署流程有明确要求,再与经过合理调参的 SGD Momentum 比较。公平比较至少需要控制模型初始化、数据顺序、更新步数和评估流程,并允许每种优化器寻找适合自己的学习率;给所有方法强行使用同一个学习率没有可比性。
训练曲线提供的线索比单个最终数字更有用:
- 损失在最初几步变成
NaN或Inf时,先检查学习率、输入和标签范围、损失实现、数值精度与异常梯度。 - 损失几乎不动时,确认参数确实需要梯度、优化器持有正确参数,并检查学习率是否过小。激活饱和和数据管线错误也会产生相似表现。
- 损失在相邻 step 间大幅上下摆动时,学习率可能过大,动量也可能放大过冲。观察梯度范数和降低学习率比立即更换优化器更有诊断价值。
- 训练损失下降而验证指标持续恶化时,问题更接近过拟合或数据分布偏差。权重衰减、数据增强、早停和模型容量需要结合任务验证。
- 改变批量大小后,梯度噪声、每个 epoch 的更新次数和调度器时间尺度都会改变。原配置中的学习率不能默认原样复用。
优化器只根据它收到的梯度和状态执行更新。损失定义错误、数据泄漏、标签噪声或评估方式不一致,都无法靠 AdamW 自动修复。把参数更新公式、训练循环顺序和实验记录对齐,才能判断一次训练失败究竟发生在哪一层。
优化器更新规则
反向传播回答某个参数在当前位置应朝哪个方向变化,优化器决定如何组合当前与历史梯度,并把方向换成有限步长。SGD 保留最直接的基线;Momentum 用历史方向缓和振荡;AdaGrad 与 RMSProp 处理坐标尺度;Adam 将两类移动平均合在一次更新中;AdamW 进一步让权重衰减避开自适应状态。
这些方法共享同一个限制:它们只看到局部梯度。模型结构、初始化、归一化、批量大小、学习率日程和数据质量都会改变梯度序列。理解优化器的状态之后,调参就可以落到具体问题上:当前步长有多大,哪个历史量在起作用,参数衰减是否进入了预期路径,以及训练曲线是否支持继续更新。
参考文献
- H. Robbins and S. Monro. A Stochastic Approximation Method. The Annals of Mathematical Statistics, 1951. DOI: 10.1214/aoms/1177729586
- I. Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance of initialization and momentum in deep learning. ICML, 2013
- J. Duchi, E. Hazan, and Y. Singer. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. Journal of Machine Learning Research, 2011
- T. Tieleman and G. Hinton. Lecture 6.5: RMSProp, Divide the Gradient by a Running Average of Its Recent Magnitude. Neural Networks for Machine Learning course notes, 2012
- D. P. Kingma and J. Ba. Adam: A Method for Stochastic Optimization. ICLR, 2015
- I. Loshchilov and F. Hutter. Decoupled Weight Decay Regularization. ICLR, 2019
- S. J. Reddi, S. Kale, and S. Kumar. On the Convergence of Adam and Beyond. ICLR, 2018
- PyTorch Contributors. torch.optim: Optimization Algorithms. PyTorch Documentation