# Taming Softmax at the Limits of Muon: A Unified Approach through Hilbert Geometry

arXiv: https://arxiv.org/abs/2610.17588 · HTML: https://mlsanity.com/paper/2610.17588/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 25 references that are in the map.
- First posted on arXiv 13 Oct 2026.

## Builds on (most cited first, 11 of 11)

- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [Scalable Optimization in the Modular Norm](https://mlsanity.com/paper/2405.14813/) (2024)
- [Memory-Efficient LLM Pretraining via Minimalist Optimizer Design](https://mlsanity.com/paper/2506.16659/) (2025)
- [Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers](https://mlsanity.com/paper/2605.18106/) (2026)
- [Non-linear Embeddings in Hilbert Simplex Geometry](https://mlsanity.com/paper/2203.11434/) (2022)

## Abstract (arXiv)

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 citeppethick2025trainingdeeplearningmodels. For scaled dot-product attention, the same principle yields the first-order query--key steepest descent problem underlying Compositional Muon citepkeigwin2026compositional. These connections motivate using Hilbert distance as a guiding principle for optimizer design in softmax modules.