MALT:轻量级曲率感知Muon优化器

MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

精选理由

MALT优化器给Muon加了轻量对角预条件,在GPT-2上比Muon更快更好,内存几乎没增加,适合大模型预训练。

AI 摘要

MALT在Muon基础上加入轻量对角预条件器,以感知损失曲率并降低对曲率各向异性的敏感度。它使用Newton-Schulz迭代正交化预条件后的动量,并通过范数嫁接控制更新幅度。进一步提出的MALTER通过自适应步长重缩放提升对随机梯度噪声的鲁棒性。在GPT-2 Small、Medium、Large预训练实验中,MALT和MALTER均优于Muon,且内存占用和耗时几乎不变。

原文 · arXiv cs.LG

MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

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.