理解 Weight Decay（一）：EMA 视角

这篇博客介绍了近期从 Exponential Moving Average (EMA) 视角理解 weight decay 的工作，主要参考 How to set AdamW’s weight decay as you scale model and dataset size [1] 和 Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training [2] 两篇文章的相关思想和结论。

Weight Decay 相关背景

Weight decay (WD) 中文名为权重衰减，顾名思义，就是在训练的过程中显式地控制模型的权重大小。记模型权重为 $w$，优化器导致的更新量为 $u$，学习率为 $\eta$，WD 为 $\lambda$，当前主流 WD 的实现方式为 decoupled 形式，可以概括为

\[w_{t} = w_{t-1} - \eta (u_t + \lambda w_{t-1}) = (1 - \eta\lambda) w_{t-1} - \eta u_t.\]

可以看到，在这种实现框架下，实际作用的权重衰减系数不是 $\lambda$，而是 $\eta \lambda$，我们后续会将其称为 independent weight decay (IWD)。从上式出发，我们也可以很直接地推导出 WD 对模型权重的限制作用。不失一般性，我们假设优化器更新量的大小是有界的，即 $\Vert u_t\Vert \leq C$，那么我们从上式可以推得权重大小的递推关系

\[\Vert w_{t}\Vert \leq (1 - \eta\lambda) \Vert w_{t-1}\Vert + \eta \Vert u_t\Vert \leq (1 - \eta\lambda) \Vert w_{t-1}\Vert + \eta C.\]

迭代此关系可以递推得到

\[\begin{aligned} \Vert w_{t}\Vert &\leq (1 - \eta \lambda)^t \Vert w_0\Vert + \eta C (1 + (1-\eta\lambda) + \cdots + (1 - \eta\lambda)^{t-1}) \\ &= (1 - \eta \lambda)^t \Vert w_0\Vert + \eta C \cdot \frac{1 - (1 - \eta\lambda)^{t}}{\eta \lambda} \\ &\to \frac{C}{\lambda} \quad (t \to \infty). \end{aligned}\]

因此，越大的 WD 意味着对模型权重大小更强的控制。从经典学习理论 [3] 的角度来说，越强的权重控制意味着更小的函数空间，更小的模型容量（model capacity、Rademacher complexity），从而能够从理论上给出更小的泛化误差（train error - test error），缓解模型潜在的过拟合问题。事实上，在传统机器学习时代，相对于模型容量，数据量往往较少，比较容易出现过拟合的问题，WD 最初的提出也正是为了缓解此类问题。

然而，在如今大模型训练的场景中，数据量和复杂度都远超以往的任务，甚至是远超模型容量，如 Token per Param/TPP 一般是几百上千，这使得过拟合暂时不再是一个重要的问题。如下图所示 [4]，不同于上一个时代的 Cifar-10 分类任务，当前 LLM 训练的 train loss 和 validation loss 是比较 match 的，不存在显著的过拟合问题。

那么这是否意味着 WD 在如今主流的大模型训练场景已经失去了作用呢？答案是否定的，WD 对当今的大模型训练效果仍然有着重要的影响。例如下图所示，Kimi 的 Muon-Moonlight [5] 相较于原始 Muon 的一大改进正是加入了 WD，在实际的长程训练中能够显著提升模型效果。

总而言之，传统的控制模型容量角度的 WD 优势解释已经不太适用于当前的大模型训练场景，理解 WD 需要新的视角和思路。近年来有不少新的理解 WD 的思路，这篇文章就介绍一下近期从 EMA 角度理解 WD 的文章 [1,2]。

Exponential Moving Average

因为这篇博客理解 WD 的方法借鉴了 EMA 的形式，所以我们先简单介绍一下 EMA 的必要内容。EMA 原本是一种常用于平滑时间序列数据的方法，定义原始数据为 ${x_t}$，EMA 的输出为 ${s_t}$，则完整的 EMA 过程可以被写为

\[\begin{aligned} & s_0 = x_0, \\ & s_t = (1 - \alpha) s_{t-1} + \alpha x_t, \quad t > 0. \end{aligned}\]

EMA 的形式相信读者已经非常熟悉，因为它已经广泛地存在于机器学习中，例如 Adam、Muon 等优化器中的动量，模型预训练中常见的权重融合等。我们这里再进一步介绍一下其中的超参 $\alpha$，这是后续理解 WD 的主要依据。毫无疑问，$\alpha$ 是 EMA 中最重要的超参数，它反映了当前输出要考虑之前多久的历史数据，确立了一个有效数据的时间窗口。具体来说，我们将 EMA 展开可以得到

\[s_t = (1 - \alpha) s_{t-1} + \alpha x_t = \cdots (1-\alpha)^t x_0 + \sum_{i=1}^t \alpha(1-\alpha)^{t-i} x_i.\]

因为 $(1-\alpha)^t + \sum_{i=1}^t \alpha(1-\alpha)^{t-i} =1$，EMA 其实就是对历史数据做了一个加权求和，越远离当前时间的历史数据获得的权重更小，$\alpha$ 就是控制了这个权重的衰减速度。从上面的式子，我们可以得到最近 $T$ 步贡献的权重和为

\[\alpha \sum_{i=0}^{T-1} (1-\alpha)^i = 1 - (1 - \alpha)^T.\]

从这里我们可以推出来 $\alpha$ 对应的时间窗口。例如，我们认为最近贡献了总权重 $a$（如 0.9）的 $T$ 步历史数据为 EMA 实际考虑的时间窗口，那么可以推出来

\[T = \frac{\log(1 - a)}{\log(1 - \alpha)} \approx -\frac{\log(1 - a)}{\alpha}.\]

由于 $a$ 是一个常数，我们完全也可以定义 $1/\alpha$ 为 EMA 对应的时间窗口常数。好的，现在我们已经有了足够的背景知识了。

从 EMA 角度理解 WD

为了和 EMA 的形式对应上，我们可以把 WD 的优化过程重写一下：

\[w_{t} = (1 - \eta\lambda) w_{t-1} - \eta u_t = (1 - \eta\lambda)w_{t-1} - \eta\lambda \cdot \frac{u_t}{\lambda}.\]

可以看到，带 WD 的优化过程也是一个标准的 EMA 过程，对应的 $\alpha = \eta \lambda$，时间窗口常数为 $1/ \eta\lambda$，这说明 IWD 对训练效果的影响至关重要。IWD 越大，越考虑近期的权重更新；IWD 越小，则能够考虑更远的权重情况。另一方面，和标准的 EMA 不同，初始化点和后续的数据点并不是来自一个时间序列：初始化点是 $w_0$，为模型的初始化权重；后续的更新点是 $u_t / \lambda$，为优化器的更新量。因此，即使 IWD 相同，不同的 $\lambda$（或不同的学习率 $\eta$）也会影响权重初始化和优化更新量的关系，从而影响模型效果。但总的来看，从 EMA 的角度来说，IWD 对结果的影响可能更为重要。

在实际训练中，我们的训练规模/步数可能会发生改变，所以一个更适用于不同训练规模的 IWD 作用度不是时间窗口的绝对大小，而是时间窗口占整个训练过程的比例。因此，记训练步数为 $S$，训练数据量为 $D$，训练 batch size 为 $B$，我们定义时间窗口比例常数为

\[\tau = \frac{1}{\eta\lambda S} = \frac{B}{\eta\lambda D}.\]

[1,2] 均在实际中观察了这个量随训练规模的关系，我们接下来对 [2] 的核心结论进行介绍。

实际观察

[2] 在不同规模的场景下，系统地探索了如何设置 lr 和 WD，它正是基于 EMA 理论的。为了规避 lr 和 WD 两个维度的 grid search，[2] 用 muP/CompleteP 来设定不同模型大小 $N$ 下的 lr，然后在不同 $N,D$ 规模下 sweep WD，从而得到最优 $\tau$/IWD 的 scaling law。[2] 发现 $\tau$ 和 TPP 有着非常显著的幂律关系，如下图所示

基于 muP 的 lr 调整和 IWD 的 scaling law，我们可以在大规模下实现对 lr 和 WD 较为科学的设置。

References

1. How to set AdamW’s weight decay as you scale model and dataset size, ICML, 2025
2. Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training, NeurIPS, 2025
3. Foundations of Machine Learning, 2018
4. Rethinking Conventional Wisdom in Machine Learning: From Generalization to Scaling, 2025
5. Muon is Scalable for LLM Training, 2025