Don't Be Lazy: CompleteP Enables Compute-Efficient Deep Transformers¶

会议: NeurIPS 2025
arXiv: 2505.01618
代码: https://github.com/EleutherAI/nanoGPT-mup/tree/completep
领域: LLM 扩展律, 超参优化
关键词: 参数化, 深度缩放, 超参转移, 计算效率, CompleteP, μP

一句话总结¶

CompleteP 参数化（α=1）是唯一同时实现深度方向超参转移和完全特征学习的方案，在深模型上相比 μP 节省 12-34% FLOPs。

研究背景与动机¶

LLM 训练中，模型规模增大意味着需要重新搜索超参（学习率、初始化方差等），成本极高。μP（maximal update parameterization）通过保证宽度方向的超参转移，实现了"小模型调参、大模型直接用"的策略。然而 μP 仅处理宽度缩放，当模型深度 L 变化时，最优超参会发生漂移，导致两个后果：（1）必须为不同深度重新调参，费时费力；（2）若不调参则接受次优训练、浪费算力。

已有工作将 μP 扩展到深度维度，引入残差缩放因子 α∈[0.5,1]，但对最优 α 值存在分歧：Yang et al. 认为 α=0.5 最佳且 HP 转移不可能，Bordelon et al. 则发现 α=1 理论上更优。本文系统比较这两个极端，最终确认 α=1 是唯一正确选择。

方法详解¶

残差块深度缩放¶

Transformer 的残差连接被参数化为：

\[h^{\ell+1} = h^{\ell} + L^{-\alpha} \mathcal{F}_\ell(h^{\ell})\]

其中 \(\mathcal{F}_\ell\) 是第 ℓ 层的 MLP 或注意力块，α 控制每层残差贡献的衰减速率。α=0 即标准参数化（SP），α=0.5 是 Yang et al. 推荐值，α=1 是本文提出的 CompleteP。

完整参数化规则（Table 1）¶

CompleteP 不仅涉及残差缩放，还需要配套调整多组超参，定义宽度乘子 \(m_N = N/N_\text{base}\)、深度乘子 \(m_L = L/L_\text{base}\)：

隐藏层初始化方差: \(\sigma^2_\text{base} \cdot m_N^{-1}\)（继承 μP）
隐藏层学习率: \(\eta_\text{base} \cdot m_N^{-1} \cdot m_L^{\alpha-1}\)（α=1 时随深度缩小）
Pre-LN 学习率: \(\eta_\text{base} \cdot m_L^{\alpha-1}\)（关键扩展，SP/μP 未做此调整）
Bias 学习率: \(\eta_\text{base} \cdot m_L^{\alpha-1}\)
残差块乘子: \(m_L^{-\alpha}\)（即 \(L^{-1}\) 当 α=1）
AdamW ε（残差块）: \(\epsilon_\text{base} \cdot m_N^{-1} \cdot m_L^{-\alpha}\)
权重衰减: \(\lambda_\text{base} \cdot m_N\)（继承 μP）

这些扩展——尤其是 LayerNorm LR 和 AdamW ε 的深度缩放——是本文的重要实践贡献。α=0.5 若不做这些调整会出现训练不稳定。

三条设计准则（Desiderata）¶

Desideratum 1: 稳定初始化。 要求所有层的隐藏表示 \(\|h^\ell\|^2/N = \Theta(1)\)，约束 α≥0.5。

Desideratum 2: 最大残差流更新。 每层参数更新对 \(h^{\ell+1}\) 的贡献应为 \(\Theta(1/L)\)，确定学习率的深度依赖为 \(\eta = \Theta(L^{1-\alpha})\)，约束 α≤1。

Desideratum 3: 完全特征学习（本文核心新贡献）。 所有层的表示 \(h^\ell\) 对任意前序参数子集都不应是"惰性"（lazy）的——即不应退化为其线性化近似。具体地，对 α<1，Taylor 展开中非线性项的阶为 \(\Theta(L^{\alpha-2})\)，当 L→∞ 时相对线性项趋于零，导致深层网络实质上退化为线性模型。只有 α=1 时非线性项与线性项同阶，保证完全的特征学习能力。

这一理论分析解释了为何 α=0.5 在小深度时看似有效、但在大深度时性能不如 α=1：浅模型的最优超参依赖于线性和非线性动态的平衡，α<1 破坏了这一平衡。

实验验证¶

深度方向 HP 转移¶

在固定 300M token 训练设置下，测试 L=2 到 L=128 的学习率和初始化标准差转移： - SP / μP / α=0.5: 最优学习率随深度显著漂移，无法稳定转移 - CompleteP (α=1): 最优学习率和 σ_init 在所有深度保持稳定，等高线呈同心结构

在 compute-optimal 设置（20 TPP、按 FLOPs 选 batch size、调优 weight decay）下，虽然对学习率的敏感性降低，但 α=1 仍然在不做额外调参的情况下一致取得更好 loss。

最优宽深比 N:L¶

在 \(P_\text{non-emb} \in \{50M, 300M, 1.5B\}\) 三个规模、每个规模 7-10 种 N:L 配置下训练：

模型规模	μP 最优 N:L	CompleteP 最优 N:L	CompleteP ≤1% 损失范围
50M	~40	~40	N:L≥~12
300M	~50	~50	N:L≥~15
1.5B	~62	~62	N:L≥11.8（μP 需≥38.7）

CompleteP 将可接受的 N:L 范围大幅扩展，使窄深模型也能保持接近最优性能。这对内存受限硬件（逐层流式推理/训练）意义重大。

FLOP 节省（vs μP）¶

配置	FLOP 节省
1.5B 最优 N:L (N=1984, L=32)	11.8%
1.5B 最深 (N=832, L=179)	34.4%
300M 最深 (N=448, L=125)	显著

深度越大，CompleteP 相对 μP 的优势越明显，因为 μP 在大深度下超参失调严重。

下游任务（零样本，1.5B 模型，20 TPP）¶

任务	μP (最优 N:L)	CompleteP (最优)	μP (L=179)	CompleteP (L=179)
HellaSwag	53.3±0.5	54.2±0.5	49.1±0.5	52.7±0.5
ARC-Easy	54.4±1.0	55.6±1.0	50.0±1.0	54.6±1.0
LAMBADA	54.3±0.7	54.9±0.7	51.8±0.7	53.3±0.7
PIQA	70.7±1.1	71.5±1.1	69.6±1.1	70.6±1.1
BoolQ	58.4±0.9	60.7±0.9	57.8±0.9	59.0±0.9
下游平均	54.3±0.3	55.2±0.3	52.0±0.3	54.3±0.3

上游增益一致传导到下游，且 CompleteP 在深模型上的优势尤为突出（54.3 vs 52.0）。

200 TPP 长训练验证¶

在 50M 和 300M 模型上进一步以 200 TPP 训练，CompleteP 在所有配置下均取得最低验证 loss，确认结论在更长训练中同样成立。

亮点与洞察¶

首次实现深度 HP 转移: 从 L=2 成功转移到 L=128，超越 LLaMA-70B/405B 的深度
理论贡献——完全特征学习: 提出 Desideratum 3，证明 α<1 导致层的非线性动态在深度极限下消失（lazy regime），只有 α=1 在数学上保证所有层保持完全非线性学习
实用贡献——完整参数化表: Table 1 给出了 Pre-LN Transformer + AdamW 全部超参的缩放公式，可直接工程落地
宽深比弹性: 打破了"N:L≈100 最优"的传统认知（Kaplan et al.），证明此前结论受 SP 深度失调的干扰
实现极简: 代码改动集中在残差乘子 \(L^{-1}\)、分层学习率和 ε 缩放，无需修改模型架构

局限性¶

最大实验规模为 1.5B non-embedding 参数，尚未在 7B+ 规模验证
仅在 SlimPajama 数据集上测试，未验证跨数据分布的鲁棒性
理论分析在固定 token 数设定下推导，直接推广到 compute-optimal 设定仍需假设
下游评估信号在小模型上相对较弱
未涉及 MoE 或并行子网络等架构变体

评分¶

⭐⭐⭐⭐⭐ (5/5)