写在前面

在上一篇文章中,我们讨论了如何用数据并行、张量并行和流水线并行,把一个大模型“拆开”放到成百上千张 GPU 上训练。但我们始终回避了一个根本问题:拆好之后,参数究竟是怎么被“更新”的?

这个问题的答案,就是优化器(Optimizer)。如果说神经网络是一张通往“智能”的复杂地图,损失函数决定了哪条路通向最优解,那么优化器就是指引模型在崎岖的损失地貌中如何一步步下山的“导航算法”

它不参与前向计算,也不定义损失函数,却掌控着每一次参数更新的方向与步长——就像登山者在浓雾中选择的下山策略,决定了你能否又快又稳地抵达谷底。

从最朴素的随机梯度下降(SGD),到集大成者 Adam,再到如今大模型时代的标配 AdamW,优化器的演变路径几乎就是深度学习工程史的一个缩影。但你是否真正想过这些问题:

  • SGD 上加一个“动量”到底解决了什么?
  • Adam 相比 SGD 最大的优势是什么,又有什么隐忧?
  • 为什么 Adam 已经如此强大,业界还要推出 AdamW?
  • 同样是“自适应”,各种 Adam 变体之间究竟有什么区别?
  • 在实际训练中,到底该选哪个优化器?

优化算法的重要性不言而喻——斯坦福大学的 CS336 课程作业中,与模型结构(Transformer)并列的训练基础设施就是损失函数优化器(明确要求使用 AdamW)。MLSys 课程的“优化器”模块也将 SGD、Adam 和 AdamW 列为三大必须掌握的“生产级优化器”。

一、起点:梯度下降与随机梯度下降(SGD)

1.1 梯度下降:最朴素的“向下走”

所有优化器的起点,都是梯度下降(Gradient Descent)

想象你站在一座浓雾笼罩的山上(损失函数的高维曲面),能见度极低,只能感知脚下局部的坡度。你想尽快下到谷底(即找到全局最小损失值),最朴素的想法就是:找坡度最陡的方向,沿着它往下走

数学上,这个“最陡的方向”就是梯度的负方向:

θ_{t+1} = θ_t - η · ∇L(θ_t)
  • θ_t 是当前参数
  • ∇L(θ_t) 是损失函数关于参数的梯度
  • η学习率(Learning Rate) ,即每一步迈出去的“步长”

完整梯度下降法要求批梯度下降每次迭代都遍历全部训练数据以计算梯度,对当今动辄 TB 级别的数据集来说完全不可行。好在可以“偷工减料”——随机抽取一小批(Mini-Batch)样本来近似估计全部数据上的真实梯度,于是,诞生了随机梯度下降(Stochastic Gradient Descent, SGD)

θ_{t+1} = θ_t - η · ∇L_B(θ_t)

其中 ∇L_B 是基于当前小批量(Batch)数据计算出的梯度。

1.2 SGD 的三大痛点

现代深度学习所依赖的高度非凸损失曲面,让朴素 SGD 暴露了三个显眼的短板:

  1. 震荡问题:在高维非凸的深度学习损失曲面中,不同参数的曲率差异极大,梯度在“狭窄峡谷”(Ravine)中会在参数不同轴反复震荡而偏移最优方向,严重拖慢收敛。
  2. 平坦区停滞:当参数陷入平稳高原时,梯度接近零,SGD 会几乎“原地踏步”,导致漫长的训练时间。
  3. 对学习率极其敏感:学习率太小,收敛如龟速;学习率太大,模型可能在最优解附近“反复横跳”甚至发散。

这三个问题的核心都在于:SGD 是“近视”的。它在每个位置仅依据“当前这一步的小批量梯度”做决策,对曾经走过的历史路径毫无记忆。为了突破这一限制,科学家们开始给 SGD 加上了“记忆”和“自适应”的能力。

二、第一道加法:动量 —— 为优化器注入“惯性”

2.1 Momentum:让梯度产生“惯性”

动量法(Momentum) 是优化器演进的第一道关键加法。它引入一个“速度”变量 v_t,让参数更新方向不仅取决于当前小批量随机梯度,还保留了过去所有时刻梯度贡献的指数移动平均:

v_t = β · v_{t-1} + (1 - β) · ∇L(θ_t)
θ_{t+1} = θ_t - η · v_t

这里的 β 是动量系数,通常取 0.9,控制历史信息的衰减速度。

直观上,这相当于给参数更新赋予“惯性”——即使当前梯度很小,只要过去方向一致,更新仍会继续推进。

动量法为优化器补上了“纵向记忆”——它能显著抑制在狭窄峡谷中的横向震荡,并加速穿越平坦高原区域。在 ResNet、VGG 等经典 CNN 上,动量法显著提升了训练的稳定性和收敛速度。甚至连 2012 年的 AlexNet 也依赖于 SGD+M 才能有效训练。

2.2 NAG:先探路,再行动

标准动量法仍存在一个潜在缺陷:它在已累积的动量和当前位置的梯度两股力共同作用下“冲过头”,然后在反向梯度下折返,表现为反复超调(Overshoot)。

Nesterov Accelerated Gradient(NAG) 对此进行了精妙的修正:它先根据当前动量预测下一步会“滚”到哪个临时位置,再在那个前瞻坐标上计算梯度:

临时位置:  θ̃_t = θ_t - η · β · v_{t-1}
计算前瞻梯度: ∇L(θ̃_t)
速度更新:  v_t = β · v_{t-1} + (1 - β) · ∇L(θ̃_t)
参数更新:  θ_{t+1} = θ_t - η · v_t

这种“先探路、再行动”的前瞻性,使 NAG 在接近极小值时会提前减速,相比标准动量更稳定、更精细。理论分析也证明,NAG 在凸优化中拥有比标准动量更好的收敛速度。Pytorch 等框架可通过 nesterov=True 直接启用。

不过在一个物理方向上做再多加法,也无法一劳永逸,因为深度学习真正的“尺度失衡”问题,动量法从未触及。

三、第二道加法:自适应学习率

在深度神经网络里,不同层的参数、甚至同一层内不同列参数的梯度在数量和尺度上也千差万别。固定全局的学习率对一部分参数刚好合适,对另一部分参数却可能导致震荡或停滞。与其用同一个学习率控制所有参数,不如为每个参数单独计算最适合它的“步幅”——这就是自适应优化器的核心思想。

3.1 AdaGrad:自适应法的“第一个脚印”

2011 年提出的 AdaGrad 首次实现了参数级学习率调整。它为每个参数维护一个累积梯度平方和 G_t

G_t = G_{t-1} + (∇L(θ_t))²
θ_{t+1} = θ_t - η · ∇L(θ_t) / √(G_t + ε)

参数被更新的频率越高、梯度的平方和越大,其有效学习率会被缩得更小,对矫正极度不平衡的特征极其有效。在处理推荐系统等特征高度稀疏的场景时,AdaGrad 优势明显。

然而它的致命缺陷在于:平滑累积的方式让 G_t 只能随时间单调增长,导致训练的中后期有效学习率过早衰减至极小,参数几乎停止更新。

3.2 RMSProp:用“近大远小”替代“单调累积”

RMSProp 在 2012 年解决了 AdaGrad 的“学习率早衰”问题。它不再傻傻地累积全部历史梯度平方,而是采用指数加权移动平均(EWMA)

v_t = β₂ · v_{t-1} + (1 - β₂) · (∇L(θ_t))²
θ_{t+1} = θ_t - η · ∇L(θ_t) / √(v_t + ε)

其中 β₂ 通常取 0.999。指数衰减赋予了优化器“近期偏好”,让自适应学习率主要被近期的梯度平方均方根(RMS)所主导,有效规避了 AdaGrad 学习率过早衰减的问题。

RMSProp 的效果显著:它允许为那些梯度变化剧烈的参数动态减小步幅,同时为平坦参数积极放大步长,至今仍是 RNN 训练的有效组件。

四、集大成者:Adam —— 动量与自适应的完美融合

4.1 Adam 的“三合一”架构

理解了动量法和自适应学习率,再来看 Adam 就会有一种“水到渠成”的豁然开朗感。2015 年提出的 Adam动量(一阶矩)自适应学习率(二阶矩)无缝融合,并加入偏差校正机制

Adam 在每一步维护两个指数移动平均:

  • 一阶矩 m_t(动量项,梯度的指数滑动平均)
  • 二阶矩 v_t(梯度平方的指数滑动平均,来自 RMSProp)
m_t = β₁ · m_{t-1} + (1 - β₁) · g_t
v_t = β₂ · v_{t-1} + (1 - β₂) · g_t²

由于 m_0v_0 被初始化为 0,在训练早期它们的估计值被系统性地拉低。为解决这一问题,Adam 引入了偏差校正(Bias Correction)

m̂_t = m_t / (1 - β₁ᵗ)
v̂_t = v_t / (1 - β₂ᵗ)

最终参数更新公式为:

θ_{t+1} = θ_t - η · m̂_t / (√(v̂_t) + ε)

Adam 的默认超参数在大量实践中被证明极为鲁棒:η = 0.001β₁ = 0.9β₂ = 0.999ε = 10⁻⁸,几乎“开箱即用”。在 BERT、GPT 等 Transformer 架构的预训练过程中,Adam 凭借其自适应特性,被证明显著优于 SGD

4.2 在语言模型中的“降维打击”:为什么 Transformer 几乎只认 Adam?

一个有趣的现象是:在计算机视觉任务中,精心调参的 SGD + Momentum 通常能达到与 Adam 相当甚至更优的最终性能;但在 Transformer 语言模型的预训练中,Adam 对 SGD 却形成了近乎“降维打击”的优势

大量研究一致指出了同样的核心结论:在语言模型训练中,除了 SGD 明显偏差较大以外,Adam 与其他主流自适应优化器(Lion、Adafactor 加动量、Signum)的性能与超参数鲁棒性都非常接近

Why?原因之一在于 Transformer 架构内部梯度尺度的极端异质性——不同参数矩阵(如 Q、K、V 的投影矩阵,FFN 的放大矩阵)在训练早期的梯度尺度可能相差几个数量级。SGD 的全局统一学习率在这种“多尺度”环境中进退两难:学习率小了,某些层几乎原地不动;大了,另一些层又会梯度爆炸。Adam 的自适应机制恰好能自动平衡各层的有效更新步长,使其成为 Transformer 训练的“必备品”。

4.3 Adam 的系统代价:3 倍的显存

但自适应优化器精妙的算法背后是需要付出系统代价的。在计算和存储依赖上,Adam 需要额外维护两个和参数同样大小的一阶与二阶矩状态,总体上Adam 的总内存占用约为模型参数的 3 倍(参数 + m + v)。对于 70B 级别的稠密 LLM,这意味着大约 420GB 的纯优化器状态,即便放在 H100 集群上也是巨大的负担。

五、AdamW:大模型时代的“灵魂修补”

5.1 L2 正则化 vs 权重衰减:为何在 Adam 中混为一谈会翻车?

在进入 AdamW 之前,需要厘清一个深层概念:权重衰减(Weight Decay)L2 正则化在标准 SGD 下是完全等价的,但在 Adam 这类自适应优化器中,两者却出现了微妙而关键的差别。

在 SGD 中,权重衰减直接在更新步骤中按比例收缩参数:θ_{t+1} = θ_t - η · (∇L(θ_t) + λ·θ_t)。这等价于在损失函数中加一个 L2 惩罚项。

但在 Adam 中,自适应学习率导致“L2 正则化作为损失项导出的梯度”与“自适应学习率分母”发生了不当耦合:对于那些自适应学习率被放大的参数,L2 正则化的实际惩罚效应被过度增强;反之则被严重削弱。结果就是 Adam 的权重衰减效果“飘忽不定”,远不如在 SGD 中那般稳定可靠。

5.2 解耦:将“防过拟合”与“参数更新”分离

2019 年,Loshchilov 和 Hutter 提出了 AdamW,只做了一件干干净净的事:把权重衰减从自适应梯度计算中彻底解耦出来

# 标准 Adam(错误耦合版本)
g_t = ∇L(θ_{t-1}) + λ · θ_{t-1}

# AdamW(正确的解耦版本)
g_t = ∇L(θ_{t-1})
▸ 使用 g_t 正常计算 Adam 的一阶/二阶矩和自适应更新
▸ 在 Adam 更新完成之后,再额外执行一步独立的 θ_t = θ_t - η · λ · θ_t

解耦后的直接收益是:权重衰减的超参数 λ 在不同任务和模型间具有更强的可迁移性;训练更稳定,尤其在大型 Transformer 模型上优势明显。研究表明,AdamW 相比标准 Adam 能取得更好的收敛效果与泛化能力,尤其在 BERT、GPT 等大规模应用中。

正是因为完成了解耦操作,AdamW 才成为现代 Transformer 训练的默认优化器。CS336 从零构建 LLM 的课程与作业,也直接选用了 AdamW 作为标准配置。

5.3 资源受限时的小提示

在资源受限、需要进行小幅度过参数扫描时,可考虑用 Lion 替代 AdamW 节省约 50% 的显存。CS336 课程在实验类作业中亦明确支持使用 Lion 以减少内存占用。

六、最新前沿:内存、Scaling Law 与符号进化

即使在 AdamW 已几乎触及算法优化的顶峰后,对“更极致优化器”的探索脚步仍在加速。当前的探索围绕三个方向进行攻坚:

6.1 系统瓶颈攻坚战:内存墙

Adam 最尖锐的问题,不是性能,而是内存。其二阶矩状态让大模型训练的内存开销成为瓶颈。因此前沿研究正在探索“更轻量级的 Adam”。

  • APOLLO(低秩放缩近似):通过纯随机投影构造一个低秩的辅助优化器状态来近似 Adam 的自适应学习率缩放规则。其 rank-1 版本(APOLLO-Mini)即可在 SGD 级内存成本上实现媲美 AdamW 的预训练性能,在 8×A100-80GB 的任务上吞吐量直接达到 AdamW 的 3 倍
  • SCALE(列归一+末层动量裁剪):通过更激进的简化策略——对梯度沿列维度归一化 + 只在梯度方差最大的输出层加一阶动量——SCALE 仅用 35-45% 的总内存就匹配甚至超越了全 Adam 性能。
  • GaLore:在 Adam 优化器的更新过程中加入低秩自适应,以更小的参数量模拟 Adam 的行为,降低对大内存的需求。
  • ZeRO 优化器状态分片:与数据并行深度融合,将 Adam 的优化器状态切分到所有 GPU 上,使单卡存储压力分散,实现近乎 O(1/N) 的优化器内存占用。

6.2 Scaling Law 的探索:超参数与模型规模的隐藏因果

前沿研究还发现,Adam 的自适应能力会随模型规模与 Batch Size 产生显著变化。在小 Batch 训练中调参准确的 SGD+M 能匹敌 Adam;随着 Batch Size 扩大,Adam 的优势开始拉大。

进一步的超参数研究发现,随着模型规模和训练 Token 数增加,Adam 的最优 β₂ 和学习率调度策略可能需要遵循特定的 Scaling Law 进行调整。研究者们正在系统挖掘这些规律,试图为万亿参数级的优化器提供批量自动缩放的“配方”。

6.3 符号式搜索进化:Lion、Sophia 与符号回归

最新的突破甚至在改变优化器的“诞生方式”——不再依赖人脑设计,而是通过程序自动搜索。

Lion(Evolved Sign Momentum) 是谷歌通过大规模程序搜索发现的新型优化器。它只保留一阶动量并采用符号运算,内存效率远高于 Adam,在同等计算资源下训练 BERT 收敛速度提升 18%,但在预训练任务中可以匹敌甚至超越 AdamW。

Sophia 则结合了对角二阶近似与 EMA 梯度平滑,在语言模型预训练上展现出了优于 AdamW 的收敛速度。

七、总结与展望

7.1 整体演进图谱

回顾这七十年优化器的演变,我们能看到两条清晰的主线在 2015 年合流:

阶段代表优化器核心突破关键缺陷
经典时代(20世纪50年代-2010)SGD随机逼近,数据效率大幅提升高维震荡与高原停滞
动量时代(2011-2012)Momentum, NAG震荡抑制与前瞻机制尺度失衡问题仍未解决
自适应前夜(2011-2013)AdaGrad, RMSProp历史平方梯度的参数级放缩学习率早衰(AdaGrad)
一阶+二阶融合(2014-2015)Adam动量 + RMSProp + 偏差校正Adam 不规范权重衰减
权重衰减解耦(2019)AdamW将权重衰减与梯度更新完全分离LLM 默认之选
自适应探索时代(2023+)Lion, Sophia, APOLLO, SCALE低内存/自动搜索,迈向极致系统效率新方法仍在业普适化迭代中

7.2 实战选择指南

在真实项目中,该如何选择?

场景推荐加速器理由
Transformer / LLM 从头预训AdamW无需争议的工业界 / 学术界共识
中小型语言模型的实验性微调AdamWLionAdamW 稳定强大,Lion 节约显存较快
计算机视觉(CNN / ResNet / ViT)SGD+MAdamWSGD+M 广泛有效、泛化好;大规模 ViT/多模态优先选用 AdamW
推荐系统 / 稀疏特征场景Adam / AdagradAdagrad 稀疏处理能力强,Adam 亦可
边缘设备 / 显存极度受限Lion / APOLLO / SCALE兼顾吞吐量且大幅减少显存占用

7.3 优化器在整个训练体系中的位置

回顾我们系列走过的路:

  1. Tokenizer → Embedding → RoPE
  2. Attention(MHA/MQA/GQA)→ FFN(SwiGLU)
  3. 预归一化/后归一化 + 残差连接
  4. MoE(专家与路由机制)
  5. KV Cache 与 PagedAttention
  6. 数据并行、张量并行、流水线并行
  7. 优化器:优化状态下降路径(本文)

如果说前面几篇文章是在组装一辆“智能赛车”的发动机、传动系统和燃料管道,那么优化器就是驾驶赛车的“AI 车手”——它决定了在油箱加满、道路崎岖的条件下,你能不能以最高效率,最快一圈冲向终点。

作者

884705373@qq.com

相关文章

QLoRA微调原理详解:与LoRA的性能与内存对比

引言:为什么大模型微调需要QLoRA? 在深...

读出全部

关于Norm的解析

可以说,如果没有残差连接和 Layer No...

读出全部