从公式到代码：拆解AdamW如何将权重衰减与梯度更新解耦

1. 优化器之争为什么AdamW比Adam更适合现代深度学习在深度学习的训练过程中优化器的选择往往决定了模型能否快速收敛到理想状态。Adam作为最流行的优化算法之一凭借其自适应学习率的特性已经成为许多开发者的默认选择。但近年来AdamW逐渐崭露头角特别是在大语言模型训练中它几乎成为了事实上的标准。我第一次接触AdamW是在训练一个图像分类模型时。当时使用Adam优化器模型在训练集上表现很好但在验证集上却出现了明显的过拟合。尝试了各种正则化方法后效果都不理想。直到切换到AdamW验证准确率才得到了显著提升。这让我意识到优化器中权重衰减的实现方式对模型泛化能力有着至关重要的影响。AdamW的核心改进在于它重新设计了权重衰减weight decay的应用方式。在传统Adam中权重衰减是通过L2正则化实现的即在损失函数中直接添加一个正则项。这种方式看似简单但实际上会干扰Adam的自适应学习率机制。而AdamW则将权重衰减从损失函数中解耦出来直接作用于参数更新步骤从而避免了这种干扰。2. 数学原理深度解析Adam与AdamW的公式对比2.1 Adam的数学实现要真正理解AdamW的创新之处我们需要先剖析Adam的工作原理。Adam结合了动量Momentum和RMSprop两种方法的优点通过维护两个移动平均值来调整每个参数的学习率一阶矩估计动量 $$m_t \beta_1 m_{t-1} (1-\beta_1)g_t$$二阶矩估计RMSprop $$v_t \beta_2 v_{t-1} (1-\beta_2)g_t^2$$其中$g_t$是当前时间步的梯度$\beta_1$和$\beta_2$是衰减率超参数通常分别设置为0.9和0.999。由于这些移动平均值在初始阶段偏向于0Adam还引入了偏差校正 $$\hat{m}_t \frac{m_t}{1-\beta_1^t}$$ $$\hat{v}_t \frac{v_t}{1-\beta_2^t}$$最终的参数更新公式为 $$\theta_{t1} \theta_t - \frac{\eta}{\sqrt{\hat{v}_t}\epsilon}\hat{m}_t$$当加入L2正则化时损失函数变为 $$L(\theta) \frac{\lambda}{2}||\theta||^2$$这会导致梯度计算时多出一个$\lambda\theta$项从而影响自适应学习率的计算。2.2 AdamW的关键改进AdamW的突破在于它改变了权重衰减的应用方式。它不再将权重衰减项混入损失函数而是直接在参数更新时添加$$\theta_{t1} \theta_t - \eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t}\epsilon} \lambda\theta_t\right)$$这种看似微小的改变带来了显著的优势权重衰减不再干扰自适应学习率机制正则化效果更加稳定和可预测超参数特别是权重衰减系数的设置更加鲁棒在实际应用中这意味着我们可以更独立地调整学习率和权重衰减强度而不必担心两者之间的相互影响。我在多个项目中发现AdamW允许使用更大的权重衰减值而不会导致训练不稳定这有助于更好地控制模型复杂度。3. 代码实现从零编写Adam和AdamW优化器理解数学原理后让我们通过代码实现来加深认识。下面我将展示如何用PyTorch手动实现这两个优化器。3.1 Adam优化器实现class MyAdam: def __init__(self, params, lr1e-3, betas(0.9, 0.999), eps1e-8, weight_decay0): self.params list(params) self.lr lr self.beta1, self.beta2 betas self.eps eps self.weight_decay weight_decay self.m [torch.zeros_like(p) for p in self.params] self.v [torch.zeros_like(p) for p in self.params] self.t 0 def step(self): self.t 1 for i, param in enumerate(self.params): if param.grad is None: continue grad param.grad if self.weight_decay ! 0: grad grad self.weight_decay * param.data self.m[i] self.beta1 * self.m[i] (1 - self.beta1) * grad self.v[i] self.beta2 * self.v[i] (1 - self.beta2) * grad.pow(2) m_hat self.m[i] / (1 - self.beta1 ** self.t) v_hat self.v[i] / (1 - self.beta2 ** self.t) param.data - self.lr * m_hat / (v_hat.sqrt() self.eps)3.2 AdamW优化器实现class MyAdamW: def __init__(self, params, lr1e-3, betas(0.9, 0.999), eps1e-8, weight_decay0): self.params list(params) self.lr lr self.beta1, self.beta2 betas self.eps eps self.weight_decay weight_decay self.m [torch.zeros_like(p) for p in self.params] self.v [torch.zeros_like(p) for p in self.params] self.t 0 def step(self): self.t 1 for i, param in enumerate(self.params): if param.grad is None: continue grad param.grad self.m[i] self.beta1 * self.m[i] (1 - self.beta1) * grad self.v[i] self.beta2 * self.v[i] (1 - self.beta2) * grad.pow(2) m_hat self.m[i] / (1 - self.beta1 ** self.t) v_hat self.v[i] / (1 - self.beta2 ** self.t) param.data - self.lr * (m_hat / (v_hat.sqrt() self.eps) self.weight_decay * param.data)关键区别在于权重衰减的应用时机和方式。在Adam中权重衰减通过修改梯度实现而在AdamW中它直接作用于参数更新步骤。这种差异在实际训练中会产生显著不同的效果。4. 实战对比Adam与AdamW在不同场景下的表现4.1 图像分类任务中的表现在CIFAR-10数据集上我对比了两种优化器的表现。使用ResNet-18架构设置相同的初始学习率(0.001)和权重衰减(0.01)训练50个epoch后指标AdamAdamW训练准确率98.2%97.8%验证准确率89.5%91.2%最终损失值0.150.12虽然Adam在训练集上略胜一筹但AdamW在验证集上表现更好显示出更强的泛化能力。这正是解耦权重衰减带来的好处——更有效的正则化。4.2 自然语言处理任务中的差异在文本分类任务上使用BERT-base模型差异更加明显训练稳定性AdamW的训练损失曲线更加平滑较少出现剧烈波动收敛速度AdamW在前期收敛更快特别是在fine-tuning阶段最终性能AdamW平均能带来1-2个百分点的准确率提升这些优势使得AdamW成为Transformer架构的首选优化器。实际上Hugging Face的Transformers库现在默认就使用AdamW。4.3 超参数敏感性测试为了全面比较两种优化器我测试了它们对不同超参数的敏感性学习率AdamW对学习率的变化更具鲁棒性权重衰减AdamW允许使用更大的衰减系数而不会破坏训练批量大小在大批量训练时AdamW的优势更加明显特别是在使用大型预训练模型时AdamW的这些特性使得调参过程更加简单可靠。不必再小心翼翼地平衡学习率和权重衰减的关系可以更独立地优化这两个超参数。