Robust Layerwise Scaling Rules by Proper Weight Decay Tuning

TL;DR

本文提出了一种针对矩阵类参数的权重衰减缩放规则 ($\lambda \propto \sqrt{d}$)，通过在AdamW训练的稳态下保持子层增益在不同模型宽度下的不变性，从而扩展了最大更新参数化($\mu$P)框架，实现了学习率和权重衰减超参数的零样本迁移。

关键定义

本文的研究建立在最大更新参数化($\mu$P)之上，并引入了对训练动态的新理解。以下是理解本文核心思想的关键概念：

  1. 最大更新参数化 (Maximal-update Parameterization, $\mu$P)：一个理论框架，旨在通过调整参数初始化和学习率，使得模型宽度变化时，参数更新的幅度保持不变。其核心思想是，对于不同类型的参数（向量类和矩阵类），应采用不同的缩放规则，从而实现超参数（尤其是学习率）在不同尺寸模型间的迁移。

  2. 优化器控制的稳态 (Optimizer-governed steady state)：与$\mu$P关注的训练早期、由初始化主导的动态不同，本文指出在现代的缩放不变架构中，训练会迅速进入一个稳态。在此状态下，参数的范数（scale）主要由优化器（如AdamW）的超参数（学习率$\eta$和权重衰减$\lambda$）共同决定，而非初始化方差。

  3. 子层增益 (Sublayer gain)：定义为一个层（sublayer）的输出激活值的均方根范数（RMS norm）与输入激活值均方根范数的比值，即 $\ \mid \mathbf{y}\ \mid _{\mathrm{rms}}/\ \mid \mathbf{x}\ \mid _{\mathrm{rms}}$。本文的核心目标是通过调整权重衰减来保持该增益在模型宽度变化时不发生改变，从而维持稳定的训练动态。

  4. 参数类别 (Parameter Classes)：根据参数尺寸随模型宽度 $d$ 的变化方式，将参数分为两类：

    • 向量类参数 (Vector-like parameters)：参数数量与 $d$ 呈线性关系，如嵌入层、RMSNorm增益、偏置项等。
    • 矩阵类参数 (Matrix-like parameters)：参数数量与 $d$ 呈二次关系，如注意力机制和前馈网络中的投影矩阵。本文提出的新权重衰减规则主要针对此类参数。

相关工作

当前，大规模模型的设计主要遵循两种思路：一是宏观的经验性缩放定律 (Empirical scaling laws)，它们描述了模型性能与模型大小、数据量和计算预算之间的幂律关系，指导“分配什么”资源；二是微观的最大更新参数化 ($\mu$P)，它通过保持参数更新幅度的宽度不变性，解决了“如何迁移”超参数（主要是学习率）的问题。

然而，$\mu$P的理论分析主要局限于训练的早期近初始化阶段。在现代使用了归一化层（如LayerNorm）的缩放不变架构中，训练很快会进入一个由优化器决定的稳态。在这个阶段，归一化层虽然保证了前向传播的尺度不变性，但却引入了反向传播的尺度敏感性：归一化层的梯度会反比于其输入的尺度。由于不同宽度的模型在稳态下的激活值尺度不同，这导致了有效学习率的宽度依赖性，从而破坏了$\mu$P的超参数迁移效果。

本文旨在解决这一具体问题：如何扩展$\mu$P框架，使其在更符合实际的、由优化器主导的训练稳态下依然有效？ 具体而言，作者寻求一个与$\mu$P学习率规则相匹配的权重衰减缩放规则，以稳定跨宽度的子层增益，最终实现学习率和权重衰减两种超参数的同时、鲁棒的“零样本”迁移。

本文方法

本文的核心方法是为$\mu$P框架补充一个针对权重衰减的缩放规则，以解决在训练稳态下超参数迁移失效的问题。

创新点

该方法建立在一个关键的理论和经验观察之上：在AdamW训练的稳态中，参数矩阵的范数不再由其初始化决定，而是稳定在与 $\sqrt{\eta/\lambda}$ 成正比的水平上。因此，为了在改变模型宽度 $d$ 时保持子层增益不变，必须合理地缩放权重衰减 $\lambda$ 来控制参数矩阵的范数。

1. 经验观察与规则推导

通过对LLaMA风格的Transformer模型进行实验，作者观察到以下两个关键现象：

图1：固定宽度下， FFN权重矩阵的统计数据。子层增益（左）和奇异值谱（中、右）的幅度均与 $\sqrt{\eta/\lambda}$ 成正比。
\[\sqrt{\frac{\eta}{\lambda}} \propto d^{-0.75}\]

这个关系确保了不同宽度模型中，关键的信号放大行为（由顶部奇异值主导）保持一致，从而实现了子层增益的宽度不变性。

2. 权重衰减缩放规则的提出

结合上述经验观察和$\mu$P为矩阵类参数设定的学习率缩放规则 $\eta_2 \propto d^{-1}$，可以推导出权重衰减 $\lambda_2$ 的缩放规则：

将 $\eta_2 \propto d^{-1}$ 代入经验公式，得到：

\[\sqrt{\frac{d^{-1}}{\lambda_2}} \propto d^{-0.75} \implies \frac{1}{\sqrt{\lambda_2 d}} \propto d^{-0.75} \implies \sqrt{\lambda_2} \propto \frac{d^{0.75}}{\sqrt{d}} = d^{0.25}\]

因此，权重衰减应遵循：

\[\lambda_2 \propto \sqrt{d}\]

3. 完整的层级缩放方案

综合$\mu$P和本文的发现，作者提出了一个完整的、针对不同参数类别的层级缩放方案：

类别 初始方差 学习率 权重衰减
向量类 $\sigma_{1}^{2}=\Theta_{d}(1)$ $\eta_{1}=\Theta_{d}(1)$ $\lambda_{1}=0$
矩阵类 $\sigma_{2}^{2}=\Theta(d^{-1})$ $\eta_{2}=\Theta(d^{-1})$ $\lambda_{2}=\Theta(\sqrt{d})$

这个方案的核心优点在于：它通过显式控制训练稳态下的参数尺度，将$\mu$P的有效性从训练早期扩展到了整个训练过程。它提供了一个理论上合理且经验上验证有效的配方，消除了为不同宽度的模型反复搜索学习率和权重衰减的巨大开销。

此外，本文还提供了一个简单的诊断方法：通过检查不同宽度模型下权重矩阵的顶部奇异值是否对齐，可以验证子层增益是否保持了不变性。

图2：在不同权重衰减缩放规则 ($\lambda\_2$)下，FFN权重矩阵的统计数据。当 $\lambda\_2 \propto \sqrt{d}$ 时，不同宽度模型（$d=256, 512, 1024$）的子层增益（左）和奇异值谱（中、右）表现出最佳的对齐。

实验结论

本文通过在LLaMA风格的Transformer模型上的一系列实验，验证了所提出的权重衰减缩放规则的有效性。

图3：最优学习率和权重衰减在不同宽度模型间的迁移。经过本文提出的规则缩放后，不同宽度模型（$d=256$ 至 $d=2048$）的损失曲面最小值点（最优超参点）高度对齐。