Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement¶

会议: NeurIPS 2026
arXiv: 2602.22681
OpenReview: https://openreview.net/forum?id=lKebtNgu8P
代码: https://github.com/SHUCHENZHU/LITE
领域: 预训练
关键词: LLM预训练、自适应优化器、黎曼ODE动力学、平坦方向增强、Hessian阻尼

一句话总结¶

基于黎曼流形上带Hessian阻尼的惯性动力学系统（RISHD），本文提出通用加速策略LITE，通过无额外优化器状态的高效子空间滤波，在平坦方向选择性放大更新步长与Hessian阻尼，显著加速Muon与SOAP的大模型预训练。

研究背景与动机¶

大语言模型预训练对计算资源有着极高需求，优化器的收敛效率直接决定了模型扩展的成本边界。当前主流的预训练基线 AdamW 仅使用坐标级对角预条件子，忽略了参数间丰富的 Hessian 交叉曲率，在高度各向异性的损失曲面上前进步伐受限。近期以 Muon、SOAP 为代表的矩阵级预条件优化器通过引入近似 Fisher 信息的几何结构，在工业级训练中展现出超越 AdamW 的优异收敛速度。

然而，现有矩阵优化器依然受制于两个核心瓶颈。一方面，预条件子诱导的参数更新幅度在不同特征方向上趋于各向同性，导致在占主导地位且决定最终损失下降的平坦方向上更新过于保守，而在极少数正曲率巨大的陡峭方向上又存在发散风险；另一方面，基于指数移动平均（EMA）的简单线性动量机制本质上施加的是各向同性的阻尼衰减，即使结合传统 Nesterov 加速，其 Hessian 阻尼也往往继承自凸优化假定，无法自适应非凸且极端病态的深度网络损失曲面。这背后的核心矛盾在于：缺乏一个将预条件矩阵与动量机制深度统一的非凸动力学分析框架，使得研究者无法系统性解耦并差异化调节陡峭震荡动力学与平坦缓慢漂移。

为了攻克这一问题，作者将连续时间惯性动力学拓展至黎曼流形，揭示预条件子构建局部度量以缓解病态性、动量充当度量空间内阻尼以抑制震荡的协同本质。核心 idea：基于统一黎曼 ODE 动力学框架，提出 LITE 策略，在保持陡峭方向稳定性的同时，显式识别并增强平坦方向的 Hessian 阻尼与学习率，从而定向激活动力学瓶颈以实现预训练大幅加速。

方法详解¶

整体框架¶

LITE（Accelerating Adaptive Optimizers in LLM Pre-training）围绕大模型各向异性损失曲面的“时间尺度分离（Time-scale Separation）”现象展开设计：陡峭方向上动力学表现为高频快速震荡但几乎不贡献损失下降，平坦方向上参数演化平缓缓慢却主导总损失降低。LITE 的核心目标在于保持陡峭方向超参数固定的前提下，精准提取平坦子空间，并在该子空间中同时放大 Hessian 阻尼项系数与更新步长。

方法的整体工作流包含三个阶段：首先利用梯度与一阶动量估计参数块的局部曲率主空间，接着通过无状态的轻量化算子划分出陡峭子空间投影与平坦子空间投影，最后依据黎曼动力学解耦公式分别对平坦与陡峭分量施加不同的动量阻尼校正与学习率缩放并更新参数。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：随机梯度 G_k 与历史动量 M_k-1"] --> B["阶段1：黎曼度量与动量解耦<br/>构造一阶动量与 Nesterov 修正项"]
    B --> C["阶段2：无状态快速子空间滤波<br/>复合 Newton-Schulz 或特征平滑滤波划分陡峭与平坦空间"]
    C --> D["阶段3：平坦方向动力学增强<br/>陡峭方向保持基线，平坦方向施加放大步长与大 Hessian 阻尼"]
    D --> E["输出：合成优化更新步 U_k 并完成参数更新"]

关键设计¶

1. 黎曼度量与动量协同的 RISHD 理论建模：揭示预条件与动量的连续动力学本质

现有研究往往割裂分析动量与预条件子，导致对二者在非凸曲面上如何相互作用缺乏统一视角。本文提出配备预条件诱导黎曼度量 $F(w)$ 的带 Hessian 阻尼惯性系统（Riemannian Inertial System with Hessian Damping, RISHD）： $$\nabla_{\dot{w}_t}^F \dot{w}_t + \alpha_t \dot{w}_t + \beta_t \operatorname{Hess}_F(w_t)\dot{w}_t + \gamma_t \operatorname{grad}_F f(w_t) = 0$$ 通过将系统解耦为切空间（Tangent Space）速度 $\dot{w}_t$ 与余切空间（Cotangent Space）动量 $m_t$ 的一阶耦合动力学，证明动量在余切空间内积累原始梯度力，而预条件子 $F(w_t)^{-1}$ 则作为双线性映射将动量拉回切空间驱动参数位移。这一视角严谨解释了现代自适应优化器“先积攒原始梯度动量、再作用逆预条件矩阵”的架构必要性。更为关键的是，连续动力学指出，Nesterov 加速中的梯度修正项 $\beta \nabla f(w_k)$ 本质是连续 Hessian 阻尼项 $\beta_t \nabla^2 f(w_t)\dot{w}_t$ 的差分离散形式，具有曲率自适应耗散动能的物理功能。

2. 平坦方向动力学增强机制（LITE）：差异化放大学习率与阻尼

针对各向异性损失曲面中平坦方向行进缓慢且受非凸鞍点困扰的痛点，LITE 对更新规则进行正交子空间分解，设 $P_k$ 为陡峭子空间投影矩阵，$Q_k = I - P_k$ 为平坦子空间投影矩阵。在平坦方向上引入放大因子 $\chi \ge 1$ 以及强化 Hessian 阻尼系数 $\beta_2 \ge \beta_1 \ge 0$： $$u_k = (I - Q_k) F(w_k)^{-1}(m_k + \beta_1 \nabla f(w_k)) + \chi Q_k F(w_k)^{-1}(m_k + \beta_2 \nabla f(w_k))$$ 在物理直觉上，该机制增大了平坦方向的驱动力系数 $\gamma_t = \eta_t(1 + \alpha\beta)$ 并赋予更大的曲率阻尼，促使动量更快穿透平坦平原与非凸鞍点；同时，陡峭方向参数完全锚定基线配置，抑制高曲率方向的离散发散风险。理论上作者基于高维 River-Valley 流形模型证明，LITE 轨迹能以指数速率迅速收敛至平坦“河流”流形 $\mathcal{R}$，且在平坦流形上的黎曼梯度积分上限由 $\mathcal{O}(1/(\chi \beta_2))$ 严格约束，从数学上保证了平坦区域内更快的遍历收敛。

3. 面向 Muon 与 SOAP 的无状态快速子空间近似工程实现

若通过常规 SVD 或特征值分解动态计算投影矩阵 $P_k$，庞大的计算与显存开销将完全抵消加速收益。为此，本文设计了硬件友好的免优化器状态滤波算法： - Muon-LITE：利用修正动量矩阵外积 $\widetilde{M}_k^\top \widetilde{M}_k$ 作为期望协方差 $\mathbb{E}[G_k^\top G_k]$ 的稳健无偏代理，利用两路极速 Newton-Schulz（NS Polar Express）迭代构造滤波算子 $T_k = \frac{1}{2}\operatorname{NS}(\widetilde{M}_k) + \frac{1}{2}\operatorname{NS}(\frac{\widetilde{M}_k}{\tau_k} - \operatorname{NS}(\widetilde{M}_k))$，通过阈值自适应反馈环调控标量因子 $l_k$ 锁定前 $d_s$ 个奇异值构成的空间，仅需在每步增加两组可高度并行化的 NS 矩阵乘法，吞吐损失仅约 1%。 - SOAP-LITE：由于 SOAP 已经显式计算了旋转基下的特征方差矩阵 $V_k$，LITE 采用软阈值过渡方案，构造平滑权重矩阵 $(P_k)_{i,j}$，在陡峭向平坦过渡的区域线性插值，杜绝硬切断带来的数值抖动，且规避了昂贵的 PyTorch top_k 调用。

损失函数 / 训练策略¶

在训练实施上，除 Transformer 内部二维权重矩阵采用 Muon-LITE 或 SOAP-LITE 优化外，Embedding、RMSNorm 与 Output 头参数沿用 AdamW。超参数选择上，经网格搜索确定陡峭空间比例 $r_s = d_s / \min(m,n) = 0.1$（SOAP-LITE 为 0.2）。对于 Muon-LITE，选取 $\chi = 2$（Norm/Emb 取 4），$\beta_1 = -0.25$，$\beta_2 = 1.0$。学习率调度支持 Cosine 与 Warmup-Stable-Decay（WSD，80% 稳定期 + 20% 线性衰减）。

实验关键数据¶

主实验¶

在 Dense 模型（LLaMA-2 架构）多尺度预训练任务中，Muon-LITE 相对 Muon 基线显著降低了验证集困惑度与终验损失；在下游 8 项零样本常识与推理任务上（ARC-E, ARC-C, PIQA, HellaSwag, BoolQ, WinoGrande, MMLU, TruthfulQA）均取得全面领先，展现出扎实的泛化能力。

模型规模 / 架构	预训练数据集	评测指标	本文 (Muon-LITE)	之前SOTA (Muon)	相对增益 / 评价
LLaMA 0.25B (10k step)	Pile	Terminal Loss	2.086	2.110	-0.024 显著更优
LLaMA 0.5B (10k step)	Pile	Terminal Loss	1.838	1.854	-0.016 优势稳定
LLaMA 1.3B (10k step)	Pile	Terminal Loss	1.638	1.652	-0.014 保持优势
LLaMA 1.3B (25k step)	Pile	Downstream AVG	45.49	43.44	+2.05 全面提升
LLaMA 1.3B (25k step)	Pile	BoolQ Acc (%)	60.52	52.45	+8.07 表现亮眼
LLaMA 1.3B (25k step)	Pile	MMLU Acc (%)	25.15	23.76	+1.39 推理增强
QwenMoE 1B (总1B/激活0.3B)	Pile (WSD调度)	Terminal Loss	显著领先	基线水准	稳定期优势随迭代扩大

消融实验¶

在 LLaMA-0.25B 上对平坦方向增强超参数和施加区域进行严谨消融，验证了联合增加学习率比率 $\chi$ 与 Hessian 阻尼 $\beta_2$ 的必要性，以及限制增强仅在平坦方向生效的合理性。

配置	终验损失 (Loss)	说明
Muon-LITE (完整方案: $\chi=2, \beta_2=1.0$)	2.086	完整模型：同时提升平坦方向步长与阻尼
变体 Muon-LITE-L (仅增大学习率 $\chi=2, \beta_2=0$)	2.096	欠缺 Hessian 阻尼调节，损失降低幅度缩小
变体 Muon-LITE-H (仅增大阻尼 $\chi=1, \beta_2=2.0$)	2.094	欠缺动力学推动力放大，收益受限
对照组：全空间均匀阻尼 ($\beta_{1,2} = 0.5$)	2.113	陡峭方向承受过大阻尼，性能劣于基线
对照组：全空间均匀阻尼 ($\beta_{1,2} = 1.0$)	2.129	陡峭方向过度震荡或受制，性能进一步恶化
原始基线 (Muon Baseline: $\chi=1, \beta_{1,2}=0$)	2.110	未进行任何子空间差异化调控的标准更新

关键发现¶

两重加速维度的互补性：仅放大学习率（LITE-L）或仅强化 Hessian 阻尼（LITE-H）均优于基线，但两者结合时收益最大，印证了连续时间动力学中驱动力 $\gamma$ 与阻尼 $\beta$ 协同演化的理论预测。
陡峭方向必须保持保守：消融实验表明，一旦将针对平坦方向设置的较大阻尼施加给陡峭方向（$\beta_{1,2} \ge 0.5$），训练损失反而劣于基线，证实了子空间解耦设计的关键价值。
长周期训练展现 2× 加速比：在训练 Token 预算从 40× 扩大至 100× 与 200× 参数量的长程实验中，Muon-LITE 在 100× Token 预算下达到 Muon 相同损失时消耗算力减少近半，具备出色的 Scaling 延展性。

亮点与洞察¶

动力学视角的范式拓展：将主流矩阵优化器（Muon、SOAP、AdamW、Lion）成功收敛进黎曼带 Hessian 阻尼惯性系统（RISHD），用连续 ODE 视角理顺了预条件子（度量）与动量（阻尼）的分工。
无状态极速 Newton-Schulz 空间滤波：利用复合函数在极小开销下通过符号函数性质过滤出前 $d_s$ 奇异子空间，不需要缓存任何高阶矩阵分解状态，吞吐损失仅约 1%，极度适合工业部署。
直击各向异性核心痛点：避开在陡峭方向上冒险，把激进加速全部押在主导损失下降的平坦方向上，思路清晰且安全性高。

局限与展望¶

超参数搜索依赖离线试探：平坦/陡峭方向维度划分比例 $r_s$、步长放大倍数 $\chi$ 与阻尼 $\beta_2$ 目前在小模型（0.25B）上确定后跨尺度泛化，缺乏训练中根据动态曲率自动演化的自适应闭环机制。
架构兼容性与适用边界：虽然在 Dense 与 MoE 上验证了有效性，但输出层（Output Head）由于对数值敏感，仍需人工排除在加速范围外；未来需进一步探索对混合精度极低位宽量化训练（如 FP4/FP8）的数值鲁棒性。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 建立了自适应优化器的黎曼 ODE 统一框架，平坦方向选择性动力学增强设计兼具物理美感与数学严谨性。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 Dense/MoE 架构、Pile/C4 数据集、Cosine/WSD 调度器，并包含 8 项下游评测与细致消融。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从直观物理图像到严格高维 River-Valley 流形收敛性证明一应俱全。
价值: ⭐⭐⭐⭐⭐ 针对万卡集群 LLM 预训练成本痛点，即插即用且吞吐损耗仅 1%，工业落地与学术参考价值极高。

配置	终验损失 (Loss)	说明
Muon-LITE (完整方案: \(\chi=2, \beta_2=1.0\))	2.086	完整模型：同时提升平坦方向步长与阻尼
变体 Muon-LITE-L (仅增大学习率 \(\chi=2, \beta_2=0\))	2.096	欠缺 Hessian 阻尼调节，损失降低幅度缩小
变体 Muon-LITE-H (仅增大阻尼 \(\chi=1, \beta_2=2.0\))	2.094	欠缺动力学推动力放大，收益受限
对照组：全空间均匀阻尼 (\(\beta_{1,2} = 0.5\))	2.113	陡峭方向承受过大阻尼，性能劣于基线
对照组：全空间均匀阻尼 (\(\beta_{1,2} = 1.0\))	2.129	陡峭方向过度震荡或受制，性能进一步恶化
原始基线 (Muon Baseline: \(\chi=1, \beta_{1,2}=0\))	2.110	未进行任何子空间差异化调控的标准更新