跳转至

Rethinking Language Model Scaling under Transferable Hypersphere Optimization

会议: NeurIPS 2026
arXiv: 2603.28743
OpenReview: https://openreview.net/forum?id=Yiph6nL9MR
代码: https://github.com/microsoft/ArchScale
领域: LLM 其他
关键词: 超球面优化, 学习率迁移律, Muon 优化器, 混合专家模型, 训练稳定性

一句话总结

针对大语言模型在二阶/矩阵预条件优化器下缺乏跨算力超参迁移律且规模化易失稳的问题,本文提出基于 Frobenius 超球面约束的 HyperP 参数化框架与 SqrtGate 路由机制,首次实现跨模型宽度、深度、数据量及 MoE 粒度的最优学习率零样本迁移,在 \(6\times 10^{21}\) FLOPs 算力下取得 1.58 倍计算效率杠杆并实现表征稳定性迁移。

研究背景与动机

神经缩放律是指导大语言模型高效预训练的核心理论支柱。为了在可承受的预算内探索架构变体与数据配方,业界普遍在小规模模型上扫描超参数,期望所观察到的最优配置能直接泛化至千亿参数与万亿 token 的前沿计算规模。然而,要揭示真实的缩放行为,缩放曲线上的每个模型都必须运行在其自身规模下的近最优超参数上。在传统 AdamW 优化体系下,尽管 \(\mu\)P 与 \(\mu\)P++ 建立了跨宽度与深度的学习率迁移理论,但这类一阶迁移律未考虑训练步数/数据量(Tokens)对学习率缩放的内生影响,更无法从结构上抑制大规模预训练中频发的 Logit 爆炸、激活离群值(Activation Outliers)与 Loss Spikes 等致命失稳现象。

工程实践中,研究者不得不引入 z-loss 正则化或精心手工编排 Weight Decay 衰减调度,但这些启发式补丁不仅脆弱,而且 Weight Decay 的最优取值与学习率、训练时长乃至模型宽度高度纠缠,严重破坏了跨尺度的超参可迁移性。近期兴起的超球面优化(Hypersphere Optimization)通过将权重矩阵约束在固定范数球面上,为消除输出 Logit 爆炸提供了结构性约束,但现有工作如 MuonH 存在两个关键缺陷:一是误认为超球面归一化自身已具备深度可迁移性而忽略了残差层累积角向漂移,导致大深度下最优学习率严重漂移;二是缺乏跨数据量与混合专家模型(MoE)细粒度路由的系统缩放法则。

面对二阶矩阵优化器与几何流形约束带来的全新动力学特性,如何构建一套兼具最优化效率与结构稳定性的完整超参迁移理论,成为决定前沿大模型能否实现可预测缩放的关键瓶颈。核心 idea:本文提出面向 Frobenius 超球面优化的 HyperP 参数化体系,从理论上证明权重衰减在超球面投影下为一阶零效应并导出严格的 Depth-\(\mu\)P 深度缩放律,揭示数据量缩放与 AdamW 共享 0.32 幂律指数,并配套提出保持 RMS 尺度不变的 MoE SqrtGate 门控,首次实现跨模型宽度、深度、数据预算及专家粒度的全局最优学习率迁移与结构稳定性保障。

方法详解

整体框架

HyperP 的核心目标是消除大语言模型预训练对高维超参联合网格搜索的依赖,使在最小代理模型(如 208M 激活参数、10.4B tokens)上调优的单个基础学习率,能够直接无缝迁移到前沿算力预算(数十亿参数、数百亿 tokens)。整个系统以 Frobenius 超球面上的 MuonH 矩阵优化器为底座,通过分析流形切空间投影消除了权重衰减维度;针对残差网络的前向与反向动力学推导深度缩放系数;通过大规模实证拟合数据规模对学习率的幂律衰减;最后针对 MoE 架构的细粒度划分提出平方根门控以抹平输出范数扰动。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据流与小尺度基础学习率"] --> B["Frobenius 超球面优化与权重衰减消除<br/>切空间投影使得 Weight Decay 沦为一阶零效应"]
    B --> C["跨维度与深度的几何迁移律<br/>各向同性保 RMS 宽度迁移 + 残差累积角向漂移 Depth-μP"]
    C --> D["数据规模幂律学习率调度<br/>通用 magic exponent 0.32 跨 token 衰减"]
    D --> E["SqrtGate 粒度无关 MoE 路由<br/>平方根门控消除细粒度输出 RMS 衰减与 Logit 尖峰"]
    E --> F["全局可迁移的最优预训练与稳定性保障"]

关键设计

1. Frobenius 超球面投影与权重衰减消除:将二维超参搜索降至一维

标准优化器需要同时在学习率 \(\eta\) 与权重衰减 \(\lambda\) 的二维平面上做繁重的网格搜索,且最优 \(\lambda\) 会随模型尺度与训练时长动态漂移。本文证明在 Frobenius 超球面约束下,权重衰减是完全多余的一阶无效操作。设权重矩阵 \(W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}\) 满足 \(\|W\|_F = c_W\),定义更新步为 \(\widetilde{W} = W + \Delta\),投影后为 \(W^+ = c_W \frac{\widetilde{W}}{\|\widetilde{W}\|_F}\)。当更新量 \(\|\Delta\|_F\) 足够小时,一阶泰勒展开表明更新量严格退化为流形切空间投影: $\(W^+ - W = \Pi_T(\Delta) + \mathcal{O}(\|\Delta\|_F^2), \quad \Pi_T(\Delta) = \Delta - \frac{\langle \Delta, W \rangle_F}{\|W\|_F^2} W\)$ 若将传统的解耦权重衰减项加入更新步,即 \(\Delta = -\eta G - \eta \lambda W\),由于径向分量 \(W\) 在切空间投影算子作用下恒等为零,即 \(\Pi_T(W) = W - \frac{\langle W, W \rangle_F}{\|W\|_F^2} W = 0\),权重衰减项在投影后对参数更新的一阶贡献完全消失,仅产生 \(\mathcal{O}(\eta^2)\) 的高阶扰动。因此,在基于 Frobenius 球面的 MuonH 优化中直接将权重衰减置零,不仅不会损失表征容量,反而将超参搜索空间压缩至单个学习率标量。

2. 宽度与深度流形迁移律:修正 MuonH 深度失稳与各向同性保留

在宽度扩展维度,根据谱范数与 Frobenius 范数夹逼引理 \(\|W\|_2 \le \|W\|_F \le \sqrt{r} \|W\|_2\)(其中 \(r = \text{rank}(W)\)),Muon 优化器对梯度更新步的奇异值正交化处理能有效避免极端的特征各向异性,维持扁平的奇异值谱分布。令 \(\|W\|_F = C \sqrt{d_{\text{out}}}\),对于典型输入 \(X\) 有 \(\|WX\|_2 \approx \frac{\|W\|_F}{\sqrt{\min(d_{\text{in}}, d_{\text{out}})}} \|X\|_2\),由此自然诱导出 \(\|WX\|_{\text{rms}} \approx C \|X\|_{\text{rms}}\),在无需像标准 \(\mu\)P 那样对隐藏层显式引入 \(1/w\) 学习率缩放的前提下天然实现宽度迁移。

更为关键的是在深度扩展维度。原 MuonH 文献声称由于参数范数已归一化,优化器天然具备跨深度迁移能力。本文从理论上推翻了这一论断:在深度残差网络 \(x_{l+1} = x_l + \alpha_L f_l(x_l; W_l)\) 中,当权重与更新步均采用固定范数归一化(\(\|U_l\|_F = c_G\))时,经局部雅可比矩阵传导,各残差分支累积到网络输出的总一阶函数摄动量为 \(\mathcal{O}(L \alpha_L \eta_l)\)。为了在大深度下保持 \(\mathcal{O}(1)\) 的函数更新尺度,学习率必须满足: $\(\eta_l = \mathcal{O}\left(\frac{1}{L \alpha_L}\right)\)$ 当采用标准深度稳定化残差缩放 \(\alpha_L = L^{-1/2}\) 时,学习率必须严格按 \(\eta_l = \mathcal{O}(L^{-1/2})\) 进行 Depth-\(\mu\)P 衰减。该结论在 Pre-Norm 与 Post-Norm 下均严格成立。忽视残差累加角向漂移会导致无 Depth-\(\mu\)P 时最优学习率随深度加深而持续下挫(由 8 层的 0.016 跌至 24 层的 0.008),唯有引入 \(\eta \propto 1/\sqrt{d}\) 才能使最优学习率在所有深度上锁定在常数区间。

3. 数据规模幂律学习率传输:跨优化器的普适魔数 0.32

超参迁移律长期面临的一个理论空白是学习率如何随总训练 token 预算 \(T\) 进行外推。本文在固定架构规模下,将数据量由 10.4B 扩展至 166.4B(跨度达 16 倍),在对数空间通过精细网格二次拟合发现最优学习率满足极高精度的幂律关系: $\(\eta^* = 24.27 \cdot T^{-0.320}\)$ 留一交叉验证(Leave-One-Out CV)的最优学习率平均绝对预测误差仅为 1.50%。令人惊叹的是,该幂律指数 0.32 与此前文献在 AdamW 上针对不同网络架构和数据集测得的数据缩放指数完全重合。这一发现强烈暗示,数据扩展指数 0.32 是基于梯度的神经网络优化在数据流形上的内在普适属性,与具体的优化器形态(一阶 AdamW 或二阶超球面 MuonH)解耦。HyperP 据此将隐藏层学习率乘子统一定义为 \(\propto 1 / (T^{0.32} \sqrt{d})\)。

4. SqrtGate 粒度不变门控:攻克细粒度 MoE 路由衰减与 Logit 暴涨

在细粒度混合专家模型中,当总激活专家数 \(k\) 增加时,传统 Softmax 门控输出 \(y_{\text{route}} = \sum_{i=1}^k g_i E_i(x)\)(\(\sum g_i = 1\))存在严峻的尺度坍塌缺陷。在超球面优化诱导的专家均方根能量均等(\(\|E_i(x)\|_{\text{rms}} = r\))且互相关较弱的合理假设下,路由输出的 RMS 满足 \(\|y_{\text{route}}\|_{\text{rms}} \approx r \sqrt{\sum_{i=1}^k g_i^2}\)。在权重分布相对均匀的常规状态下(\(g_i \approx 1/k\)),输出范数缩水为 \(\|y_{\text{route}}\|_{\text{rms}} \approx r / \sqrt{k}\),仅在坍塌为独热路由(Top-1)的极端退化情况下才能保持为 \(r\)。这意味着随着专家细粒度划分增大,路由信号被严重压低,引发后续层表征紊乱与路由 Logit 爆炸。

针对这一几何缺陷,本文提出 SqrtGate(平方根门控),将路由权重形式重构为: $\(y_{\text{route}}' = \sum_{i=1}^k \sqrt{g_i} E_i(x), \quad \sum_{i=1}^k g_i = 1\)$ 在相同弱相关假设下,输出范数严格满足: $\(\|y_{\text{route}}'\|_{\text{rms}} \approx r \sqrt{\sum_{i=1}^k (\sqrt{g_i})^2} = r \sqrt{\sum_{i=1}^k g_i} = r\)$ SqrtGate 在任意路由分布下均从数学上保证了输出 RMS 对激活专家粒度 \(k\) 的严格不变性。结合超球面约束对线性投影输出的有界性保证(\(\|Wx\|_{\text{rms}} \le C \sqrt{d_{\text{in}} / d_{\text{out}}} \|x\|_{\text{rms}}\)),SqrtGate 将路由器的 z-value(Logit 平方和对数度量)峰值直接压低了 5 倍(由超过 190 骤降并稳定在 40 以下),彻底消除了细粒度专家扩展中的数值不稳定。

损失函数 / 训练策略

  1. 参数化配置(HyperP):矩阵类隐层权重采用 MuonH 优化器,学习率乘子设为 \(1 / (T^{0.32} \sqrt{d})\),初始化标准差为 \(1/\sqrt{d_{\text{in}}}\),残差分支乘子取 \(1/\sqrt{2d}\),权重衰减严格置为 0;LM Head 投影矩阵采用 AdamH,向量参数与 Embedding 采用 AdamW,学习率乘子统一保持 \(1/\sqrt{d}\),无需权重衰减。
  2. MoE 辅助负载均衡损失:采用全局批次 Switch 均衡损失 \(\mathcal{L}_{\text{aux}} = \gamma \cdot N \cdot \sum_{i=1}^N f_i \cdot P_i\)。得益于超球面优化对 Logit 幅值的内在边界约束,辅助损失不会干扰主干语言建模,本文将惩罚权重 \(\gamma\) 激进提升至 0.1(常规方案仅用 \(10^{-2}\) 或因掉点而寻求免辅助损失),在实现近乎完全均衡的专家分流(Mean MaxVio 降至 0.086)的同时,取得了最低的验证集损失。
  3. 退火调度与批次设置:遵循无 Warmup 的线性衰减策略(终点为峰值的 10%),全局训练上下文长度设为 4K,批大小固定为 2M tokens(处于临界批大小安全区内)。

实验关键数据

主实验

主实验在 SlimPajama 语料上评估,采用精确计入自注意力和输出头的 Chinchilla FLOPs 统计法,对比基线 Muon(配置 \(\mu\)P++ 与 \(1/w\) 权重衰减缩放)、无 HyperP 的原版 MuonH、以及搭载 HyperP 的稠密与 MoE 模型。计算效率杠杆(CEL)衡量基线达到相同验证损失所需算力与本文方法实际算力之比(\(\rho = C_{\text{base}} / C^*\))。

深度 \(d\) 总算力 (FLOPs) 稠密 Muon 基线 (nats) 原版 MuonH (nats) MuonH+HyperP (nats) HyperP CEL 优势
8 \(2.14 \times 10^{19}\) 2.4777 2.4845 2.4804 \(0.99\times\)
12 \(1.49 \times 10^{20}\) 2.2257 2.2099 2.2192 \(1.04\times\)
16 \(6.59 \times 10^{20}\) 2.0671 2.0500 2.0526 \(1.16\times\)
20 \(2.19 \times 10^{21}\) 1.9591 1.9558 1.9311 \(1.35\times\)
24 \(5.96 \times 10^{21}\) 1.8785 1.9015 1.8365 \(1.58\times\)

在最大算力规模 \(5.96 \times 10^{21}\) FLOPs 下,MuonH+HyperP 取得了 1.58 倍的计算效率杠杆,其相对于基线的优势随算力规模单调递增;而未采用 HyperP 的原版 MuonH 因深度学习率漂移,在大规模下 CEL 锐减至 \(0.70\times\)。此外,稀疏度 \(S=8, k=8\) 的 Transformer-Next-MoE 模型在相同超参迁移下,相较稠密基线实现了高达 \(3.38\times\) 的算力杠杆收益。

消融实验

消融实验重点验证 Depth-\(\mu\)P 的深度保真性、超球面优化消除 Weight Decay 的无损性、以及 SqrtGate 对专家粒度的平滑作用。

配置与消融项 评估维度 核心指标 / 最优 \(\eta^*\) 关键表现与说明
Muon (最佳 \(\lambda=10^{-3}\)) 优化器对比 (\(d=8\)) 最优 Loss: 2.479 (\(\eta^*=0.0222\)) 对 \(\lambda\) 极度敏感,偏离最优值导致性能显著滑坡
MuonH (\(\lambda=0\), HyperP) 权重衰减消除 (\(d=8\)) 最优 Loss: 2.475 (\(\eta^*=0.0155\)) 彻底去除 \(\lambda\) 仍保持略优性能,搜索维度减半
无 Depth-\(\mu\)P 深度扩展 (\(d=8 \to 24\)) \(\eta^*: 0.016 \to 0.008\) 最优学习率随深度腰斩,损失曲面随深度变陡
包含 Depth-\(\mu\)P 深度扩展 (\(d=8 \to 24\)) \(\eta^*: 0.014 \sim 0.016\) 最优学习率完全锁定在平稳常数区间,实现零样本迁移
标准门控 (\(k=64\)) MoE 粒度扩展 最优 Loss: 2.3244 粒度超过 \(k=16\) 后收益饱和衰退,路由信号缩水
SqrtGate (\(k=64\)) MoE 粒度扩展 最优 Loss: 2.3154 在所有 \(k\) 下均持续降低 Loss,且峰值性能延展至 \(k=32\)
辅助损失 \(\gamma=10^{-3}\) 负载均衡权重扫描 Loss: 2.334, MaxVio: 0.848 负载极不均衡,部分专家存在严重的利用率违规
辅助损失 \(\gamma=10^{-1}\) 负载均衡权重扫描 Loss: 2.332, MaxVio: 0.086 负载均衡逼近完美分布,且未造成任何语言建模质量损伤

关键发现

  • 超参迁移的零样本收益放大:由 208M 模型确定的单组基础学习率在向大尺度推进时,HyperP 的算力节省效益不是恒定的,而是随算力呈现单调放大趋势(从 8 层的 \(0.99\times\) 放大至 24 层的 \(1.58\times\)),表明前沿大模型预训练能从该迁移律中捕获极显著的算力边际收益。
  • 稳定性指标随规模非递增:全程监测的 6 项稳定性指标(注意力与 MoE 路由的 z-value、输出 RMS、激活离群值比例)在参数量从 913M 扩展至 13.3B 过程中全部保持有界甚至单调下降,打破了“模型越大越易出现数值离群点与尖峰”的工程魔咒。
  • 超参搜索点数敏感度不对称:实证表明最优 Loss 是学习率的二阶不变量,仅需 5 个粗粒度扫描点即足以将拟合 Loss 相对误差压至 0.04%(绝对值 \(\sim 0.001\) nats),大幅削减了小尺度代理实验的算力开销。

亮点与洞察

  • 切空间几何消除权重衰减:巧妙利用 Frobenius 球面的微元投影几何性质,证明了权重衰减在切空间投影中精确恒等为零,以极简的数学论证消除了预训练中最为棘手的联合调参痛点。
  • 跨优化器的数据缩放魔数:在全新的二阶流形优化器上复现了原本在 AdamW 上发现的 0.32 数据扩展幂律指数,揭示了深度学习优化中可能独立于具体更新规则的深层普适几何结构。
  • SqrtGate 的优雅自洽:针对 MoE 细粒度扩展中的信号衰减,不采用启发式的增益放大层,而是从弱相关向量和的范数出发引入平方根归一化,以极简修改达成了跨专家粒度的能量守恒与数值稳定。

局限与展望

  • 算力最优假说的局限性:作者直接借用了 Chinchilla 的 Token/Param 黄金比例计算各尺度的理论 FLOPs,但不同数据清洗质量与领域混合分布下的最优数据配比可能存在漂移。
  • 数据缩放指数缺乏纯理论解析:尽管 0.32 指数在两类完全不同的优化器上展现出高度一致性,但目前仍建立在经验拟合基础之上,尚缺乏关于数据流形内在维度的完整数学推导。
  • 模型架构覆盖面有待拓宽:验证主要集中在标准的 Transformer-Next 与 Transformer-Next-MoE 架构上,尚未覆盖 Mamba 等混合状态空间模型(SSM)或递归线性注意力架构。

相关工作与启发

  • vs 标准 \(\mu\)P / \(\mu\)P++:\(\mu\)P 族主要面向一阶 AdamW 优化器,通过特征更新尺度的极限分析给出宽度与深度缩放规则,但未闭环数据量(Tokens)对学习率的影响,且无法杜绝大尺度训练中的数值失稳;HyperP 面向 Frobenius 超球面 MuonH,消除了权重衰减,统一了模型规模与数据量两个维度的缩放法则。
  • vs 原版 MuonH:原版 MuonH 忽视了多层残差连接在网络深层累积的角向扰动,错误断言超球面归一化具备内在深度迁移性;HyperP 给出严密的流形雅可比扰动分析,证伪了该假设并补充了必不可少的 Depth-\(\mu\)P 修正。
  • vs 免辅助损失 MoE 负载均衡:DeepSeek 等工作因发现辅助损失恶化语言建模表现而转向无损失路由;HyperP 证明在超球面的有界 Logit 庇护下,传统辅助损失即使放大 100 倍(\(\gamma=0.1\))也不会损害主任务,反而能兼得最优困惑度与极致负载均衡。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次建立超球面优化的全维度迁移律,证伪原版 MuonH 深度假设并提出 SqrtGate
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 208M 到 13.3B 参数、跨越 4 个数量级的算力验证,全流程监控 6 大数值稳定性指标
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严密,实验设计逻辑环环相扣,图表详实且可复现性高
  • 价值: ⭐⭐⭐⭐⭐ 为新一代基于 Muon 的大模型预训练提供了开箱即用的工业级缩放方案与稳定性范式