前面的文章已经把损失函数沿计算图反向求导,PyTorch 用一行 loss.backward() 将结果累积到各个参数的 grad 中。训练循环紧接着还有一行 optimizer.step()。模型参数真正发生变化,就在这一步。

梯度只描述当前位置附近,损失增长最快的方向。优化器要把它变成有限长度的更新,还要处理小批量采样带来的波动、不同参数的尺度差异,以及此前若干步留下的信息。SGD 直接使用当前梯度;Momentum 记住近期方向;RMSProp 为每个坐标估计梯度尺度;Adam 同时维护方向和尺度;AdamW 又把权重衰减从自适应更新中分离出来。

本文从同一条参数更新公式出发,依次拆开这些状态。目标是读懂每个超参数在更新中处于什么位置,以及训练异常时应该先检查什么。示意实验只解释算法的更新几何,不比较模型精度,也不据此判断哪种优化器普遍更好。

梯度估计与参数更新

设训练集包含 NN 个样本,模型参数为 θ\theta,平均经验损失写成

L(θ)=1Ni=1Ni(θ).L(\theta)=\frac{1}{N}\sum_{i=1}^{N}\ell_i(\theta).

完整梯度需要遍历全部样本。神经网络训练通常从数据中取一个大小为 BB 的小批量 Bt\mathcal{B}_t,在第 tt 步计算

gt=1BiBtθi(θt1).g_t = \frac{1}{B}\sum_{i\in\mathcal{B}_t} \nabla_\theta \ell_i(\theta_{t-1}).

当样本按相同概率随机抽取时,gtg_t 的期望对应完整梯度;某一个批次仍会偏离它。批量越小,这种波动通常越明显。数据分布、采样方式、数据增强和损失缩放也会改变实际看到的梯度,不能只用批量大小解释所有噪声。

优化器根据当前梯度与内部状态产生更新向量 utu_t,再用学习率 ηt\eta_t 控制步长:

ut=Optimizer(θt1,gt,statet1),θt=θt1ηtut.u_t=\operatorname{Optimizer}(\theta_{t-1},g_t,\text{state}_{t-1}), \qquad \theta_t=\theta_{t-1}-\eta_t u_t.

一次这样的更新称为一个 step 或 iteration。一个 epoch 表示训练数据大致被遍历一遍,通常包含许多 step。学习率调度器究竟按 step 还是按 epoch 更新,必须与具体实现保持一致;混淆两者会让衰减速度相差一个 epoch 所含的批次数。

SGD

随机梯度下降(Stochastic Gradient Descent,SGD)的基本形式没有额外状态:

θt=θt1ηtgt.\theta_t=\theta_{t-1}-\eta_t g_t.

负梯度是当前点在欧氏几何下最陡的下降方向,这个结论只描述足够小的局部移动。学习率把方向变成实际步长。ηt\eta_t 太大时,更新可能越过低损失区域并持续振荡;太小时,方向即使正确,有限训练预算内也走不了多远。

不同方向的曲率差异会放大学习率的矛盾。考虑一个只有两个参数的二次目标:

L(x,y)=12(x2+12y2),L(x,y)=(x,12y).L(x,y)=\frac{1}{2}\left(x^2+12y^2\right), \qquad \nabla L(x,y)=(x,12y).

yy 方向比 xx 方向陡 12 倍。固定学习率下,SGD 的两个坐标分别乘上 1η1-\eta112η1-12\eta。为了让 yy 方向稳定,需要满足 0<η<1/60<\eta<1/6;这时 xx 方向可能下降得很慢。图中的 SGD 使用 η=0.12\eta=0.12,所以 yy 会交替越过谷底,xx 则缓慢接近零。

狭长椭圆等高线上的三条优化轨迹从左上角同一点出发,SGD 先沿陡峭方向来回跨越谷底再缓慢横移,Momentum 以逐渐收窄的振荡接近中心,Adam 经过坐标缩放后沿弧线靠近中心
同一二次目标上的 72 次确定性更新。三种方法分别使用图中标注的稳定超参数,轨迹用于观察更新几何;不同学习率没有构成受控的性能比较。作者计算并绘制

真实训练中的小批量还会让轨迹围绕平均下降方向波动。波动可能帮助参数离开狭窄区域,也可能使训练在低损失附近迟迟不能稳定。随着训练接近收敛点,常见做法是降低学习率,让更新幅度随之缩小。SGD 本身不会决定何时衰减,这部分由训练配置或学习率调度器负责。

Momentum

Momentum 为每个参数保存一个速度 vtv_t。一种常见写法是

vt=μvt1+gt,θt=θt1ηtvt,v_t=\mu v_{t-1}+g_t, \qquad \theta_t=\theta_{t-1}-\eta_t v_t,

其中 v0=0v_0=0μ[0,1)\mu\in[0,1) 为动量系数。展开递推式可以看到,当前速度是历史梯度的指数加权和:

vt=gt+μgt1+μ2gt2+.v_t = g_t+\mu g_{t-1}+\mu^2g_{t-2}+\cdots.

若某个方向的梯度连续同号,这些项会累积,参数沿该方向走得更快。若梯度在狭长谷底两侧反复变号,旧项与新项会部分抵消,横跨谷底的振荡逐渐减弱。上图中的橙色轨迹仍有过冲,但振幅随更新缩小,并且比同图的 SGD 更快穿过平缓的 xx 方向。

μ\mu 与学习率共同决定速度的尺度。增大动量后沿用原学习率,实际更新可能明显变大,因此两者需要一起调节。不同教材和框架还会把 (1μ)(1-\mu) 放在梯度项前,或者用速度的相反数记号;公式外观会变化,历史方向被指数加权这一点不变。

Nesterov Momentum 在计算当前梯度时考虑速度将要到达的位置,使更新能较早感知前方曲率。PyTorch 通过 nesterov=True 提供这一变体,但它仍然依赖合适的学习率、动量与初始化。是否启用应以目标架构的训练配方和验证结果为准。

AdaGrad 与 RMSProp

SGD 和 Momentum 对所有坐标使用同一个学习率。若某些参数的梯度长期很大,另一些参数只偶尔收到稀疏梯度,一个全局尺度很难同时照顾它们。AdaGrad 为每个坐标累计梯度平方:

st=st1+gt2,θt=θt1ηgtst+ϵ.s_t=s_{t-1}+g_t^2, \qquad \theta_t = \theta_{t-1} -\eta\frac{g_t}{\sqrt{s_t}+\epsilon}.

平方与除法都按元素进行。经常出现大梯度的坐标会积累较大的 sts_t,后续步长随之变小;很少更新的坐标保留相对较大的步长。这一性质适合高维稀疏特征。代价是 sts_t 只增不减,训练足够久后,所有有效步长都可能变得很小。

RMSProp 将无限累积改成梯度平方的指数移动平均:

rt=ρrt1+(1ρ)gt2,θt=θt1ηgtrt+ϵ.r_t = \rho r_{t-1}+(1-\rho)g_t^2, \qquad \theta_t = \theta_{t-1} -\eta\frac{g_t}{\sqrt{r_t}+\epsilon}.

较早的梯度会按 ρ\rho 的幂次逐渐淡出,分母因而反映近期尺度。ϵ\epsilon 用于避免除零并改善数值稳定性,它在低精度训练和极小梯度下可能实际影响更新,不能只把它看作形式上的常数。

AdaGrad 与 RMSProp 解决的是坐标尺度问题,分子仍以当前梯度为主。若还希望平滑梯度方向,需要再加入动量,或者使用同时维护一阶与二阶状态的 Adam。

Adam

Adam 为每个参数维护两个与参数形状相同的状态:mtm_t 是梯度的一阶矩估计,vtv_t 是梯度平方的二阶原始矩估计。

mt=β1mt1+(1β1)gt,m_t = \beta_1m_{t-1}+(1-\beta_1)g_t, vt=β2vt1+(1β2)gt2.v_t = \beta_2v_{t-1}+(1-\beta_2)g_t^2.

m0m_0v0v_0 都初始化为零,这会让训练早期的移动平均偏向零。假设梯度暂时保持为常量 gg,递推后有

mt=(1β1t)g.m_t=(1-\beta_1^t)g.

β1\beta_1 接近 1 且 tt 很小时,1β1t1-\beta_1^t 明显小于 1。Adam 用下面两项消除这部分初始化偏差:

m^t=mt1β1t,v^t=vt1β2t.\widehat{m}_t=\frac{m_t}{1-\beta_1^t}, \qquad \widehat{v}_t=\frac{v_t}{1-\beta_2^t}.

最终更新为

θt=θt1ηm^tv^t+ϵ.\theta_t = \theta_{t-1} -\eta \frac{\widehat{m}_t}{\sqrt{\widehat{v}_t}+\epsilon}.

分子汇总近期方向,分母按坐标缩放更新。若一个坐标的梯度稳定同号,m^t\widehat{m}_t 保留较强信号;若符号频繁变化,一阶平均趋近零。梯度整体乘上常数时,分子与分母的尺度大体相消,但 ϵ\epsilon、历史状态和有限精度使这种尺度不敏感性并不严格成立。

Adam 论文给出的常用起点是 η=103\eta=10^{-3}β1=0.9\beta_1=0.9β2=0.999\beta_2=0.999ϵ=108\epsilon=10^{-8},PyTorch 也沿用这组默认值。它们适合用作基线,不代表任何数据和模型上的最佳配置。较大的 β2\beta_2 会让二阶状态记住更长历史;训练总步数很少、梯度分布快速改变或梯度十分稀疏时,这段记忆可能需要重新评估。

Adam 的自适应缩放也没有消除学习率问题。学习率过大仍会越过低损失区域,过小仍会进展缓慢。原始 Adam 的收敛论证后来还出现了反例,AMSGrad 通过保留二阶状态的历史最大值,为特定在线凸优化条件提供了修正。深度网络属于非凸问题,实践中仍应依靠训练曲线、验证集和成熟配方判断。

AdamW

训练通常还会约束参数规模。给损失增加 L2 正则项后,目标变为

Lλ(θ)=L(θ)+λ2θ22,L_\lambda(\theta) = L(\theta)+\frac{\lambda}{2}\lVert\theta\rVert_2^2,

对应梯度为 gt+λθt1g_t+\lambda\theta_{t-1}。在最基本的 SGD 中,把这项放进更新可以整理为

θt=θt1η(gt+λθt1)=(1ηλ)θt1ηgt.\theta_t = \theta_{t-1} -\eta(g_t+\lambda\theta_{t-1}) = (1-\eta\lambda)\theta_{t-1}-\eta g_t.

参数每一步都乘上 1ηλ1-\eta\lambda,这就是乘法形式的权重衰减。在 SGD 里,L2 正则与权重衰减经过系数换算可以得到相同更新。

Adam 会先用二阶状态对每个坐标缩放梯度。若把 λθ\lambda\theta 直接加到 gtg_t,正则项会进入 mtm_tvtv_t,随后再被各坐标不同的分母缩放。参数收缩因而依赖梯度历史,已经不等同于统一乘上 1ηλ1-\eta\lambda

AdamW 将两条计算分开。mtm_tvtv_t 只根据数据损失的梯度更新,参数衰减单独执行:

θ~t=(1ηλ)θt1,\widetilde{\theta}_t =(1-\eta\lambda)\theta_{t-1}, θt=θ~tηm^tv^t+ϵ.\theta_t = \widetilde{\theta}_t -\eta \frac{\widehat{m}_t}{\sqrt{\widehat{v}_t}+\epsilon}.
左右两幅流程图比较 Adam 中的 L2 正则与 AdamW 权重衰减,左侧的 lambda theta 先加入梯度并进入一阶矩和二阶矩,右侧数据梯度单独更新矩估计,参数沿另一条支路乘以一减 eta lambda 后再与自适应更新合并
在 Adam 中直接加入 L2 梯度时,正则项会参与两类移动平均;AdamW 让参数收缩绕开自适应状态。作者绘制

“解耦”描述的是权重衰减与梯度预条件器的关系,它没有让 λ\lambda 失去超参数含义。学习率变化时,每步收缩因子也会变化。许多训练配方还会把偏置和归一化层的缩放参数放进 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速度 vtv_t全局学习率指数累积历史方向
AdaGrad平方和 sts_t1/(st+ϵ)1/(\sqrt{s_t}+\epsilon)分母持续增长,适合稀疏梯度
RMSProp平方移动平均 rtr_t1/(rt+ϵ)1/(\sqrt{r_t}+\epsilon)根据近期梯度尺度调整各坐标
Adammtm_tvtv_t一阶矩除以二阶矩平方根同时平滑方向并调整坐标尺度
AdamWmtm_tvtv_t与 Adam 相同自适应更新外单独收缩参数

这也解释了内存差异。纯 SGD 不保存与参数等大的优化器状态;Momentum 通常多一个状态张量;Adam 和 AdamW 通常多两个。大模型训练时,参数、梯度、优化器状态和低精度训练使用的主权重需要一起计算,优化器选择因此也会影响可容纳的模型规模。

优化器选择与诊断

成熟模型通常已有经过验证的训练配方。经典卷积网络常见 SGD、Momentum 与学习率衰减的组合,许多 Transformer 配方使用 AdamW、预热和后续衰减。沿用对应实现的基线,再改变一个因素观察结果,通常比从优化器名称开始猜测更可靠。

没有现成配方时,可以先用 AdamW 建立容易运行的基线,同时记录训练损失、验证指标、学习率和梯度范数。若任务对最终泛化、显存或既有部署流程有明确要求,再与经过合理调参的 SGD Momentum 比较。公平比较至少需要控制模型初始化、数据顺序、更新步数和评估流程,并允许每种优化器寻找适合自己的学习率;给所有方法强行使用同一个学习率没有可比性。

训练曲线提供的线索比单个最终数字更有用:

  • 损失在最初几步变成 NaNInf 时,先检查学习率、输入和标签范围、损失实现、数值精度与异常梯度。
  • 损失几乎不动时,确认参数确实需要梯度、优化器持有正确参数,并检查学习率是否过小。激活饱和和数据管线错误也会产生相似表现。
  • 损失在相邻 step 间大幅上下摆动时,学习率可能过大,动量也可能放大过冲。观察梯度范数和降低学习率比立即更换优化器更有诊断价值。
  • 训练损失下降而验证指标持续恶化时,问题更接近过拟合或数据分布偏差。权重衰减、数据增强、早停和模型容量需要结合任务验证。
  • 改变批量大小后,梯度噪声、每个 epoch 的更新次数和调度器时间尺度都会改变。原配置中的学习率不能默认原样复用。

优化器只根据它收到的梯度和状态执行更新。损失定义错误、数据泄漏、标签噪声或评估方式不一致,都无法靠 AdamW 自动修复。把参数更新公式、训练循环顺序和实验记录对齐,才能判断一次训练失败究竟发生在哪一层。

优化器更新规则

反向传播回答某个参数在当前位置应朝哪个方向变化,优化器决定如何组合当前与历史梯度,并把方向换成有限步长。SGD 保留最直接的基线;Momentum 用历史方向缓和振荡;AdaGrad 与 RMSProp 处理坐标尺度;Adam 将两类移动平均合在一次更新中;AdamW 进一步让权重衰减避开自适应状态。

这些方法共享同一个限制:它们只看到局部梯度。模型结构、初始化、归一化、批量大小、学习率日程和数据质量都会改变梯度序列。理解优化器的状态之后,调参就可以落到具体问题上:当前步长有多大,哪个历史量在起作用,参数衰减是否进入了预期路径,以及训练曲线是否支持继续更新。

参考文献

  1. H. Robbins and S. Monro. A Stochastic Approximation Method. The Annals of Mathematical Statistics, 1951. DOI: 10.1214/aoms/1177729586
  2. I. Sutskever, J. Martens, G. Dahl, and G. Hinton. On the importance of initialization and momentum in deep learning. ICML, 2013
  3. J. Duchi, E. Hazan, and Y. Singer. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. Journal of Machine Learning Research, 2011
  4. 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
  5. D. P. Kingma and J. Ba. Adam: A Method for Stochastic Optimization. ICLR, 2015
  6. I. Loshchilov and F. Hutter. Decoupled Weight Decay Regularization. ICLR, 2019
  7. S. J. Reddi, S. Kale, and S. Kumar. On the Convergence of Adam and Beyond. ICLR, 2018
  8. PyTorch Contributors. torch.optim: Optimization Algorithms. PyTorch Documentation