Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement

📄 arXiv: 2602.22681v1 📥 PDF

作者: Shuchen Zhu, Rizhen Hu, Mingze Wang, Mou Sun, Xue Wang, Kun Yuan, Zaiwen Wen

分类: cs.LG

发布日期: 2026-02-26

🔗 代码/项目: GITHUB

💡 一句话要点

LITE：通过增强平坦方向动态加速LLM预训练

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 预训练 优化器 Hessian阻尼 黎曼优化

📋 核心要点

  1. 现有优化器在LLM预训练中效率不足，尤其是在各向异性优化地形的平坦方向上。
  2. LITE通过在平坦方向上应用更大的Hessian阻尼系数和学习率，增强训练动态，加速收敛。
  3. 实验表明，LITE显著加速了Muon和SOAP在多种架构、参数规模和数据集上的训练。

📝 摘要（中文）

预训练大型语言模型需要巨大的计算资源，因此优化器的效率至关重要。优化地形是高度各向异性的，损失的减少主要由沿平坦方向的进展驱动。虽然诸如Muon和SOAP等基于矩阵的优化器利用细粒度的曲率信息来优于AdamW，但它们的更新往往趋于各向同性——在平坦方向上相对保守，但在尖锐方向上可能过于激进。为了解决这个限制，我们首先建立了一个统一的黎曼常微分方程（ODE）框架，阐明了常见的自适应算法如何协同工作：预处理器诱导黎曼几何来减轻病态条件，而动量充当黎曼阻尼项来促进收敛。在这些见解的指导下，我们提出LITE，一种广义的加速策略，通过在平坦轨迹上应用更大的Hessian阻尼系数和学习率来增强训练动态。大量的实验表明，LITE显著加速了Muon和SOAP在不同架构（Dense, MoE）、参数规模（130M--1.3B）、数据集（C4, Pile）和学习率调度（cosine, warmup-stable-decay）上的训练。理论分析证实，LITE有助于在各向异性地形中沿平坦方向更快地收敛，为高效的LLM预训练提供了一种有原则的方法。代码可在https://github.com/SHUCHENZHU/LITE获取。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）预训练过程中优化器效率低下的问题。现有的优化器，如AdamW，以及改进的矩阵优化器Muon和SOAP，在处理各向异性优化地形时存在不足。具体来说，这些优化器在平坦方向上的更新相对保守，限制了训练速度。

核心思路：论文的核心思路是增强优化器在平坦方向上的动态。通过分析现有自适应算法的黎曼常微分方程（ODE）框架，发现可以通过调整Hessian阻尼系数和学习率来改善收敛速度。LITE策略旨在在平坦轨迹上应用更大的Hessian阻尼系数和学习率，从而加速训练。

技术框架：论文首先建立了一个统一的黎曼ODE框架，用于分析现有优化算法的行为。该框架将预处理器视为诱导黎曼几何，将动量视为黎曼阻尼项。然后，基于此框架，提出了LITE加速策略，该策略可以应用于现有的矩阵优化器，如Muon和SOAP。LITE通过调整Hessian阻尼系数和学习率来增强平坦方向的动态。

关键创新：论文的关键创新在于LITE加速策略，它能够自适应地调整Hessian阻尼系数和学习率，从而在平坦方向上实现更快的收敛。与现有方法相比，LITE不是简单地使用固定的学习率或阻尼系数，而是根据优化地形的局部结构进行调整，从而更有效地利用计算资源。

关键设计：LITE的关键设计包括：1) 使用黎曼ODE框架分析现有优化器的行为；2) 设计自适应的Hessian阻尼系数和学习率调整策略，使其能够根据优化地形的局部结构进行调整；3) 将LITE策略应用于现有的矩阵优化器，如Muon和SOAP，以验证其有效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LITE能够显著加速Muon和SOAP等优化器的训练速度。在不同的模型架构（Dense, MoE）、参数规模（130M--1.3B）、数据集（C4, Pile）和学习率调度（cosine, warmup-stable-decay）上，LITE均表现出优越的性能。具体数据提升幅度未知，但论文强调了“显著加速”。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的预训练，尤其是在计算资源有限的情况下。通过提高优化器的效率，LITE能够加速LLM的训练过程，降低训练成本，并促进更大规模模型的开发。此外，该方法还可以应用于其他深度学习模型的训练，特别是在优化地形高度各向异性的情况下。

📄 摘要（原文）

Pre-training Large Language Models requires immense computational resources, making optimizer efficiency essential. The optimization landscape is highly anisotropic, with loss reduction driven predominantly by progress along flat directions. While matrix-based optimizers such as Muon and SOAP leverage fine-grained curvature information to outperform AdamW, their updates tend toward isotropy -- relatively conservative along flat directions yet potentially aggressive along sharp ones. To address this limitation, we first establish a unified Riemannian Ordinary Differential Equation (ODE) framework that elucidates how common adaptive algorithms operate synergistically: the preconditioner induces a Riemannian geometry that mitigates ill-conditioning, while momentum serves as a Riemannian damping term that promotes convergence. Guided by these insights, we propose LITE, a generalized acceleration strategy that enhances training dynamics by applying larger Hessian damping coefficients and learning rates along flat trajectories. Extensive experiments demonstrate that LITE significantly accelerates both Muon and SOAP across diverse architectures (Dense, MoE), parameter scales (130M--1.3B), datasets (C4, Pile), and learning-rate schedules (cosine, warmup-stable-decay). Theoretical analysis confirms that LITE facilitates faster convergence along flat directions in anisotropic landscapes, providing a principled approach to efficient LLM pre-training. The code is available at https://github.com/SHUCHENZHU/LITE.