Intrinsic Muon: Spectral Optimization on Riemannian Matrix Manifolds¶

会议: NeurIPS 2026
arXiv: 2605.09238
代码: https://github.com/1bang118/manifold-intrinsic-muon
领域: 优化/理论
关键词: 黎曼流形优化, 谱范数优化, Muon, 低秩适配 (LoRA), 线性最大化预言机 (LMO)

一句话总结¶

针对 Muon 等谱范数优化器难以直接推广到流形约束的问题，本文提出内在 Muon（iMuon），通过黎曼度量将酉不变范数典则提升至切空间，彻底消除了商流形对称性破坏与切空间范数耦合难题，并在固定秩、SPD、Stiefel 与 Grassmann 等多种矩阵流形上实现了精确闭式更新与与条件数无关的收敛保证。

研究背景与动机¶

以 Muon 为代表的矩阵范数约束优化器通过在线性最大化预言机（LMO）中施加谱范数球约束，在大型语言模型预训练与 LoRA 微调中展现出显著超越常规一阶优化器的泛化能力与训练效率。然而，现代深度表示学习中大量参数本质上具有流形约束结构——例如 LoRA 权重直接栖息在固定秩流形上，正交与酉约束定义了 Stiefel 和 Grassmann 流形，而协方差与度量学习则生活在对称正定（SPD）矩阵流形上。现存的矩阵谱优化理论几乎全部建立在非约束欧氏空间中，难以直接拓展到黎曼矩阵流形场景。

将欧氏谱优化生硬迁移到流形切空间会遭遇两个根本性理论阻碍：其一是商流形上的对称性破缺（O1），欧氏谱范数并不具备商几何等价类群变换的不变性，例如在 LoRA 分解 \(X=BA\) 中，任意可逆矩阵规范变换 \((BN^{-1}, NA)\) 保持乘积不变，却会导致不同的逐因子欧氏 Muon 步长；其二是切空间约束与外围范数球边界的强几何耦合（O2），强制在切空间内求解外围谱范数极值将破坏极因子分解（polar factor）形式，导致绝大多数矩阵流形上根本不存在闭式解。以往的权宜方案（如 Spectron 依赖动态启发式幅度缩放、Riemannion 依赖低秩截断逼近、MCSD 丢弃切空间直接向流形外投影）均无法从几何机理上同时化解这两大难题。

本文的切入视角非常纯粹：既然黎曼度量张量本身诱导了切空间上的局部内积与 Frobenius 范数，那么任何酉不变矩阵范数都能够通过该度量唯一、典则地“提升”为切空间内在范数。核心 idea：利用黎曼度量 operator 将外围酉不变范数典则提升至切空间形成内在约束，通过度量的左右分解与缩放切空间的奇异值不变性消除两重约束耦合，在各主流流形上推导出具有商对称不变性的精确闭式 iMuon 更新。

方法详解¶

整体框架¶

iMuon 统一了黎曼矩阵流形上的酉不变范数线性最大化预言机（Intrinsic LMO）。给定流形 \(\mathcal{M}\) 上的黎曼度量 \(g_x(\xi, \zeta) = \langle G_x \xi, \zeta \rangle\)（其中 \(G_x \succ 0\) 为自伴正定度量算子），针对任意酉不变范数 \(\varphi\)（如谱范数、Frobenius 范数、核范数或 Schatten-\(p\) 范数），定义切向量 \(\xi \in T_x\mathcal{M}\) 的内在范数为 \(\|\xi\|_{\varphi, x} := \varphi(G_x^{1/2} \xi)\)。优化器每一轮在切空间内求解内在约束极值方向 \(\xi^*\)，并沿黎曼退缩（retraction）更新参数。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：当前流形参数与欧氏梯度"] --> B["内在范数构建：度量算子提取与几何提升"]
    B --> C["内在 LMO 解耦与闭式求解：左右度量缩放变换"]
    C --> D["流形几何退缩更新：保证参数落回目标矩阵流形"]
    D --> E["输出：新一代流形参数"]

关键设计¶

1. 内在范数构建与商对称性保持：消除商流形表示冗余引起的步长畸变

针对商流形上代表元选取破坏欧氏范数的痛点，iMuon 构造了内在范数 \(\|\xi\|_{\varphi, x} = \varphi(G_x^{1/2}\xi)\)。对于具有左右乘积结构度量 \(g_x(\xi, \zeta) = \text{tr}(\xi^\top P_x \zeta Q_x)\)（其中 \(P_x, Q_x \succ 0\)）的商流形 \(\mathcal{M} = \overline{\mathcal{M}}/\mathcal{G}\)，度量算子分解为 \(G_x^{1/2}\xi = P_x^{1/2}\xi Q_x^{1/2}\)。论文在理论上证明了群作用变换 \(h \in \mathcal{G}\) 诱导的正交等距性使内在范数满足：

\[\varphi(G_{\tilde{x}}^{1/2}\tilde{\xi}) = \varphi(G_x^{1/2}\xi)\]

这意味着无论是可行集边界还是最优点更新，在同胚等价类内部完全严格同构，消除了以往算法需引入启发式因子重缩放的弊端。

2. 奇异值不变性与约束解耦：切空间与谱范数双重约束的闭式化简

为了解决切空间从属与范数边界双重约束耦合导致的无解析解难题，iMuon 引入了缩放切空间概念 \(\mathcal{Z}_x := \{G_x^{1/2}\eta : \eta \in T_x\mathcal{M}\}\)。当 \(\mathcal{Z}_x\) 具备奇异值（SV）不变性（即对任意属于该空间的矩阵，改变其非零奇异值后仍停留在空间内）时，利用 von Neumann 迹不等式，切空间约束在极值点处退化为平凡约束，内在优化方向由以下两步精确闭式给出：

\[\xi^* = G_x^{-1/2} Z^*, \quad Z^* = \tau \text{Ortho}(H)\]

其中 \(H := G_x^{1/2}(\text{grad} f)\) 为缩放后的梯度，\(\text{Ortho}(H) = UV^\top\) 代表奇异值正交极因子分解。对于切空间呈现乘积形式的流形，进一步利用 \(\ell_\infty\) 积范数解耦为彼此独立的闭式子问题。

3. 多矩阵流形闭式实例化：覆盖固定秩、SPD、Stiefel 与 Grassmann 流形

在四大核心矩阵流形上，iMuon 均导出了只需欧氏反向传播梯度的简洁更新式： - 固定秩流形 \(\mathcal{M}_r\)（LoRA 场景）：对分解 \(X=BA\)，采用 \(GL(r)\) 不变度量，因子分别计算 \(\dot{B}^* = \tau \text{Ortho}(\nabla_B f (AA^\top)^{-1/2}) (AA^\top)^{-1/2}\) 与 \(\dot{A}^* = \tau (B^\top B)^{-1/2} \text{Ortho}((B^\top B)^{-1/2} \nabla_A f)\)。 - 对称正定流形 \(S_{++}^n\)（协方差学习）：在仿射不变度量下，闭式方向为 \(\xi^* = \tau X^{1/2} \text{sign}(X^{1/2}(\nabla_X f)X^{1/2}) X^{1/2}\)。 - Stiefel 流形 \(\text{St}(m, r)\) 与 Grassmann 流形 \(\text{Gr}(m, r)\)：将切向量分解为列空间反对称部分与正交补法向部分，各自闭式正交化后合成，保证更新方向严格驻留在真实切空间内。

损失函数 / 训练策略¶

算法以常数步长 \(\eta = \sqrt{2\Delta_0 / (L C_\varphi)} / (\tau \sqrt{T})\) 或衰减步长 \(\eta_t = \mathcal{O}(1/\sqrt{t})\) 执行更新。收敛界由切空间约束几何常数 \(C_\varphi\) 决定，其定义为 \(C_\varphi := \sup_{x \in \mathcal{M}} \max\{\|\xi\|_x^2 : \varphi(G_x^{1/2}\xi) \leq 1\}\)。在谱范数下，各流形的 \(C_\varphi\) 分别为：固定秩流形 \(\mathcal{M}_r\) 为 \(2r\)，SPD 流形为 \(n\)，Grassmann 流形为 \(\min(m-r, r)\)。该常数纯粹取决于流形固有维度与秩，完全摆脱了对权重因子条件数与范数大小的依赖。

实验关键数据¶

主实验¶

论文在两类主流场景上进行了严谨评测：固定秩流形上的 LLM LoRA 微调（涵盖 GPT-2 Medium 在 E2E NLG 生成任务，以及 4-bit 量化 Mistral-7B 在 GLUE 9 项全任务上的评测，均采用无动量设置以纯粹对比 LMO 更新方向的质量），以及 SPD 流形上的 CIFAR-100 协方差原型分类。

任务 / 数据集	指标	iMuon (本文)	Factor-wise Muon	Riemannion	Scaled GD	SGD / 基线
E2E NLG (GPT-2, \(r=4\))	BLEU	70.74	70.02	70.02	69.20	66.60
E2E NLG (GPT-2, \(r=4\))	NIST	8.88	8.81	8.78	8.71	8.54
E2E NLG (GPT-2, \(r=4\))	ROUGE-L	72.14	71.68	71.99	70.90	68.20
E2E NLG (GPT-2, \(r=4\))	CIDEr	2.54	2.53	2.52	2.48	2.32
GLUE Avg. (Mistral-7B)	平均分	86.25	81.68	82.65	80.36	71.21
GLUE MRPC (Mistral-7B)	F1 / Acc	85.54	82.35	82.60	81.62	70.10
GLUE WNLI (Mistral-7B)	Accuracy	74.65	45.07	60.56	56.34	49.30
GLUE RTE (Mistral-7B)	Accuracy	88.09	84.12	82.67	54.15	50.90

消融实验¶

在 CIFAR-100 上的 SPD 原型矩阵分类任务中，对比了三种酉不变范数（Frobenius 范数、谱范数、核范数）下外围欧氏 LMO 与本文内在黎曼 LMO 的表现：

范数约束类型	欧氏 LMO 变体	内在 iMuon 变体 (本文)	准确率对比 (Test Acc ± Std)	相对增益
Frobenius 范数	EGD (0.523 ± 0.026)	RGD (黎曼自然梯度)	0.548 ± 0.002	+2.5%
谱范数 (Spectral)	Muon (0.407 ± 0.016)	iMuon-Spectral	0.475 ± 0.009	+6.8%
核范数 (Nuclear)	NuMuon (0.422 ± 0.032)	iMuon-Nuclear	0.546 ± 0.000	+12.4%

关键发现¶

在固定秩 LoRA 微调中，iMuon 在 GLUE 平均分上大幅超越因子式欧氏 Muon 达 4.57 个点，并在小样本与低资源任务（如 WNLI 提升近 30 个点、RTE 提升约 4 个点）中展现出压倒性优势，验证了商流形不变性对梯度噪声和因子范数失衡的强抗扰性。
在 SPD 流形实验中，当约束范数从 Frobenius 向谱范数、核范数逐级强化时，外围欧氏方法与内在度量方法的性能差距持续放大（从 2.5% 扩大到 12.4%），证明更新方向的秩越低、几何侵略性越强，越依赖黎曼度量张量对流形局部曲率的精确矫正。

亮点与洞察¶

几何统一视角的建立：发现通过度量张量二次型开方映射 \(G_x^{1/2}\xi\) 能够将任意酉不变外围范数无缝提升为流形内在范数，将经典黎曼梯度下降（Frobenius 特例）与新兴矩阵谱优化（谱范数/核范数）统摄于同一套测地凸/非凸优化理论框架之下。
理论与工程的精妙契合：利用缩放切空间的奇异值不变性绕过了拉格朗日乘子系统中的流形非线性边界求交难题，使得多类重要矩阵流形上的更新可以直接退化为 SVD 或正交极分解，工程实现完全无需显式计算高阶克里斯托费尔符号或黎曼曲率张量。

局限与展望¶

算力开销依赖 SVD 或 Newton-Schulz 迭代：每次迭代需要对缩放梯度做极分解，当矩阵规模巨大且秩数 \(r\) 较高时计算开销高于标量 SGD。
动量机制与方差缩减仍待深入：当前主理论以标准梯度为主，实际大规模训练中若融入重球动量或 Adam 式二阶自适应分母，如何保持严格的流形测地线平行移动仍有待进一步几何形式化。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从黎曼内在度量视角完美解决了谱范数优化器在矩阵流形上的对称性与闭式求解两大核心冲突。
实验充分度: ⭐⭐⭐⭐☆ 覆盖了 LLM LoRA 微调及多类代表性矩阵流形，实验数据真实详实，若补充数十亿参数模型预训练验证会更完满。
写作质量: ⭐⭐⭐⭐⭐ 数学证明严丝合缝，几何机理与图表阐释清晰，各流形实例化公式推导极为完备。
价值: ⭐⭐⭐⭐⭐ 为参数高效微调（PEFT）和非欧深层模型优化提供了坚实且优雅的底层优化器理论基石。