On the Nonlinearity of Learning Rate Scaling for LLM Training¶
会议: NeurIPS 2026
arXiv: 2606.29158
OpenReview: https://openreview.net/forum?id=3fy5G5iNJf
论文: https://arxiv.org/abs/2606.29158
代码: https://github.com/neurips2026lrtransfer/LR-Transfer-Trajectory
领域: 预训练
关键词: 学习率缩放定律, 有效学习率, 超参数迁移, AdamW 动力学, 权重范数松弛
一句话总结¶
本文系统评估了大模型训练中经典对数线性学习率缩放定律的失效机制，揭示出最优学习率在扩展规模时存在系统性向上弯曲，证明改用基于归一化权重方向的有效学习率 \(\eta_{\mathrm{eff}}\) 配合沿数据量 \(D\) 轴外推能大幅消除非线性并降低计算开销，同时从 AdamW 瞬态权重范数演化的数学机制上解释了非线性根源。
研究背景与动机¶
在万亿 token 级别训练超大规模语言模型时，直接在目标规模下对学习率等超参数进行网格搜索的算力开销是不可承受的。为此，学术界与工业界广泛采用超参数迁移（Hyperparameter Transfer）范式——在低成本的小模型或小数据量预算下探索最优学习率，再通过参数化缩放定律外推至目标训练规模。以 \(\mu\mathrm{P}\) 为代表的理论参数化方案主要覆盖模型宽度扩展，但在权重范数动态漂移或沿数据量轴缩放时预测精度受限；以经验幂律为代表的研究则普遍假设最优学习率 \(\eta^*(D, N)\) 与数据规模 \(D\) 及模型参数量 \(N\) 严格服从对数线性关系 \(\log \eta^* = a \log D + b \log N + c\)。
然而，这一对数线性假说在更宽广的训练视界（Horizon）下从未得到充分检验，且已有文献往往忽略了小规模拟合本身的算力约束对最终外推精度的实质影响。当训练规模从数十亿 token 扩展到上百亿乃至千亿 token 时，经验最优学习率往往显著偏离对数线性拟合线，直接导致外推给出的学习率过小，迫使下游需要耗费巨大的额外计算量（Extra Compute Ratio, ECR）才能弥合验证集损失差距。更深层的矛盾在于：现代 Transformer 由于 LayerNorm 等归一化层具备尺度不变性，网络输出和损失实质上仅取决于参数的方向而非绝对模长，经典优化器中名义学习率与实际参数更新轨迹之间存在隐式的动力学解耦。
针对这一根本痛点，本文在统一的计算预算约束下系统评测了学习率在模型轴与数据轴上的外推行为，提出以归一化权重空间步长定义的有效学习率（Effective Learning Rate, \(\eta_{\mathrm{eff}}\)）代替名义学习率。核心 idea：最优名义学习率随训练规模扩展出现的向上弯曲源于 AdamW 权重范数在小学习率下远离稳态平衡的慢速瞬态演化，改用表征权重方向角位移的有效学习率 \(\eta_{\mathrm{eff}}\) 并优先沿数据量 \(D\) 轴外推，能够彻底恢复稳定的对数线性缩放规律，在节省 80% 以上扫参算力的同时仅引入微量额外训练代价。
方法详解¶
整体框架¶
本文旨在建立准确且计算高效的大模型训练学习率外推框架。整个研究 pipeline 包含四个阶段：密集网格轨迹采集、双轴（数据轴 \(D\) vs 模型轴 \(N\)）与双参数化（名义学习率 \(\eta\) vs 有效学习率 \(\eta_{\mathrm{eff}}\)）缩放定律拟合、受限算力下的分布外泛化与额外计算量评测，以及基于 AdamW 权重范数动力学和显式控制优化器 AdamH 的机理解释。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["密集网格预训练轨迹采集<br/>GPT-2 (22M-707M), FineWeb (5B-100B)"] --> B["有效学习率参数化与稳态聚合<br/>计算归一化权重空间角位移步长"]
B --> C["算力约束双轴外推评测<br/>固定预算比较 D 轴与 N 轴拟合准确度"]
C --> D["AdamW 权重范数松弛机制分析<br/>平衡态与预平衡态的两阶段动力学推导"]
D --> E["AdamH 显式控制超球面优化验证<br/>约束更新范数与权重范数彻底消除曲率"]在数据采集阶段，作者基于 nanoGPT 框架在 FineWeb-100B 数据集上系统训练了参数量覆盖 22M、64M、124M、208M、417M、707M（隐藏层维度 256 至 2048）的 12 层 GPT-2 模型，数据步长从 5B 扩展到 100B tokens（以 2.5B 为增量），并在每个配置下对学习率进行以 2 为底对数步长为 1 的大范围网格扫描（\(\log_2 \eta \in [-16, -7]\)）。通过在 WSD（Warmup-Steady-Decay）调度期间周期性保存检查点并追加统一的 5000 步衰减阶段，高效构建出覆盖稠密损失轨迹和优化统计量的开源基准数据集。
关键设计¶
1. 有效学习率参数化与稳态时间聚合
现代 Transformer 架构中的注意力投影矩阵和前馈层通常位于 LayerNorm 之后，具有尺度不变性 \(\mathcal{L}(w) = \mathcal{L}(\alpha w), \forall \alpha > 0\)。名义学习率 \(\eta\) 无法如实反映权重在超球面上的角位移。为此，本文定义有效学习率为归一化权重方向 \(\hat{w}_t = w_t / \|w_t\|_2\) 的单步位移：
实测显示在固定学习率下，除初始 warmup 外，\(\eta_{\mathrm{eff}}(t)\) 在整个稳定训练阶段演化极其平缓。因此，作者在除 embedding 以外的所有权重矩阵上计算 \(\eta_{\mathrm{eff}}\)，并通过时间平均将其标量化为 \(\overline{\eta}_{\mathrm{eff}}(\eta; D, N) = \frac{1}{T}\sum_{t=0}^{T-1}\eta_{\mathrm{eff}}(t; \eta, D, N)\)。对于给定的训练规模 \((D, N)\)，通过对验证损失曲线进行三次多项式拟合分别确定经验最优的名义学习率 \(\eta^*(D, N)\) 与最优有效学习率 \(\eta_{\mathrm{eff}}^*(D, N)\)。
2. 统一算力预算下的双轴外推基准评测
为了客观衡量获取缩放定律所需的实际前置成本，本文提出了受限算力外推评估协议。设目标规模训练基准算力为 \(C(D_{\mathrm{test}}, N_{\mathrm{test}}) = 6 N_{\mathrm{test}} D_{\mathrm{test}}\)，给定总算力比例约束 \(\beta \in (0, 1)\)，对比两条正交的外推路径： - 数据轴（\(D\) 轴）外推：使用完整目标模型 \(N_{\mathrm{test}}\)（707M），但在较小的数据子集 \(\{\alpha_k D_{\mathrm{test}}\}_{k=1}^K\) 上训练以拟合定律 \(\log \eta^* = a_N \log D + b_N\)（或其有效学习率版本）； - 模型轴（\(N\) 轴）外推：在全量目标数据 \(D_{\mathrm{test}}\)（如 100B）上运行较小参数量模型 \(\{\alpha_k N_{\mathrm{test}}\}_{k=1}^K\) 来拟合定律 \(\log \eta^* = c_D \log N + d_D\)。
在满足 \(\sum_{k=1}^K \alpha_k = \beta\) 的相同算力约束下，本文引入分布外决定系数 \(R^2_{\mathrm{OOD}}(\beta)\) 以及额外计算量开销比率 \(\mathrm{ECR}(\beta) = \sum \Delta C / \sum C\)，将统计拟合偏差直接映射为工程中弥合性能损失所需要的真实算力浪费。
3. AdamW 权重范数松弛动力学解释非线性曲率
理论分析指出，最优名义学习率在对数-对数尺度上的向上弯曲本质上是由 AdamW 的瞬态松弛效应导致的。当考虑 AdamW 更新步 \(w_{t+1} = (1 - \eta \lambda) w_t - \eta u_t\)（其中 \(u_t\) 为标准化梯度动量，模长近似常数 \(\|u_t\|_2 \approx U\)）时，权重范数并非瞬时恒定，而是以特征衰减时间 \((2\eta \lambda)^{-1}\) 指数收敛向平衡态模长 \(W_{+\infty} = U \sqrt{\frac{\eta}{2\lambda}\frac{1+\beta_1}{1-\beta_1}}\)：
这导致了两个截然不同的物理区间： - 平衡态区间（\(t \gg (2\eta\lambda)^{-1}\)）：指数项衰减殆尽，\(\eta_{\mathrm{eff}} \propto \sqrt{\eta}\)，在对数空间中斜率为 \(\log \eta = 2 \log \eta_{\mathrm{eff}} + \text{const}\)； - 预平衡态区间（\(t \ll (2\eta\lambda)^{-1}\)）：初始权重模长占主导，\(\eta_{\mathrm{eff}} \approx \eta U / W_0 \propto \eta\)，在对数空间中斜率为 \(\log \eta = \log \eta_{\mathrm{eff}} + \text{const}\)。
随着训练步数 \(T\)（数据规模 \(D\)）扩大，理论上最优的有效学习率单调下降，从而迫使最优名义学习率 \(\eta^*\) 变得极小；而较小的 \(\eta^*\) 会导致松弛时间 \((2\eta\lambda)^{-1}\) 剧增（例如当 \(\eta=2^{-16}, \lambda=0.1\) 时需 32.7 万步才达稳态），使训练被困在预平衡态区间。这种在长视界下从斜率 2 向斜率 1 的渐进过渡，直接在对数-对数图上塑造出随规模扩展而向上翘起的非线性曲率。
4. 基于 AdamH 优化器的显式约束验证
为验证上述动力学假设，消除优化器隐式演化带来的干扰，作者引入了 AdamH 优化器作为对比实验。AdamH 将权重与更新向量同时显式约束在固定范数的超球面上：每步计算 Adam 更新向量 \(u_t\) 后，将临时更新设为 \(\tilde{W}_{t+1} = W_t - \eta R \cdot \mathrm{Normalize}(u_t)\)，再重投影至半径 \(R = \|W_0\|_2\)。在 AdamH 下，\(\eta_{\mathrm{eff}}\) 受到优化器的严格显式控制，彻底消除了权重模长漂移的瞬态相位。实验证明，在 AdamH 下最优学习率与数据规模展现出几近完美的对数线性拟合关系（Pearson 相关系数 \(r = -0.9994\)），强有力地证实了理论推导。
损失函数 / 训练策略¶
语言模型训练采用标准自回归交叉熵损失。所有实验采用 12 层 GPT-2 解码器架构，注意力头维度固定为 128，隐藏层维度在 \(\{256, 512, 768, 1024, 1536, 2048\}\) 间缩放。优化器使用 AdamW，设置 \(\beta_1 = 0.95, \beta_2 = 0.95, \epsilon = 10^{-8}\)，权重衰减 \(\lambda = 0.1\)。Embedding 层使用独立的 AdamW 参数优化（学习率固定为 0.0036，\(\beta_1=0.9, \beta_2=0.95\)）。学习率调度器采用 WSD，包含 1000 步线性 warmup，随后保持稳定值，并在指定训练终点前执行 5000 步退火衰减，最终在独立的 10M token 验证集上测定损失。
实验关键数据¶
主实验¶
在目标规模 \((D_{\mathrm{test}}, N_{\mathrm{test}}) = (100\mathrm{B}, 707\mathrm{M})\) 下，评估不同算力预算比例 \(\beta\)、不同外推轴向（数据轴 \(D\) vs 模型轴 \(N\)）以及不同参数化形式（有效学习率 \(\eta_{\mathrm{eff}}\) vs 名义学习率 \(\eta\)）的预测误差、验证集损失差距与额外计算量开销：
| 算力预算比例 \(\beta\) (%) | 外推策略设定 | 对数学习率预测误差 | 验证集损失差 (\(10^{-4}\)) | 额外计算量比例 ECR (%) |
|---|---|---|---|---|
| 13.00% | \(D\) 轴, \(\eta_{\mathrm{eff}}\) | +0.1235 | 4.77 | 1.21 |
| 13.00% | \(D\) 轴, \(\eta\) | -0.8144 | 66.42 | 16.96 |
| 17.00% | \(D\) 轴, \(\eta_{\mathrm{eff}}\) | +0.1148 | 4.12 | 1.04 |
| 17.00% | \(D\) 轴, \(\eta\) | -0.7295 | 53.21 | 13.72 |
| 20.64% | \(D\) 轴, \(\eta_{\mathrm{eff}}\) | +0.1271 | 5.06 | 1.29 |
| 20.64% | \(N\) 轴, \(\eta_{\mathrm{eff}}\) | -0.1903 | 11.04 | 2.21 |
| 20.64% | \(D\) 轴, \(\eta\) | -0.5736 | 32.82 | 8.65 |
| 20.64% | \(N\) 轴, \(\eta\) | -0.5461 | 29.74 | 7.87 |
| 37.63% | \(D\) 轴, \(\eta_{\mathrm{eff}}\) | +0.0935 | 2.73 | 0.67 |
| 37.63% | \(N\) 轴, \(\eta_{\mathrm{eff}}\) | -0.1320 | 5.35 | 1.11 |
| 37.63% | \(D\) 轴, \(\eta\) | -0.2061 | 4.21 | 1.20 |
| 37.63% | \(N\) 轴, \(\eta\) | -0.2328 | 5.37 | 1.52 |
| 55.10% | \(D\) 轴, \(\eta_{\mathrm{eff}}\) | +0.0983 | 3.02 | 0.75 |
| 55.10% | \(N\) 轴, \(\eta_{\mathrm{eff}}\) | -0.1357 | 5.64 | 1.17 |
| 55.10% | \(D\) 轴, \(\eta\) | -0.2286 | 5.18 | 1.47 |
| 55.10% | \(N\) 轴, \(\eta\) | -0.2179 | 4.71 | 1.34 |
消融实验¶
针对显式控制有效学习率的 AdamH 优化器，与标准 AdamW 在 \(N=64\mathrm{M}, D=100\mathrm{B}\) 设定下对比各预算比例下的外推线性度与算力代价：
| 优化器与拟合设定 | 线性相关系数 \(r\) | 预算 \(\beta=13\%\) 时的 ECR (%) | 预算 \(\beta=20\%\) 时的 ECR (%) | 说明 |
|---|---|---|---|---|
| AdamH (\(D\) 轴拟合) | -0.9994 | < 1.0% | < 0.8% | 严格受限超球面，完全消除瞬态模长效应 |
| AdamW (\(D\) 轴, \(\eta_{\mathrm{eff}}\)) | -0.9950 | 1.21% | 1.29% | 间接统计有效学习率，线性度高度逼近 |
| AdamW (\(D\) 轴, \(\eta\)) | -0.9620 | 16.96% | 8.65% | 存在系统性向上翘曲，低预算下严重低估学习率 |
| AdamW (\(N\) 轴, \(\eta\)) | -0.9580 | 不可用 (无法覆盖) | 7.87% | 模型轴外推误差波动更大，缺乏极小算力区间支持 |
关键发现¶
- 向上曲率导致严重低估：采用名义学习率对数线性外推时，预测误差均为负值（如 \(\beta=13\%\) 时为 -0.8144），这意味着传统幂律在长视界训练中会大幅低估最优学习率，导致训练欠收敛；而有效学习率 \(\eta_{\mathrm{eff}}\) 展现出更出色的对称性和抗偏能力（误差仅在 +0.1 左右）。
- 数据轴外推全面优于模型轴：在相同受限算力 \(\beta\) 下，沿数据量 \(D\) 轴扩展比沿模型规模 \(N\) 轴扩展具有更高的分布外决定系数 \(R^2_{\mathrm{OOD}}\) 和更低的额外算力消耗（如 \(\beta=20.64\%\) 下，\(D\) 轴 \(\eta_{\mathrm{eff}}\) 的 ECR 为 1.29%，显著优于 \(N\) 轴的 2.21%）。
- 泛化误差与算力代价呈严格线性映射：实验测定额外计算量比率与分布外泛化误差之间满足严密的比例关系 \(\mathrm{ECR} \propto 1 - R^2_{\mathrm{OOD}}\)，说明学习率外推精度哪怕只有轻微劣化，工程上也会按比例直接放大补齐性能所需消耗的 GPU 显卡时。
亮点与洞察¶
- 揭示了学习率缩放定律的局部有效性：打破了过去文献普遍将「最优学习率关于模型和数据规模呈全局对数线性」奉为信条的认知，指出其仅在小规模区间局部近似成立，大尺度下的向上翘曲是优化器动力学必然导致的结构性现象。
- 确立了有效学习率作为超参数迁移的真正第一性量：由于 LayerNorm 的尺度不变性，网络更新的实质是方向的变化；证明了以超球面单位步长 \(\eta_{\mathrm{eff}}\) 作为拟合对象，比直接拟合名义学习率 \(\eta\) 具有本质上的几何优越性。
- 给出了清晰的工程指导原则：在大模型预训练预算规划中，优先固定目标大模型、用缩短 token 数（数据量 \(D\) 轴）进行小成本探索并外推有效学习率，只需消耗约 13% 的基准算力即可锁定最优超参数，相比动辄全量扫参节约了 80% 以上的前置算力。
局限与展望¶
- 模型架构与优化器覆盖面受限：主干实验集中在 12 层 GPT-2 标准架构与 AdamW 优化器上，尚未推广到混合专家架构（MoE）、具有 RoPE/RMSNorm 的 LLaMA 风格模型，以及新兴的 Muon/SOAP 等二阶矩阵优化器。
- 最优值估计依赖三次曲线拟合：经验最优学习率由三次多项式在扫参点上拟合极小值得到，当真值落在网格边界附近时，拟合多项式自身可能引入微小的人为偏差。
- 实验最大规模上限：当前受限于学术算力，模型规模验证至 707M、数据量至 100B tokens，尚未探究 70B+ 参数级别和 10T+ tokens 极端训练体制下是否会出现新的二阶相变现象。
相关工作与启发¶
- vs \(\mu\mathrm{P}\) (Maximal Update Parameterization, Yang et al.)：\(\mu\mathrm{P}\) 侧重于模型宽度维度的零样本超参数迁移，但难以直接覆盖数据量 \(D\) 轴的扩展，且当权重范数在训练中漂移时假设往往被打破；本文与 \(\mu\mathrm{P}\) 形成有力互补，直接攻克了跨 token 视界的外推难题。
- vs Step Law (Li et al., 2025) & Bjorck et al. (2025)：以往数据-模型联合对数线性定律假定学习率遵循刚性幂律 \(D^a N^b\)；本文明确指出了该假定在长视界下的系统性曲率失效，并通过有效学习率提供了无须复杂经验修正的闭环理论解释。
- vs 权重范数旋转平衡理论 (Kosson et al., 2024; Wen et al., 2025)：继承了权重范数演化与超球面步长的理论工具，并创造性地将其应用于超参数外推预测这一实用工程场景，打通了理论机制与实践算力节省之间的桥梁。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统证伪了学习率全局对数线性假设，揭示了尺度扩展下的向上弯曲现象并给出了优美的动力学解释。
- 实验充分度: ⭐⭐⭐⭐⭐ 耗时两月消耗 16 卡 H200 算力，进行了全网格穷举扫描与开源，涵盖数百个完整的训练与损失衰减轨迹。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨简洁，两阶段渐进动力学与工程实验对照鲜明，论证逻辑环环相扣。
- 价值: ⭐⭐⭐⭐⭐ 为前沿大模型预训练提供了立即可落地的低成本超参数迁移范式，对工业界万卡集群训练规划极具指导意义。