Move on Muon: A Hamiltonian Probability Gradient Flow Perspective of Muon Optimizer¶

会议: NeurIPS 2026
arXiv: 2605.23871
论文: https://arxiv.org/abs/2605.23871
领域: 优化/理论
关键词: 矩阵优化器, Muon, 哈密顿动力学, 概率测度梯度流, 动量加速

一句话总结¶

将硬谱正交化 Muon 优化器正则化为核范数 Fenchel 对偶平滑诱导的镜像步，并提升至矩阵概率测度空间建立阻尼哈密顿动力学流与精确耗散恒等式，确立了有限粒子 McKean-Vlasov 极限、混沌传播保证及指数收敛率。

研究背景与动机¶

在大规模神经网络与大语言模型的预训练中，权重矩阵的内部代数与谱几何结构对训练稳定性和收敛速度至关重要。传统一阶优化算法（如 SGD、AdamW）本质上将模型参数视为平坦欧几里得空间中的独立标量坐标，依赖逐坐标缩放来调整更新幅度。这种坐标级视角无法利用矩阵特有的奇异值分布和奇异子空间旋转特征。近期提出的 Muon 优化器通过对动量矩阵进行正交化投影，仅保留其左右奇异向量而将奇异值统一截断为一，在多项大模型训练中展现出显著优于 AdamW 的收敛效率。

然而，Muon 在理论分析上面临基础性的数学屏障。理想 Muon 的极分解正交化算子在秩亏矩阵处不连续且不可微，导致其在连续时间极限下退化为复杂的微分包含式，破坏了经典常微分方程与偏微分方程的适定性。更重要的是，现有研究主要将 Muon 解释为单矩阵参数在算子范数单位球上的线性最小化预言机（LMO）或谱信赖域方法，缺少能够解释动量如何与非欧几何协同加速的变分框架；同时，在过参数化均值场与集成学习视角下，如何将 Muon 的谱几何提升至概率测度空间，并刻画粒子交互演化规律，依然是一片理论空白。

为了打通从离散正交化步到连续概率动力学的数学路径，本文构造了一种保持奇异向量方向并平滑奇异值饱和的正则化 Muon 映射。核心 idea：将 Muon 动量正交化识别为核范数光滑 Fenchel 对偶函数生成的镜像步，将动量视为对偶坐标，进而将动力学提升至矩阵概率测度相空间构建阻尼哈密顿流，以动能-势能对偶耗散机制证明目标函数差距的指数收敛性与均值场混沌传播。

方法详解¶

整体框架¶

本文首先将单矩阵参数更新建立在光滑 Fenchel 对偶对 $(\Psi_\varepsilon, \Phi_\varepsilon)$ 之上，将动量更新与正交化位移重构为精确的近端镜像迭代。随后，将该框架提升至概率测度空间 $\mathcal{P}(\mathcal{X})$，考虑形如 $\mathcal{J}(\rho) = \mathcal{R}(\int F(W) d\rho(W))$ 的非线性泛函。在此基础上，引入时间步长 $h$ 与动量阻尼系数 $\beta_h = 1 - \gamma h + o(h)$ 的惯性时间缩放，推导出相互作用粒子系统的特征常微分方程组及其对应的相空间 McKean-Vlasov 连续性演化偏微分方程。最终，将连续演化识别为总能量单调衰减的阻尼哈密顿概率流，并拓展至 Transformer 混合专家模型（MoE）的乘积矩阵空间。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：目标泛函 J(ρ)=R(∫F dρ)<br/>与相空间初值 Law(W_0, P_0)"] --> B["1. 核范数平滑与谱镜像对偶<br/>构造 (Ψ_ε, Φ_ε) 对偶对并诱导光滑 Orth_ε 算子"]
    B --> C["2. 测度空间提升与粒子相互作用系统<br/>均值场内积配对与粒子间第一变分力场耦合"]
    C --> D["3. 惯性时间缩放与 McKean-Vlasov 连续性方程<br/>取 η=h, β=1-γh 导出位置-动量相空间连续极限"]
    D --> E["4. 阻尼哈密顿流结构与精确能量耗散<br/>总能量 H=K+γJ 单调递减与动能-速度耦合"]
    E --> F["5. 修正李雅普诺夫泛函与指数收敛性证明<br/>引入对齐项 C_t 桥接 PL 条件与谱动能耗散"]
    F --> G["6. 乘积空间拓展与混合专家架构适配<br/>块分离镜像势能支持门控路由器与多头注意力优化"]

关键设计¶

1. 核范数平滑与谱镜像对偶：消除极分解不连续性并构建对偶势能

理想 Muon 依赖硬极分解算子 $\text{Orth}(P) = U V^\top$（其中 $P = U \Sigma V^\top$ 为奇异值分解）。该算子在秩亏矩阵处对应核范数次微分 $\partial \|P\|_{\text{nuc}}$ 的多值包含项，引发连续极限的不连续性。本文在标量层级引入凸共轭函数对 $\psi_\varepsilon(a) = \sqrt{a^2 + \varepsilon^2} - \varepsilon$ 与 $\phi_\varepsilon(b) = \varepsilon (1 - \sqrt{1 - b^2}) + \iota_{[-1, 1]}(b)$（$\varepsilon > 0$），并通过奇异值算子谱提升定义矩阵对偶对： $$\Psi_\varepsilon(P) = \sum_{i=1}^q \left(\sqrt{\sigma_i(P)^2 + \varepsilon^2} - \varepsilon\right), \quad \Phi_\varepsilon(G) = \varepsilon \sum_{i=1}^q \left(1 - \sqrt{1 - \sigma_i(G)^2}\right) + \iota_{\{\|G\|_{\text{op}} \leq 1\}}(G)$$ 由此诱导的光滑正交化映射 $\text{Orth}_\varepsilon(P) = \nabla \Psi_\varepsilon(P) = U \text{diag}\left(\frac{\sigma_i(P)}{\sqrt{\sigma_i(P)^2 + \varepsilon^2}}\right) V^\top$ 处处单值且满足 $1/\varepsilon$-Lipschitz 连续性。在此框架下，离散更新步骤等价于关于势函数 $\Phi_\varepsilon$ 的 Bregman 散度近端步，动量矩阵 $P$ 严格扮演对偶动量坐标的角色。

2. 测度空间提升与粒子相互作用系统：从单参数优化迈向均值场动力学

为了刻画过参数化神经网络中的神经元或矩阵块群体行为，本文将优化目标推广至概率测度泛函 $\mathcal{J}(\rho) = \mathcal{R}(\int_\mathcal{X} F(W) d\rho(W))$。在由 $N$ 个粒子组成的经验分布 $\rho_W^N = \frac{1}{N}\sum_{i=1}^N \delta_{W_i}$ 上，引入均值场平均内积 $\langle U, V \rangle_{\text{avg}} = \frac{1}{N}\sum_{i=1}^N \langle U_i, V_i \rangle_F$。该几何结构确保经验目标的平均梯度严格与 Wasserstein 第一变分测度力场 $a_i(W) = \mathcal{R}'(F^N(W)) \nabla F(W_i)$ 对齐。粒子更新形式为： $$P_{i, k+1} = \beta P_{i, k} + (1 - \beta) a_i(W_k), \quad W_{i, k+1} = W_{i, k} - \eta \text{Orth}_\varepsilon(P_{i, k+1})$$ 不同粒子之间仅通过标量均值 $F^N(W)$ 发生全局弱耦合，而在动量正交化与位移推进时保持严格的块分离性。

3. 惯性时间缩放与 McKean-Vlasov 连续性方程：推导相空间联合概率演化

若在离散步长趋近于零时保持动量衰减率 $\beta$ 固定，动量将迅速衰减至过阻尼极限并丢失二阶惯性加速效应。因此，本文设定惯性连续时间尺度：$\eta_h = h$，$\beta_h = 1 - \gamma h + r_h$（其中 $r_h = O(h^2)$）。在此尺度下，$N$ 粒子系统一致收敛至非线性特征 ODE 组：$\dot{W}_i(t) = -\text{Orth}_\varepsilon(P_i(t))$，$\dot{P}_i(t) = \gamma (a_i(W(t)) - P_i(t))$。在粒子数 $N \to \infty$ 的统计极限下，状态变量联合分布 $\mu_t = \text{Law}(W_t, P_t)$ 严格满足 McKean-Vlasov 连续性偏微分方程： $$\partial_t \mu_t + \nabla_W \cdot \left(-\text{Orth}_\varepsilon(P) \mu_t\right) + \nabla_P \cdot \left(\gamma (a_t(W) - P) \mu_t\right) = 0$$ 并且在此极限下，粒子系统展现出混沌传播特性，即任意有限个粒子的边际联合分布在有限时间跨度内以 $O(1/N)$ 的均方速率强收敛至均值场解的张量积。

4. 阻尼哈密顿流结构与精确能量耗散：揭示动能与势能的非单调转换机制

定义包含正规化谱动能与目标势能的测度空间总哈密顿量： $$\mathscr{H}_{\varepsilon, \gamma}(\mu) = \int_\mathcal{Z} \Psi_\varepsilon(P) d\mu(W, P) + \gamma \mathcal{R}\left(\int_\mathcal{Z} F(W) d\mu(W, P)\right)$$ 上述 McKean-Vlasov 方程在数学上严格等价于一个含摩擦阻尼项的哈密顿动力学方程。沿该系统的解轨迹，力场传递与速度输运交叉项发生精确对消，导出总能量的精确微分耗散恒等式： $$\frac{d}{dt} \mathscr{H}_{\varepsilon, \gamma}(\mu_t) = -\gamma \int_\mathcal{Z} \langle P, \text{Orth}_\varepsilon(P) \rangle_F d\mu_t(W, P) = -\gamma \int_\mathcal{Z} \sum_{i=1}^q \frac{\sigma_i(P)^2}{\sqrt{\sigma_i(P)^2 + \varepsilon^2}} d\mu_t \leq 0$$ 由于耗散项等价于动能与对偶速度的 Fenchel 强耦合，耗散率非负且当且仅当动量 $P=0$ 时为零。这从力学视角严谨阐明了 Muon 优化过程中目标函数本身无需单调递减的原因：动量充当动能蓄水池，能量在势能与动能之间震荡互换，而系统总机械能则保持严格的单调下耗散。

5. 修正李雅普诺夫泛函与指数收敛性证明：平衡动量对齐与曲率余项

虽然哈密顿量保证了动力学的物理稳定性，但纯粹的动量耗散无法直接控制目标函数差距 $U_t = \mathcal{J}(\rho_t) - \mathcal{J}_\star$。为此，本文设计了包含受控对齐项的修正李雅普诺夫泛函 $\mathcal{L}_t = \mathcal{H}_t - \alpha C_t$，其中 $C_t = \int \langle a_t(W), P \rangle_F d\mu_t$ 度量动量与下降力场的瞬时对齐程度。在动量有界区域动力学矫顽性、泛函 Polyak-Łojasiewicz（PL）条件以及二阶曲率有界的假设下，选取微小权重 $\alpha > 0$ 使得 $\mathcal{L}_t \approx \mathcal{H}_t$，通过对齐项微分方程引入强迫力模长 $A_t$，最终在连续与离散时间下分别建立了目标函数差距的指数收敛保证： $$\mathcal{J}(\rho_t) - \mathcal{J}_\star \leq \frac{e^{-c_{\alpha, r} t}}{\gamma(1 - \alpha M_C)} \left(\mathcal{H}_0 - \alpha C_0\right), \quad \mathcal{J}(\rho_k) - \mathcal{J}_\star \leq \frac{(1 - h c_\star)^k}{\gamma(1 - \alpha M_C)} \left(\mathcal{H}_0 - \alpha C_0\right)$$

6. 乘积空间拓展与混合专家架构适配：兼容注意力机制与平滑路由分发

将单一矩阵空间推广至有限笛卡尔积空间 $\Theta = \Theta_{\text{exp}} \times \Theta_{\text{gate}} = \prod_{b=1}^B \mathbb{R}^{m_b \times n_b}$，使正规化镜映射能够独立逐块作用于多头注意力的查询、键、值、输出矩阵投影 $(Q, K, V, O)$、前馈网络权重 $(W_1, W_2)$ 以及路由门控网络参数 $\phi$。针对输入相关的专家路由，将专家输出与平滑门控统一编码进增广希尔伯特特征空间 $F_{\text{soft}}(\omega, \phi) = (e^{s_\phi} \psi_\omega, e^{s_\phi})$，使得 Softmax 归一化路由与非归一化平滑门控均能完美嵌入通用概率泛函框架，完整继承阻尼哈密顿流分析的所有收敛与适定性结论。

实验关键数据¶

主实验¶

论文在两类确定性合成动力学基准上进行了粒子系统轨迹验证：实验 1 为单矩阵块的线性矩阵均值匹配任务（$F(W)=W$，目标极小值 $J_\star = 0$）；实验 2 为双矩阵乘积参数空间上的非线性教师-学生神经网络回归任务（粒子参数 $\theta_i = (A_i, B_i) \in \mathbb{R}^{p \times r} \times \mathbb{R}^{r \times d}$，带有 $\tanh$ 激活函数和未共享力场）。对比更新策略包括标准欧几里得动量梯度下降（GD + momentum）、硬极分解 Muon、5 次 5 阶 Newton-Schulz 迭代近似 Muon 以及本文提出的平滑正则化 Muon。

任务配置	参数空间与粒子规模	迭代步数	GD + 动量	硬极分解 Muon	Newton-Schulz Muon	最优正则化 Muon (本文)	最优 $\varepsilon$
矩阵均值匹配 (1)	$\mathbb{R}^{16 \times 8}, (M, N)=(1, 10)$	10,000	$2.4 \times 10^{-29}$	$1.1 \times 10^{-3}$	$1.9 \times 10^{-3}$	$\mathbf{2.0 \times 10^{-32}}$	$3 \times 10^{-2}$
矩阵均值匹配 (2)	$\mathbb{R}^{16 \times 8}, (M, N)=(4, 32)$	10,000	$8.5 \times 10^{-31}$	$1.1 \times 10^{-3}$	$1.4 \times 10^{-3}$	$\mathbf{1.5 \times 10^{-33}}$	$10^{-2}$
乘积空间教师-学生 (1)	$\mathbb{R}^{4 \times 6} \times \mathbb{R}^{6 \times 10}, (M, N)=(3, 12)$	2,000	$1.1 \times 10^{-2}$	$4.4 \times 10^{-6}$	$2.8 \times 10^{-6}$	$\mathbf{1.3 \times 10^{-7}}$	$10^{-5}$
乘积空间教师-学生 (2)	$\mathbb{R}^{4 \times 6} \times \mathbb{R}^{6 \times 10}, (M, N)=(10, 10)$	2,000	$4.5 \times 10^{-3}$	$5.4 \times 10^{-6}$	$4.6 \times 10^{-6}$	$\mathbf{6.0 \times 10^{-7}}$	$10^{-5}$

消融实验¶

消融实验重点考查正则化平滑超参数 $\varepsilon$ 在连续谱系中的过渡规律（在实验 2 的过参数化场景 $(M, N)=(3, 12)$ 下评估 2,000 步时的最终损失与动力学特性）：

平滑参数 $\varepsilon$ 配置	最终目标损失 $J_N$	动态行为特征与理论对应状态
$\varepsilon = 10^{-1}$ (强平滑)	$\sim 10^{-3}$	过度平滑，行为退化为阻尼欧几里得动量更新，丧失谱加速优势
$\varepsilon = 10^{-3}$ (中平滑)	$\sim 5.0 \times 10^{-7}$	具备稳定谱正交化加速，同时保留较好的数值下界收敛能力
$\varepsilon = 10^{-5}$ (最优平滑)	$\mathbf{1.3 \times 10^{-7}}$	完美平衡谱特征利用与极值点邻域收敛半径，能量单调耗散至深谷
$\varepsilon \downarrow 0$ (硬极分解极限)	$4.4 \times 10^{-6}$	有限步长下正交化步长在最优解附近无法自适应衰减，形成 $\sim 10^{-6}$ 残差平台

关键发现¶

消除极分解有限步残差平台：硬极分解与 Newton-Schulz 算子在动量非零时始终输出范数恒为 1 的归一化更新方向。在固定有限步长 $h=0.01$ 下，这种无法动态衰减的步长会导致轨迹在极小点附近震荡，卡在 $10^{-3} \sim 10^{-6}$ 的误差平台上；而正则化算子在动量较小时近似满足 $\text{Orth}_\varepsilon(P) \approx P/\varepsilon$，具备步长自收缩特性，使优化过程平滑趋于数值零点。
谱正交化加速在非凸乘积空间的巨大优势：在非线性的教师-学生两层网络中，欧几里得动量梯度下降仅能达到 $10^{-2} \sim 10^{-3}$ 阶损失，而基于谱正交化的 Muon 系列方法能将误差压低 4 个数量级以上，证实了矩阵奇异子空间旋转更新在多层矩阵架构中的内在优越性。
哈密顿能量单调递减的经验自洽性：实验日志监控表明，尽管在局部复杂曲率区域目标函数 $J_N$ 偶有微幅抖动，但总哈密顿能量 $K_t + \gamma J_N$ 在对数刻度下严格单调衰减，有力印证了能量耗散恒等式。

亮点与洞察¶

变分视角的根本统一：首次阐明了 Muon 优化器在算子范数约束下并非启发式的硬投影技巧，而是正规化核范数 Fenchel 对偶势能所生成的连续镜像步，使非平滑启发式算法回归至经典凸几何与对偶分析的统一大厦。
物理相空间中的二阶阻尼流诠释：跳脱出将动量优化器简化为一阶过阻尼 Wasserstein 测度流的传统范式，建立了带有真实动能与摩擦因子的相空间二阶微分系统，将深度网络训练的“震荡加速”具象化为动能与势能间的相空间耗散。
混合专家（MoE）路由优化的均值场理论拓荒：通过增广特征空间技巧将离散/连续路由器与专家网络权重同时纳入乘积矩阵泛函框架，为分析大规模 MoE 架构的联合动力学与路由坍缩提供了全新的偏微分方程工具。

局限与展望¶

理论收敛假设在实际大模型中的验证难度：指数收敛率依赖于全局或轨迹局部的 PL 条件、有界动量以及受控曲率假设。在实际千万亿参数的大语言模型预训练非凸非平滑地形中，严格验证这些假设极其困难。
硬 Top-k 离散路由的理论断裂：当前理论要求路由得分算子具备二次连续可微性（$C^2$ 平滑），无法直接覆盖目前工业界广泛采用的不可微硬性 Top-k 稀疏路由（需要引入差异化包含或次微分松弛）。
随机小批量噪声与非确定性分析的缺失：论文主要构建了全批次确定性均值场极限，尚未将随机梯度噪声（SGD-Muon）与有限步 Newton-Schulz 多项式截断误差纳入随机哈密顿微分方程（SDE）体系中深入探讨。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立 Muon 优化器的平滑镜像流、阻尼哈密顿概率动力学与 McKean-Vlasov 均值场理论。
实验充分度: ⭐⭐⭐⭐ 包含严密的二维线性与非线性动力学对比及消融，但缺少全尺寸真实 Transformer 大模型预训练实验。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨精炼，从凸对偶、粒子系统到测度流的逻辑链路层层递进、自洽完备。
价值: ⭐⭐⭐⭐⭐ 为当前大模型前沿广泛应用的矩阵谱优化器提供了坚固扎实的理论基石与深刻的物理学直觉。

任务配置	参数空间与粒子规模	迭代步数	GD + 动量	硬极分解 Muon	Newton-Schulz Muon	最优正则化 Muon (本文)	最优 \(\varepsilon\)
矩阵均值匹配 (1)	\(\mathbb{R}^{16 \times 8}, (M, N)=(1, 10)\)	10,000	\(2.4 \times 10^{-29}\)	\(1.1 \times 10^{-3}\)	\(1.9 \times 10^{-3}\)	\(\mathbf{2.0 \times 10^{-32}}\)	\(3 \times 10^{-2}\)
矩阵均值匹配 (2)	\(\mathbb{R}^{16 \times 8}, (M, N)=(4, 32)\)	10,000	\(8.5 \times 10^{-31}\)	\(1.1 \times 10^{-3}\)	\(1.4 \times 10^{-3}\)	\(\mathbf{1.5 \times 10^{-33}}\)	\(10^{-2}\)
乘积空间教师-学生 (1)	\(\mathbb{R}^{4 \times 6} \times \mathbb{R}^{6 \times 10}, (M, N)=(3, 12)\)	2,000	\(1.1 \times 10^{-2}\)	\(4.4 \times 10^{-6}\)	\(2.8 \times 10^{-6}\)	\(\mathbf{1.3 \times 10^{-7}}\)	\(10^{-5}\)
乘积空间教师-学生 (2)	\(\mathbb{R}^{4 \times 6} \times \mathbb{R}^{6 \times 10}, (M, N)=(10, 10)\)	2,000	\(4.5 \times 10^{-3}\)	\(5.4 \times 10^{-6}\)	\(4.6 \times 10^{-6}\)	\(\mathbf{6.0 \times 10^{-7}}\)	\(10^{-5}\)

平滑参数 \(\varepsilon\) 配置	最终目标损失 \(J_N\)	动态行为特征与理论对应状态
\(\varepsilon = 10^{-1}\) (强平滑)	\(\sim 10^{-3}\)	过度平滑，行为退化为阻尼欧几里得动量更新，丧失谱加速优势
\(\varepsilon = 10^{-3}\) (中平滑)	\(\sim 5.0 \times 10^{-7}\)	具备稳定谱正交化加速，同时保留较好的数值下界收敛能力
\(\varepsilon = 10^{-5}\) (最优平滑)	\(\mathbf{1.3 \times 10^{-7}}\)	完美平衡谱特征利用与极值点邻域收敛半径，能量单调耗散至深谷
\(\varepsilon \downarrow 0\) (硬极分解极限)	\(4.4 \times 10^{-6}\)	有限步长下正交化步长在最优解附近无法自适应衰减，形成 \(\sim 10^{-6}\) 残差平台