Adam Optimizer:从梯度下降到自适应学习率
训练神经网络时,我们经常会看到下面这行代码:
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
它很短,却承担着一个关键任务:根据反向传播得到的梯度,不断更新模型中的权重和偏置。
Adam 的全称是 Adaptive Moment Estimation,可以翻译为“自适应矩估计”。它同时利用梯度的移动平均和梯度平方的移动平均,为不同参数自动调整更新方向与步长。
一句话概括:
Adam 在动量方法的基础上记录梯度方向,又根据历史梯度大小为每个参数调整学习率,因此通常比普通 SGD 更容易上手。
1. 优化器到底负责什么
神经网络训练通常包含三个核心步骤:
optimizer.zero_grad()
loss.backward()
optimizer.step()
它们的职责并不相同:
| 代码 | 作用 |
|---|---|
optimizer.zero_grad() |
清空上一轮保存在参数上的梯度 |
loss.backward() |
根据计算图和链式法则计算梯度 |
optimizer.step() |
根据梯度和优化算法更新参数 |
需要特别注意:
loss.backward()只负责计算梯度,真正修改参数的是optimizer.step()。
对于普通梯度下降,参数更新公式为:
$$
\theta_t=\theta_{t-1}-\alpha g_t
$$
其中:
- $\theta_t$ 是第 $t$ 步更新后的参数。
- $\alpha$ 是学习率。
- $g_t=\nabla_\theta L_t$ 是当前损失对参数的梯度。
Adam 仍然沿着梯度方向更新参数,但不会直接使用当前梯度,而是先对梯度进行两类统计。
2. 为什么不总是直接使用 SGD
SGD 简单、可靠,而且在很多任务上有很好的泛化能力。不过,普通 SGD 也会遇到一些困难。
2.1 不同方向的梯度尺度差异很大
假设损失曲面像一条狭长山谷:
- 水平方向比较平缓。
- 垂直方向非常陡峭。
如果所有参数共用相同学习率,SGD 可能会在陡峭方向来回震荡,在平缓方向前进得很慢。
2.2 小批量梯度存在噪声
深度学习通常使用 mini-batch 训练。每个 batch 包含的样本不同,因此每一步梯度都会略有变化。
如果只看当前梯度,更新方向可能左右摇摆。
2.3 不同参数需要的有效步长可能不同
某些参数经常获得较大梯度,另一些参数的梯度可能很小或很稀疏。
统一学习率无法针对每个参数的历史梯度尺度进行调整。
Adam 的设计就是同时处理“方向不稳定”和“尺度不一致”这两个问题。
3. Adam 的核心思想
Adam 可以理解为把两种思想组合在一起:
- 用梯度的指数移动平均估计更新方向。
- 用梯度平方的指数移动平均估计梯度尺度。
第一部分类似 Momentum,第二部分与 RMSProp 的思路接近。
3.1 一阶矩:记录梯度方向
Adam 使用下面的公式更新一阶矩:
$$
m_t=\beta_1m_{t-1}+(1-\beta_1)g_t
$$
$m_t$ 可以理解为近期梯度的加权平均。
如果连续多个 batch 的梯度方向一致,$m_t$ 会积累这个方向;如果某一步梯度突然反向,历史信息会减小这次波动的影响。
这有点像一个带惯性的球:不会因为某一步坡度轻微变化就立即完全改变方向。
PyTorch 中 $\beta_1$ 默认通常取:
$$
\beta_1=0.9
$$
这表示新的 $m_t$ 保留 90% 的历史估计,并加入 10% 的当前梯度。
3.2 二阶矩:记录梯度大小
Adam 还会更新梯度平方的指数移动平均:
$$
v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2
$$
这里的平方是逐元素计算。
$v_t$ 记录的是近期梯度的尺度:
- 某个参数长期梯度较大,$v_t$ 就会较大,实际更新步长会被压小。
- 某个参数长期梯度较小,$v_t$ 就会较小,更新不会被过度缩小。
PyTorch 中 $\beta_2$ 默认通常取:
$$
\beta_2=0.999
$$
它比 $\beta_1$ 更接近 1,意味着梯度平方的统计通常会保留更长的历史。
4. Adam 的完整更新公式
假设当前参数是 $\theta_{t-1}$,当前梯度为:
$$
g_t=\nabla_\theta L_t
$$
Adam 的一次更新可以拆成下面几步。
第一步:更新一阶矩
$$
m_t=\beta_1m_{t-1}+(1-\beta_1)g_t
$$
第二步:更新二阶矩
$$
v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2
$$
第三步:进行偏差修正
$$
\hat{m}_t=\frac{m_t}{1-\beta_1^t}
$$
$$
\hat{v}_t=\frac{v_t}{1-\beta_2^t}
$$
第四步:更新参数
$$
\theta_t=
\theta_{t-1}
-\alpha\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}
$$
其中:
- $\alpha$ 是基础学习率。
- $\beta_1$ 控制一阶矩的平滑程度。
- $\beta_2$ 控制二阶矩的平滑程度。
- $\epsilon$ 是很小的正数,用来防止分母为 0 并提升数值稳定性。
- $\hat{m}_t$ 和 $\hat{v}_t$ 是偏差修正后的估计。
5. 为什么需要偏差修正
Adam 通常把一阶矩和二阶矩初始化为 0:
$$
m_0=0,\quad v_0=0
$$
在第一步训练时:
$$
m_1=(1-\beta_1)g_1
$$
如果 $\beta_1=0.9$,那么:
$$
m_1=0.1g_1
$$
它明显比真实梯度小,因为移动平均刚刚开始,还没有积累足够的历史信息。
二阶矩同样存在这个问题。如果 $\beta_2=0.999$:
$$
v_1=0.001g_1^2
$$
所以 Adam 使用 $1-\beta_1^t$ 和 $1-\beta_2^t$ 进行修正。
第一步的一阶矩修正后:
$$
\hat{m}_1=\frac{0.1g_1}{1-0.9}=g_1
$$
第一步的二阶矩修正后:
$$
\hat{v}_1=\frac{0.001g_1^2}{1-0.999}=g_1^2
$$
随着 $t$ 增大,$\beta_1^t$ 和 $\beta_2^t$ 会逐渐接近 0,偏差修正的影响也会逐渐减小。
6. 手工计算两步 Adam 更新
为了看清每个变量的作用,下面只优化一个参数。
假设:
$$
\theta_0=1
$$
并使用:
$$
\alpha=0.001,\quad
\beta_1=0.9,\quad
\beta_2=0.999
$$
为了简化手算,暂时忽略非常小的 $\epsilon$。
6.1 第一步:梯度 $g_1=2$
一阶矩:
$$
m_1=0.9\times0+0.1\times2=0.2
$$
二阶矩:
$$
v_1=0.999\times0+0.001\times2^2=0.004
$$
偏差修正:
$$
\hat{m}_1=\frac{0.2}{1-0.9}=2
$$
$$
\hat{v}_1=\frac{0.004}{1-0.999}=4
$$
参数更新量:
$$
\Delta\theta_1
=0.001\frac{2}{\sqrt{4}}
=0.001
$$
所以:
$$
\theta_1=1-0.001=0.999
$$
6.2 第二步:梯度 $g_2=1$
一阶矩:
$$
m_2=0.9\times0.2+0.1\times1=0.28
$$
二阶矩:
$$
v_2=0.999\times0.004+0.001\times1^2=0.004996
$$
偏差修正:
$$
\hat{m}_2=\frac{0.28}{1-0.9^2}
\approx1.4737
$$
$$
\hat{v}_2=\frac{0.004996}{1-0.999^2}
\approx2.4992
$$
参数更新量大约为:
$$
\Delta\theta_2
=0.001\frac{1.4737}{\sqrt{2.4992}}
\approx0.000932
$$
因此:
$$
\theta_2\approx0.999-0.000932=0.998068
$$
可以发现,虽然第二步梯度从 2 变成了 1,但 Adam 的更新并不是简单地把学习率乘以当前梯度。历史方向和历史尺度都会影响这一步。
7. Adam 与其他优化器的区别
| 优化器 | 是否记录方向 | 是否调整参数尺度 | 常见特点 |
|---|---|---|---|
| SGD | 否 | 否 | 简单,更新完全由当前梯度决定 |
| SGD + Momentum | 是 | 否 | 减少震荡,加快稳定方向上的前进 |
| RMSProp | 否 | 是 | 根据梯度平方缩放每个参数的步长 |
| Adam | 是 | 是 | 同时使用一阶矩和二阶矩 |
| AdamW | 是 | 是 | 在 Adam 基础上解耦权重衰减 |
Adam 并不是在所有任务上都一定优于 SGD。
- Adam 往往在训练前期下降更快,对学习率不那么敏感。
- SGD 配合 Momentum 在部分视觉任务中可能获得更好的最终泛化效果。
- 稀疏梯度、大模型预训练、Transformer 等任务经常使用 Adam 或 AdamW。
选择优化器时,不能只看训练损失下降速度,也要观察验证集效果。
8. Adam 的超参数
PyTorch 中的常见写法是:
optimizer = torch.optim.Adam(
model.parameters(),
lr=1e-3,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=0.0,
)
8.1 lr
lr 是基础学习率,也是最值得优先调整的参数。
常见起点:
1e-3 = 0.001
如果训练不稳定、损失剧烈波动,可以尝试:
3e-4、1e-4
如果损失下降非常慢,也可以在确保稳定的前提下适当增大学习率。
“自适应学习率”并不代表 Adam 不需要设置 lr。Adam 只是根据历史梯度对基础学习率进行逐参数缩放。
8.2 betas
betas=(0.9, 0.999)
beta1控制梯度方向的移动平均。beta2控制梯度平方的移动平均。
大多数普通任务可以先使用默认值。
如果梯度变化很快,减小 beta2 会让二阶矩更快响应近期变化,但也会让估计更容易波动。
8.3 eps
eps=1e-8
eps 主要用于数值稳定。通常不需要调整。
在低精度训练或特定数值不稳定场景中,有时会使用更大的 eps,但应结合实际损失和梯度情况判断。
8.4 weight_decay
权重衰减用于限制参数过度增大,帮助缓解过拟合。
如果希望使用解耦式权重衰减,通常更推荐 AdamW:
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
weight_decay=1e-2,
)
9. Adam 和 AdamW 有什么区别
传统 L2 正则化会把参数惩罚项加入损失函数:
$$
L’=L+\frac{\lambda}{2}\lVert\theta\rVert_2^2
$$
它会让梯度中多出一项:
$$
\nabla_\theta L’=\nabla_\theta L+\lambda\theta
$$
对于普通 SGD,L2 正则化与权重衰减在常见设置下可以得到等价形式。
但 Adam 会对梯度做自适应缩放。如果把 $\lambda\theta$ 直接混入梯度,这个正则化项也会被一阶矩和二阶矩处理,效果不再等同于直接衰减参数。
AdamW 将权重衰减从梯度更新中分离出来,可以直观理解为在 Adam 更新之外单独缩小参数:
$$
\theta_t
=\theta_{t-1}
-\text{AdamUpdate}
-\alpha\lambda\theta_{t-1}
$$
实际使用时可以先记住:
需要权重衰减时,尤其在 Transformer 等模型中,通常优先考虑
torch.optim.AdamW。
10. PyTorch 使用 Adam 完整训练一个 MLP
下面使用 Adam 训练一个 XOR 分类器。
import torch
from torch import nn
torch.manual_seed(42)
X = torch.tensor([
[0.0, 0.0],
[0.0, 1.0],
[1.0, 0.0],
[1.0, 1.0],
])
y = torch.tensor([0, 1, 1, 0])
class XORNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 8),
nn.Tanh(),
nn.Linear(8, 2),
)
def forward(self, x):
return self.net(x)
model = XORNet()
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(
model.parameters(),
lr=0.01,
)
for epoch in range(1000):
model.train()
logits = model(X)
loss = loss_fn(logits, y)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
if epoch % 200 == 0:
print(f"epoch={epoch}, loss={loss.item():.6f}")
model.eval()
with torch.no_grad():
logits = model(X)
probabilities = torch.softmax(logits, dim=1)
predictions = logits.argmax(dim=1)
print("predictions:", predictions.tolist())
print("probabilities:\n", probabilities)
训练成功后,预测类别应接近:
predictions: [0, 1, 1, 0]
10.1 一次训练迭代发生了什么
logits = model(X)
执行前向传播,得到两个类别的 logits。
loss = loss_fn(logits, y)
根据真实标签计算交叉熵损失。
optimizer.zero_grad(set_to_none=True)
清空旧梯度。set_to_none=True 会把梯度设为 None,通常可以减少不必要的内存写入。
loss.backward()
Autograd 计算每个参数的梯度,并写入参数的 .grad。
optimizer.step()
Adam 读取每个参数的梯度,更新该参数对应的 $m_t$、$v_t$ 和步数 $t$,然后修改参数值。
11. 优化器内部保存了什么状态
SGD 在不使用 Momentum 时,几乎只需要读取当前梯度。
Adam 则需要为每个可训练参数保存:
- 一阶矩 $m_t$。
- 二阶矩 $v_t$。
- 当前更新步数 $t$。
因此 Adam 会比普通 SGD 使用更多优化器状态内存。
如果模型参数本身占用 $P$ 单位内存,仅从概念上看,Adam 还需要额外保存与参数形状相同的一阶矩和二阶矩,优化器状态可能带来大约额外 $2P$ 的开销;实际训练还要计算参数、梯度、混合精度副本等其他内存。
可以检查优化器状态:
# 至少执行一次 optimizer.step() 后才会创建大部分状态
state = optimizer.state_dict()
print(state.keys())
print(state["param_groups"])
print(state["state"])
保存断点时,除了模型参数,还应该保存优化器状态:
torch.save({
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": epoch,
}, "checkpoint.pth")
恢复训练:
checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
start_epoch = checkpoint["epoch"] + 1
如果只加载模型参数、不加载优化器状态,Adam 之前积累的一阶矩、二阶矩和步数都会丢失,恢复后的更新轨迹可能发生变化。
12. 为不同参数设置不同学习率
PyTorch 优化器支持参数组。
假设一个模型包含特征提取器和分类头,可以让分类头使用更大的学习率:
optimizer = torch.optim.AdamW([
{
"params": model.backbone.parameters(),
"lr": 1e-4,
},
{
"params": model.classifier.parameters(),
"lr": 1e-3,
},
], weight_decay=1e-2)
这种写法常用于迁移学习和微调:
- 预训练层已经学到较好表示,使用较小学习率。
- 新加入的分类头需要从头学习,使用较大学习率。
还可以对偏置和归一化参数关闭权重衰减:
decay_params = []
no_decay_params = []
for name, param in model.named_parameters():
if not param.requires_grad:
continue
if name.endswith("bias") or "norm" in name.lower():
no_decay_params.append(param)
else:
decay_params.append(param)
optimizer = torch.optim.AdamW([
{"params": decay_params, "weight_decay": 1e-2},
{"params": no_decay_params, "weight_decay": 0.0},
], lr=1e-3)
13. Adam 仍然可以配合学习率调度器
Adam 会为不同参数调整有效步长,但基础学习率仍然非常重要。
因此 Adam 也常与学习率调度器一起使用。
例如每 20 个 epoch 把学习率乘以 0.5:
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=20,
gamma=0.5,
)
for epoch in range(100):
train_one_epoch(model, optimizer)
scheduler.step()
Transformer 训练中还常见 warmup:训练开始时先逐渐增大学习率,避免模型在参数和优化器状态尚不稳定时迈出过大的步子。
14. 常见错误
14.1 忘记清空梯度
错误写法:
loss.backward()
optimizer.step()
如果循环中一直不调用 zero_grad(),梯度会不断累积。
普通训练建议写成:
optimizer.zero_grad()
loss.backward()
optimizer.step()
14.2 在训练循环中反复创建优化器
错误写法:
for epoch in range(100):
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# training...
每次重新创建 Adam,都会丢失之前的一阶矩和二阶矩。
优化器通常应该在训练循环外创建一次。
14.3 认为 Adam 不需要调学习率
Adam 的自适应缩放不能代替基础学习率调节。
如果出现下面情况,仍然要检查 lr:
- 损失变成
NaN。 - 损失剧烈震荡。
- 参数或梯度突然变得非常大。
- 训练长期没有明显进展。
14.4 忘记保存优化器状态
如果要继续训练,应同时保存:
model.state_dict()optimizer.state_dict()- 当前 epoch 或 step
- 学习率调度器状态(如果使用)
14.5 把训练损失下降快等同于泛化更好
Adam 可能很快降低训练损失,但最终模型是否更好仍然要看验证集指标。
15. Adam 什么时候适合使用
Adam 适合作为这些场景的起点:
- 希望快速得到一个可用基线。
- 梯度尺度差异较大。
- 梯度较稀疏。
- 训练 Transformer、语言模型或复杂多任务网络。
- 不确定 SGD 的学习率和 Momentum 应如何配置。
可以进一步考虑 AdamW 的情况:
- 需要明确使用权重衰减。
- 训练 Transformer、ViT、BERT 等模型。
- 希望将自适应梯度更新与参数衰减分离。
可能考虑 SGD + Momentum 的情况:
- 在特定计算机视觉任务中追求最终泛化效果。
- 已经有成熟的学习率、Momentum 和调度策略。
- 优化器状态内存比较紧张。
16. 总结
Adam 的核心可以浓缩成四步:
计算当前梯度 g
-> 更新梯度移动平均 m
-> 更新梯度平方移动平均 v
-> 偏差修正后更新参数
从直觉上看:
- $m_t$ 判断最近应该往哪个方向走。
- $v_t$ 判断每个参数方向上的历史梯度有多大。
- 偏差修正解决训练初期统计量偏向 0 的问题。
- 学习率 $\alpha$ 仍然控制整体更新尺度。
在 PyTorch 中,最核心的训练模板仍然是:
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
Adam 并不是一个不用调参、永远最优的优化器,但它通常稳定、易用,是深度学习实验非常实用的起点。
参考资料
- Diederik P. Kingma, Jimmy Ba, Adam: A Method for Stochastic Optimization: https://arxiv.org/abs/1412.6980
- Ilya Loshchilov, Frank Hutter, Decoupled Weight Decay Regularization: https://arxiv.org/abs/1711.05101
- PyTorch Documentation,
torch.optim.Adam: https://docs.pytorch.org/docs/stable/generated/torch.optim.Adam.html - PyTorch Documentation,
torch.optim.AdamW: https://docs.pytorch.org/docs/stable/generated/torch.optim.AdamW.html - PyTorch Tutorials, Optimizing Model Parameters: https://docs.pytorch.org/tutorials/beginner/basics/optimization_tutorial.html
