Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining

2026-08-28

AI Summary

为解决LLM预训练中梯度噪声和病态曲率问题，提出了一种带球约束的曲率条件多尺度动量方法，加速平坦方向上的训练。

Abstract

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.

Problem

• LLM Pretraining
• Ill-Conditioned Loss Landscape
• Gradient Normalization
• Flat Directions
• Curvature

Innovation

• curvature-conditioned
• multiscale momentum
• sphere constraints
• flat directions
• parameter inflation