通过平坦方向动力学增强 LLM 预训练加速
机器学习
2026-02-27 v1
摘要
预训练大型语言模型需要巨大的计算资源，因此优化器效率至关重要。优化 landscape 高度各向异性，损失减少主要由沿平坦方向的进展所驱动。虽然矩阵基优化器如 Muon 和 SOAP 利用精细曲率信息在 AdamW 方面取得优势，但其更新倾向于各向同性——在平坦方向上相对保守，同时在锐利方向上可能过于激进。为解决这一局限，我们首先建立统一的黎曼常微分方程 (ODE) 框架，阐明常见自适应算法如何协同工作：预条件器诱导黎曼几何以缓解不良条件，而动量则作为黎曼阻尼项以促进收敛。基于这些见解，我们提出 LITE，一种通用的加速策略，通过沿平坦轨迹施加更大的 Hessian 阻尼系数和学习率来增强训练动力学。广泛实验表明，LITE 在 diverse architectures (Dense、MoE)、parameter scales (130M--1.3B)、datasets (C4、Pile) 和 learning-rate schedules (cosine、warmup-stable-decay) 上显著加快了 Muon 和 SOAP 的速度。理论分析确认，LITE 在各向异性 landscape 中沿平坦方向实现更快收敛，为高效 LLM 预训练提供了原则方法。代码已公开于 https://github.com/SHUCHENZHU/LITE。
关键词
引用
@article{arxiv.2602.22681,
title = {Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement},
author = {Shuchen Zhu and Rizhen Hu and Mingze Wang and Mou Sun and Xue Wang and Kun Yuan and Zaiwen Wen},
journal= {arXiv preprint arXiv:2602.22681},
year = {2026}
}