论文

本征 μ 子:黎曼矩阵流形的谱优化

Intrinsic Muon: Spectral Optimization on Riemannian Matrix Manifolds

模型训练参数高效训练

摘要

Muon 和相关的范数约束矩阵优化器已成为大规模学习问题的核心。它们被公式化为无约束欧几里得空间中环境矩阵范数球上的线性最大化预言(LMO)。然而,这些并不能完全推广到流形值参数,例如低秩分解、正交性约束或对称正定 (SPD) 矩阵。天真地将 Muon LMO 限制在切线空间(i)打破了商对称性,并且(ii)将切线空间约束与环境范数界限耦合起来,从而阻碍了各种感兴趣流形上的封闭形式解。我们通过一次观察解决了这两个问题:每个黎曼度量规范地将酉不变的欧几里得范数提升为每个切空间上的内在范数,并且由此产生的内在范数约束的 LMO 是对称性保持的。在此基础上,我们引入了内在 Muon (iMuon),这是一个统一的框架,可为任何酉不变范数(包括谱范数、Frobenius 范数和核范数)的固定秩、SPD、Stiefel 和 Grassmann 流形生成闭合形式更新。我们为确定性和随机 iMuon 建立了收敛保证,其速率常数仅取决于流形维度。值得注意的是,在固定秩流形上,该常数仅取决于秩,从而使速率独立于因子条件,并消除了先前工作所需的运行时因子重新缩放。 LLM 的 LoRA 微调、图像分类和子空间学习的实验说明了该方法的有效性。