针对大模型训练中平坦方向收敛慢的问题,提出分尺度动量与球约束方法加速优化。
Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining
- 在平坦方向使用分尺度动量,慢衰减降噪,快衰减适应曲率
- 结合球约束防止参数膨胀和学习率过快衰减
- 适用于多种架构与规模模型,显著提升训练速度
预训练占大语言模型训练总计算成本的很大比例。然而,噪声主导的梯度和高度病态的损失曲面带来了严峻挑战。尽管现代自适应优化器如AdamW和Muon在大规模预训练中表现优异,但其依赖梯度归一化,对病态曲率的缓解有限。沿平坦方向(小特征值对应方向)的优化进展缓慢,而这一方向主导了最终损失下降。为提升平坦方向的训练动态,本文提出一种曲率感知的多尺度动量方法并引入球约束,实现大模型预训练中的稳定加速。该多尺度动量仅作用于平坦方向,结合慢衰减分量降噪与快衰减分量快速适应曲率,发挥互补优势。关键在于采用球约束技术,防止因简单组合导致的参数膨胀和有效学习率过快衰减。大量实验表明,该方法在多种架构(密集型、MoE)和模型规模(0.12B–2.3B参数)下显著加速了Muon。理论上验证了加速效果,并揭示了平坦方向多尺度动量的设计原理。
原文摘要 · Abstract (English)
Pretraining accounts for a large fraction of the total computational cost in LLM training. However, noise-dominant gradients and the highly ill-conditioned loss landscape bring severe challenges. Although modern adaptive optimizers such as AdamW and Muon have achieved great success in large-scale pretraining, their reliance on gradient normalization offers limited mitigation of the ill-conditioned curvature. The progress along flat directions (eigen-directions of small eigenvalues), which dominates the final loss reduction, remains relatively slow. To enhance training dynamics along flat directions, we propose a curvature-conditioned multiscale momentum method with sphere constraints, delivering steady acceleration in LLM pretraining. This multiscale momentum, applied only along flat directions, pairs a slow-decay component for noise reduction with a fast-decay component for rapid curvature adaptation, harnessing their complementary strengths. Crucially, we employ a sphere constraint technique to prevent parameter inflation and excessively rapid effective learning rate decay that would otherwise arise from a naive combination. Extensive experiments show that the proposed method significantly accelerates Muon across diverse architectures (dense, MoE) and model sizes (0.12B--2.3B parameters). Theoretically, we verify the acceleration effect and provide insight into the design principles underlying the flat-direction multiscale momentum.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。