arXiv:2608.05088cs.LG2026-08被引 1

MALT通过轻量对角预条件提升Muon对损失曲率的适应性,训练更稳定高效。

MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

论文配图:MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning
图 1 · 摘自论文原文
  • 用轻量双侧对角预条件近似捕捉损失曲率几何
  • 在GPT-2小中大模型上优于Muon,内存和时间开销几乎不变
  • 适合追求高鲁棒性与低开销的语言模型预训练场景

Muon最近作为AdamW的替代方案在语言模型预训练中崭露头角,通过Newton-Schulz迭代正交化动量矩阵。尽管它缓解了梯度各向异性问题,但未显式考虑损失曲面的曲率几何,仍可能受曲率各向异性影响。为此,我们提出MALT(Muon Augmented by Lightweight Two-sided Preconditioning),利用轻量对角预条件器降低Muon对曲率各向异性的敏感性。具体地,MALT采用计算与内存开销低的双侧对角预条件器,近似捕获损失曲面的曲率几何,再以Newton-Schulz迭代正交化预条件后的动量,并映射回更新方向,同时通过范数嫁接控制更新幅度。为进一步提升对随机梯度噪声的鲁棒性,我们引入MALTER(MALT with Adaptive stEpsize Rescaling)。在随机非凸设置下,我们为MALT提供了收敛保证。在GPT-2 Small、Medium和Large模型上的预训练实验表明,所提方法优于Muon,同时保持几乎相同的内存占用和实际运行时间。

原文摘要 · Abstract (English)

Muon has recently emerged as a promising alternative to AdamW for language model pretraining by orthogonalizing momentum matrices using Newton-Schulz iterations. Although Muon mitigates gradient anisotropy, it does not explicitly account for the curvature geometry of the loss landscape and may therefore remain sensitive to curvature anisotropy. We bridge this gap by proposing MALT (Muon Augmented by Lightweight Two-sided Preconditioning), which uses lightweight diagonal preconditioners to reduce the sensitivity of Muon to curvature anisotropy. Specifically, MALT uses two-sided diagonal preconditioners with low memory and computational overhead to approximately capture the curvature geometry of the loss landscape. It orthogonalizes the preconditioned momentum using Newton-Schulz iterations and maps the result back to define the update direction, while norm grafting controls the update magnitude. To improve the robustness of MALT to stochastic gradient noise, we further propose MALTER (MALT with Adaptive stEpsize Rescaling). Convergence guarantees are provided for MALT in the stochastic non-convex setting. Experiments on GPT-2 Small, Medium, and Large pretraining show that the proposed methods outperform Muon while maintaining nearly the same memory footprint and wall-clock time.

优化器语言模型曲率感知

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。