宽度-深度联合缩放下 µP 的谱条件 Spectral Condition for μP under Width-Depth Scaling
用统一谱框架解决宽度-深度联合缩放下的超参数迁移问题
前置知识
Maximal Update Parameterization (µP)
µP 是一种通过重新参数化超参数（如学习率、初始化方差）使其随模型尺寸变化的方法，目的是在模型宽度缩放时保持特征学习的尺度不变性，同时最大化参数更新对特征的贡献。核心思想是：当网络宽度 n 增大时，标准参数化（SP）会导致特征更新爆炸或消失，而 µP 通过精心设计的缩放规则让每层的权重范数 $\|W_l\|_R = \Theta(1)$ 和更新范数 $\|\Delta W_l\|_R = \Theta(1)$ 保持恒定，从而实现稳定的训练动态。
本文的核心贡献就是将 µP 从单纯的宽度缩放扩展到宽度-深度联合缩放，理解 µP 的基本原理是理解本文的关键前提。
RMS 算子范数
对于矩阵 $A \in \mathbb{R}^{m \times n}$，RMS 算子范数定义为 $\|A\|_R := \max_{v \neq 0} \frac{\|Av\|_R}{\|v\|_R} = \sqrt{\frac{m}{n}} \|A\|_2$，其中 $\|A\|_2$ 是谱范数（最大奇异值）。与普通谱范数不同，RMS 算子范数在宽度缩放时自动归一化向量维度，使得不同宽度的矩阵可以直接比较其'有效强度'。本文采用 RMS 算子范数而非谱范数来表述谱条件，主要是为了符号简洁。
论文的整个理论框架都建立在 RMS 算子范数之上，谱条件的核心不等式都用该范数表达，不理解它就无法理解 Condition 3.1 的含义。
残差网络中的块深度 k
在残差网络 $h_l(x) = h_{l-1}(x) + \alpha_l W_l^{(2)} W_l^{(1)} h_{l-1}(x)$ 中，k 表示每个残差分支中包含的线性变换层数。k=1 对应单层残差分支，k=2 对应两层残差分支（如 Transformer 中的 attention + FFN 结构）。论文证明 k 的取值决定了深度缩放规则：k=1 产生较宽松的 Depth-µP 缩放，而 k≥2 由于出现高阶更新项，需要更严格的 CompleteP 缩放。
k 值是本文理论的核心变量，它解释了为什么之前不同架构和优化器的 µP 结论不一致——本质上是因为它们的残差分支深度不同。
谱条件（Spectral Condition）
谱条件是一组关于权重范数和更新范数的缩放约束。在宽度缩放下，谱条件要求 $\|W_l\|_R = \Theta(1)$ 且 $\|\Delta W_l\|_R = \Theta(1)$。在宽度-深度联合缩放下，本文提出的 Condition 3.1 将这些约束细化：隐藏层需要 $\alpha_l \|W_l^{(2)}\|_R \|W_l^{(1)}\|_R = \Theta(1/L)$，且更新约束包括一阶项 $\alpha_l \|\Delta W_l^{(2)}\|_R \|W_l^{(1)}\|_R = \Theta(1/L)$ 和二阶项 $\alpha_l \|\Delta W_l^{(2)}\|_R \|\Delta W_l^{(1)}\|_R = \Theta(1/L)$。
谱条件是连接理论分析与实际超参数实现的桥梁，通过它可以从统一的数学约束推导出各种优化器的具体 HP 缩放规则。
超参数迁移（HP Transfer）
µP 的一个关键实用价值是：在小模型上调好的最优超参数可以直接迁移到大模型。具体做法是先确定基础超参数 $\eta_{base}$、$\sigma^2_{base}$、$\alpha_{base}$，然后根据宽度缩放比 $r_n = n/n_{base}$ 和深度缩放比 $r_L = L/L_{base}$ 计算目标模型的超参数。例如隐藏层学习率设为 $\eta_l = \eta_{base}/\sqrt{r_n}$。这大大降低了大模型训练的调参成本。
论文的实验核心验证就是 HP 迁移能力——µP 能否让小模型上调好的学习率在大模型上依然最优。
标准参数化（SP）
标准参数化是神经网络训练中的默认缩放方式，通常使用与模型尺寸无关的固定超参数。在宽度缩放时，SP 会导致隐藏层特征的范数随宽度变化（通常爆炸或消失），使得训练不稳定且最优学习率随模型尺寸漂移。SP 是本文的基线对照组。
理解 SP 的缺陷才能理解 µP 的价值——论文的实验反复对比 SP 与 µP，展示 SP 在宽度-深度联合缩放下的失败。
研究动机
生成式基础模型正在同时增大宽度和深度（如 GPT-4、Llama 等百亿参数模型），但现有 µP 理论在宽度-深度联合缩放场景下存在严重不足。首先，已有扩展方法（如 Depth-µP、CompleteP）各自依赖特定的架构假设和优化器选择，结论高度碎片化：例如 Depth-µP 适用于 k=1 的残差块，CompleteP 适用于 Transformer 风格的多变换残差分支，但两者的理论推导分别依赖于 Tensor Programs 或动力学均场论等技术性极强的工具，普通研究者难以理解和扩展。其次，不同论文给出的 µP 公式看似矛盾（如残差乘子 $\alpha_l$ 的缩放分别是 $\Theta(1/\sqrt{L})$ 和 $\Theta(1/L)$），缺乏统一的解释框架来说明这些差异的根源。最后，当模型参数达到数十亿级别时，特征学习动态变得不稳定，超参数调优成本极高，亟需一个简单、统一且实用的理论指导。
本文的目标是本文的目标是建立一个简单统一的谱框架，为宽度-深度联合缩放下的 µP 提供可操作的理论基础。具体而言，作者希望：(1) 用基本的线性代数和概率论工具（而非 Tensor Programs 等复杂理论）推导出统一的谱条件；(2) 解释为什么 k=1 和 k≥2 的残差块需要不同的深度缩放规则；(3) 将谱条件映射到具体优化器的超参数实现方案；(4) 在 GPT-2 风格的 Transformer 上验证其实用性。
与已有工作不同的是，本文的独特切入角度是'谱视角'（spectral perspective）。之前 Yang 等人（2022）已经在宽度缩放场景下建立了谱条件 $\|W_l\|_R = \Theta(1), \|\Delta W_l\|_R = \Theta(1)$ 的统一框架，但该条件仅适用于固定深度。本文将这一谱视角扩展到宽度-深度联合缩放场景，通过分析残差分支中一阶和二阶更新项的累积效应，揭示了残差块深度 k 如何决定深度缩放规则。这个视角的优势在于：(1) 只需要基本的线性代数运算（子可加性、子可乘性、随机矩阵谱范数估计），大大降低了理解门槛；(2) 自然统一了之前分散的 µP 结论，将 Depth-µP 和 CompleteP 纳入同一个框架的两个特例；(3) 可以直接推导出任意优化器的 HP 缩放规则，无需重复复杂的理论推导。
核心方法
本文的方法路线可以概括为：建立模型 → 推导谱条件 → 映射到优化器实现 → 实验验证。首先，作者建立了包含 k 层线性变换的深度残差网络模型 $h_l(x) = h_{l-1}(x) + \alpha_l W_l^{(2)} W_l^{(1)} h_{l-1}(x)$，其中宽度 n 和深度 L 同时趋向无穷。然后，通过分析前向传播的特征稳定性（初始条件）和单步梯度更新的特征变化（更新条件），推导出一组关于权重范数和更新范数的缩放约束（Condition 3.1）。关键的直觉是：每个残差块对最终特征的贡献需要均摊为 $\Theta(1/L)$，而当残差分支包含 k≥2 层时，更新展开中会出现二阶项 $\Delta W^{(2)} \Delta W^{(1)}$，这要求更严格的缩放约束。最后，通过将谱条件中的 RMS 算子范数约束映射到具体优化器的更新规则，推导出每个优化器的 µP 超参数缩放方案。
本文的核心创新是揭示了残差块深度 k 对深度缩放规则的决定性影响，以及二阶更新项在其中的关键作用。具体而言，当分析残差网络的单步特征更新 $\Delta h_l(x)$ 时，可以将其展开为零阶、一阶和二阶项：零阶项 $\epsilon_0$ 是前序层更新的传播，一阶项 $\epsilon_1$ 包含 $\Delta W^{(2)} W^{(1)}$ 和 $W^{(2)} \Delta W^{(1)}$ 这类恰好更新一个权重的项，二阶项 $\epsilon_2$ 包含 $\Delta W^{(2)} \Delta W^{(1)}$ 这类同时更新两个权重的项。对于 k=1 的残差块，不存在二阶项，因此约束较宽松，得到 $\alpha_l = \Theta(1/\sqrt{L})$ 的 Depth-µP 缩放。对于 k≥2 的残差块，二阶项的存在将约束收紧为 $\alpha_l = \Theta(1/L)$ 的 CompleteP 缩放。这个发现统一了之前看似矛盾的 µP 结论：它们不是矛盾的，而是分别适用于不同深度的残差块。
方法步骤详情
方法分为以下步骤：(1) 模型建立：定义 k=2 的深度残差网络 $h_0(x) = \alpha_0 W_0 x$，$h_l(x) = h_{l-1}(x) + \alpha_l W_l^{(2)} W_l^{(1)} h_{l-1}(x)$，$h_{L+1}(x) = \alpha_{L+1} W_{L+1} h_L(x)$，权重初始化为 i.i.d. 高斯分布，宽度 $n, n_l = \Theta(n)$，深度 $L \to \infty$。(2) 推导初始条件：从输入层到输出层逐层分析前向传播稳定性，要求 $\alpha_0 \|W_0\|_R = \Theta(1)$，$\alpha_{L+1} \|W_{L+1}\|_R = \Theta(1)$，隐藏层初步要求 $\alpha_l \|W_l^{(2)}\|_R \|W_l^{(1)}\|_R = O(1/\sqrt{L})$。(3) 推导更新条件：展开单步更新为零阶、一阶、二阶项，要求一阶项 $\alpha_l \|\Delta W_l^{(2)}\|_R \|W_l^{(1)}\|_R = \Theta(1/L)$，二阶项 $\alpha_l \|\Delta W_l^{(2)}\|_R \|\Delta W_l^{(1)}\|_R = \Theta(1/L)$。(4) 联合求解：将一阶更新条件相乘再除以二阶更新条件，得到最终初始条件 $\alpha_l \|W_l^{(1)}\|_R \|W_l^{(2)}\|_R = \Theta(1/L)$，将初步的 $O(1/\sqrt{L})$ 收紧为 $\Theta(1/L)$。(5) 优化器映射：以 Muon-Kimi 为例，其更新范数满足 $\|\Delta W_l\|_R = \Theta(\eta_l \sqrt{n_{in}})$（隐藏层），代入谱条件得到学习率 $\eta_l = \Theta(1/\sqrt{n_{in}})$，结合宽度缩放比得到 $\eta_l = \eta_{base}/\sqrt{r_n}$。
技术新颖性
本文的技术新颖性体现在多个层面。首先，在理论层面，这是首次用基本线性代数工具（而非 Tensor Programs 或动力学均场论）推导出宽度-深度联合缩放的 µP 条件，大大降低了理论门槛。具体来说，整个推导仅依赖子可加性 $\|A+B\|_R \leq \|A\|_R + \|B\|_R$、子可乘性 $\|AB\|_R \leq \|A\|_R \|B\|_R$ 和高斯随机矩阵的谱范数估计 $\|A\|_2 = \Theta(\sigma(\sqrt{m}+\sqrt{n}))$。其次，在洞察层面，通过'更新阶数'视角统一了 Depth-µP（k=1）和 CompleteP（k≥2）两种之前独立发展的方法，揭示了它们的差异源于残差分支深度而非理论框架的不同。第三，在实用性层面，论文给出了一个简洁的实现规则：对于大多数现代优化器（除 SGD 外），只需在宽度缩放 µP 的基础上添加隐藏残差乘子 $\alpha_l = \Theta(1/L)$ 即可。这是因为归一化和预处理操作消除了梯度中的深度因子，使得优化器特定的 HP 规则与宽度缩放 µP 相同。
实验结果
论文在 GPT-2 风格的 Transformer 语言模型上进行了系统实验，使用 OpenWebText 数据集、GPT-2 tokenizer、最大序列长度 1024。基础模型配置为 $(n_{base}, L_{base}) = (256, 4)$，宽度最大缩放到 4096，深度最大缩放到 256 层。实验使用 Muon-Kimi-AdamW、Muon-AdamW、Shampoo-AdamW 和 Sophia 四种优化器。核心发现包括：(1) 特征学习稳定性：在 10 步坐标检查测试中，SP 的特征范数 $\|h_L\|_R$ 随宽度和深度快速增长（图 1a,b 中从 $2^{10}$ 到 $2^{12}$），而 µP 始终保持稳定的尺度不变特征（约 $2^{5}$ 到 $2^{9}$ 的范围内）。(2) HP 迁移能力：在 300M tokens 训练中，SP 的最优学习率随宽度显著漂移（图 1c），而 µP 的最优基础学习率几乎不变。(3) 深度缩放对比：Condition 3.1 (k≥2) 实现了稳定的学习率迁移，而 Condition B.1 (k=1) 随深度增加导致最优学习率漂移（图 2），验证了残差块深度 k 的决定性作用。(4) LayerNorm 消融：移除 LayerNorm 后，SP 训练在深度缩放时变得不稳定，而 µP 在 L=256 时仍然稳定，说明 SP 的表面稳定部分依赖于归一化层的掩盖效应。(5) 优化器通用性：在所有四种优化器上，µP 都展现出一致的优势。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 特征学习稳定性（宽度缩放，10步） | 输出特征 RMS 范数 $\|h_L\|_R$ 的稳定性 | µP (Condition 3.1)：特征范数在宽度 128-4096 范围内保持 $\Theta(1)$ 恒定 | SP：特征范数随宽度快速增长，从约 $2^5$ 增长到超过 $2^{12}$ | µP 实现了尺度不变的特征学习，SP 导致特征爆炸 |
| 特征学习稳定性（深度缩放，10步） | 输出特征 RMS 范数 $\|h_L\|_R$ 的稳定性 | µP (Condition 3.1)：特征范数在深度 4-256 范围内保持稳定 | SP：特征范数随深度快速增长 | µP 在深度方向同样保持稳定，SP 在深度缩放时失败 |
| HP 迁移（宽度缩放） | 最优基础学习率 $\eta_{base}$ 的跨宽度一致性 | µP：最优 $\eta_{base}$ 在宽度 256-4096 范围内几乎不变 | SP：最优学习率随宽度显著漂移 | µP 实现了可靠的宽度方向 HP 迁移 |
| HP 迁移（深度缩放） | 最优基础学习率 $\eta_{base}$ 的跨深度一致性及验证损失 | µP (Condition 3.1)：最优 $\eta_{base}$ 稳定，且损失随深度降低（约 3.5-4.0） | SP：学习率部分稳定但损失较高；Condition B.1：学习率随深度漂移 | µP (k≥2) 同时实现稳定迁移和更低损失 |
| 残差块深度 k 的影响验证 | 深度缩放下的 HP 迁移稳定性 | Condition 3.1 (k≥2)：Muon-Kimi-AdamW 和 Muon-AdamW 均实现稳定迁移 | Condition B.1 (k=1)：两种优化器均出现学习率漂移 | k≥2 的 CompleteP 缩放是 Transformer 架构的正确选择 |
局限与改进
本文存在以下局限性：(1) 理论推导基于线性残差 MLP 的单步更新简化假设，虽然作者在附录 G 中引入了一些自然假设来推广到多步、非线性和多样本场景，但这些假设的严格验证仍不充分。(2) 实验主要在 GPT-2 风格的 Transformer 上进行，模型规模最大约到中等大小（宽度 4096、深度 256），尚未在真正的大规模预训练（如 7B+ 参数）中验证。(3) 论文发现 SP 在 Muon-Kimi-AdamW 下的深度缩放表现看似尚可（图 1d），作者将此归因于 LayerNorm 等归一化组件的掩盖效应，这说明实际架构中的归一化层可能部分缓解了缩放问题，但论文未深入分析这种交互作用。(4) 谱条件的推导依赖于'非相消'和'对齐'假设（即权重矩阵与特征向量之间不存在特殊的相消对齐），虽然这是神经网络理论中的常见假设，但在实际训练中的有效性需要更多验证。(5) 论文未讨论 µP 与学习率调度（如 warmup、cosine decay）的交互，实际训练中这些因素可能影响 HP 迁移效果。
独立分析的弱点
本文的弱点包括：(1) 实验规模有限——所有实验的最大模型宽度为 4096、深度为 256，这对于验证大模型缩放理论来说仍然偏小。建议在未来工作中扩展到 7B-70B 参数规模的预训练实验，以验证 µP 在真正大规模场景下的有效性。(2) 缺乏与现有方法的直接对比——论文没有与 CompleteP 和 Depth-µP 的原始实现进行直接对比实验，而是通过 Condition B.1 作为替代。建议补充与这些方法原始论文中的实验配置的直接对比。(3) 理论假设的局限——非相消假设在宽度-深度联合缩放下的有效性缺乏严格的数学证明，仅在附录中通过实验验证。建议发展更严格的理论工具来支撑这些假设。(4) 优化器覆盖不完整——虽然推导了 9 种优化器的 µP 方案，但实验部分主要集中在 Muon-Kimi-AdamW，其他优化器的结果仅在附录中简要展示。建议在正文中增加更多优化器的实验对比。(5) 未讨论 µP 与梯度裁剪、学习率 warmup 等常见训练技巧的交互，这些在实际大规模训练中至关重要。
未来方向
基于本文的谱框架，未来研究可以在多个方向展开：(1) 将框架扩展到更复杂的架构，如混合专家模型（MoE）、状态空间模型（SSM）等，这些架构的残差结构可能与标准 Transformer 不同。(2) 研究 µP 与训练稳定性技术（如梯度裁剪、QK-Norm、Pre-Norm vs Post-Norm）的联合优化，论文已经观察到 LayerNorm 会掩盖 SP 的缩放问题，这提示归一化与参数化的交互值得深入研究。(3) 在真正的大规模预训练中验证 µP 的实用性，特别是在多模态模型和长上下文场景下。(4) 将谱框架扩展到微调和迁移学习场景，研究预训练模型的 µP 缩放规则是否仍然适用。(5) 探索 µP 与架构搜索的结合，利用 µP 的 HP 迁移能力来降低架构搜索的计算成本。(6) 发展更严格的理论工具来支撑联合缩放下非相消假设的有效性。
复现评估
论文的复现条件较好。作者在 GitHub 上开源了代码（https://github.com/ML-GSAI/Width-Depth-muP），提供了完整的实验配置。实验使用 OpenWebText 数据集（公开可用）和 GPT-2 tokenizer（标准工具），基础模型配置为 $(n_{base}, L_{base}) = (256, 4)$，训练 300M tokens，batch size 240。使用的学习率调度为线性 warmup + cosine decay。论文的理论推导基于基本线性代数，不依赖复杂的数值方法，理论结果可以直接复现。然而，完整的超参数搜索实验（特别是图 1c,d 中的学习率扫描）可能需要较多计算资源，因为需要在多个宽度和深度配置上训练多个模型。总体而言，复现难度中等，主要挑战在于计算资源而非技术细节。
论文图表
