Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement

📄 arXiv: 2602.22681 📥 PDF

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

分类: cs.LG

发布日期: 2026-02-28

💡 一句话要点

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

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

关键词: 大型语言模型 预训练 优化器 Hessian阻尼 黎曼优化 加速训练 平坦方向 自适应学习率

📋 核心要点

  1. 现有基于矩阵的优化器（如Muon和SOAP）在平坦方向上更新保守，限制了LLM预训练的效率。
  2. 论文提出LITE，通过在平坦方向上应用更大的Hessian阻尼系数和学习率，增强训练动态，加速收敛。
  3. 实验表明，LITE显著加速了Muon和SOAP在多种架构、参数规模和数据集上的预训练过程。

📝 摘要（中文）

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

🔬 方法详解

问题定义：大型语言模型（LLM）的预训练需要大量的计算资源，优化器的效率至关重要。现有的优化器，尤其是基于矩阵的优化器（如Muon和SOAP），虽然利用了曲率信息，但在平坦方向上的更新过于保守，导致收敛速度慢。这种保守性源于对尖锐方向的过度规避，从而限制了在平坦方向上的有效探索。

核心思路：论文的核心思路是增强优化器在平坦方向上的动态。通过增加在平坦方向上的Hessian阻尼系数和学习率，LITE鼓励优化器更积极地探索这些方向，从而加速收敛。这种方法基于一个关键观察：损失的减少主要发生在平坦方向上，因此更有效地利用这些方向可以显著提高预训练效率。

技术框架：LITE方法建立在一个统一的黎曼常微分方程（ODE）框架之上。该框架将常见的自适应算法（如AdamW、Muon和SOAP）视为黎曼流形上的ODE求解器。预处理器（如AdamW中的对角缩放）诱导黎曼几何，以减轻病态条件，而动量项充当黎曼阻尼，促进收敛。LITE通过调整黎曼阻尼系数和学习率，专门针对平坦方向进行优化。

关键创新：LITE的关键创新在于其广义的加速策略，该策略通过自适应地调整Hessian阻尼系数和学习率来增强训练动态。与现有方法不同，LITE不是简单地增加全局学习率，而是有选择地增强平坦方向上的更新。这种方法能够更有效地利用计算资源，并避免在尖锐方向上引入不稳定性。

关键设计：LITE的关键设计包括：1) 使用Hessian信息的估计来识别平坦方向；2) 增加在这些方向上的Hessian阻尼系数，以减少噪声和提高稳定性；3) 增加在这些方向上的学习率，以加速收敛。具体的实现细节可能取决于所使用的基础优化器（如Muon或SOAP），但核心思想是相同的：增强平坦方向上的动态。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LITE显著加速了Muon和SOAP在不同架构（Dense，MoE）、参数规模（1.3亿-13亿）、数据集（C4，Pile）和学习率调度（cosine，warmup-stable-decay）上的表现。具体性能提升数据未知，但摘要强调了“显著加速”，表明LITE在多种设置下均能有效提高预训练效率。

🎯 应用场景

LITE方法可广泛应用于加速各种大型语言模型的预训练过程，尤其是在计算资源受限的情况下。通过提高训练效率，LITE可以降低LLM的开发成本，并促进更大规模、更复杂模型的开发。此外，LITE还可以应用于其他深度学习任务，例如图像识别和自然语言处理，只要这些任务的优化地形具有各向异性。

📄 摘要（原文）

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 atthis https URL.