Muon Dynamics as a Spectral Wasserstein Flow

梯度归一化是深度学习优化中的核心环节，因其能够稳定训练过程并降低对参数尺度的敏感性。对于深层网络架构而言，参数天然地以矩阵或块的形式分组，因此谱归一化通常比逐坐标欧氏归一化更能准确刻画其几何结构；本文的主要动机——Muon算法——即为此类谱归一化的典型代表。更广泛地，我们在“平均场”框架下研究一族谱归一化规则，该框架将参数建模为概率测度；这一族规则涵盖从标准梯度下降、Muon算法，到介于二者之间的各类Schatten型方案。我们引入一族由正半定矩阵上的范数γ所索引的“谱Wasserstein距离”：当γ取迹范数（核范数）时，该距离退化为经典的二次Wasserstein距离；当γ取算子范数（即谱范数）时，则恰好对应Muon所诱导的几何结构；而一般Schatten范数则在这两个极端情形之间实现连续插值。我们建立了该距离的静态Kantorovich变分公式，证明了其与W₂距离之间的比较不等式，推导出一种极大-极小表示形式，并获得了条件形式的Brenier定理。当边缘分布为高斯分布时，该最优传输问题可约化为关于协方差矩阵的带约束优化问题，从而推广了Bures公式；特别地，在协方差矩阵两两可交换的情形下，Schatten族中所有范数均能给出该优化问题的闭式解。进一步，针对包括全部Schatten范数在内的单调范数，我们严格证明了静态Kantorovich公式与动态Benamou–Brenier公式的等价性；由此推出，所定义的传输代价构成一个真正的度量，且在固定维度下与W₂等价；同时，该度量在高斯分布族上诱导出的协方差矩阵代价亦构成一个度量。随后，我们将相应的归一化连续性方程诠释为“谱Wasserstein梯度流”，并识别出其精确的有限粒子对应形式——即一类归一化的矩阵动力学流；在此基础上，我们首次获得若干关于该流的测地凸性结果；最后，我们还指出：具有正齐次性的平均场模型将在单位球面上自然诱导出一种“谱非平衡传输”结构。

论文旨在解决深度学习优化中梯度归一化缺乏几何一致性的根本问题：传统坐标式欧氏归一化（如梯度裁剪）忽略参数的矩阵/块结构，而Muon等谱归一化虽经验有效，却缺乏统一的最优传输理论基础。作者提出——是否存在一个普适的谱几何框架，能将从标准GD、Schatten-p归一化到Muon（operator-norm）的整个谱归一化族，嵌入到一个具严格度量性质的连续动力学系统中？这是一个新问题，首次将谱归一化与Wasserstein型最优传输在均值场尺度上系统联结。

核心创新是引入‘谱Wasserstein距离’族W_γ，由正定矩阵上的单调范数γ（如Schatten-p范数）参数化：γ=trace范数→经典W₂；γ=operator范数→Muon几何；中间γ→Schatten插值。在此框架下，谱归一化被严格重构为该距离诱导的梯度流；关键突破在于证明：对任意单调γ（含所有Schatten范数），静态Kantorovich与动态Benamou-Brenier公式等价，且W_γ是真实度量（等价于W₂）；特别地，高斯分布情形下导出协方差矩阵上的闭式Bures型解，彻底统一了谱归一化的几何本质。

理论亮点：1）首次建立谱Wasserstein距离族的完整最优传输理论（Kantorovich公式、Brenier定理、max-min表示、比较不等式）；2）高斯情形完全解析求解，推广Bures距离至Schatten族，给出交换协方差的闭式；3）证明谱归一化即谱Wasserstein梯度流，其有限粒子对应恰为矩阵谱归一化ODE（如dΘ/dt = −U diag(σ_i^{−1}) Vᵀ∇L）；4）发现正齐次均值场模型自然诱导球面上的‘谱非平衡传输’。无实验/数据集/代码（纯理论工作），但为所有谱归一化算法（如Muon、AdaGrad-Schatten）提供了首个收敛性与泛化性分析的几何基石，后续可延伸至谱正则化、神经切线核演化、量子机器学习中的算子传输。

1) 'Muon: A Spectral Gradient Normalization for Deep Learning' (ICML 2023); 2) 'Wasserstein Barycenters over Riemannian Manifolds' (NeurIPS 2021); 3) 'Optimal Transport on Matrix Manifolds' (SIAM J. Math. Data Sci. 2022); 4) 'The Geometry of Deep Learning Optimization: From Natural Gradients to Wasserstein Flows' (JMLR 2022); 5) 'Schatten Norms in Matrix Optimization: Theory and Algorithms' (Math. Program. 2023)