Adam 与 AdamW 优化器对比：weight_decay 在 PyTorch 2.0 中的 2 种实现差异

Adam与AdamW优化器对比：深入解析PyTorch中的weight_decay实现差异

在深度学习模型训练过程中，优化器的选择对模型性能有着至关重要的影响。Adam和AdamW作为两种广泛使用的优化算法，它们在处理权重衰减(weight_decay)时存在关键性差异，这些差异往往被大多数使用者忽视。本文将深入剖析这两种优化器在PyTorch框架中的具体实现差异，并通过实验验证它们对模型训练的实际影响。

1. 权重衰减的本质与作用

权重衰减(weight_decay)是深度学习中一种重要的正则化技术，其核心思想是通过在损失函数中添加模型权重的L2范数惩罚项，防止模型参数值过大，从而降低过拟合风险。数学表达式为：

L = L_original + λ/2 * ||w||²

其中λ就是weight_decay参数。传统SGD优化器中，weight_decay的实现非常直接——在每次参数更新时，会从当前权重中减去λ*w：

w = w - lr * (grad + λ * w)

然而，在自适应优化器如Adam中，weight_decay的实现方式变得复杂且存在争议。这主要是因为自适应优化器本身已经包含了每个参数的自适应学习率，直接添加weight_decay可能会干扰这种自适应机制。

2. Adam优化器中的weight_decay实现

PyTorch中标准Adam优化器的weight_decay实现方式被称为"L2正则化"或"衰减式weight_decay"。其参数更新过程可以分解为以下步骤：

1. 计算梯度：g = ∇L(θ)
2. 计算一阶矩估计：m = β1*m + (1-β1)*g
3. 计算二阶矩估计：v = β2*v + (1-β2)*g²
4. 偏差校正：m_hat = m/(1-β1^t), v_hat = v/(1-β2^t)
5. 参数更新：θ = θ - ηm_hat/(√v_hat + ε) - ηλ*θ

关键点在于最后一步：weight_decay项(λ*θ)是与自适应学习率项(m_hat/(√v_hat + ε))分开应用的。这种实现方式存在一个潜在问题——当梯度很小但权重本身较大时，weight_decay项可能会主导更新过程，导致训练不稳定。

以下是PyTorch中Adam优化器的典型使用方式：

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=0.01)

3. AdamW优化器的创新设计

AdamW（Adam with Weight decay correction）是Adam的一个变种，由Ilya Loshchilov和Frank Hutter在2017年提出。它修正了原始Adam中weight_decay的实现方式，将权重衰减与自适应学习率机制解耦。其更新步骤如下：

1. 计算梯度：g = ∇L(θ)
2. 计算一阶矩估计：m = β1*m + (1-β1)*g
3. 计算二阶矩估计：v = β2*v + (1-β2)*g²
4. 偏差校正：m_hat = m/(1-β1^t), v_hat = v/(1-β2^t)
5. 参数更新：θ = θ*(1 - ηλ) - ηm_hat/(√v_hat + ε)

关键区别在于：AdamW将weight_decay项(λ*θ)从自适应学习率项中分离出来，直接应用于参数本身，而不是梯度。这种实现更接近传统SGD中的weight_decay行为，在实践中通常能获得更好的泛化性能。

PyTorch中的AdamW使用方式：

optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=0.01)

4. 实验对比：Adam与AdamW的实际表现差异

为了直观展示这两种优化器的差异，我们设计了一个简单的线性回归实验。我们创建一个带有噪声的线性数据集，并分别使用Adam和AdamW进行训练，观察参数更新的轨迹。

4.1 实验设置

import torch import matplotlib.pyplot as plt # 生成带噪声的线性数据 torch.manual_seed(42) X = torch.randn(100, 1) true_w, true_b = 2.0, -1.0 y = true_w * X + true_b + 0.1 * torch.randn(X.size()) # 模型定义 class LinearModel(torch.nn.Module): def __init__(self): super().__init__() self.linear = torch.nn.Linear(1, 1, bias=True) def forward(self, x): return self.linear(x) # 训练函数 def train_with_optimizer(optimizer_class, weight_decay): model = LinearModel() criterion = torch.nn.MSELoss() optimizer = optimizer_class(model.parameters(), lr=0.1, weight_decay=weight_decay) weights = [] for epoch in range(100): optimizer.zero_grad() outputs = model(X) loss = criterion(outputs, y) loss.backward() optimizer.step() weights.append(model.linear.weight.item()) return weights # 对比实验 weight_decay = 0.1 adam_weights = train_with_optimizer(torch.optim.Adam, weight_decay) adamw_weights = train_with_optimizer(torch.optim.AdamW, weight_decay) # 可视化结果 plt.figure(figsize=(10, 6)) plt.plot(adam_weights, label=f'Adam (wd={weight_decay})') plt.plot(adamw_weights, label=f'AdamW (wd={weight_decay})') plt.axhline(y=true_w, color='r', linestyle='--', label='True weight') plt.xlabel('Epoch') plt.ylabel('Weight value') plt.title('Weight Trajectory: Adam vs AdamW') plt.legend() plt.grid(True) plt.show()

4.2 结果分析

从实验结果可以观察到几个关键现象：

1. 收敛速度：AdamW通常比Adam收敛得更快，特别是在使用较大weight_decay值时
2. 最终值：AdamW的参数值往往比Adam更接近真实值
3. 稳定性：AdamW的更新轨迹通常更加平滑，震荡较小

这些差异源于两种优化器处理weight_decay的根本方式不同。Adam将weight_decay应用于梯度空间，而AdamW直接应用于参数空间，后者更符合weight_decay作为正则化项的本质目的。

5. 实际应用中的选择建议

根据理论分析和实践经验，我们总结出以下建议：

关键参数设置指南：

• 学习率：AdamW通常可以使用比Adam稍大的学习率
• weight_decay：对于AdamW，常用范围是0.01-0.1；对于Adam，常用范围是0.0001-0.001
• beta参数：保持默认(0.9, 0.999)在大多数情况下效果良好

注意：当从Adam切换到AdamW时，通常需要重新调整weight_decay值，因为它们的作用强度不同。建议从一个较小的值开始，如0.01，然后根据验证集表现进行调整。

6. 高级话题：weight_decay与模型泛化

weight_decay对模型泛化能力的影响不仅取决于优化器类型，还与模型架构、数据特性等因素相关。一些最新研究发现：

1. 分层weight_decay：对不同层使用不同的weight_decay值可能提升性能
2. 动态weight_decay：在训练过程中调整weight_decay强度可能优于固定值
3. 与批归一化的交互：weight_decay对带有BN层的模型影响较小

以下是一个实现分层weight_decay的示例代码：

# 为不同层设置不同的weight_decay optimizer = torch.optim.AdamW([ {'params': model.conv_layers.parameters(), 'weight_decay': 0.01}, {'params': model.fc_layers.parameters(), 'weight_decay': 0.001}, {'params': model.bias_parameters(), 'weight_decay': 0.0} # 通常不对偏置项应用weight_decay ], lr=0.001)

在实际项目中，我发现合理使用AdamW配合适当的学习率调度（如余弦退火）往往能取得最佳效果。特别是在计算机视觉任务中，这种组合显著提升了模型的泛化能力，同时减少了调参的复杂度。