arXiv：cs.AI· Aditya Somasundaram, Charles Guille-Escuret, Alexander Moreno, Zhengzhong Liu, Eric Xing·· 6 小时前
用 Hilbert 几何统一 Muon 在 softmax 极限下的优化：RowNorm 与 Compositional Muon 的数学框架
Taming Softmax at the Limits of Muon: A Unified Approach through Hilbert Geometry
AI 导读
Muon 通过线性层输出的 RMS 位移控制权重更新，但该代理指标对 softmax 失真，RMS 激活变化无法忠实反映输出分布变化。作者提出用 Hilbert 距离约束输出分布变化，可界定交叉熵损失变化并限制概率重分配，对 softmax 表现为预激活扰动的变分半范数，从而把概率约束转为参数约束。
正文
Abstract:Muon controls weight updates through the RMS movement of a linear layer's output. This implicitly treats the linear output as a proxy for the full module, including the nonlinearity that follows. This proxy can be misleading for softmax, where the RMS shift in activations does not faithfully measure the change in the output distribution. We propose Hilbert distance as a natural metric for bounding changes in these output distributions. Hilbert distance bounds changes in cross-entropy loss and limits probability redistribution. For softmax, it takes the form of the variation seminorm of the preactivation perturbation, allowing probability constraints to be expressed as parameter constraints. Formulating the steepest descent problem with a Hilbert distance constraint brings two recently proposed methods under a unified mathematical framework. For the LM head, solving a row-wise approximation to the resulting diameter-constrained problem leads us to RowNorm \citep{pethick2025trainingdeeplearningmodels}. For scaled dot-product attention, the same principle yields the first-order query--key steepest descent problem underlying Compositional Muon \citep{keigwin2026compositional}. These connections motivate using Hilbert distance as a guiding principle for optimizer design in softmax modules.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.17588 [cs.LG] |
| (or arXiv:2610.17588v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.17588 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Aditya Somasundaram [view email]
[v1]
Tue, 13 Oct 2026 02:07:50 UTC (189 KB)
来源：arXiv：cs.AI · arxiv.org