为什么LLaMA和Stable Diffusion都选择AdamW？深度解析优化器进化之路

在训练一个包含数十亿参数的大模型时，优化器的选择往往决定了整个项目的成败。2023年开源社区最耀眼的两个明星——Meta的LLaMA系列和Stability AI的Stable Diffusion，都不约而同地在官方实现中采用了AdamW优化器。这绝非偶然，而是深度学习优化器演进历程中的一个必然选择。

1. 从Adam到AdamW：优化器的关键进化

1.1 Adam优化器的核心机制

Adam（Adaptive Moment Estimation）自2014年提出以来，迅速成为深度学习领域的默认优化器。它巧妙融合了两种经典思想：

• 动量（Momentum）：通过梯度的一阶矩估计（移动平均）保持参数更新的惯性
• RMSProp：通过梯度的二阶矩估计（平方的移动平均）实现参数自适应的学习率

# Adam的典型实现伪代码
m = beta1 * m + (1 - beta1) * grad  # 一阶动量
v = beta2 * v + (1 - beta2) * grad**2  # 二阶动量
m_hat = m / (1 - beta1**t)  # 偏差修正
v_hat = v / (1 - beta2**t)
param -= lr * m_hat / (sqrt(v_hat) + eps)

这种设计使得Adam具备以下优势：

• 自动适应不同参数的特性
• 对初始学习率选择相对鲁棒
• 在稀疏梯度场景表现优异

1.2 Adam的致命缺陷：权重衰减的误用

然而在实践中，研究者们逐渐发现Adam存在一个关键问题——它对L2正则化（权重衰减）的处理方式存在理论缺陷。传统SGD中，权重衰减直接作用于参数：

θ = θ - η*(∇L + λθ)

但在Adam中，权重衰减项会被自适应学习率机制二次调整，导致：

优化器	权重衰减效果	实际表现
SGD	直接作用于参数	稳定可靠
Adam	被自适应学习率缩放	过弱或过强

这种耦合会导致模型要么欠正则化（过拟合），要么学习受阻（欠拟合），在大模型训练中尤为明显。

2. AdamW的革命性设计

2.1 解耦的核心思想

2017年ICLR论文《Decoupled Weight Decay Regularization》提出了AdamW优化器，其核心创新在于：

"将权重衰减从梯度更新中完全解耦，使其独立于自适应学习率机制"

这种设计带来了三个关键改进：

1. 真正的L2正则化效果
2. 更稳定的训练动态
3. 更好的泛化性能

2.2 数学形式的精妙差异

对比Adam和AdamW的更新公式：

# Adam更新（耦合权重衰减）
param -= lr * (m_hat/(sqrt(v_hat)+eps) + wd*param)

# AdamW更新（解耦权重衰减） 
param -= lr * m_hat/(sqrt(v_hat)+eps)
param -= lr * wd * param  # 独立操作

虽然看似微小，但这种解耦在大规模训练中会产生显著差异：

• 权重衰减不再受自适应学习率影响
• 正则化强度在不同参数层间保持一致
• 超参数（尤其是wd）变得更容易调节

3. 大模型为何偏爱AdamW

3.1 LLaMA的训练实践

Meta在训练LLaMA-2时特别强调了优化器选择：

• 使用AdamW而非原始Adam
• 设置β1=0.9，β2=0.95
• 权重衰减固定为0.1
• 学习率预热2000步

这种配置下，模型展现出：

• 更稳定的训练损失曲线
• 更少的梯度爆炸/消失现象
• 更好的最终zero-shot性能

3.2 Stable Diffusion的优化需求

扩散模型的训练面临独特挑战：

1. 时间步间的梯度差异极大
2. 需要平衡生成质量和多样性
3. 长时训练容易过拟合

AdamW通过以下方式应对这些挑战：

• 自适应学习率处理不同时间步的梯度变化
• 精确的权重衰减控制模型复杂度
• 解耦设计避免训练后期震荡

4. 实战中的AdamW调参技巧

4.1 超参数设置指南

基于主流大模型实践，推荐配置：

参数	推荐值	调整建议
β1	0.9	通常固定
β2	0.95-0.999	越大越稳定
权重衰减	0.01-0.1	大模型用较小值
学习率	1e-5-1e-4	配合预热策略
ε	1e-8	除非梯度特别小

4.2 实际训练中的注意事项

1. 学习率预热：大模型前1-2%的训练步骤应采用线性/余弦预热
2. 梯度裁剪：即使使用AdamW，仍建议设置梯度范数阈值（如1.0）
3. 权重衰减排除：对LayerNorm/bias参数应禁用权重衰减
4. 混合精度：配合AMP使用时需监控梯度缩放情况

# 典型AdamW实现示例（PyTorch）
optimizer = AdamW(
    model.parameters(),
    lr=5e-5,
    betas=(0.9, 0.999),
    weight_decay=0.01,
    eps=1e-8
)
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=2000,
    num_training_steps=100000
)

5. 超越AdamW：优化器的未来方向

虽然AdamW目前是大模型训练的事实标准，但研究社区仍在探索更优方案：

• Lion优化器：Google提出的符号动量方法，在部分任务上表现更好
• Sophia：引入曲率信息的二阶优化方法
• Adan：通过嵌套动量加速收敛

不过在实际项目中，当面对数十亿参数的大模型训练时，AdamW依然是那个最稳妥的选择——它可能不是每个指标上的最优解，但几乎永远不会成为系统中最薄弱的环节。