Muon Dynamics as a Spectral Wasserstein Flow¶

会议: NeurIPS 2026
arXiv: 2604.04891
OpenReview: https://openreview.net/forum?id=sqPfZgXVES
代码: https://github.com/gpeyre/muon-dynamics
领域: 优化/理论
关键词: Muon优化器, 谱Wasserstein距离, 均值场动力学, 梯度流, 线性极小化预言机

一句话总结¶

本文为以 Muon 为代表的谱归一化优化器构建了严格的无限宽均值场数学框架，通过在正半定锥上引入单调谱度规建立谱 Wasserstein 距离，证明了归一化均值场 PDE 本质上正是该传输度规下的最大斜率梯度流。

研究背景与动机¶

在大规模深度神经网络与大语言模型的预训练中，梯度归一化与谱归一化优化器（以 Muon 为典型代表）展现出优于标准欧氏梯度下降（SGD/Adam）的训练稳定性和收敛表现。Muon 等优化器通过利用动量梯度的极分解或 Newton-Schulz 迭代计算正交极因子更新，替换了参数空间的标准欧氏几何，引入了一种适配矩阵参数块的谱几何。然而，现有关于此类优化器的理论分析主要局限于有限维、离散时间或启发式预条件解释，缺乏一个能够覆盖现代超参数化超宽网络、在无穷宽连续极限下自洽的变分与几何理论。

传统的过参数化两层网络分析依赖于均值场（Mean-Field）理论，将网络参数表示为参数空间上的概率测度，并借助经典 Wasserstein 空间（\(\mathsf{W}_{2}\)）的梯度流刻画神经元连续演化。然而，经典 \(\mathsf{W}_{2}\) 几何的切空间速度动作由速度场的动能（即位移梯度的 \(L^2\) 范数或协方差的迹）所主导，天然诱导的是各向同性的欧氏最速下降，完全无法解释 Muon 等算法在更新时施加的谱范数（算子范数）约束及其对参数协方差谱结构的动态调整。

为了填补这一几何与优化鸿沟，本文的核心切入点是打破“传输代价必须基于点对点欧氏距离”的传统设定，转而在全局位移协方差矩阵上定义矩阵范数规范。本文的核心 idea 是：通过在半正定矩阵锥上引入单调谱度规 \(\gamma\)，构建耦合级谱 Wasserstein 传输距离 \(\mathsf{W}_\gamma\)，证明其与 Benamou-Brenier 动力学公式等价，并将 Muon 均值场偏微分方程精准刻画为测度空间上该传输度规下的最大斜率梯度流。

方法详解¶

整体框架¶

本文构建的数学体系遵循“矩阵谱归一化 \(\to\) 测度线性极小化预言机（Measure-LMO）\(\to\) 连续均值场连续性方程 \(\to\) 静态谱 Wasserstein 距离 \(\to\) 动态 Benamou-Brenier 变分等价与测度梯度流”的自底向上路线。

在整体流程中，对于目标泛函 \(\mathcal{F}(\mu)\)，经典欧氏梯度流所对应的第一变分给出 Wasserstein 力场 \(g_{\mu}(x)=\nabla_x \frac{\delta\mathcal{F}}{\delta\mu}(x)\)。在谱几何下，力场 \(g_\mu\) 不再直接充当速度，而是作为余切向量输入测度级线性极小化预言机 \(\mathcal{J}_\gamma^\mu\)；该预言机基于力场二阶矩协方差 \(S_g\) 求解有限维矩阵半正定优化问题，输出经过谱加权修正的真实速度场 \(v_t(x)\)，进而驱动粒子流演化；同时，静态传输通过凸对偶与极大极小定理诱导了鲁棒最优传输距离，与动力学动量作用量完美闭环。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["经验粒子或测度状态 μ<br/>与目标泛函 ℱ(μ)"] --> B["谱规范与矩阵线性极小化预言机<br/>定义单调度规 γ 与参数块下降方向"]
    B --> C["有限维约化的测度预言机<br/>利用力场协方差 S_g 计算最优谱速度场"]
    C --> D["归一化均值场连续性偏微分方程<br/>∂_t μ_t + div(μ_t v_t) = 0 驱动粒子演化"]
    D --> E["静态谱Wasserstein鲁棒传输几何<br/>极大极小定理与代价鲁棒测度 W_γ"]
    E --> F["动态Benamou-Brenier等价与测度梯度流<br/>最大斜率曲线与能量耗散恒等式闭环"]

关键设计¶

1. 谱规范与矩阵线性极小化预言机：将块参数更新方向与谱范数严格对齐

传统欧氏梯度下降采用固定的二次正则项，导致更新步长和方向对矩阵参数奇异值的各向异性极度敏感。针对矩阵块参数优化，本文在半正定对称矩阵锥 \(\mathbb{S}_{+}^d\) 上定义保偏序的单调规范 \(\gamma\)（即当 \(0 \preceq S \preceq S'\) 时 \(\gamma(S) \leq \gamma(S')\)），并在矩阵空间 \(\mathbb{R}^{n \times d}\) 上诱导矩阵范数 \(\mathrm{N}(V) \coloneqq \gamma(V^\top V)^{1/2}\)。基于此范数，梯度下降的最优下降方向被定义为线性极小化预言机（LMO）：

\[\mathrm{J}_\gamma(G) \in \operatorname*{argmin}_{V} \left\{ \langle G, V \rangle_F + \frac{1}{2}\mathrm{N}(V)^2 \right\}\]

当选用 Schatten \(p\)-度规 \(\gamma_p(S) = \|\lambda(S)\|_{\ell_p}\) 时，其对偶指数呈现丰富内插：当 \(p=1\) 时（迹范数）\(\mathrm{J}_{\gamma_1}(G) = -G\)，精准退化为经典梯度流；而当 \(p=\infty\) 时（算子范数/最大特征值），\(\mathrm{J}_{\gamma_\infty}(G) = -\|G\|_{S_1} U W^\top = -\operatorname{tr}((G^\top G)^{1/2}) G (G^\top G)^{\dagger/2}\)，正好对应了 Muon 优化器所依赖的按核范数缩放的局部极分解更新。这种将离散极小化算子转化为变分预言机的机制，构成了连接离散算法与连续介质力学的基石。

2. 有限维约化的测度预言机：消除无穷维力场选择的计算阻碍

当模型宽度推向无穷大时，网络状态表示为概率测度 \(\mu \in \mathcal{P}_2(\mathbb{R}^d)\)，速度场 \(v \in L^2(\mu; \mathbb{R}^d)\) 的动作范数定义为 \(\mathcal{N}_\mu(v)^2 \coloneqq \gamma\left(\int v(x)v(x)^\top d\mu(x)\right)\)。由于对无穷维泛函直接求解测度级预言机 \(\mathcal{J}_\gamma^\mu(g) \in \operatorname*{argmin}_v \left\{ \int g \cdot v \, d\mu + \frac{1}{2}\mathcal{N}_\mu(v)^2 \right\}\) 极为抽象，本文证明了该无穷维优化问题能够完全约化为在力场协方差 \(S_g \coloneqq \int g(x)g(x)^\top d\mu(x)\) 的值域上的紧致有限维矩阵极值问题。

该约化机制通过极分解构造算子投影，证明最优速度场必具有形式 \(\mathcal{J}_\gamma^\mu(g)(x) = A_g^\top S_g^{\dagger/2} g(x)\)，其中 \(A_g\) 满足有限维凸规划。特别地，在 Muon 对应的算子范数极值点（\(p=\infty\)），显式速度选择器与速度协方差满足：

\[\mathcal{J}_{\gamma_\infty}^\mu(g)(x) = -\operatorname{tr}(S_g^{1/2}) S_g^{\dagger/2} g(x), \qquad \int v v^\top d\mu = \operatorname{tr}(S_g^{1/2})^2 P_g\]

其中 \(P_g\) 为 \(S_g\) 值域上的正交投影算子。该机制揭示出 Muon 动力学在无穷维均值场中的本质效应：它并非将粒子的微观位移限制在单一优势方向，而是在整个活跃力场子空间上强行将速度协方差彻底“各向同性白化/等权化”，实现了矩阵正交化思想在测度空间的宏观映射。

3. 静态谱Wasserstein鲁棒传输几何：通过极大极小对偶确立距离结构

为了探究归一化连续性方程是否由某个度量空间的梯度流驱动，必须首先确立概率测度空间上的底层距离。本文提出了静态谱 Wasserstein 传输代价：

\[\mathsf{W}_\gamma(\mu, \nu)^2 \coloneqq \inf_{\pi \in \Pi(\mu, \nu)} \gamma\left( \int_{\mathbb{R}^d \times \mathbb{R}^d} (y-x)(y-x)^\top d\pi(x, y) \right)\]

与经典 Wasserstein 距离直接积分欧氏距离平方不同，\(\mathsf{W}_\gamma\) 先在耦合 \(\pi\) 下对位移外积张量求期望得到全局位移协方差 \(\Sigma_\pi\)，再对其作用谱度规 \(\gamma\)。为了突破静态形式下三角不等式与测度性质证明的障碍，本文利用单调正半定规范的凸共轭紧致表示集 \(\mathcal{K}_\gamma \subset \mathbb{S}_+^d\)，应用 Sion 极大极小定理将极小与极大算子对调，得到了极具深刻物理意义的鲁棒代价表达：

\[\mathsf{W}_\gamma(\mu, \nu)^2 = \max_{Q \in \mathcal{K}_\gamma} \inf_{\pi \in \Pi(\mu, \nu)} \int (y-x)^\top Q (y-x) d\pi(x, y) = \max_{Q \in \mathcal{K}_\gamma} \mathsf{W}^Q(\mu, \nu)^2\]

该设计表明，\(\mathsf{W}_\gamma\) 本质上是对抗性地在允许的二次型半正定度量矩阵 \(Q \in \mathcal{K}_\gamma\) 中选取最劣传输距离。对于 Muon（\(p=\infty\)），\(\mathcal{K}_{\gamma_\infty} = \{Q \succeq 0: \operatorname{tr}(Q) = 1\}\)，最坏方向正是位移协方差最大特征值对应的投影方向。

4. 动态Benamou-Brenier等价与测度梯度流：闭环均值场能量耗散动力学

确立了静态度量后，最关键的桥梁是证明连续时间流动的微观动能消耗与宏观最优静态位移一致。本文在单调规范假设下，引入连续性方程约束下的速度动力学动作：

\[\mathsf{W}_\gamma^{\mathrm{BB}}(\mu_0, \mu_1)^2 \coloneqq \inf_{(\mu_t, v_t)} \int_0^1 \gamma\left( \int v_t(x) v_t(x)^\top d\mu_t(x) \right) dt\]

利用叠加热核轨迹分解与条件 Jensen 不等式，严格证明了静态距离与动态作用量的等价性 \(\mathsf{W}_\gamma^{\mathrm{BB}}(\mu_0, \mu_1) = \mathsf{W}_\gamma(\mu_0, \mu_1)\)，确立了常速测地线插值 \(\mu_t = ((1-t)x + ty)_\# \pi^*\)。在此基础上，通过构造一阶输运展开式与对偶范数关系，证明了测度演化曲线 \(\mu_t\) 正是能量泛函 \(\mathcal{F}\) 在度量空间 \((\mathcal{P}_2(\mathbb{R}^d), \mathsf{W}_\gamma)\) 下的最大斜率曲线（Curve of Maximal Slope），几乎处处满足精准能量耗散恒等式：

\[\frac{d}{dt}\mathcal{F}(\mu_t) = -\mathcal{N}_{\mu_t}^*(g_t)^2 = -|\mu'|_\gamma(t)^2 = -|\partial\mathcal{F}|_\gamma(\mu_t)^2\]

从而将 Muon 优化器在连续均值场下的演化过程彻底纳入 Ambrosio-Gigli-Savaré 度量梯度流理论大厦。

损失函数 / 训练策略¶

在均值场参数化架构下，两层感知机预测器为 \(H_\mu(z) = \int_{\mathbb{R}^d} \phi(z, x) d\mu(x)\)，经验风险为 \(\mathcal{F}(\mu) = \frac{1}{2} \int |H_\mu(z) - y^*(z)|^2 d\rho(z)\)。其连续时间演化方程由如下非线性输运 PDE 刻画：

\[\partial_t \mu_t + \operatorname{div}\left( \mu_t \mathcal{J}_\gamma^{\mu_t}(g_{\mu_t}) \right) = 0, \qquad g_{\mu_t}(x) = \nabla_x \frac{\delta \mathcal{F}}{\delta \mu}(x)\]

在数值模拟实现中，离散粒子系统采用显式 Euler-forward 时间离散更新 \(X_{k+1} = X_k + \tau \mathrm{J}_\gamma(\nabla \mathrm{F}(X_k))\)，并对粒子矩阵步长的 Frobenius 范数施加截断保护以防止数值发散。

实验关键数据¶

主实验¶

论文通过两层神经网络拟合、高斯线性网络、多头浅层注意力机制（Shallow Multi-head Attention）以及最大均值差异（MMD）粒子梯度流等多个基准实验，系统对比了欧氏测度流（\(p=1\)）与 Muon 谱测度流（\(p=\infty\)）在均值场极限下的动力学特征。

表 1 汇总了数值实验采用的基准架构与离散时间时钟规格参数（来自论文 Table 1 及 Appendix G）：

实验任务	粒子数 / 测度规模 \(n\)	演化步数 \((p=1, p=\infty)\)	名义步长 \(\Delta t_1\)	名义步长 \(\Delta t_\infty\)	步长截断阈值 (Step Cap)	坐标截断保护
两层 ReLU MLP 拟合	1920 神经元	(2000, 2000)	75	50	0.20	\(\pm 30\)
浅层多头注意力机制	200 头 (\(d=48\))	(960, 480)	100	180	0.45	\(\pm 10\)
MMD 能量距离粒子流	200 粒子	(6000, 6000)	5.0	3.2	无截断 (None)	无

消融与动力学特性分析¶

针对高斯线性双层网络在不同 Schatten 规范 \(p \in \{1, 2, 4, \infty\}\) 下的模态降维动力学，论文研究了协方差秩亏（Rank Loss）与几何轨迹的演化对比：

规范几何配置	谱指标参数 \(q\)	模态耦合状态	秩亏临界残差阈值 \(R_{p, 1}\)	协方差动态特征
\(p=1\) (经典迹范数 / \(\mathsf{W}_2\))	\(q=2\)	模态间完全解耦，独立演化	\(R_{1,1} \to \infty\)（无有限时间秩亏）	各向同性收缩，协方差保持正定
\(p=2\) (Frobenius 中间态)	\(q=4/3\)	全局非线性谱范数耦合	\(R_{2, 1} = 0.400\)	产生优先衰减方向，模态相互挤压
\(p=4\) (高阶 Schatten 态)	\(q=8/7\)	强谱范数耦合与收敛偏置	\(R_{4, 1} \simeq 0.252\)	次要奇异值加速塌缩
\(p=\infty\) (算子范数 / Muon)	\(q=1\)	强全局耦合，速度协方差投影白化	\(R_{\infty, 1} = 0.200\)	极易发生有限时间秩亏，随后贴合边界滑移

关键发现¶

谱白化驱动的各向同性横向发散：在 MMD 优化实验中，算子范数（\(p=\infty\)）诱导的粒子流在演化初期展现出强烈的横向协调扩展，这一现象归因于 Muon 选择器将活跃力场子空间上的速度协方差强行拉平，赋予微弱横向扰动相对较大的放大倍率（\(\propto \operatorname{tr}(S_g^{1/2})/\sqrt{\lambda_i}\)）。
高斯线性网络中的模态共享约束：对于 \(p > 1\) 的谱几何，不同神经元特征模态之间通过全局归一化产生强相互作用耦合；当目标残差较大时，收缩模态会在有限时间内迅速发生秩亏（特征值衰减为 0），随后沿正半定锥边界继续滑移。
注意力权重的稀疏性响应差异：在浅层注意力实验中，高低温注意力稀疏度分布（有效 token 数 \(N_{\text{eff}} \approx 16\) 与 \(N_{\text{eff}} \approx 2.1\)）对 \(p=1\) 与 \(p=\infty\) 产生显著不同的损失收敛曲率，验证了谱归一化优化器对 Softmax 激活极值特性的内在几何自适应。

亮点与洞察¶

将离散矩阵启发式升华至无限维几何：论文没有停留在对 Muon 离散更新技巧的工程性修补，而是从连续介质和最优传输的第一性原理出发，确立了矩阵谱优化对应的度量空间 \((\mathcal{P}_2, \mathsf{W}_\gamma)\)。
Sion 极大极小定理对非标准传输代价的妙用：对于非点对点相加的矩阵外积协方差代价，通过引入凸对偶集 \(\mathcal{K}_\gamma\) 构造鲁棒二次型，不仅巧妙化解了动态与静态等价性证明中的凸性困境，还直接衍生出正半定锥上的 Bures 距离推广。
均值场偏微分方程的闭式高斯降维：证明了仿射力场在谱流驱动下仍保持高斯族闭合，使具有无穷多神经元的非线性演化在特定两层线性模型中能精准降维为低维常微分方程，为理论分析提供了绝佳玩具模型。

局限与展望¶

作者承认的理论局限：本文梯度流定理（Theorem 5.1）基于既有曲线上的验证性假设（如力场一致 Lipschitz 与二阶输运展开），尚未证明非光滑激活函数（如包含扭结点的 ReLU）下的全局 PDE 解的存在唯一性，亦未涵盖动量项与随机小批量（Minibatch）噪声对测度流的扰动。
实际实现的动力学偏移：工程实现中的 Muon 通常使用 5 次 Newton-Schulz 迭代近似极分解，这在数学上对应于多项式谱映射而非严格符号投影，且离散更新中加入了较大的动量缓冲（Phase-space Liouville 方程），与无动量纯测度梯度流存在一定间隙。
未来改进方向：构建结合动量与阻尼项的二阶谱 Wasserstein 动量动力学（Kramers / Langevin 谱型流）；发展针对有限样本随机采样的非渐近测度集中不等式。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立谱归一化优化器与测度空间谱 Wasserstein 几何的严格理论对应。
实验充分度: ⭐⭐⭐⭐☆ 包含 MMD、MLP、高斯闭合及多头注意力的多维度几何诊断，实验重在揭示几何机制而非刷榜。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密精巧，物理直觉清晰，概念演进环环相扣。
价值: ⭐⭐⭐⭐⭐ 为超大规模大模型中 Muon 系列矩阵优化器提供了首个坚实的微分几何与均值场分析理论基础。