Paper deep dive
Taming Softmax at the Limits of Muon: A Unified Approach through Hilbert Geometry
Aditya Somasundaram, Charles Guille-Escuret, Alexander Moreno, Zhengzhong Liu, Eric Xing
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.17588v1
- Canonical: https://arxiv.org/abs/2610.17588v1
Trouble viewing inline? Open PDF directly →
Full Text
82,295 characters extracted from source content.
Expand or collapse full text
Taming Softmax at the Limits of Muon: A Unified Approach through Hilbert Geometry Aditya Somasundaram Charles Guille-Escuret Alexander Moreno Affiliation: Zhengzhong Liu, Eric Xing Affiliation: Institute of Foundation Models, Affiliation: Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) Email: aditya.somasundaram@mbzuai.ac.ae 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 (Pethick et al., 2025). For scaled dot-product attention, the same principle yields the first-order query–key steepest descent problem underlying Compositional Muon (Keigwin et al., 2026). These connections motivate using Hilbert distance as a guiding principle for optimizer design in softmax modules. 1 Introduction An important consideration for stable training is the extent to which a single update changes the network’s intermediate representations. Each intermediate representation is both the output of one layer and the input to the next. This dependence is especially important in deep sequential models, where changing one representation alters the inputs processed by every layer that follows. Without proper control, these representation shifts can accumulate across layers and destabilize training. Muon formalizes this idea for linear layers by controlling the worst-case RMS movement of the layer’s output over inputs of bounded size (Jordan et al., 2024; Bernstein, 2025a). This requirement leads to a suitably scaled spectral norm constraint on the weight update. This spectral constraint is not arbitrary: it follows from choosing RMS to measure the functional change produced by the update. More broadly, it illustrates how the geometry of an optimizer depends on the notion of change assigned to a module. The Modula project and Scion generalize this perspective by advocating update geometries that reflect the function of each module (Bernstein, 2025b; Pethick et al., 2025). Muon’s standard update geometry, however, is derived for the linear map and does not explicitly adapt to architectural operations such as normalization and nonlinear activations. As a result, Muon does not yet provide a single optimization rule for every parameter group in the model. In standard Muon training recipes, the LM head, embeddings, and scalar or vector parameters are typically optimized with AdamW (Jordan et al., 2024). Large-scale Muon training can also require additional stabilization. For example, Kimi K2 introduced MuonClip to address excessive attention-score growth by rescaling query and key weights after an update (Team et al., 2026). We focus on how optimizer design changes when a linear map is followed by softmax. The linear layer’s output is an intermediate state, and its RMS displacement does not faithfully measure the change in the probability distribution produced by softmax. We instead directly constrain the change in the output probability distribution. We propose using Hilbert’s projective metric, or Hilbert distance, for this purpose. Our contributions are threefold: • We formulate the steepest descent problem using Hilbert distance to constrain changes in softmax output distributions. This single formulation brings two seemingly independent optimization methods under a unified mathematical framework. • For the LM head, we show that this probability constraint becomes a row diameter constraint. Its exact dual scales quadratically with vocabulary size. A tighter common-ball constraint instead leads to RowNorm in closed form (Pethick et al., 2025). • For standard attention, we show that this probability constraint produces a spectral constraint on the first-order change in the QK product, recovering the steepest descent problem posed in Compositional Muon (Keigwin et al., 2026). Beyond unifying these recent methods, we provide a guiding principle for optimizer design for any neural network module that outputs a probability distribution. We first expose the mismatch between RMS and softmax (Section 2). Next, we introduce Hilbert geometry and establish its suitability for measuring changes in softmax output distributions (Section 3). We subsequently formulate the corresponding steepest descent problems for the LM head and attention. For the LM head, a tighter approximation leads to RowNorm in closed form. For attention, our formulation recovers the QK steepest descent problem studied by Compositional Muon (Sections 4 and 5). 2 Preliminaries on steepest descent Steepest descent chooses an update that gives the greatest first-order decrease in the loss while remaining small under a chosen measure of size. For a minibatch gradient G=∇WℒG= _WL, an update ΔW W changes the loss to first order by ⟨G,ΔW⟩F G, W _F. Let W(ΔW)D_W( W) measure the size of this update at the current weights W. We consider minΔW⟨G,ΔW⟩Fsubject toW(ΔW)≤ρ, _ W G, W _F to _W( W)≤ρ, (1) where ρ>0ρ>0 sets the permitted update size. We leave WD_W unspecified for now, allowing its choice to reflect the module being optimized. The importance of the underlying choice of metric cannot be overstated. Choosing the Frobenius norm measures the update directly in parameter space and gives the SGD direction (Robbins and Monro, 1951). Muon takes a functional perspective instead: for a linear layer, it controls the worst-case change in the output produced by the update (Bernstein, 2025a). Measuring this change with RMS leads to a spectral norm constraint on the weight update, whose steepest descent solution gives the Muon direction. The Modula project generalizes this perspective by advocating update geometries that reflect the function of each module (Bernstein, 2025b). Around this principle, Scion develops a systematic optimization framework, formulating updates as Linear Minimization Oracles under ‖⋅‖p→q · _p→ q operator norm constraints (Pethick et al., 2025). Here, the norms ‖⋅‖p · _p and ‖⋅‖q · _q describe the input and output geometries of the module. The shared ethos of these works is to build optimizers whose geometry reflects the function of each module. This functional perspective becomes more delicate when a linear layer is followed by a nonlinearity. The linear output is then only an intermediate state in the computation of the full module. Controlling the RMS movement of this intermediate state does not account for the geometry introduced by the subsequent nonlinearity. For a softmax module, the linear output is the preactivation vector, which softmax maps to a probability distribution on the simplex. We therefore ask whether RMS displacement of the preactivation faithfully measures changes in the output distribution. Figure 1: A common RMS bound permits very different changes in softmax probabilities. All cases share the same centered z∈ℝ32z ^32 and initial distribution p=softmax(z)p=softmax(z). The centered perturbations produce q=softmax(z+Δz)q=softmax(z+ z) and all satisfy RMS(Δz)≤1RMS( z)≤ 1. The pairs (RMS(Δz),ℋ(p,q))(RMS( z),H(p,q)) are (1,4)(1,4) for blue, (1,2)(1,2) for purple, and (1/2,4)(1/2,4) for red, where ℋH denotes Hilbert distance. Blue substantially reshapes the distribution, while purple barely changes it. Red matches blue’s Hilbert distance at half its RMS. 2.1 Geometric mismatch between Muon and softmax We argue that Muon is inadequate for modules involving softmax. Muon’s spectral geometry can allocate updates poorly for such modules. The distinction lies between the nature of RMS and softmax: RMS measures how much the preactivation vector moves on average, while softmax measures how much one coordinate gains relative to others. A simple example. Consider an update where all the rows are equal, ΔW=b⊤ W=1b . In this case, Δz=ΔWx=(b⊤x) z= Wx=(b x)1, a scaled ones vector. Such a perturbation to the preactivation vector does not alter the probability distribution, despite incurring an arbitrarily large RMS norm. This particular case is not problematic in practice because softmax modules receive row-centered gradients, causing Muon to consider only centered directions; see Section 6. It nevertheless illustrates a broader geometric mismatch between RMS bounds and softmax that can produce practical pathologies. Figure 1 shows that this discrepancy persists under centered perturbations. Unlike the preceding example, all three cases share the same initial distribution and use centered perturbations satisfying a common RMS bound. The blue and purple updates have equal RMS, yet the blue perturbation substantially reshapes the distribution while the purple perturbation barely changes it. The red perturbation has half their RMS but still causes a substantial probability redistribution. This mismatch compels direct control of the output distribution. Designing an effective update rule for softmax modules requires choosing a suitable metric on the output distribution and formulating the induced steepest descent problem. This raises a fundamental question: What is the appropriate metric for perturbations of the softmax probabilities? 3 Hilbert geometry of softmax We turn to Hilbert’s projective metric, or simply Hilbert distance (Birkhoff, 1957; De La Harpe, 1993). For strictly positive vectors x,y∈ℝ>0nx,y _>0^n, it is defined by ℋ(x,y)=logmaxixi/yiminjxj/yj.H(x,y)= _ix_i/y_i _jx_j/y_j. We apply this metric to strictly positive probability distributions on the simplex =x∈ℝ>0n:∑ixi=1P=\x _>0^n: _ix_i=1\ (Nielsen and Sun, 2018; Nielsen and Sun, 2023). For p,q∈p,q , Hilbert distance can equivalently be written as ℋ(p,q)=maxi,j|logqi/qjpi/pj|.H(p,q)= _i,j | q_i/q_jp_i/p_j |. (2) Thus, Hilbert distance measures the largest change in pairwise odds. Hilbert’s projective metric has several useful geometric properties on positive cones and probability simplices (Lemmens and Nussbaum, 2012; Papadopoulos and Troyanov, 2014). Bounding the Hilbert distance directly controls the change in cross-entropy loss and limits probability redistribution in total variation. For softmax, the distance also admits a closed-form expression in terms of the preactivation update: it coincides with the range between the largest and smallest coordinates (the variation seminorm). Consequently, uniform shifts leave the distance invariant, mirroring the translation invariance of the softmax itself. Motivated by these geometric properties, we adopt the Hilbert distance to quantify changes in output distributions and formulate the resulting steepest descent problem. Hilbert distance bounds pairwise odds ratios. Throughout, let z∈ℝnz ^n denote the preactivation supplied to softmax, and let p=softmax(z)p=softmax(z) be the resulting output distribution. A perturbation Δz z produces the updated distribution q=softmax(z+Δz)q=softmax(z+ z). By Equation 2, bounding Hilbert distance uniformly bounds the odds ratio for every pair of outcomes: ℋ(p,q)≤ε⟺e−ε≤qi/qjpi/pj≤eε∀i,j.H(p,q)≤ e^- ≤ q_i/q_jp_i/p_j≤ e ∀ i,j. In particular, for softmax, Hilbert distance has the closed form ℋ(softmax(z),softmax(z+Δz))=maxi,j|Δzi−Δzj|=maxiΔzi−miniΔzi=‖Δz‖var.H (softmax(z),softmax(z+ z) )= _i,j z_i- z_j = _i z_i- _i z_i= z _var. (3) The variation seminorm ‖Δz‖var z _var, also called the oscillation osc(Δz) ( z), is the difference between the largest and smallest entries of Δz z. It vanishes precisely when Δz z is constant, matching the translation invariance of softmax. Hilbert distance bounds probability redistribution. Bounding pairwise odds ratios also limits the redistribution of probability mass. We measure this redistribution using total variation, TV(p,q)=12∑i|pi−qi|=maxA⊆1,…,n|∑i∈A(pi−qi)|.TV(p,q)= 12 _i p_i-q_i = _A \1,…,n\ _i∈ A(p_i-q_i) . Thus, TV(p,q)≤δTV(p,q)≤δ guarantees that no set of outcomes can gain or lose more than δ probability mass. Hilbert distance gives a sharp bound on total variation (Cohen and Fausti, 2024, Theorem 5.1), TV(p,q)≤tanh(ℋ(p,q)4).TV(p,q)≤ \! ( H(p,q)4 ). The bound is sharp even when the preactivation perturbation is fixed. For every Δz∈ℝn z ^n, supz∈ℝnTV(softmax(z),softmax(z+Δz))=tanh(‖Δz‖var4). _z ^nTV (softmax(z),softmax(z+ z) )= \! ( z _var4 ). The variation seminorm therefore determines the largest probability redistribution that a given preactivation update can produce across all starting distributions. Appendix A.2 proves this sharpness statement. Hilbert distance bounds change in cross-entropy (CE) loss. Let p and q be the prediction distributions before and after an update. For a target outcome i, the CE loss is ℓi(p)=−logpi _i(p)=- p_i, so its change is Δℓi=ℓi(q)−ℓi(p)=logpiqi _i= _i(q)- _i(p)= p_iq_i. Hilbert distance is the range of these changes, ℋ(p,q)=maxiΔℓi−miniΔℓiH(p,q)= _i _i- _i _i. It therefore measures how differently the update affects the loss across possible targets. Moreover, ∑ipi=∑iqi=1 _ip_i= _iq_i=1 implies minipiqi≤1≤maxipiqi _i p_iq_i≤ 1≤ _i p_iq_i. Taking logarithms gives miniΔℓi≤0≤maxiΔℓi _i _i≤ 0≤ _i _i. Consequently, ℋ(p,q)≤εH(p,q)≤ guarantees that the change in CE loss is at most ε , |Δℓi|≤maxjΔℓj−minjΔℓj≤ℋ(p,q)∀i. _i ≤ _j _j- _j _j (p,q) ∀ i. Steepest descent problem. Together, these results show that Hilbert distance controls pairwise odds, probability redistribution, and changes in cross-entropy loss. Having established these properties, we now formulate the steepest descent problem using Hilbert distance to measure the size of an update. We return to Equation 1 and express the update constraint in terms of this distance. Let W denote the module’s parameters and let pW(x)p_W(x) be the output distribution produced at input x. An update ΔW W changes this distribution from pW(x)p_W(x) to pW+ΔW(x)p_W+ W(x). Requiring these distributions to remain within Hilbert distance ε for every input in a prescribed set X gives minΔW⟨∇Wℒ,ΔW⟩Fsubject tosupx∈ℋ(pW(x),pW+ΔW(x))≤ε. _ W _WL, W _F to _x H (p_W(x),p_W+ W(x) )≤ . (4) 4 The LM Head For the LM head, bounding Hilbert distance over the hidden state space is equivalent to bounding the diameter of the update rows. The resulting steepest descent problem has an exact dual whose size grows quadratically with the vocabulary. A tighter common-ball approximation avoids this coupling and leads to RowNorm in closed form. We first specialize Equation 4 to the LM head. For a hidden state h∈ℝdh ^d, the LM head produces the logit vector z=Wh∈ℝVz=Wh ^V and the token distribution p=softmax(z)p=softmax(z). We consider an untied head without bias, trained with CE and no other loss term depending on W. Throughout this section, we compare the head before and after an update at the same hidden state. Hilbert distance induces a row diameter constraint. Let ΔW W be a weight update with rows Δwi⊤ w_i . The update changes the logit for token i by Δwi⊤h w_i h. For two tokens i and j, their logit difference therefore changes by (Δwi−Δwj)⊤h( w_i- w_j) h. Equation 3 gives ℋ(pW(h),pW+ΔW(h))=‖ΔWh‖var=maxi,j|(Δwi−Δwj)⊤h|.H (p_W(h),p_W+ W(h) )= Wh _var= _i,j ( w_i- w_j) h . For a fixed pair of rows, the largest value of this inner product over ‖h‖2≤R h _2≤ R is R‖Δwi−Δwj‖2R w_i- w_j _2. Taking the largest value over all row pairs gives sup‖h‖2≤Rℋ(pW(h),pW+ΔW(h))=Rmaxi,j‖Δwi−Δwj‖2=RD(ΔW). _ h _2≤ RH (p_W(h),p_W+ W(h) )=R _i,j w_i- w_j _2=RD( W). (5) The maximum pairwise distance is precisely the diameter D(ΔW)D( W) of the row cloud Δw1,…,ΔwV\ w_1,…, w_V\ (Rudin, 1976, Definition 3.9). For this input ball, the general constraint therefore specializes to W(ΔW)=RD(ΔW)D_W( W)=RD( W). Choosing R=dR= d gives the RMS unit ball. Thus, requiring a Hilbert distance of at most ε throughout this ball is equivalent to imposing D(ΔW)≤ρheadD( W)≤ _head, where ρhead=ε/R _head= /R. This single constraint on the update rows controls both the predicted distribution and the cross-entropy loss. By Section 3, no group of tokens can gain or lose more than tanh(ε/4) ( /4) probability mass, and the CE loss changes by at most ε for any target token. Steepest descent under the row diameter constraint. We seek the update with the greatest first-order loss decrease under the row diameter bound. Let G=∇WℒG= _WL have rows gi⊤g_i . Using the row diameter constraint in Equation 5, we obtain minΔW⟨G,ΔW⟩Fsubject toD(ΔW)≤ρhead. _ W G, W _F to D( W)≤ _head. (6) Since the loss depends on W only through the softmax output, its gradient satisfies ⊤G=01 G=0. Adding the same vector to every update row produces a common shift of the logits and therefore leaves both the output distribution and the loss unchanged. Because pairwise row differences are also unchanged, the diameter constraint places no restriction on the common translation of the update rows. Appendix A.3 gives an exact dual characterization of this problem. The dual is a vector-valued optimal transport problem (Ciosmak, 2021). In our setting, this transport problem introduces a d-dimensional flow variable for each of the (V2) V2 token pairs. This quadratic growth with vocabulary size makes solving the flow problem computationally expensive. We proceed on a quest to find a cheaper approximation. A common-ball approximation leads to RowNorm. We ask for a tighter constraint, requiring all update rows to lie in a common closed Euclidean ball ℬ(ζ,ρhead/2)B(ζ, _head/2). With this restriction, Equation 6 becomes minΔW,ζ⟨G,ΔW⟩Fsubject toΔwi∈ℬ(ζ,ρhead/2)∀i. _ W,ζ G, W _F to w_i (ζ, _head/2) ∀ i. (7) Since the gradient rows sum to zero, translating every update row by the same vector leaves the objective unchanged. Using this freedom, we invoke a change of variable ΔW^=ΔW−ζ⊤ W= W-1ζ , which expresses each row relative to the ball’s center. Each row Δw^i w_i then lies in ℬ(0,ρhead/2)B(0, _head/2). Equation 7 therefore has the same minimum as minΔW^⟨G,ΔW^⟩Fsubject toΔw^i∈ℬ(0,ρhead/2)∀i. _ W G, W _F to w_i (0, _head/2) ∀ i. (8) Conversely, any ΔW W satisfying this constraint gives a feasible common-ball update by choosing ΔW=ΔW W= W and ζ=0ζ=0. A natural curiosity would be the suboptimality gap introduced by the tighter ball constraint. Let Δ⋆ _ and Δball _ball denote the greatest first-order decreases under the diameter and common-ball constraints, respectively. These are the negatives of the minimum objective values in Equations 6 and 7. For G≠0G≠ 0 and ρhead>0 _head>0, let k=rank(G)k=rank(G). Jung’s theorem (Jung, 1901) bounds the radius needed to enclose a set of rows with a given diameter. Applying it to the two optimization problems bounds the relative suboptimality gap by 0≤Δ⋆−ΔballΔ⋆≤1−k+12k.0≤ _ - _ball _ ≤ 1- k+12k. (9) The bound is sharp. Appendix A.4 provides its proof and elaborates on the equivalence brought about by the change of variable. To find an update with this guarantee, we solve Equation 8 one row at a time. Each row has its own constraint and contributes independently to the objective. For a feasible row Δw^i w_i, the Cauchy–Schwarz inequality and the radius constraint give ⟨gi,Δw^i⟩≥−‖gi‖2‖Δw^i‖2≥−ρhead2‖gi‖2. g_i, w_i ≥- g_i _2 w_i _2≥- _head2 g_i _2. For gi≠0g_i≠ 0, equality is attained by choosing Δw^i w_i opposite to gig_i with norm ρhead/2 _head/2. For a zero gradient row, we choose Δw^i=0 w_i=0. Let RN normalize each nonzero row to unit Euclidean norm and leave zero rows unchanged. Choosing ΔW=ΔW W= W and ζ=0ζ=0, we obtain one minimizing update for the common-ball problem, ΔWRN=−ρhead2(G). W_RN=- _head2 RN(G). (10) This update is feasible for the original diameter problem and attains the common-ball optimum. Its first-order decrease is Δball _ball, so Equation 9 applies directly to RowNorm. A practical version subtracts the mean row from ΔWRN W_RN to obtain projected RowNorm. This common translation preserves the objective, the row diameter, and the output probabilities. Our Hilbert distance argument therefore independently recovers the direction derived by Lau and Su (2026) from the symmetries of the LM head. Appendix A.5 gives an implementation with momentum and numerical stabilization. 5 The attention module We extend our principle from the LM head to attention. Attention presents a more intricate setting: its softmax probabilities arise from bilinear scores that jointly depend on the query and key weights. Bounding the Hilbert distance between attention distributions leads us to the QK constraint underlying Compositional Muon. Attention distributions. Attention uses softmax to distribute probability across the positions visible to each query. We study how this distribution changes when the query and key weights are updated, holding the input hidden states h1,…,hT∈ℝdh_1,…,h_T ^d fixed. For a single head with query and key dimension dhd_h, the weights WQ,WK∈ℝd×dhW_Q,W_K ^d× d_h determine the attention scores through M=WQWK⊤∈ℝd×d,sij=hi⊤Mhjdh.M=W_QW_K ^d× d, s_ij= h_i Mh_j d_h. We consider scaled dot-product attention without rotary position transforms or query–key normalization. Under causal masking, let pijp_ij denote the probability assigned by query i to position j∈i=1,…,ij _i=\1,…,i\. We measure Hilbert distance on this visible set, where all probabilities are positive. Hilbert distance bounds changes in score differences. Let ΔM M denote the change in the QK product, so that the updated product is M+ΔM+ M. At position i, write pi=(pij)j∈ip_i=(p_ij)_j _i for the original attention distribution and pi+p_i^+ for the distribution after this change. Holding the hidden states fixed, the score sijs_ij changes by Δsij=hi⊤(M+ΔM)hj−hi⊤Mhjdh=hi⊤ΔMhjdh. s_ij= h_i (M+ M)h_j-h_i Mh_j d_h= h_i Mh_j d_h. Within row i, the score difference between positions j and k therefore changes by hi⊤ΔM(hj−hk)/dhh_i M(h_j-h_k)/ d_h. Applying Equation 3 within each attention row gives ℋ(pi,pi+)=maxj,k∈i|Δsij−Δsik|=1dhmaxj,k∈i|hi⊤ΔM(hj−hk)|.H(p_i,p_i^+)= _j,k _i s_ij- s_ik = 1 d_h _j,k _i h_i M(h_j-h_k) . (11) Intuitively, the probability constraint, specific to the attention mechanism, asks how much an update can change a query’s preference between two keys. The induced spectral norm constraint. Assume that every hidden state satisfies ‖hℓ‖2≤R h_ _2≤ R, with R=dR= d corresponding to the RMS unit ball. The hidden state hih_i has norm at most R, while the difference hj−hkh_j-h_k has norm at most 2R2R. Consequently, |hi⊤ΔM(hj−hk)|≤‖hi‖2‖ΔM‖op‖hj−hk‖2≤2R2‖ΔM‖op. h_i M(h_j-h_k) ≤ h_i _2 M _op h_j-h_k _2≤ 2R^2 M _op. For causal sequences with T≥3T≥ 3, the worst-case Hilbert distance over this input set is sup‖hℓ‖2≤R1≤ℓ≤Tmax1≤i≤Tℋ(pi,pi+)=2R2dh‖ΔM‖op. _ subarrayc h_ _2≤ R\\ 1≤ ≤ T subarray _1≤ i≤ TH(p_i,p_i^+)= 2R^2 d_h M _op. (12) Requiring a Hilbert distance of at most ε for every attention row throughout this input set is therefore equivalent to ‖ΔM‖op≤ρQK,ρQK=εdh2R2. M _op≤ _QK, _QK= d_h2R^2. (13) Why does the same Hilbert distance principle give a row diameter for the LM head and a spectral norm for attention? In the LM head, each comparison involves a single hidden state, whereas attention involves a query input and a difference between key inputs. Under the stated norm bounds, these two vectors can be aligned independently to attain the spectral norm of ΔM M. The bound is attained when hih_i is R times a unit leading left singular vector of ΔM M, and hj−hkh_j-h_k is 2R2R times the corresponding unit right singular vector. Steepest descent over query and key updates. Let ΔWQ W_Q and ΔWK W_K denote updates to the query and key weights WQW_Q and WKW_K, respectively. Expanding the updated product gives ΔM=(WQ+ΔWQ)(WK+ΔWK)⊤−WQWK⊤=ΔWQWK⊤+WQΔWK⊤⏟ΔMlin+ΔWQΔWK⊤. M=(W_Q+ W_Q)(W_K+ W_K) -W_QW_K = W_QW_K +W_Q W_K _ M_lin+ W_Q W_K . The first two terms give the first-order change ΔMlin(ΔWQ,ΔWK) M_lin( W_Q, W_K), while ΔWQΔWK⊤ W_Q W_K is second order in the updates and is therefore omitted. Applying Equation 13 to this first-order change gives ‖ΔWQWK⊤+WQΔWK⊤‖op≤ρQK. W_QW_K +W_Q W_K _op≤ _QK. We assume that these weights affect the loss only through the attention probabilities. Write GQ=∇WQℒG_Q= _W_QL and GK=∇WKℒG_K= _W_KL for their gradients. Constructing our steepest descent problem under this Hilbert distance bound, we obtain minΔWQ,ΔWK _ W_Q, W_K ⟨GQ,ΔWQ⟩F+⟨GK,ΔWK⟩F G_Q, W_Q _F+ G_K, W_K _F (14) subject to to ‖ΔWQWK⊤+WQΔWK⊤‖op≤ρQK. W_QW_K +W_Q W_K _op≤ _QK. Equation 14 is precisely the first-order QK steepest descent problem posed by Keigwin et al. (2026, Section 1.2). Their derivation controls changes in the composed bilinear map through the spectral norm of its first-order perturbation. Our starting point is a bound on the Hilbert distance between attention distributions. Translating this probability bound through softmax and the bilinear scores recovers the same constraint on the first-order QK product change. Our Hilbert distance derivation therefore provides a natural prelude to Compositional Muon, motivating its QK constraint through the stability of the attention distributions. The LM head and attention affect optimization through both forward probability distributions and the gradients propagated to the hidden states. Our main analysis bounds changes in the former and derives the corresponding steepest descent problems. Holding the inputs and incoming gradients fixed, Appendix B bounds the resulting perturbations of the latter. 6 Discussion Can Muon be adapted to softmax simply by removing the common shifts that softmax ignores? Because softmax is invariant to common shifts, one might try adapting Muon by centering the update rows before measuring their spectral norm. We show that the resulting steepest descent problem has the same minimum and that the Muon update is already row centered. Translation symmetry alone therefore does not alter Muon, motivating the stronger relative control provided by Hilbert distance. Centering leaves Muon unchanged. Consider a linear module with n outputs and a parameter matrix whose influence on the loss is entirely through the softmax output distribution. Let P=In−⊤/nP=I_n-11 /n denote the orthogonal projection that subtracts the mean row. Since gradient rows sum to zero, PG=GPG=G. Since P is an orthogonal projection, ‖PΔW‖op≤‖P‖op‖ΔW‖op≤‖ΔW‖op P W _op≤ P _op W _op≤ W _op. Thus, every update satisfying ‖ΔW‖op≤ρ W _op≤ρ also satisfies ‖PΔW‖op≤ρ P W _op≤ρ. Conversely, take any update ΔW W satisfying ‖PΔW‖op≤ρ P W _op≤ρ. The centered update PΔWP W has spectral norm at most ρ. It also preserves the objective, since ⟨G,PΔW⟩F=⟨PG,ΔW⟩F=⟨G,ΔW⟩F G,P W _F= PG, W _F= G, W _F. The additional freedom therefore leaves the minimum unchanged, min‖PΔW‖op≤ρ⟨G,ΔW⟩F=min‖ΔW‖op≤ρ⟨G,ΔW⟩F. _ P W _op≤ρ G, W _F= _ W _op≤ρ G, W _F. The Muon update is already centered. Indeed, if G=AΣB⊤G=A B is the compact SVD of a nonzero gradient, then every left singular vector satisfies ⊤aj=(⊤Gbj)/σj=01 a_j=(1 Gb_j)/ _j=0, giving ΔWMuon=−ρAB⊤,PΔWMuon=−ρ(PA)B⊤=ΔWMuon. W_Muon=-ρ AB , P W_Muon=-ρ(PA)B = W_Muon. Accounting for translation symmetry alone therefore recovers the standard Muon solution. Our Hilbert distance approach goes further by bounding the largest change in the difference between any two preactivation coordinates. For a centered preactivation perturbation Δz∈ℝn z ^n, with p=softmax(z)p=softmax(z) and q=softmax(z+Δz)q=softmax(z+ z), we have 2‖Δz‖RMS≤ℋ(p,q)≤2n‖Δz‖RMS2 z _RMS (p,q)≤ 2n\, z _RMS. Since the upper bound is sharp, a fixed bound on centered RMS permits Hilbert distance to grow with the output width. Appendix A.1 proves these bounds and discusses their sharpness. Limitations. The main objective of this work is to provide a unifying mathematical foundation for softmax modules. Through Hilbert distance, we independently recover RowNorm for the LM head and the first-order QK constraint underlying Compositional Muon. These mathematical connections naturally raise the question of what this perspective offers in practice. We leave that question to future experimental work and make no empirical claim of improved training speed or final model quality. The work of Lau and Su (2026) on RowNorm and that of Keigwin et al. (2026) on Compositional Muon report improvements in language-model pretraining and include ablations of their optimizer choices. These experiments provide existing empirical evidence for the practical relevance of the methods connected by our framework. A second limitation arises from the relative nature of Hilbert distance. Because it controls pairwise odds, it can be sensitive to large relative changes in low probability outcomes even when little probability mass moves. For example, suppose an update raises the probability of one outcome from 10−910^-9 to 10−610^-6 while leaving the rest of the distribution nearly unchanged. The probability of this outcome increases by a factor of 10001000, producing a large Hilbert distance despite the small movement of probability mass, TV(p,q)=9.99×10−7,ℋ(p,q)≈6.91.TV(p,q)=9.99× 10^-7,H(p,q)≈ 6.91. A small Hilbert distance constraint can therefore exclude updates that move very little probability mass. Future work could investigate metrics that preserve the useful guarantees developed here while being less sensitive to low probability outcomes. 7 Related work From Adam to module specific optimizers. Adam has long served as the default optimizer across deep learning (Kingma and Ba, 2014). Muon has recently challenged this convention (Jordan et al., 2024). For a linear layer, Muon controls the worst-case change in its output as measured by RMS, resulting in a bound on the spectral norm of the weight update. The resulting steepest descent problem forms the mathematical foundation of the Muon update (Bernstein, 2025a). The Modula project generalizes this perspective and advocates constructing optimizers whose geometry reflects the function of each module (Bernstein, 2025b; Large et al., 2024). Modular duality and Scion provide systematic ways to construct updates from such geometries (Bernstein and Newhouse, 2024; Pethick et al., 2025). We extend this principle to modules involving softmax. Geometry on probability distributions. A related line of work connects parameter updates to the probability distributions they induce. Natural gradient equips a family of probability distributions with the Fisher geometry and chooses updates according to the resulting notion of steepest descent (Amari, 1998; Martens, 2020). TRPO makes this connection explicit by constraining the expected KL divergence between the action distributions induced by the original and updated policies (Schulman et al., 2015). Hilbert’s projective metric offers a different geometry for positive vectors. Beginning with the contraction theory of positive operators, this geometry has subsequently appeared in matrix scaling, clustering, nonlinear embeddings, and the study of probability simplices (Birkhoff, 1957; De La Harpe, 1993; Cuturi, 2013; Nielsen and Sun, 2018; Nielsen and Sun, 2023). More recently, Cohen and Fausti (2024) established a sharp relationship between Hilbert distance and total variation. We draw on this literature to study the parameter geometry induced by Hilbert distance on the probability distributions produced by softmax. Optimizers for the LM head and attention. The LM head and attention have each motivated specialized optimization methods. For the LM head, Scion (Pethick et al., 2025) includes RowNorm and SCALE (Glentis et al., 2026) proposes a momentum-based variant; Lau and Su (2026) derive projected RowNorm from its permutation and shared-shift symmetries. Attention has motivated a different collection of interventions. Query and key normalization controls the scale of attention scores, while spectral reparameterization addresses attention entropy collapse (Henry et al., 2020; Zhai et al., 2023). MuonClip monitors the maximum attention logits and conditionally rescales the query and key projection weights after a Muon update (Team et al., 2026). Compositional Muon goes further by constraining the first-order change in the product of the query and key matrices (Keigwin et al., 2026). Although row normalization for the LM head and product control for attention appear to be independent design choices, we show that they share common ground in bounding changes in the probability distributions produced by softmax. 8 Conclusion We develop a unified framework for optimizing composed softmax modules by bounding changes in their output probabilities and deriving the corresponding parameter constraints. This addresses a limitation of Muon’s RMS control, which can conceal substantial probability redistribution when applied to the intermediate linear output. We propose using Hilbert distance to constrain these changes in the output distribution directly. Hilbert distance respects the common-shift invariance of softmax and controls both the redistribution of probability mass and changes in cross-entropy loss. Its representation as the variation seminorm of the preactivation perturbation lets us translate a common constraint on the output distribution into a parameter geometry for each module. For the LM head, this gives a row diameter constraint. A tighter ball approximation leads to RowNorm, with centering giving its projected form. For attention, the same principle recovers the first-order QK steepest descent problem underlying Compositional Muon. A single probability constraint thus connects the separately developed RowNorm update and Compositional Muon’s QK optimization problem. We offer this framework as a contribution to the broader Modula perspective, demonstrating how the output geometry of composed softmax modules can guide optimizer design. We hope these connections spark new ideas for optimizers and architectures built around composed modules. Hilbert distance provides a starting point for softmax, and we encourage researchers to ask which output geometries could play that role for other compositions. AI use statement The idea of bounding changes in softmax output probabilities originated with the authors. The choice of Hilbert distance was found by studying many existing distances in the literature. AI-assisted mathematical exploration also informed the selection of Hilbert distance. AI tools were extensively used to survey literature, identify related work, check mathematical arguments, and verify correctness. In particular, the authors initially obtained a weaker approximation bound for RowNorm relative to the diameter-constrained optimum. AI tools subsequently identified Jung’s theorem and its applicability in tightening this bound. The initial manuscript was written by the authors. AI tools were subsequently used to improve wording, technical exposition, and presentation, and to provide critical feedback on the manuscript. The authors reviewed all AI assisted editorial changes before incorporating them. We, the authors, take responsibility for the final content of the paper, including its mathematical claims, proofs, and references. Ethics statement This work develops a unified framework for optimizing modules which involve softmax. The research presented here does not involve human participants, personal data, or the collection or release of datasets. The framework is general purpose, and all ethical considerations in downstream applications depend on the training data, models, and deployment contexts in which it is used. Reproducibility statement The assumptions underlying our theoretical results are stated alongside the derivations in the main text. Appendix A provides the sharpness argument for the total variation bound, the exact dual formulation of the diameter-constrained problem, and the proof of the sharp RowNorm approximation guarantee. Appendix B derives the forward and backward perturbation bounds for the LM head and attention. Appendix A.5 specifies the practical RowNorm update, including centering, momentum, and numerical stabilization. To reproduce Figure 1, Appendix C gives deterministic formulas for the initial logits and all three updates, together with their RMS norms, Hilbert distances, and total variation values. Author Contributions If you’d like to, you may include a section for author contributions as is done in many journals. This is optional and at the discretion of the authors. Acknowledgments Use unnumbered third level headings for the acknowledgments. All acknowledgments, including those to funding agencies, go at the end of the paper. References Amari (1998) S. Amari Natural gradient works efficiently in learning. Neural computation 10 (2), p. 251–276. Cited by: §7. Bernstein and Newhouse (2024) J. Bernstein and L. Newhouse Modular duality in deep learning. arXiv preprint arXiv:2410.21265. Cited by: §7. Bernstein (2025a) J. Bernstein Deriving muon. External Links: Link Cited by: §1, §2, §7. Bernstein (2025b) J. Bernstein The modula docs. External Links: Link Cited by: §1, §2, §7. Birkhoff (1957) G. Birkhoff Extensions of jentzsch’s theorem. Transactions of the American Mathematical Society 85 (1), p. 219–227. Cited by: §3, §7. Ciosmak (2021) K. J. Ciosmak Optimal transport of vector measures. Calculus of Variations and Partial Differential Equations 60 (6), p. 230. Cited by: §A.3, §4. Cohen and Fausti (2024) S. N. Cohen and E. Fausti Hyperbolic contractivity and the hilbert metric on probability measures. External Links: 2309.02413, Link Cited by: §A.2, Appendix B, §3, §7. Cuturi (2013) M. Cuturi Sinkhorn distances: lightspeed computation of optimal transport. Advances in neural information processing systems 26. Cited by: §7. De La Harpe (1993) P. De La Harpe On hilbert’s metric for simplices. Geometric group theory 1, p. 97–119. Cited by: §3, §7. Glentis et al. (2026) A. Glentis, J. Li, A. Han, and M. Hong Memory-efficient llm pretraining via minimalist optimizer design. External Links: 2506.16659, Link Cited by: §7. Henry et al. (2020) A. Henry, P. R. Dachapally, S. S. Pawar, and Y. Chen Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, p. 4246–4253. Cited by: §7. Jordan et al. (2024) K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, and J. Bernstein Muon: an optimizer for hidden layers in neural networks. External Links: Link Cited by: §1, §1, §7. Jung (1901) H. Jung Ueber die kleinste kugel, die eine räumliche figur einschliesst.. Journal für die reine und angewandte Mathematik (Crelles Journal) 1901 (123), p. 241–257. Cited by: §A.4.2, §4. Keigwin et al. (2026) B. Keigwin, L. Yang, Y. Zhang, D. Pai, and A. DeWulf Towards compositional steepest descent. External Links: Link Cited by: 3rd item, §5, §6, §7, Abstract. Kingma and Ba (2014) D. P. Kingma and J. Ba Adam: a method for stochastic optimization. arXiv preprint arXiv:1412.6980. Cited by: §7. Large et al. (2024) T. Large, Y. Liu, M. Huh, H. Bahng, P. Isola, and J. Bernstein Scalable optimization in the modular norm. External Links: 2405.14813, Link Cited by: §7. Lau and Su (2026) T. T. Lau and W. Su Symmetry-compatible principle for optimizer design: embeddings, lm heads, swiglu mlps, and moe routers. External Links: 2605.18106, Link Cited by: §4, §6, §7. Lemmens and Nussbaum (2012) B. Lemmens and R. D. Nussbaum Nonlinear perron-frobenius theory. Cambridge University Press. Cited by: §3. Martens (2020) J. Martens New insights and perspectives on the natural gradient method. Journal of Machine Learning Research 21 (146), p. 1–76. Cited by: §7. Nielsen and Sun (2018) F. Nielsen and K. Sun Clustering in hilbert’s projective geometry: the case studies of the probability simplex and the elliptope of correlation matrices. In Geometric Structures of Information, p. 297–331. External Links: ISBN 9783030025205, ISSN 1860-4870, Link, Document Cited by: §3, §7. Nielsen and Sun (2023) F. Nielsen and K. Sun Non-linear embeddings in hilbert simplex geometry. In Topological, Algebraic and Geometric Learning Workshops 2023, p. 254–266. Cited by: §3, §7. Papadopoulos and Troyanov (2014) A. Papadopoulos and M. Troyanov Handbook of hilbert geometry. European Mathematical Society. Cited by: §3. Pethick et al. (2025) T. Pethick, W. Xie, K. Antonakopoulos, Z. Zhu, A. Silveti-Falls, and V. Cevher Training deep learning models with norm-constrained lmos. External Links: 2502.07529, Link Cited by: 2nd item, §1, §2, §7, §7, Abstract. Robbins and Monro (1951) H. Robbins and S. Monro A stochastic approximation method. The annals of mathematical statistics, p. 400–407. Cited by: §2. Rudin (1976) W. Rudin Principles of mathematical analysis. 3 edition, International Series in Pure and Applied Mathematics, McGraw-Hill. External Links: Link Cited by: §4. Schulman et al. (2015) J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz Trust region policy optimization. In International conference on machine learning, p. 1889–1897. Cited by: §7. Team et al. (2026) K. Team, Y. Bai, Y. Bao, Y. Charles, C. Chen, G. Chen, H. Chen, H. Chen, J. Chen, N. Chen, R. Chen, Y. Chen, Y. Chen, Y. Chen, Z. Chen, J. Cui, H. Ding, M. Dong, A. Du, C. Du, D. Du, Y. Du, Y. Fan, Y. Feng, K. Fu, B. Gao, C. Gao, H. Gao, P. Gao, T. Gao, Y. Ge, S. Geng, Q. Gu, X. Gu, L. Guan, H. Guo, J. Guo, X. Hao, T. He, W. He, W. He, Y. He, C. Hong, H. Hu, Y. Hu, Z. Hu, W. Huang, Z. Huang, Z. Huang, T. Jiang, Z. Jiang, X. Jin, Y. Kang, G. Lai, C. Li, F. Li, H. Li, M. Li, W. Li, Y. Li, Y. Li, Y. Li, Z. Li, Z. Li, H. Lin, X. Lin, Z. Lin, C. Liu, C. Liu, H. Liu, J. Liu, J. Liu, L. Liu, S. Liu, T. Y. Liu, T. Liu, W. Liu, Y. Liu, Y. Liu, Y. Liu, Y. Liu, Z. Liu, E. Lu, H. Lu, L. Lu, Y. Luo, S. Ma, X. Ma, Y. Ma, S. Mao, J. Mei, X. Men, Y. Miao, S. Pan, Y. Peng, R. Qin, Z. Qin, B. Qu, Z. Shang, L. Shi, S. Shi, F. Song, J. Su, Z. Su, L. Sui, X. Sun, F. Sung, Y. Tai, H. Tang, J. Tao, Q. Teng, C. Tian, C. Wang, D. Wang, F. Wang, H. Wang, H. Wang, J. Wang, J. Wang, J. Wang, S. Wang, S. Wang, S. Wang, X. Wang, Y. Wang, Y. Wang, Y. Wang, Y. Wang, Y. Wang, Z. Wang, Z. Wang, Z. Wang, Z. Wang, C. Wei, Q. Wei, H. Wu, W. Wu, X. Wu, Y. Wu, C. Xiao, J. Xie, X. Xie, W. Xiong, B. Xu, J. Xu, L. H. Xu, L. Xu, S. Xu, W. Xu, X. Xu, Y. Xu, Z. Xu, J. Xu, J. Xu, J. Yan, Y. Yan, H. Yang, X. Yang, Y. Yang, Y. Yang, Z. Yang, Z. Yang, Z. Yang, H. Yao, X. Yao, W. Ye, Z. Ye, B. Yin, L. Yu, E. Yuan, H. Yuan, M. Yuan, S. Yuan, H. Zhan, D. Zhang, H. Zhang, W. Zhang, X. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Z. Zhang, H. Zhao, Y. Zhao, Z. Zhao, H. Zheng, S. Zheng, L. Zhong, J. Zhou, X. Zhou, Z. Zhou, J. Zhu, Z. Zhu, W. Zhuang, and X. Zu Kimi k2: open agentic intelligence. External Links: 2507.20534, Link Cited by: §1, §7. Vandenberghe and Boyd (2004) L. Vandenberghe and S. Boyd Convex optimization. Vol. 1, Cambridge university press Cambridge. Cited by: §A.3. Zhai et al. (2023) S. Zhai, T. Likhomanenko, E. Littwin, D. Busbridge, J. Ramapuram, Y. Zhang, J. Gu, and J. M. Susskind Stabilizing transformer training by preventing attention entropy collapse. In International conference on machine learning, p. 40770–40803. Cited by: §7. Appendix A Proofs and further connections A.1 Comparing centered RMS and Hilbert distance This subsection proves the comparison used in Section 6. For centered preactivation perturbations, Hilbert distance bounds RMS independently of the output width, whereas the reverse bound necessarily grows with the width. Let z∈ℝnz ^n be a preactivation vector and let Δz∈ℝn z ^n satisfy ⊤Δz=01 z=0. Define p=softmax(z)p=softmax(z), q=softmax(z+Δz)q=softmax(z+ z), and ‖Δz‖RMS=‖Δz‖2/n z _RMS= z _2/ n. We prove 2‖Δz‖RMS≤ℋ(p,q)≤2n‖Δz‖RMS.2 z _RMS (p,q)≤ 2n\, z _RMS. (15) The softmax identity gives ℋ(p,q)=maxi,j|Δzi−Δzj|H(p,q)= _i,j z_i- z_j . Applying the Cauchy–Schwarz inequality gives |Δzi−Δzj|≤2Δzi2+Δzj2≤2∑k=1nΔzk2=2n‖Δz‖RMS,i≠j. z_i- z_j ≤ 2 z_i^2+ z_j^2≤ 2 _k=1^n z_k^2= 2n z _RMS, i≠ j. Taking the maximum over i and j proves the upper bound. Equality is attained by Δz=(a,−a,0,…,0) z=(a,-a,0,…,0) for any a>0a>0, so the factor 2n 2n is sharp. For the lower bound, let m=miniΔzim= _i z_i, M=maxiΔziM= _i z_i, and c=(m+M)/2c=(m+M)/2. Since ∑iΔzi=0 _i z_i=0, ‖Δz‖RMS2≤‖Δz‖RMS2+c2=1n∑i(Δzi−c)2≤maxi(Δzi−c)2=(M−m2)2=ℋ(p,q)24. z _RMS^2≤ z _RMS^2+c^2= 1n _i( z_i-c)^2≤ _i( z_i-c)^2= ( M-m2 )^2= H(p,q)^24. Dropping the nonnegative term c2c^2 proves the lower bound. When n is even, equality is attained by taking half the coordinates equal to a and half equal to −a-a. Equation 15 reveals an asymmetry between the two measures. Hilbert distance bounds centered RMS with a constant independent of n, but a fixed RMS bound permits Hilbert distance to grow as n n. A similar dependence appears in the standard comparison between the spectral and Frobenius norms: for a matrix A of rank r, ‖A‖op≤‖A‖F≤r‖A‖op. A _op≤ A _F≤ r\, A _op. In both comparisons, an aggregate Euclidean measure can conceal an extremal change whose magnitude grows with dimension or rank. Appendix A.2 shows that, for softmax, this separation can persist at the level of the output distribution: centered RMS can vanish while total variation approaches 1/21/2. A.2 Probability redistribution under softmax This subsection establishes two consequences of the Hilbert geometry developed in the main text. First, we show that the total variation bound is sharp for every fixed preactivation perturbation. We then construct centered perturbations whose RMS tends to zero while their probability redistribution remains bounded away from zero as the output width grows. The sharp total variation bound. Recall that the variation seminorm is γ=‖Δz‖var=maxℓΔzℓ−minℓΔzℓγ= z _var= _ z_ - _ z_ . For p=softmax(z)p=softmax(z) and q=softmax(z+Δz)q=softmax(z+ z), Equation 3 and the sharp bound of Cohen and Fausti (2024, Theorem 5.1) give TV(p,q)≤tanh(γ/4)TV(p,q)≤ (γ/4). To approach equality when γ>0γ>0, choose i,ji,j with Δzi−Δzj=γ z_i- z_j=γ and set zi=−γ/2z_i=-γ/2, zj=0z_j=0. The difference between these preactivation coordinates therefore changes from −γ/2-γ/2 to γ/2γ/2. Let all other preactivations tend to −∞-∞, concentrating both distributions on i,ji,j. The probability gained by coordinate i then gives TV(p,q)⟶eγ/21+eγ/2−11+eγ/2=tanh(γ4).TV(p,q) e^γ/21+e^γ/2- 11+e^γ/2= \! ( γ4 ). This approaches the upper bound through strictly positive softmax distributions, proving the claimed supremum. For γ=0γ=0, the update is a common shift and both sides are zero. Vanishing RMS with nonvanishing probability redistribution. We now show that centering alone does not make RMS a uniform measure of probability change. Take n≥3n≥ 3, initial preactivation vector z=0z=0, and an=log(n−1)a_n= (n-1). The initial probabilities and centered update are p=n,Δz=an(1−1n,−1n,…,−1n)⊤,⊤Δz=0.p= 1n, z=a_n (1- 1n,- 1n,…,- 1n ) , 1 z=0. By softmax’s common-shift invariance, q=softmax(Δz)q=softmax( z) satisfies q1=eanean+n−1=12,qi=1ean+n−1=12(n−1)(i>1).q_1= e^a_ne^a_n+n-1= 12, q_i= 1e^a_n+n-1= 12(n-1) (i>1). Yet, as n→∞n→∞, ‖Δz‖RMS=n−1nlog(n−1)⟶0,TV(p,q)=12−1n⟶12. z _RMS= n-1n (n-1) 0, (p,q)= 12- 1n 12. Even after centering, vanishing RMS does not ensure vanishing probability redistribution uniformly across widths. Meanwhile, ℋ(p,q)=log(n−1)H(p,q)= (n-1) captures the growing relative change. A.3 The exact dual of the diameter constraint This subsection derives the exact dual of the row diameter problem in Equation 6. The dual represents each gradient row through antisymmetric vectors assigned to pairs of vocabulary rows and minimizes their total Euclidean norm. We first establish the balance constraints and the dual identity, then relate the resulting problem to vector transport. Let G∈ℝV×dG ^V× d have rows gi⊤g_i , where V≥2V≥ 2 and ∑igi=0 _ig_i=0. For ρhead>0 _head>0, consider minΔW⟨G,ΔW⟩Fsubject to‖Δwi−Δwj‖2≤ρhead∀i<j. _ W G, W _F to w_i- w_j _2≤ _head ∀ i<j. (16) Pairwise dual variables. For every pair i<ji<j, introduce a vector φij∈ℝd _ij ^d. Extend this collection by setting φji=−φij _ji=- _ij and φii=0 _i=0. The dual feasibility conditions are ∑j=1Vφij=gi,i=1,…,V. _j=1^V _ij=g_i, i=1,…,V. (17) Antisymmetry implies ∑i∑jφij=∑i<j(φij+φji)=0. _i _j _ij= _i<j( _ij+ _ji)=0. Hence, the condition ∑igi=0 _ig_i=0 is necessary for Equation 17. It is also sufficient. For example, setting φiV=gi _iV=g_i for i<Vi<V and every other independently indexed variable to zero satisfies all balance equations because gV=−∑i<Vgig_V=- _i<Vg_i. The exact dual identity. The primal problem in Equation 16 satisfies minD(ΔW)≤ρhead⟨G,ΔW⟩F=−ρheadmin∑i<jΦ∑jφij=gi∀i‖φij‖2. _D( W)≤ _head G, W _F=- _head _ subarrayc \\ _j _ij=g_i\;∀ i subarray _i<j _ij _2. (18) Both minima are attained. Consequently, the greatest first order decrease under the diameter constraint is Δ⋆=ρheadmin∑i<jΦ∑jφij=gi∀i‖φij‖2. _ = _head _ subarrayc \\ _j _ij=g_i\;∀ i subarray _i<j _ij _2. Proof. We first establish attainment of the primal minimum. Given any feasible ΔW W, subtract ΔwV w_V from every update row. This translation preserves all pairwise differences and changes the objective by −⟨∑igi,ΔwV⟩=0.- _ig_i, w_V =0. We may therefore restrict attention to representatives satisfying ΔwV=0 w_V=0. Feasibility then gives ‖Δwi‖2≤ρhead w_i _2≤ _head for every i<Vi<V. The resulting feasible set is nonempty and compact, so the continuous objective attains a finite minimum. To derive the dual, introduce an auxiliary variable ξij _ij ffor each pair i<ji<j and impose ξij=Δwi−Δwj,‖ξij‖2≤ρhead,i<j. _ij= w_i- w_j, _ij _2≤ _head, i<j. Associate the multiplier φij _ij with the equality constraint ξij−Δwi+Δwj=0 _ij- w_i+ w_j=0. The Lagrangian is Λ(ΔW,ξ,Φ) ( W,ξ, ) =⟨G,ΔW⟩F+∑i<j⟨φij,ξij−Δwi+Δwj⟩ = G, W _F+ _i<j _ij, _ij- w_i+ w_j =∑i=1V⟨gi−∑j=1Vφij,Δwi⟩+∑i<j⟨φij,ξij⟩. = _i=1^V g_i- _j=1^V _ij, w_i + _i<j _ij, _ij . The rows of ΔW W are unrestricted in the inner minimization. Its infimum is therefore finite only if ∑jφij=gi _j _ij=g_i for all i. Under these balance equations, the terms containing ΔW W vanish. Minimizing over the Euclidean balls gives infξ‖ξij‖2≤ρhead∀i<j∑i<j⟨φij,ξij⟩=−ρhead∑i<j‖φij‖2. _ subarraycξ\\ _ij _2≤ _head\;∀ i<j subarray _i<j _ij, _ij =- _head _i<j _ij _2. For φij≠0 _ij≠ 0, the infimum for that pair is attained at ξij=−ρheadφij‖φij‖2; _ij=- _head _ij _ij _2; when φij=0 _ij=0, any feasible ξij _ij is optimal. The Lagrange dual therefore maximizes −ρhead∑i<j‖φij‖2- _head _i<j _ij _2 subject to Equation 17. Since −ρhead<0- _head<0, this is equivalent to the minimization on the right hand side of Equation 18. The reformulated primal problem is convex. Moreover, ΔW=0 W=0 and ξij=0 _ij=0 satisfy every equality constraint and lie strictly inside every Euclidean ball. Slater’s condition therefore gives strong duality and attainment of the dual optimum (Vandenberghe and Boyd, 2004, Section 5.2.3). This proves Equation 18. ∎ Connection to vector transport. Let =1,…,VX=\1,…,V\ carry the discrete metric d(i,j)=1d(i,j)=1 for i≠ji≠ j, and define the vector measure μ=∑i=1Vgiδi.μ= _i=1^Vg_i _i. The condition ∑igi=0 _ig_i=0 gives μ()=0μ(X)=0. Equation 17 represents this measure through an antisymmetric vector flow on the complete graph, with transport cost ∑i<j‖φij‖2. _i<j _ij _2. The primal variables provide the corresponding Lipschitz functions. In particular, define ui=−Δwiρhead.u_i=- w_i _head. The diameter constraint implies ‖ui−uj‖2≤1=d(i,j)∀i≠j, u_i-u_j _2≤ 1=d(i,j) ∀ i≠ j, so u:→ℝdu:X ^d is 11 Lipschitz. Equation 18 is therefore the finite discrete form of vector transport duality (Ciosmak, 2021, Theorem 2). This dual formulation characterizes the optimal value without requiring the primal minimizer to be unique. It contains one vector in ℝdR^d for each of the (V2) V2 unordered pairs of vocabulary rows. The resulting quadratic growth in the number of dual variables motivates the common ball approximation studied next. A.4 The RowNorm approximation The main text obtains RowNorm by requiring all update rows to lie in a common ball and then moving that ball to the origin. We first justify why this change of variable preserves the minimum. We then bound the suboptimality introduced by replacing the diameter constraint with the ball constraint. A.4.1 Centering the enclosing ball preserves the optimum Fix a gradient G∈ℝV×dG ^V× d with rows gi⊤g_i satisfying ∑igi=0 _ig_i=0, and let ρhead>0 _head>0. Recall that ℬ(ζ,ρhead/2)B(ζ, _head/2) denotes the closed Euclidean ball with center ζ and radius ρhead/2 _head/2. In Equation 7, we choose both the update ΔW W and the center ζ, minΔW,ζ⟨G,ΔW⟩Fsubject toΔwi∈ℬ(ζ,ρhead/2)∀i. _ W,ζ G, W _F to w_i (ζ, _head/2) ∀ i. Any center that encloses all the update rows is allowed. In particular, the center is not required to equal the mean row, and the mean row remains unrestricted. Since the gradient rows sum to zero, translating every update row by the same vector leaves the objective unchanged. Using this freedom, we invoke a change of variable ΔW^=ΔW−ζ⊤. W= W-1ζ . For any feasible pair (ΔW,ζ)( W,ζ), the translated rows satisfy ‖Δw^i‖2=‖Δwi−ζ‖2≤ρhead2∀i. w_i _2= w_i-ζ _2≤ _head2 ∀ i. Thus, every row now lies in a ball centered at the origin. The objective is preserved because ⟨G,ΔW⟩F=⟨G,ΔW^⟩F+⟨∑igi,ζ⟩⏟=0=⟨G,ΔW^⟩F. G, W _F= G, W _F+ _ig_i,ζ _=0= G, W _F. After this change of variable, neither the row constraints nor the objective depends on ζ. Equation 7 therefore has the same minimum as minΔW^⟨G,ΔW^⟩Fsubject toΔw^i∈ℬ(0,ρhead/2)∀i, _ W G, W _F to w_i (0, _head/2) ∀ i, which is Equation 8. To complete the equivalence, we also need to return from this problem to the original one. Given any feasible ΔW W, choose any ζ∈ℝdζ ^d and set ΔW=ΔW^+ζ⊤. W= W+1ζ . The resulting rows lie in ℬ(ζ,ρhead/2)B(ζ, _head/2), and the same objective identity shows that their value is unchanged. Choosing ζ=0ζ=0 simply gives ΔW=ΔW W= W. The two problems therefore attain exactly the same objective values. The feasible set for the problem centered at the origin is a finite product of closed, bounded balls. It is nonempty and compact, so the continuous linear objective attains its minimum. The reverse construction then provides a minimizing pair for the original problem as well. Consequently, every minimizer of the problem centered at the origin gives a minimizer of the common-ball problem, and every minimizing pair of the common ball problem translates back to one. This change of variable centers the enclosing ball. The rows inside it can still have a nonzero mean. The RowNorm update in Equation 10 is therefore already one minimizing choice for the common ball problem, with ζ=0ζ=0. Subtracting its mean row is optional and preserves the objective by the same translation argument. Moving the ball to the origin introduces no additional suboptimality. The gap considered next comes entirely from requiring the update rows to fit inside a ball of radius ρhead/2 _head/2 instead of imposing only the original diameter constraint. A.4.2 Approximation guarantee and sharpness Moving the ball to the origin preserves the minimum. We now bound the first-order decrease lost by replacing the diameter constraint with the common-ball constraint. Assume G≠0G≠ 0, ⊤G=01 G=0, and ρhead>0 _head>0. Let Δ⋆ _ and Δball _ball denote the greatest first-order decreases under these two constraints, respectively. Since RowNorm attains the common-ball optimum, Δball=ρhead2∑i‖gi‖2>0. _ball= _head2 _i g_i _2>0. Bounding the gap. Jung’s theorem bounds the radius needed to enclose a set using its diameter and dimension. Here we can work in the space spanned by the gradient rows, =spang1,…,gVS=span\g_1,…,g_V\. By definition of matrix rank, the dimension of this row space is k=rank(G)k=rank(G). To see why this space suffices, consider the part of an update row perpendicular to S. It has zero inner product with every gradient row and therefore contributes nothing to the objective. Removing these parts from all rows is an orthogonal projection, which cannot increase pairwise distances. Every feasible update can thus be replaced by one whose rows lie in S, with the same objective value and no larger diameter. Jung’s theorem (Jung, 1901) guarantees that any such rows with diameter at most ρhead _head lie in a ball of radius at most rk=ρheadk2(k+1).r_k= _head k2(k+1). Let ζ be the center of such a ball, so that ‖Δwi−ζ‖2≤rk w_i-ζ _2≤ r_k for every row. We only need to know that this center exists. As in the preceding subsection, ∑igi=0 _ig_i=0 allows us to subtract ζ without changing the objective. Cauchy–Schwarz then gives −⟨G,ΔW⟩F=−∑i⟨gi,Δwi−ζ⟩≤rk∑i‖gi‖2.- G, W _F=- _i g_i, w_i-ζ ≤ r_k _i g_i _2. This bounds Δ⋆ _ because every feasible update can be brought into S without changing its objective value. Every common-ball update also satisfies the diameter constraint, so 0<Δball≤Δ⋆≤2k+1Δball.0< _ball≤ _ ≤ 2kk+1\, _ball. Rearranging proves the relative suboptimality bound in Equation 9. Sharpness. A regular simplex gives a gradient for which equality holds. Choose g1,…,gk+1g_1,…,g_k+1 as the unit vertices of a regular k-dimensional simplex centered at the origin in ℝdR^d. For a vocabulary with V>k+1V>k+1, set gi=0g_i=0 for every i>k+1i>k+1. The simplex rows satisfy ∑i=1k+1gi=0,‖gi‖2=1,⟨gi,gj⟩=−1k(i≠j). _i=1^k+1g_i=0, g_i _2=1, g_i,g_j =- 1k (i≠ j). The full gradient therefore satisfies ⊤G=01 G=0 and rank(G)=krank(G)=k. Consider the update ΔW=−rkG W=-r_kG. For every i>k+1i>k+1, this sets Δwi=−rkgi=0 w_i=-r_kg_i=0. Distinct simplex update rows have distance ‖Δwi−Δwj‖2=rk2(k+1)k=ρhead. w_i- w_j _2=r_k 2(k+1)k= _head. Each simplex update row is also at distance rk≤ρheadr_k≤ _head from every zero update row. The entire update is therefore feasible. It achieves a decrease of rk(k+1)r_k(k+1), matching the upper bound on Δ⋆ _ . Since Δball=ρhead(k+1)/2 _ball= _head(k+1)/2, the relative gap is Δ⋆−ΔballΔ⋆=1−ρhead2rk=1−k+12k. _ - _ball _ =1- _head2r_k=1- k+12k. This proves sharpness for every permissible rank 1≤k≤min(d,V−1)1≤ k≤ (d,V-1). RowNorm therefore attains the optimal first-order decrease when k=1k=1 and retains at least 1/21/ 2 of it for every rank. The next subsection describes the practical update with momentum and numerical stabilization. A.5 Practical RowNorm In training, we apply row normalization to a first-moment estimate of the gradient. Let P=IV−⊤/VP=I_V-11 /V denote the projection that subtracts the mean row of a matrix. Starting from B−1=0B_-1=0, for t=0,1,…t=0,1,… let Bt=βBt−1+(1−β)Gt,B^t=Bt1−βt+1,B~t=PB^t,B_t=β B_t-1+(1-β)G_t, B_t= B_t1-β^t+1, B_t=P B_t, where 0≤β<10≤β<1. For numerical stability, define δnum RN_ _num by dividing row b~t,i⊤ b_t,i by ‖b~t,i‖2+δnum b_t,i _2+ _num, with δnum>0 _num>0. With learning rate ηt≥0 _t≥ 0, the update is ΔWt=−ηtPδnum(B~t),Wt+1=Wt+ΔWt. W_t=- _tP\, RN_ _num( B_t), W_t+1=W_t+ W_t. (19) Under our assumptions, the momentum estimate is already centered in exact arithmetic, so the first projection only removes any accumulated numerical drift. Row normalization can introduce a nonzero mean, which the final projection removes. Every normalized row has norm at most one, and centering preserves pairwise differences, so D(ΔWt)≤2ηtD( W_t)≤ 2 _t. Consequently, 0≤ηt≤ε/(2R)0≤ _t≤ /(2R) ensures the required Hilbert distance bound for every ‖h‖2≤R h _2≤ R. The computation takes O(Vd)O(Vd) time, uses one first-moment buffer, and implements each projection by subtracting the mean row. Momentum and numerical stabilization preserve the diameter bound on the applied update. The forward and backward estimates therefore remain applicable under their stated assumptions. Appendix B Forward and backward perturbation bounds A spectral constraint, as used by Muon, controls perturbations in both the forward and backward computations of a linear layer. For a fixed input h and incoming gradient r, ‖ΔWh‖2≤‖ΔW‖op‖h‖2,‖ΔW⊤r‖2≤‖ΔW‖op‖r‖2. Wh _2≤ W _op h _2, W r _2≤ W _op r _2. We establish analogous bounds for our softmax modules below, comparing each module before and after an update at the same inputs. Write τ=tanh(ε/4)τ= ( /4). The inequality relating Hilbert distance and total variation gives TV(p,q)≤τTV(p,q)≤τ whenever ℋ(p,q)≤εH(p,q)≤ (Cohen and Fausti, 2024, Theorem 5.1). B.1 The LM head The diameter constraint controls the head’s probability output and also bounds the change in the gradient passed to its hidden state. Let p=pW(h)p=p_W(h) and q=pW+ΔW(h)q=p_W+ W(h) at a fixed input ‖h‖2≤R h _2≤ R. The forward bound is ℋ(p,q)≤RD(ΔW)≤ε.H(p,q)≤ RD( W)≤ . For a target token c, denote the gradients passed to h by λh _h and λh+ _h^+. Cross-entropy gives λh=W⊤(p−ec),λh+=(W+ΔW)⊤(q−ec), _h=W (p-e_c), _h^+=(W+ W) (q-e_c), where ece_c is the target’s one-hot vector. Writing wj⊤w_j and Δwj⊤ w_j for the rows of W and ΔW W, subtraction gives λh+−λh=∑jqj(Δwj−Δwc)+∑j(qj−pj)wj. _h^+- _h= _jq_j( w_j- w_c)+ _j(q_j-p_j)w_j. The first sum averages differences between update rows, so its norm is at most D(ΔW)D( W). The second sum accounts for the redistributed probability. Because p and q each have unit mass, the positive and negative coefficients qj−pjq_j-p_j have the same total mass, namely TV(p,q)TV(p,q). After normalization, they form two weighted averages of the head rows. Their distance cannot exceed D(W)D(W), giving ‖λh+−λh‖2≤D(ΔW)+D(W)TV(p,q)≤εR+τD(W). _h^+- _h _2≤ D( W)+D(W)TV(p,q)≤ R+τ D(W). The current head appears through D(W)D(W) because it transmits the changed probability gradient. B.2 Attention For attention, we again seek bounds on the change in the output distributions and the resulting change in the gradients passed to the hidden states. The backward calculation now includes both the query and key uses of each hidden state. We keep the inputs ‖hℓ‖2≤R h_ _2≤ R and visible sets iJ_i fixed. Let M=WQWK⊤M=W_QW_K , and let ΔM M denote its change. Equation 12 gives maxiℋ(pi,pi+)≤2R2dh‖ΔM‖op≤ε. _iH(p_i,p_i^+)≤ 2R^2 d_h M _op≤ . The total variation bound recalled above then gives TV(pi,pi+)≤τTV(p_i,p_i^+)≤τ, where τ=tanh(ε/4)τ= ( /4). To connect this probability bound to backpropagation, let rir_i denote the incoming gradient on pip_i. We hold these incoming gradients fixed when comparing the module before and after the QK update. Differentiating softmax gives the gradient with respect to the scores, aij=pij(rij−μi),μi=∑j∈ipijrij.a_ij=p_ij(r_ij- _i), _i= _j _ip_ijr_ij. Each score sij=hi⊤Mhj/dhs_ij=h_i Mh_j/ d_h sends this gradient to both its query input and its key input. Denote the corresponding contributions by λiQ _i^Q and λijK _ij^K: λiQ=M∑j∈iaijhjdh,λijK=aijM⊤hidh. _i^Q= M _j _ia_ijh_j d_h, _ij^K= a_ijM h_i d_h. The total gradient passed to hℓh_ adds its query contribution and its key contributions from all queries that can see it, λℓ=λℓQ+∑i:ℓ∈iλiℓK. _ = _ ^Q+ _i: _i _i ^K. There are two sources of change in these expressions: the QK product changes, and so does the score gradient. Writing Δai=ai+−ai a_i=a_i^+-a_i and Δλ=λ+−λ λ=λ^+-λ, we obtain ΔλiQ _i^Q =ΔM∑j∈iaij+hj+M∑j∈iΔaijhjdh, = M _j _ia_ij^+h_j+M _j _i a_ijh_j d_h, ΔλijK _ij^K =aij+ΔM⊤hi+ΔaijM⊤hidh. = a_ij^+ M h_i+ a_ijM h_i d_h. Since every hidden state has norm at most R, the remaining quantities to bound are ‖ai+‖1 a_i^+ _1 and ‖Δai‖1 a_i _1. The first accompanies the product change, and the second accompanies the current product. We establish these estimates for one row, temporarily omitting the index i. Write aj(p)=pj(rj−μ)a_j(p)=p_j(r_j-μ) with μ=∑jpjrjμ= _jp_jr_j. Recall that osc(r)=maxjrj−minjrj (r)= _jr_j- _jr_j. For the size of a(p)a(p), Cauchy–Schwarz reduces the calculation to the variance of the entries of r under p. With m=(maxjrj+minjrj)/2m=( _jr_j+ _jr_j)/2 denoting the midpoint of their range, ‖a(p)‖12≤∑jpj(rj−μ)2=∑jpj(rj−m)2−(μ−m)2≤osc(r)24. a(p) _1^2≤ _jp_j(r_j-μ)^2= _jp_j(r_j-m)^2-(μ-m)^2≤ (r)^24. For the change in a, we use the probability redistribution measured by total variation. Any change from p to q can be decomposed into transfers from coordinates that lose probability to coordinates that gain it. The total mass transferred in this way is TV(p,q)TV(p,q). It therefore suffices to bound the effect of one transfer and add these bounds. Suppose an amount α moves from coordinate k to coordinate j, giving q=p+α(ej−ek)q=p+α(e_j-e_k). Here ej,eke_j,e_k are coordinate vectors and 0≤α≤pk0≤α≤ p_k. Subtracting the score gradients at fixed r gives a(q)−a(p)=α[(rj−μ)ej−(rk−μ)ek−(rj−rk)q].a(q)-a(p)=α [(r_j-μ)e_j-(r_k-μ)e_k-(r_j-r_k)q ]. Because q has unit mass, ‖a(q)−a(p)‖1≤α(|rj−μ|+|rk−μ|+|rj−rk|)≤2αosc(r). a(q)-a(p) _1≤α (|r_j-μ|+|r_k-μ|+|r_j-r_k| )≤ 2α\, (r). The last step uses the fact that the three absolute differences sum to twice the range of rj,rk,μr_j,r_k,μ. All three values lie within the range of r, since μ is a weighted average of its entries. Adding the bounds over the transfers proves ‖a(q)−a(p)‖1≤2osc(r)TV(p,q) a(q)-a(p) _1≤ 2 (r)TV(p,q). Returning to the attention rows, the two estimates become ‖ai+‖1≤osc(ri)2,‖Δai‖1≤2osc(ri)τ. a_i^+ _1≤ (r_i)2, a_i _1≤ 2 (r_i)τ. Apply these to the query and key expressions above, using ‖hℓ‖2≤R h_ _2≤ R and the triangle inequality: ∑ℓ‖Δλℓ‖2 _ _ _2 ≤2Rdh∑i(‖ΔM‖op‖ai+‖1+‖M‖op‖Δai‖1) ≤ 2R d_h _i ( M _op a_i^+ _1+ M _op a_i _1 ) ≤Rdh(‖ΔM‖op+4‖M‖opτ)∑iosc(ri). ≤ R d_h ( M _op+4 M _opτ ) _i (r_i). The factor of two in the first line accounts for the query and key contributions. Both are included when a position attends to itself. For fixed M and incoming gradients, the bound tends to zero as ε tends to zero under the stated spectral constraint. For finite weight updates, ΔM M includes the quadratic term ΔWQΔWK⊤ W_Q W_K . The comparison concerns the QK probability module and does not include changes in the incoming gradients or the value pathway. Appendix C Reproducing Figure 1 Figure 1 holds the initial distribution fixed and varies only the centered update. We construct the initial logits by centering a quadratic profile over i=1,…,32i=1,…,32: ai=−(i−16.54)2,zi=ai−132∑j=132aj,p=softmax(z).a_i=- ( i-16.54 )^2, z_i=a_i- 132 _j=1^32a_j, p=softmax(z). The three updates below either shift a central block together or perturb individual preactivation coordinates within it. Here, ⋅1_\·\ is the indicator function, equal to one when its condition holds and zero otherwise: Δziblue [rgb]0,0,1 z_i^blue =213≤i≤16−217≤i≤20, =21_\13≤ i≤ 16\-21_\17≤ i≤ 20\, Δzipurple [rgb]0.5,0,0.5 z_i^purple =29≤i≤24−1, =21_\9≤ i≤ 24\-1, Δzired [rgb]1,0,0 z_i^red =2i=16−2i=19. =21_\i=16\-21_\i=19\. Each update has zero mean and produces q=softmax(z+Δz)q=softmax(z+ z). Using RMS(Δz)=132∑i=132(Δzi)2,ℋ(p,q)=maxiΔzi−miniΔzi,RMS( z)= 132 _i=1^32( z_i)^2, (p,q)= _i z_i- _i z_i, we obtain the following values, together with the total variation TV(p,q)=12∑i|pi−qi|TV(p,q)= 12 _i|p_i-q_i|: BluePurpleRedRMS(Δz)111/2ℋ(p,q)424TV(p,q)≈0.514≈0.004≈0.430. array[]c|c& [rgb]0,0,1Blue& [rgb]0.5,0,0.5Purple& [rgb]1,0,0Red\\ ( z)&1&1&1/2\\ H(p,q)&4&2&4\\ TV(p,q)&≈ 0.514&≈ 0.004&≈ 0.430 array. Purple adds a common increment to coordinates 99–2424, which initially carry more than 99.5%99.5\% of the probability mass. The pairwise odds among coordinates 9,…,249,…,24 remain unchanged, and only a small amount of probability mass moves from the tails into this block. Blue and red instead alter these ratios, producing much larger redistribution under the same RMS bound.