Abstract
Gradient normalization is central in deep-learning optimization because it stabilizes training and reduces sensitivity to scale. For deep architectures, parameters are naturally grouped into matrices or blocks, so spectral normalizations are often more faithful than coordinatewise Euclidean ones; Muon is the main motivating example of this paper. More broadly, we study a family of spectral normalization rules, ranging from ordinary gradient descent to Muon and intermediate Schatten-type schemes, in a mean-field regime where parameters are modeled by probability measures. We introduce a family of Spectral Wasserstein distances indexed by a norm gamma on positive semidefinite matrices. The trace norm recovers the classical quadratic Wasserstein distance, the operator norm recovers the Muon geometry, and intermediate Schatten norms interpolate between them. We develop the static Kantorovich formulation, prove comparison bounds with W2, derive a max-min representation, and obtain a conditional Brenier theorem. For Gaussian marginals, the problem reduces to a constrained optimization on covariance matrices, extending the Bures formula and yielding a closed form for commuting covariances in the Schatten family. For monotone norms, including all Schatten cases, we prove the equivalence between the static and dynamic Benamou-Brenier formulations, deduce that the resulting transport cost is a genuine metric equivalent to W2 in fixed dimension, and show that the induced Gaussian covariance cost is also a metric. We then interpret the associated normalized continuity equation as a Spectral Wasserstein gradient flow, identify its exact finite-particle counterpart as a normalized matrix flow, obtain first geodesic-convexity results, and show how positively homogeneous mean-field models induce a spectral unbalanced transport on the sphere.
Abstract (translated)
梯度归一化在深度学习优化中至关重要，因为它能稳定训练过程并降低对尺度变化的敏感性。对于深度架构，参数自然按矩阵或分块分组，因此谱归一化通常比逐坐标欧几里得归一化更忠实；本文的主要动机案例Muon即体现了这一特性。更广泛地，我们在参数被建模为概率测度的平均场框架下，研究了一类谱归一化规则，涵盖从普通梯度下降到Muon及中间施瓦茨型方案的连续谱系。我们引入了一族以正半定矩阵上的范数γ为指标的谱Wasserstein距离。迹范数恢复经典二次Wasserstein距离，算子范数恢复Muon几何，中间施瓦茨范数则实现两者间的插值。我们发展了静态康托罗维奇公式，证明了与W₂的比较界，推导了极大极小表示，并获得了条件布伦耶定理。对于高斯边缘分布，问题简化为协方差矩阵上的约束优化，推广了伯斯公式，并给出了施瓦茨族中对易协方差的闭式解。对于单调范数（包括所有施瓦茨情形），我们证明了静态与动态贝纳穆-布伦耶公式的等价性，推导出所得输运成本是与W₂在固定维度下等价的真度量，并表明诱导的高斯协方差成本也是度量。随后，我们将相关的归一化连续性方程解释为谱Wasserstein梯度流，将其精确有限粒子对应识别为归一化矩阵流，获得了首个测地凸性结果，并展示了正齐次平均场模型如何在球面上诱导谱非平衡输运。
URL
https://arxiv.org/abs/2604.04891
