通过增强平坦方向动力学,显著加速大模型预训练过程。
Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement
- 提出LITE方法,在平坦方向使用更大动量和学习率以提升训练效率。
- 在130M到1.3B参数模型上,相比Muon和SOAP提速超20%。
- 适用于各种架构与数据集,尤其适合追求高效预训练的研究者。
大语言模型预训练需要大量计算资源,优化器效率至关重要。优化空间具有高度各向异性,损失下降主要依赖于平坦方向的进展。尽管基于矩阵的优化器如Muon和SOAP利用精细曲率信息优于AdamW,但其更新趋向各向同性——在平坦方向保守,在尖锐方向可能激进。为此,我们建立统一的黎曼常微分方程框架,揭示常见自适应算法的协同机制:预条件器引入黎曼几何以缓解病态,动量则作为黎曼阻尼项促进收敛。基于此,提出LITE,一种通用加速策略,通过在平坦轨迹上施加更大的海森动量系数和学习率来增强训练动态。大量实验表明,LITE在多种架构(密集、MoE)、参数规模(130M–1.3B)、数据集(C4、Pile)及学习率调度(余弦、预热-稳定-衰减)下均显著加速Muon和SOAP。理论分析证实,LITE能有效加快各向异性景观中平坦方向的收敛速度,为高效大模型预训练提供原理性方案。代码已开源:https://github.com/SHUCHENZHU/LITE。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。