AdamW/zh

Article
Topic area	optimization
Prerequisites	Adam, Stochastic gradient descent, Weight decay

概述

AdamW 是一种用于训练神经网络的随机优化算法，它将权重衰减与 Adam 的基于梯度的更新解耦。该算法由 Loshchilov 和 Hutter 于 2017 年提出，修正了自适应优化器中长期存在的一个实现缺陷：在标准 Adam 中，向损失函数添加 L2 惩罚项并不会产生真正的权重衰减，因为该惩罚项会被逐参数的自适应学习率重新缩放。AdamW 通过在自适应更新之后对参数施加一个独立的、固定比率的收缩，恢复了 Hanson and Pratt 最初的表述。这一改动只需要寥寥几行代码，却能持续改善泛化能力，并已成为基于 Transformer 的模型的默认优化器，包括 BERT、GPT 风格的语言模型以及 Vision Transformer。

动机：在 Adam 中 L2 正则化并非权重衰减

对于普通的 SGD，在损失函数中添加 L2 惩罚项 $ \tfrac{\lambda}{2}\|\theta\|^2 $，在数学上等价于在每一步将参数乘以 $ (1 - \eta\lambda) $，其中 $ \eta $ 是学习率。这两种表述——L2 正则化与权重衰减——是一致的。

这种等价性在自适应方法中不成立。Adam 会按照对二阶矩 $ \hat{v}_t $ 的估计来缩放每一个梯度分量，因此附加到梯度上的 L2 项 $ \lambda\theta $ 在被应用之前会被除以 $ \sqrt{\hat{v}_t}+\epsilon $。历史梯度较大（即条件较好的方向）的参数所受到的正则化反而比梯度较小的参数更少，这与权重衰减本应起到的作用恰好相反。Loshchilov 和 Hutter 证明，这种耦合导致 Adam 在图像分类基准上的泛化表现不及带动量的 SGD，而将权重衰减解耦能够弥合大部分差距。

算法

设 $ \theta_t $ 表示第 $ t $ 步的参数向量，$ g_t = \nabla_\theta f_t(\theta_{t-1}) $ 是在第 $ t $ 个小批量上损失函数的随机梯度，$ \eta_t $ 是（可能按调度变化的）学习率。AdamW 以衰减率 $ \beta_1, \beta_2 \in [0,1) $ 维护梯度及梯度平方的指数滑动平均：

$ {\displaystyle m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t} $

$ {\displaystyle v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2} $

经过偏差修正后的估计为

$ {\displaystyle \hat{m}_t = \frac{m_t}{1-\beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1-\beta_2^t}.} $

随后参数的更新公式为

$ {\displaystyle \theta_t = \theta_{t-1} - \eta_t \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda\,\theta_{t-1} \right).} $

关键之处在于 $ \lambda\,\theta_{t-1} $ 这一项被施加在自适应分母之外。相比之下，原始的 Adam-加-L2 更新会把 $ \lambda\theta_{t-1} $ 折叠进 $ g_t $，得到 $ \theta_t = \theta_{t-1} - \eta_t (\hat{m}_t + \lambda \theta_{t-1}\cdot\text{scaling})/(\sqrt{\hat{v}_t}+\epsilon) $，此时权重衰减项与梯度共享同一个逐参数的自适应缩放因子。

在大多数实现中，默认超参数为 $ \beta_1 = 0.9 $、$ \beta_2 = 0.999 $、$ \epsilon = 10^{-8} $，其中 $ \lambda $ 在 Transformer 预训练时通常取 $ [10^{-2}, 10^{-1}] $，在微调时通常取 $ [10^{-4}, 10^{-2}] $。

解耦权重衰减在实践中的表现

解耦带来了两个实际后果。首先，最优权重衰减系数 $ \lambda $ 现在在很大程度上独立于学习率 $ \eta $，从而简化了超参数调优——在原始 Adam 中，改变学习率实际上也会改变正则化强度，迫使开发者必须联合扫参。其次，AdamW 的最优 $ \lambda $ 通常比在 Adam 中可用的 L2 系数大一到两个数量级，因为自适应缩放不再对其进行衰减。

一个常见的细节问题是：在使用学习率调度时是否应将 $ \lambda $ 也乘以 $ \eta_t $。原论文将更新写作 $ \theta_t = \theta_{t-1} - \eta_t \hat{m}_t/(\sqrt{\hat{v}_t}+\epsilon) - \eta_t \lambda \theta_{t-1} $，即权重衰减也按调度进行缩放。某些实现（尤其是 PyTorch 的一个早期版本）则直接施加 $ \lambda\theta_{t-1} $ 而不乘以 $ \eta_t $；如今这被普遍视为一个 bug，当前的 PyTorch、JAX 与 TensorFlow 实现都遵循原论文的约定。

变体与扩展

若干优化器对 AdamW 进行了扩展或修改：

• Lion（EvoLved Sign Momentum，Chen 等人 2023）——用符号算子代替二阶矩估计，保留解耦的权重衰减；显存占用约为 AdamW 的一半。
• 带 梯度裁剪的 AdamW ——全局范数或逐层裁剪是大语言模型预训练的标准做法，用于抑制损失尖峰。
• LAMB（Layer-wise Adaptive Moments）——在 AdamW 之上加入逐层归一化以适配超大批大小（32k+），曾用于刷新 BERT 预训练耗时纪录。
• AdaFactor ——通过对二阶矩矩阵进行因式分解来节省显存；支持解耦衰减。
• Adan 与 Sophia ——受二阶方法启发的优化器，均沿用了解耦衰减的设计。

FP16/bfloat16 训练时代引入了额外的超参数：在实践中，混合精度下 $ \epsilon $ 经常被提高到 $ 10^{-6} $ 或 $ 10^{-5} $，以避免 $ \sqrt{\hat{v}_t}+\epsilon $ 的下溢。

与 Adam 和 SGD 的比较

经验上，AdamW 弥合了最初促使从业者在视觉任务中偏好带动量 SGD 的泛化差距。在 ImageNet 上，调优良好的 AdamW 训练的 ResNet-50 所达到的准确率与 SGD+动量相差不到 0.1–0.3 个百分点，而朴素的 Adam-加-L2 通常要落后 1–2 个百分点。对于 Transformer 而言，AdamW 几乎是通用选择：注意力子层与前馈网络子层之间的梯度动态范围很大，需要逐参数的自适应缩放来处理；解耦衰减则避免了嵌入层和层归一化参数被欠正则化。

一条实用的经验法则：如果你从零开始训练模型，并且架构中含有 LayerNorm 或 RMSNorm，请使用 AdamW。如果你是在对预训练模型进行微调，请配合较小的 $ \lambda $ 使用 AdamW，并可能采用较小的 $ \beta_2 $（例如 0.95），以避免陈旧的二阶矩估计冲淡预训练权重。

局限性

AdamW 继承了 Adam 的显存开销：每个参数需要额外存储两个张量（$ m_t $、$ v_t $），优化器状态相比朴素 SGD 增加到三倍。对于参数量达数十亿的模型而言，这是主要开销，因而催生了分片变体（如 ZeRO）和将优化器状态量化的 8-bit AdamW。

解耦衰减并非万能。它假设 $ \lambda $ 遵循一个固定的调度；循环或热重启调度（Loshchilov 与 Hutter 的 SGDR）与二阶矩缓冲之间会产生非平凡的相互作用，目前的最佳实践仍是先在前几百到几千步内对学习率进行预热，然后再施加完整的权重衰减。在数据量较少的场景中，AdamW 对 $ \beta_2 $ 仍然敏感，因为此时二阶矩估计噪声较大；在强化学习和持续学习中，$ \beta_2 = 0.95 $ 或 $ 0.98 $ 等取值十分常见。

最后，L2 正则化与权重衰减之间的等价性，在 AdamW 中不成立，在 Adam 中同样不成立——它们如今是两种不同的正则化器；在报告 “weight decay” 时若不说明它指的是 AdamW 中的 $ \lambda $ 还是某种L2 损失项，就会成为再现性错误的常见来源。