arXiv cs.AI· Aditya Somasundaram, Charles Guille-Escuret, Alexander Moreno, Zhengzhong Liu, Eric Xing·· 3 天前AI 评分22
用 Hilbert 几何统一 Muon 在 softmax 极限下的优化方法
Taming Softmax at the Limits of Muon: A Unified Approach through Hilbert Geometry
AI 导读
针对 Muon 以线性层输出的 RMS 变化近似整个模块、在 softmax 上失真的问题，研究者提出用 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