Paper deep dive
MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads
Linkai Ma, Xinyu Luo, Mengbo Wang, Ananth Grama, Petros Drineas, Brian Bullins
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 10/5/2026, 6:46:40 AM
Summary
The paper introduces MuonIO, an optimization method that applies a principled, norm-aware update rule to both input embedding tables and output language model heads in large language models. Unlike standard Muon implementations that use AdamW for these layers, MuonIO uses the 1->2 operator norm for embeddings and the 2->infinity operator norm for language model heads, leveraging their transposition symmetry to apply a single normalized-vector update. This approach reduces I/O optimizer state memory by 50% and update FLOPs by ~46% while improving validation perplexity.
Entities (10)
Relation Signals (8)
MuonIO → appliesnormto → 1->2 Operator Norm
confidence 95% · for the embedding table... draw on the 1 -> 2 operator norm
MuonIO → appliesnormto → 2->infinity Operator Norm
confidence 95% · for the language model head... motivate the use of the 2->infinity operator norm
1->2 Operator Norm → associatedwith → Embedding Table
confidence 95% · for the embedding table... draw on the 1 -> 2 operator norm
2->infinity Operator Norm → associatedwith → Language Model Head
confidence 95% · for the language model head... motivate the use of the 2->infinity operator norm
MuonIO → replaces → AdamW
confidence 95% · Standard Muon implementations... exclude the input... and output... layers... for which they use AdamW instead. We present MuonIO... for both of these layers.
MuonIO → evaluatedon → LLaMA
confidence 90% · 1B LLaMA pretraining on C4
MuonIO → reducesflopsby → 46%
confidence 90% · I/O update FLOPs by ~46% compared to Muon
MuonIO → reducesmemoryby → 50%
confidence 90% · MuonIO reducing I/O optimizer state memory by 50%
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:The Muon optimizer derives its update rule for hidden linear layers by solving a local linearization of the loss penalized by the spectral norm, motivated by an RMS-stability argument for dense linear layers. Standard Muon implementations, however, exclude the input (embedding table) and output (language model head) layers from this principled treatment, for which they use AdamW instead. We present MuonIO, a single Muon-style update for both of these layers. For the language model head $\mathbf{L} \in \mathbb{R}^{V \times d}$, we motivate the use of the $2\to\infty$ operator norm, due to the Lipschitz continuity of the softmax output geometry, while for the embedding table $\mathbf{E} \in \mathbb{R}^{d \times V}$, we draw on the $1 \to 2$ operator norm, based on the one-hot input geometry identified by Bernstein & Newhouse (2025). The identity $\lVert\mathbf{L}\rVert_{2\to\infty}=\lVert\mathbf{L}^\top\rVert_{1\to2}$ then puts both matrices in the same vocabulary-oriented geometry: MuonIO applies a single normalized-vector rule, which appears as column normalization for $\mathbf{E}$ and row normalization for $\mathbf{L}$. Empirical evaluations demonstrate the effectiveness of our approach, with MuonIO reducing I/O optimizer state memory by 50% and I/O update FLOPs by $\sim$46% compared to Muon for 1B LLaMA pretraining on C4, while also improving validation perplexity.
Tags
Links
- Source: https://arxiv.org/abs/2610.02705v1
- Canonical: https://arxiv.org/abs/2610.02705v1
Trouble viewing inline? Open PDF directly →
Full Text
61,177 characters extracted from source content.
Expand or collapse full text
MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads Linkai Ma †thanks: Equal contribution. Xinyu Luo11footnotemark: 1 Mengbo Wang Ananth Grama Petros Drineas Brian Bullins Affiliation: Department of Computer Science, Purdue University Affiliation: ma856, luo466, wang4887, ayg, pdrineas, bbullins@purdue.edu Abstract The Muon optimizer derives its update rule for hidden linear layers by solving a local linearization of the loss penalized by the spectral norm, motivated by an RMS-stability argument for dense linear layers. Standard Muon implementations, however, exclude the input (embedding table) and output (language model head) layers from this principled treatment, for which they use AdamW instead. We present MuonIO, a single Muon-style update for both of these layers. For the language model head ∈ℝV×dL ^V× d, we motivate the use of the 2→∞2→∞ operator norm, due to the Lipschitz continuity of the softmax output geometry, while for the embedding table ∈ℝd×VE ^d× V, we draw on the 1→21→ 2 operator norm, based on the one-hot input geometry identified by Bernstein and Newhouse (2025). The identity ‖2→∞=‖⊤‖1→2\|L\|_2→∞=\|L \|_1→ 2 then puts both matrices in the same vocabulary-oriented geometry: MuonIO applies a single normalized-vector rule, which appears as column normalization for E and row normalization for L. Empirical evaluations demonstrate the effectiveness of our approach, with MuonIO reducing I/O optimizer state memory by 50% and I/O update FLOPs by ∼ 46% compared to Muon for 1B LLaMA pretraining on C4, while also improving validation perplexity. 1 Introduction Gradient-based optimization of neural networks has long been dominated by Adam (Kingma and Ba, 2014) and its descendants. The Muon optimizer (Jordan et al., 2024) revisited a foundational question: what is the right geometry for determining the direction of a weight update? Building on ideas from modular duality (Bernstein and Newhouse, 2024; Bernstein and Newhouse, 2025), Muon chooses its update by minimizing a local first-order approximation of the loss while penalizing the update in the spectral norm geometry. This geometry is appropriate for hidden linear layers (Yang et al., 2023) and yields the polar factor update, which is computed in practice by a Newton–Schulz iteration rather than by an explicit SVD (Jordan et al., 2024; Bernstein and Newhouse, 2025). Muon’s practical impact has been rapid: it achieved a 1.35×1.35× speedup over the previous NanoGPT training record (Jordan et al., 2024), has been used in large-scale models such as Moonlight (Liu et al., 2025) and Kimi K2 (Team et al., 2025), and has inspired variants such as HTMuon (Pang et al., 2026), NorMuon (Li et al., 2026), Polar Express (Amsel et al., 2026), AdaMuon (Si et al., 2025), Scion (Pethick et al., 2025), Dion (Ahn et al., 2025b), and Dion2 (Ahn et al., 2025a). Despite this progress, standard Muon implementations still exclude two important layer types from the spectral norm update: the input embedding layer (embed) and the output language model head (lm_head). Instead, both are typically routed to AdamW, and as Jordan et al. (2024) explicitly note, this exception is not fully explained by the existing theory: “That such dynamics are also different for the output layer does not seem to follow from the theory, and is instead driven by empirics.” Similarly, Bernstein’s derivation note (Bernstein, 2025) observes that the dense layer stability argument “separates the Linear layer from, say, an Embedding layer that takes one-hot inputs.” While prior work has identified the norm induced by one-hot embedding lookups (Bernstein and Newhouse, 2024; Bernstein and Newhouse, 2025), and column/row normalization updates have been proposed on other grounds (Pethick et al., 2025; Glentis et al., 2026; Lau and Su, 2026), there has yet to be established a principled norm-aware approach that handles both vocabulary-facing layers with a single update. 1.1 Our Contribution In this work, we take the embedding layer geometry (Bernstein and Newhouse, 2024; Bernstein and Newhouse, 2025) as a starting point and present MuonIO, a single Muon-style update for the two vocabulary-facing matrices. Prior work identifies the 1→RMS1 norm and its column-normalized direction for embed, which, under our ℓ2 _2 convention, is the 1→21→ 2 norm up to a fixed dimension factor. For lm_head, we instead study the change in the probability distribution after softmax. A tight total-variation stability bound identifies the 2→∞2→∞ norm of the weight perturbation. These operator norms are related by transposition: ‖2→∞=‖⊤‖1→2\|L\|_2→∞=\|L \|_1→ 2. Once both parameters are stored in vocabulary-major form, Corollary 2.3 gives the same normalized-vector update for both. In weight-tied models where =⊤L=E (Press and Wolf, 2017; Inan et al., 2016), it acts directly on the single shared parameter. Optimizing embed and lm_head with AdamW also carries concrete computation and memory cost. These two layers are among the largest parameter matrices in a language model: each has VdVd entries, with V being the vocabulary size and d the hidden dimension. AdamW maintains two optimizer state moments per parameter, whereas MuonIO keeps only a single momentum buffer. In terms of computation, AdamW takes 13Vd13Vd FLOPS per matrix while MuonIO takes 7Vd+3V7Vd+3V. Routing these two layers to AdamW therefore roughly doubles the optimizer state memory and arithmetic cost on the model’s biggest tensors. Further details are included in Appendix C. We evaluate MuonIO by pretraining LLaMA-family models at the 60M, 130M and 1B scales. Replacing the AdamW treatment of the two I/O layers with one geometry-matched update halves their persistent optimizer state and, at the same time, yields improved validation perplexity over a strong Muon baseline that already tunes separate learning rates for the input, output, hidden, and gain/bias parameters. 1.2 Related Work The work on modular duality (Bernstein and Newhouse, 2025) gives a column normalization update for embed, but does not provide a corresponding treatment of lm_head. Other optimizer frameworks have nevertheless proposed row-normalized output layer updates, and recent work has applied vocabulary-wise normalization specifically to language model heads, motivated by linear-minimization oracles, empirical gradient variance, and equivariance (Pethick et al., 2025; Glentis et al., 2026; Lau and Su, 2026), while Ember (Shah, 2026) treats both layers specially using factorized second moments. Another related line of work observes, both empirically and theoretically, that embed benefits from a learning rate distinct from the hidden layers (Hayou and Liu, 2025; He et al., 2026; Kalra and Barkeshli, 2026). Additionally, Nair (2025) establishes tight Lipschitz bounds for the softmax function under matched ℓp _p norms. For a more detailed discussion, see Appendix A. 2 Background and Notation 2.1 Notation Scalars are denoted by lowercase letters (e.g., x); vectors by bold lowercase letters (e.g., x); matrices by bold uppercase letters (e.g., W). For a matrix ∈ℝm×nA ^m× n we write ‖F\|A\|_F for the Frobenius norm, ‖2\|A\|_2 for the spectral norm, ‖∗\|A\|_* for the nuclear norm, and ⟨,⟩=tr(⊤) ,B = tr (A B ) for the Frobenius inner product. We write :,tA_:,t for the t-th column of A and t,:A_t,: for its t-th row. We write kI_k for the k×k× k identity matrix (dropping the subscript when the dimension is clear), and [V][V] for 1,2,…,V\1,2,…,V\. We write ΔV:=p∈ℝ≥0V:∑t=1Vpt=1 ^V:= \p ^V_≥ 0: _t=1^Vp_t=1 \ for the probability simplex over V outcomes. 2.2 Steepest Descent Under a Norm Let ℒ()L(W) be a loss function and let k=∇ℒ(k)G_k= (W_k) denote the gradient at iterate kW_k. For a given norm ∥⋅∥\|·\| on the weight matrices, the steepest descent step of size λ>0λ>0 minimizes a penalized local linearization. The following generic result is Bernstein and Newhouse (2024, Proposition 1), stated for matrix spaces under the Frobenius inner product. We use the reciprocal of their sharpness parameter so that λ is a step size. Proposition 2.1 (Penalized steepest descent). Let ∈ℝm×nG ^m× n and let ∥⋅∥\|·\| be a norm on ℝm×nR^m× n, with dual norm ‖†:=max‖≤1⟨,⟩\|G\| := _\|T\|≤ 1 ,T . For λ>0λ>0, argminΔ⟨,Δ⟩+12λ‖Δ‖2=−λ‖†:∈argmax‖=1⟨,⟩. *argmin_ \ , + 12λ\| \|^2 \= \-λ\|G\| T:T∈ *argmax_\|T\|=1 ,T \. (1) Gradient descent (Frobenius geometry). Applying Proposition 2.1 with ∥⋅∥=∥⋅∥F\|·\|=\|·\|_F yields Δ⋆=−λ =- , the ordinary gradient descent step. Muon (spectral norm geometry). Taking ∥⋅∥=∥⋅∥2\|·\|=\|·\|_2 and letting =⊤G=U V be the thin SVD of the gradient, Proposition 2.1 gives the following choice of solution: Δ⋆=−λ‖∗polar(),polar():=⊤. =-λ\,\|G\|_*polar(G), (G):=UV . (2) 2.3 Operator Norms For p,q∈[1,∞]p,q∈[1,∞] and ∈ℝm×nA ^m× n, the p→qp→ q operator norm is ‖p→q:=max‖p=1‖q.\|A\|_p→ q:= _\|x\|_p=1\|Ax\|_q. For example, the spectral norm satisfies ‖2=‖2→2\|A\|_2=\|A\|_2→ 2. Bernstein and Newhouse (2024, Proposition 8) establish the following induced-norm identities. Proposition 2.2. Let ∈ℝm×nA ^m× n, let p∈[1,∞]p∈[1,∞], and let p∗p^* be its Hölder conjugate, so that 1/p+1/p∗=11/p+1/p^*=1. Then ∥1→p=maxj∈[n]∥:,j∥p,∥p→∞=maxi∈[m]∥i,:∥p∗.\|A\|_1→ p= _j∈[n]\|A_:,j\|_p, \|A\|_p→∞= _i∈[m]\|A_i,:\|_p^*. (3) In particular, ∥1→2=maxj∈[n]∥:,j∥2,∥2→∞=maxi∈[m]∥i,:∥2.\|A\|_1→ 2= _j∈[n]\|A_:,j\|_2, \|A\|_2→∞= _i∈[m]\|A_i,:\|_2. (4) Combining Propositions 2.1 and 2.2 immediately gives the following identities. Corollary 2.3. Let ∈ℝm×nG ^m× n. The dual norms of ∥⋅∥1→2\|·\|_1→ 2 and ∥⋅∥2→∞\|·\|_2→∞ are, respectively, ∥1→2†=∑j=1n∥:,j∥2,∥2→∞†=∑i=1m∥i,:∥2. \|G\|_1→ 2 = _j=1^n\|G_:,j\|_2, \|G\|_2→∞ = _i=1^m\|G_i,:\|_2. (5) Define row and column normalization by (RowNorm())i,::=i,:/∥i,:∥2,i,:≠,,otherwise.ColNorm():=RowNorm(⊤)⊤. (RowNorm(G) )_i,::= casesG_i,:/\|G_i,:\|_2,&G_i,: 0,\\ 0,&otherwise. cases (G):=RowNorm(G ) . (6) Then the 2→∞2→∞-penalized problem has the row-normalized solution Δ⋆=−λ‖2→∞†RowNorm(), =-λ\|G\|_2→∞ RowNorm(G), (7) whereas the 1→21→ 2-penalized problem has the column-normalized solution Δ⋆=−λ‖1→2†ColNorm(). =-λ\|G\|_1→ 2 ColNorm(G). (8) 2.4 Layer Types and Their Geometries A standard transformer language model contains three qualitatively distinct layer types, each with a different input/output geometry. We denote V as the vocabulary size and d as the hidden dimension. Hidden linear layers. A weight matrix ∈ℝnout×ninW ^n_out× n_in maps a dense activation x into y, both of which are expected to be dense vectors. Given a perturbation of the weight Δ , the bound ‖Δ‖2≤‖Δ‖2‖2\| \|_2≤\| \|_2\|x\|_2 identifies the spectral norm as the natural penalty (Bernstein, 2025), motivating the Muon polar factor update. Embedding layers. embed (∈ℝd×VE ^d× V) maps token indices to dense vectors via the lookup =y=Ex, where ∈ℝVx ^V is one-hot. Subsection 3.1 reviews the resulting norm from prior work. Language model head. lm_head (∈ℝV×dL ^V× d) maps a hidden state ∈ℝdy ^d to logits =∈ℝVz=Ly ^V and then to probabilities =softmax()∈ΔVp=softmax(z)∈ ^V, where ΔV ^V is the probability simplex. Given a perturbation Δ , a good measure of the perturbation is not simply ‖Δ‖2\| \|_2, since its effect must be propagated to the probability distribution. Instead, we measure that change by total variation, with the analysis in Subsection 3.2. 3 MuonIO for Vocabulary-Facing Layers The input embedding and output language model head have different storage orientations, but each has one vector associated with every vocabulary item. We first identify the appropriate norm for each layer, then orient both matrices along their vocabulary axis and give a single MuonIO update. Appendix B contains the proof of Propositions 3.3, 3.5, 3.7 and Remark 3.8. 3.1 Embedding geometry from prior work The embedding table ∈ℝd×VE ^d× V acts as a lookup: for a one-hot token indicator ∈0,1Vx∈\0,1\^V, the forward pass returns Ex, one column of E. Consequently, max∈0,1V,∑ixi=1∥2=maxt∈[V]∥:,t∥2=∥1→2. _x∈\0,1\^V,\, _ix_i=1\|Ex\|_2= _t∈[V]\|E_:,t\|_2=\|E\|_1→ 2. (9) The final equality is the p=2p=2 case of Proposition 2.2. Thus the 1→21→ 2 operator norm, rather than the spectral norm used for dense hidden layers, measures the largest change to any token embedding. Bernstein and Newhouse (2024) identify the 1→RMS1 operator norm for embedding layers, which is the same as 1→21→ 2 operator norm up to a scaling factor of d d. We use that established geometry as one half of MuonIO, and our new ingredient is the total variation justification for the matching output layer geometry developed next. 3.2 Language-model-head geometry The language model head ∈ℝV×dL ^V× d maps a hidden state ∈ℝdy ^d to a probability distribution (as opposed to a dense activation vector) over the vocabulary via =softmax()∈ΔVp=softmax(Ly)∈ ^V. We thus identify a matched penalty norm by asking: which norm on Δ controls the worst case change in the output distribution induced by the weight perturbation? 3.2.1 Total variation stability and the 2→∞2→∞ norm A canonical measure of the difference between probability distributions is the total variation distance, ‖′−‖TV=12‖′−‖1\|p -p\|_TV= 12\|p -p\|_1. We show that the 2→∞2→∞ operator norm controls this quantity. Lemma 3.1 (Jacobian of softmax). Let =softmax()∈ΔVp=softmax(z)∈ ^V. The Jacobian of softmaxsoftmax at z is J()=∂softmax()∂=diag()−⊤.J(z)= ∂\,softmax(z) =diag(p)-pp . (10) Proof. Writing i=ei/∑s=1Vesp_i=e^z_i/ _s=1^Ve^z_s, direct differentiation gives ∂i/∂j=iδij−ij _i/ _j=p_i _ij-p_ip_j. ∎ Lemma 3.2 (∞→1∞→ 1 norm of the softmax Jacobian). For every ∈ℝVz ^V, ‖J()‖∞→1≤1.\|J(z)\|_∞→ 1≤ 1. (11) Proof. Fix ∈ℝVu ^V with ‖∞≤1\|u\|_∞≤ 1, and let =softmax()p=softmax(z). The i-th coordinate of J()J(z)u is i(i−m)p_i(u_i-m), where m:=⊤=∑iim:=p u= _ip_iu_i. Hence ‖J()‖1=∑i=1Vi|i−m|.\|J(z)u\|_1= _i=1^Vp_i|u_i-m|. Introduce a random variable X that takes value iu_i with probability ip_i, so m=[X]m=E[X] and ‖J()‖1=[|X−[X]|]\|J(z)u\|_1=E[|X-E[X]|]. By Jensen’s inequality, [|X−[X]|]≤Var(X).E[|X-E[X]|]≤ Var(X). Since ‖∞≤1\|u\|_∞≤ 1, all values of X lie in [−1,1][-1,1]. Popoviciu’s inequality on variances (Popoviciu, 1935) states that Var(X)≤(b−a)2/4Var(X)≤(b-a)^2/4 for any random variable supported on [a,b][a,b], so with a=−1a=-1 and b=1b=1 this gives Var(X)≤1Var(X)≤ 1. Therefore ‖J()‖1≤1\|J(z)u\|_1≤ 1, and taking the supremum over all such u proves the claim. ∎ Proposition 3.3 (Tightness of Lemma 3.2). For every V≥2V≥ 2, inequality (11) is attained: there exist ∈ℝVz ^V and ∈±1Vu∈\± 1\^V such that ‖J()‖1=1\|J(z)u\|_1=1. Lemma 3.4 (Global Lipschitz bound for softmax). For all ,′∈ℝVz,z ^V, ‖softmax(′)−softmax()‖1≤‖′−‖∞.\|softmax(z )-softmax(z)\|_1≤\|z -z\|_∞. (12) Proof. Let γ(t)=+t(′−)γ(t)=z+t(z -z) for t∈[0,1]t∈[0,1]. By the fundamental theorem of calculus, softmax(′)−softmax()=∫01J(γ(t))(′−)t.softmax(z )-softmax(z)= _0^1J(γ(t))(z -z)\,dt. (13) Taking norms, applying Lemma 3.2 at each t, and integrating gives ‖softmax(′)−softmax()‖1≤∫01‖J(γ(t))‖∞→1‖′−‖∞t≤‖′−‖∞.∎\|softmax(z )-softmax(z)\|_1≤ _0^1\|J(γ(t))\|_∞→ 1\,\|z -z\|_∞\,dt≤\|z -z\|_∞. (14) Proposition 3.5 (Tightness of Lemma 3.4). The constant 11 in Lemma 3.4 cannot be improved. Namely, for every V≥2V≥ 2, sup′≠‖softmax(′)−softmax()‖1‖′−‖∞=1, _z \; \|softmax(z )-softmax(z)\|_1\|z -z\|_∞=1, (15) although the supremum is not attained by any finite pair ′≠z . Lemma 3.6 (TV stability of the lm_head). For every hidden state ∈ℝdy ^d and every perturbation Δ∈ℝV×d ^V× d, ‖softmax((+Δ))−softmax()‖TV≤12‖Δ‖2→∞‖2.\|softmax((L+ )y)-softmax(Ly)\|_TV≤ 12\| \|_2→∞\,\|y\|_2. (16) Proof. Set =z=Ly and ′=(+Δ)z =(L+ )y, so ′−=(Δ)z -z=( )y. Applying Lemma 3.4 and then the definition of the 2→∞2→∞ norm, ‖softmax(′)−softmax()‖1≤‖(Δ)‖∞≤‖Δ‖2→∞‖2.\|softmax(z )-softmax(z)\|_1≤\|( )y\|_∞≤\| \|_2→∞\,\|y\|_2. (17) Dividing by 22 converts the ℓ1 _1 distance to total variation. ∎ Proposition 3.7 (Tightness of Lemma 3.6). The constant 12 12 in Lemma 3.6 cannot be improved: for every V≥2V≥ 2 and d≥1d≥ 1, sup,Δ≠,≠‖softmax((+Δ))−softmax()‖TV12‖Δ‖2→∞‖2=1, _ subarraycL,\, 0,\,y 0 subarray\; \|softmax((L+ )y)-softmax(Ly)\|_TV 12\| \|_2→∞\,\|y\|_2=1, (18) although the supremum is not attained. Lemma 3.6 shows that ‖Δ‖2→∞\| \|_2→∞ gives sharp uniform control of worst-case TV perturbation when ‖2\|y\|_2 is held fixed. The 2→∞2→∞ operator norm is therefore the natural penalty for lm_head. 3.3 Transposition symmetry The two layer geometries are the same after orienting the vocabulary axis consistently. This follows from the relation between the norms identified in Sections 3.1 and 3.2. The observation is a direct consequence of the dual norm property, and it makes the common MuonIO update transparent. For completeness, we include the proof in Appendix B. Remark 3.8 (Transposition identity). For any ∈ℝm×nA ^m× n and any p,q∈[1,∞]p,q∈[1,∞] with Hölder conjugates determined by 1/p+1/p∗=11/p+1/p^*=1 and 1/q+1/q∗=11/q+1/q^*=1, ‖p→q=‖⊤‖q∗→p∗.\|A\|_p→ q\;=\;\|A \|_q^*→ p^*. (19) Specializing to p=2p=2, q=∞q=∞ (so p∗=2p^*=2 and q∗=1q^*=1) gives ‖2→∞=‖⊤‖1→2,\|A\|_2→∞\;=\;\|A \|_1→ 2, (20) i.e., the lm_head norm of L equals the embed norm of ⊤L . Consequence under weight tying. In language models that tie the embedding and output weights (=⊤L=E ) (Press and Wolf, 2017; Inan et al., 2016), Remark 3.8 reduces to the tautology ‖⊤‖2→∞=‖1→2\|E \|_2→∞=\|E\|_1→ 2. By Corollary 2.3, column normalization of E is exactly row normalization of =⊤L=E , so MuonIO acts once on the shared parameters. The next subsection expresses the same rule for untied weights. 3.4 One MuonIO update Put either vocabulary-facing matrix into the canonical row-oriented form =⊤∈ℝV×d,for the embedding table,∈ℝV×d,for the language-model head.W= casesE ^V× d,&for the embedding table,\\ L ^V× d,&for the language-model head. cases (21) By (20), both layer penalties become the same 2→∞2→∞ norm in this orientation. By Proposition 2.2, this is the maximum 2-norm of a row. Thus both layers use the same local problem Δ⋆∈argminΔ⟨k,Δ⟩+12λ‖Δ‖2→∞2,k=∇ℒ(k). ∈ _ \; _k, + 12λ\| \|_2→∞^2, _k= (W_k). (22) Corollary 2.3 now applies directly, as the unsmoothed solution row-normalizes kG_k and scales it by λ‖k‖2→∞†λ\|G_k\|_2→∞ . In the original storage orientation, the same rule normalizes columns of E and rows of L. Algorithm 1 substitutes the momentum estimate for the raw gradient and adds the usual numerical safeguard as well as an additional scaling factor d d so that each vector in ℝdR^d has 2-norm O(d)O( d). Algorithm 1 IO for embedding table or language model head input row-oriented 0∈ℝV×dW_0 ^V× d from (21), step size λ, momentum β, loss ℒL, ε>0 >0 initialize 0←M_0 0, k←0k← 0 1: while kW_k not converged do 2: Obtain ~k G_k such that [~k]=∇ℒ(k)E[ G_k]= (W_k) 3: k+1←βk+(1−β)~kM_k+1← _k+(1-β) G_k 4: (k+1)j,:←(k+1)j,:max∥(k+1)j,:∥2,ε(U_k+1)_j,:← (M_k+1)_j,: \\|(M_k+1)_j,:\|_2, \ for every j∈[V]j∈[V] 5: sk+1←∑j=1V∥(k+1)j,:∥2s_k+1← _j=1^V\|(M_k+1)_j,:\|_2 6: k+1←k−λsk+1dk+1W_k+1 _k-λ s_k+1 d\,U_k+1 7: k←k+1k← k+1 8: end while 9: return kW_k 4 Experiments Our Algorithm 1 acts as a modular component that can be naturally attached to any Muon-style optimizer that treats the I/O layers with AdamW. To demonstrate the effectiveness of our approach, we incorporate IO into Muon (Jordan et al., 2024) and Polar Express (Amsel et al., 2026), yielding the corresponding optimizers MuonIO and Polar Express IO. We also include the recent method Ember (Shah, 2026) as an additional baseline, as it also provides alternative treatment for the I/O layers. We evaluate these algorithms along three axes: validation perplexity, optimizer arithmetic cost and optimizer state memory. We additionally report various norms of the gradient for the I/O layers to offer a complementary view of the training dynamics. The central claim is that our geometry-matched update for the I/O layers improves pretraining perplexity while saving roughly 50%50\% of the computation and persistant memory costs compared to AdamW. 4.1 Experimental Setup We evaluate all optimizers on autoregressive language model pretraining using the English C4 corpus (Raffel et al., 2020), a cleaned Common Crawl dataset released under the Open Data Commons Attribution License. We train LLaMA-family models from scratch at the 60M, 130M, and 1B scales (Touvron et al., 2023; Zhao et al., 2024), using untied input embeddings and language model heads. Text is tokenized with the T5-base tokenizer, and individual examples are truncated or padded to 256 tokens. Padding tokens are excluded from the training and evaluation losses. All experiments use bfloat16 precision and an effective batch size of 512 sequences, with gradient accumulation adjusted according to the per-device batch size and GPU count. We use nominal training budgets (including padding positions) of approximately 1.15 billion tokens for the 60M and 130M models and 6.87 billion tokens for the 1B models. Learning rates increase linearly during the first 10% of training, followed by cosine decay to 10% of their peak values. Weight decay and gradient clipping are disabled. We evaluate on the held-out C4 validation split using approximately 10 million non-padding tokens per evaluation. Our primary metric is final validation perplexity, computed by exponentiating the average validation loss, with lower values indicating better language-modeling performance. We organize the six optimizer configurations into two groups. The first consists of four configurations that use standard Muon for the hidden weight matrices while varying the treatment of the I/O layers. The second examines whether our I/O update extends to other Muon-style optimizers, using Polar Express as a representative example. Across all configurations, auxiliary parameters, including normalization gains and biases when present, are optimized with AdamW. Table 1 summarizes the six configurations, where ηa,ηh,ηi _a, _h, _i, and ηo _o denote learning rates. Within each configuration, repeated learning rate symbols indicate that the corresponding parameter groups share a common learning rate during hyperparameter tuning. Table 1: Optimizer configurations used in our comparisons. Each entry specifies the optimizer and learning rate for the corresponding parameter group. Repeated learning rate symbols indicate that the same learning rate is shared across parameter groups. For each configuration, all distinct learning rates shown in the corresponding row are jointly tuned. Configuration Embedding Hidden 2D LM Head Gains & Biases Plain Muon (baseline 1) AdamW ηa _a Muon ηm _m AdamW ηa _a AdamW ηa _a Muon Tuned (baseline 2) AdamW ηi _i Muon ηm _m AdamW ηo _o AdamW ηa _a Ember Ember ηi _i Muon ηm _m Ember ηo _o AdamW ηa _a MuonIO (ours) IO ηi _i Muon ηm _m IO ηo _o AdamW ηa _a Polar Express Tuned AdamW ηi _i Polar Express ηh _h AdamW ηo _o AdamW ηa _a Polar Express IO (ours) IO ηi _i Polar Express ηh _h IO ηo _o AdamW ηa _a 4.2 Optimizer Performances Table 2 reports final validation perplexity at all three model scales. Table 2: Final validation perplexity on C4 (lower is better). The upper block uses standard Muon hidden updates, while the lower block uses Polar Express. Bold values indicate the lowest mean perplexity within each block at each model scale. For 60M and 130M, we report the mean and sample standard deviation over the three of five seeds with the lowest final validation perplexity. For 1B, we report the mean and sample standard deviation over three seeds. Configuration 60M 130M 1B Plain Muon (baseline 1) 36.963±0.05436.963± 0.054 28.513±0.02628.513± 0.026 15.116±0.000315.116± 0.0003 Muon Tuned (baseline 2) 28.502±0.01528.502± 0.015 24.238±0.00824.238± 0.008 14.527±0.00714.527± 0.007 Ember 28.835±0.04328.835± 0.043 24.074±0.00424.074± 0.004 14.516±0.03214.516± 0.032 MuonIO (ours) 28.366±0.02528.366± 0.025 23.617±0.02523.617± 0.025 14.374±0.13214.374± 0.132 Polar Express Tuned 28.385±0.00928.385± 0.009 24.055±0.01024.055± 0.010 14.361±0.00214.361± 0.002 Polar Express IO (ours) 28.430±0.02528.430± 0.025 23.646±0.01223.646± 0.012 14.160±0.00814.160± 0.008 With standard Muon hidden updates, MuonIO achieves the lowest mean validation perplexity among the four configurations at every model scale, including improvements over the strong baseline Muon Tuned which separately tuned learning rates for all four components of the model. With Polar Express hidden updates, our algorithm lowers mean perplexity at 130M and 1B, while slightly increasing it at 60M. These results show that the improvements at the two larger scales also extend to the Polar Express setting. Figure 1 complements the final perplexity comparison with validation loss trajectories for the four configurations that used Muon (top 4) for hidden 2D parameters at the 130M scale. Figure 1: Validation loss trajectories on C4 at the 130M scale for Plain Muon (baseline 1), Muon Tuned (baseline 2), Ember and MuonIO. All four configurations use standard Muon hidden updates. The inset highlights the 0.7–0.9B training token range. 4.3 IO versus AdamW: exact optimizer cost Table 3 compares Algorithm 1’s exact per-update cost with AdamW on the same V×dV× d matrix. It needs roughly half the FLOPs and, because it stores only momentum, half the optimizer state. Appendix C gives the counting convention and the operation-by-operation derivation. Table 3: FLOPs and persistent optimizer state counts for one V×dV× d I/O matrix and one optimizer update. Derivations are in Appendix C. I/O-layer optimizer FLOPs Optimizer State AdamW 13Vd13Vd 2Vd2Vd entries IO 7Vd+3V7Vd+3V VdVd entries 4.4 Gradient Norm Dynamics To examine how different optimizers affect the training dynamics of the I/O layers, Figure 2 tracks the gradients of these layers during training of the 130M model. We compare three configurations: MuonIO, Muon Tuned, and Ember, all of which use standard Muon updates for the hidden weight matrices and differ only in their treatment of the I/O layers. We consider the gradient both as a matrix and as a flattened vector. From the matrix perspective, we measure its Schatten-22 and Schatten-∞ norms, corresponding to the Frobenius and spectral norms, respectively. From the vector perspective, we consider the ℓ2 _2 and ℓ∞ _∞ norms. Since the vectorized ℓ2 _2 norm is identical to the Frobenius norm, this yields three distinct quantities: the Frobenius norm, spectral norm, and maximum absolute entry. These norms capture complementary aspects of the gradient magnitude and serve as diagnostics of training progress on the I/O layers. Figure 2: Gradient magnitudes during 130M model training. Rows show the input embedding table (top) and language model head (bottom); columns show the Frobenius norm, spectral norm, and maximum absolute gradient entry. Curves correspond to MuonIO (orange, solid), Ember (blue, dashed), and Muon Tuned (green, dash-dotted). The horizontal axis counts microbatch steps, with eight microbatches per optimizer update. Each curve represents one run using the best set of hyperparamters of the corresponding optimizer. Across all three metrics, MuonIO drives the I/O gradients down more rapidly than the competing configurations. For the embedding table, MuonIO already exhibits substantially smaller gradient norms at the earliest recorded measurements and quickly approaches near-zero values, whereas Muon Tuned shows an initial increase before decaying and Ember starts with considerably larger gradients. The difference is even more pronounced for the language-model head: MuonIO rapidly reduces all three gradient measures within the first few thousand steps, while Muon Tuned decays more gradually and Ember maintains substantially larger gradients throughout much of the displayed interval. The maximum absolute entry metric shows the clearest separation, with Ember exhibiting large and persistent individual gradient coordinates while MuonIO quickly suppresses them. Overall, these trajectories indicate that, under MuonIO, the embedding and language model head gradients contract substantially earlier in training, particularly for the language-model head. 5 Conclusion We have presented MuonIO, one Muon-style update for the two layer types that the original Muon recipe and many implementations route to AdamW: the input embedding table and the output language model head. For the embedding table ∈ℝd×VE ^d× V, we adopt the established 1→RMS1 geometry induced by one-hot lookup (Bernstein and Newhouse, 2024; Bernstein and Newhouse, 2025), or equivalently 1→21→ 2 up to the fixed d−1/2d^-1/2 RMS factor. For the output head ∈ℝV×dL ^V× d, our total variation stability analysis of softmax identifies the 2→∞2→∞ geometry. The resulting algorithm MuonIO is also cheaper than AdamW for vocabulary-facing layers in both optimizer state memory and FLOPs. Our numerical experiments further show that MuonIO improves final loss and perplexity. 6 Limitations and future work Our analysis of lm_head measures perturbations of the output distribution using total variation and controls them through the 2→∞2→∞ norm. Although this formulation provides a sharp uniform bound and leads to a computationally cheap row-normalization update, it does not explicitly account for the invariance of softmax to constant shifts of all logits. Future work could investigate shift-invariant formulations as well as other measures of distributional change, such as KL divergence, and evaluate whether their optimization benefits justify any additional computational cost. References Ahn et al. (2025a) K. Ahn, N. Amsel, and J. Langford Dion2: a simple method to shrink matrix in muon. arXiv preprint arXiv:2512.16928. Cited by: §1. Ahn et al. (2025b) K. Ahn, B. Xu, N. Abreu, Y. Fan, G. Magakyan, P. Sharma, Z. Zhan, and J. Langford Dion: distributed orthonormalized updates. arXiv preprint arXiv:2504.05295. Cited by: §1. Amsel et al. (2026) N. Amsel, D. Persson, C. Musco, and R. M. Gower The polar express: optimal matrix sign methods and their application to the muon algorithm. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §1, §4. Bernstein and Newhouse (2024) J. Bernstein and L. Newhouse Old optimizer, new norm: an anthology. In OPT 2024: Optimization for Machine Learning, Note: arXiv:2409.20325 External Links: Link Cited by: §1.1, §1, §1, §2.2, §2.3, §3.1, §5. Bernstein and Newhouse (2025) J. Bernstein and L. Newhouse Modular duality in deep learning. In International Conference on Machine Learning, p. 3920–3930. Cited by: Appendix A, §1.1, §1.2, §1, §1, §5, Abstract. Bernstein (2025) J. Bernstein Deriving muon. External Links: Link Cited by: §1, §2.4. Glentis et al. (2026) A. Glentis, J. Li, A. Han, and M. Hong Memory-efficient LLM pretraining via minimalist optimizer design. In International Conference on Machine Learning, Note: arXiv:2506.16659 External Links: Link Cited by: Appendix A, §1.2, §1. Hayou and Liu (2025) S. Hayou and L. Liu Optimal embedding learning rate in llms: the effect of vocabulary size. arXiv preprint arXiv:2506.15025. Cited by: §1.2. He et al. (2026) D. He, S. Tu, K. Wang, L. Yin, and S. Liu One LR doesn’t fit all: heavy-tail guided layerwise learning rates for LLMs. In ICLR 2026 2nd Workshop on Deep Generative Model in Machine Learning: Theory, Principle and Efficacy, External Links: Link Cited by: §1.2. Inan et al. (2016) H. Inan, K. Khosravi, and R. Socher Tying word vectors and word classifiers: a loss framework for language modeling. arXiv preprint arXiv:1611.01462. Cited by: §1.1, §3.3. 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: Appendix A, §1, §1, §4. Kalra and Barkeshli (2026) D. S. Kalra and M. Barkeshli Quantifying hyperparameter transfer and the importance of embedding layer learning rate. arXiv preprint arXiv:2605.21486. Cited by: §1.2. Kingma and Ba (2014) D. P. Kingma and J. Ba Adam: a method for stochastic optimization. arXiv preprint arXiv:1412.6980. Cited by: §1. Lau and Su (2026) T. T. Lau and W. Su Symmetry-compatible principle for optimizer design: embeddings, LM heads, SwiGLU MLPs, and MoE routers. arXiv preprint arXiv:2605.18106. External Links: Link Cited by: Appendix A, §1.2, §1. Li et al. (2026) Z. Li, L. Liu, C. Liang, W. Chen, and T. Zhao NorMuon: making muon more efficient and scalable. In Forty-third International Conference on Machine Learning, External Links: Link Cited by: §1. Liu et al. (2025) J. Liu, J. Su, X. Yao, Z. Jiang, G. Lai, Y. Du, Y. Qin, W. Xu, E. Lu, J. Yan, et al. Muon is scalable for llm training. arXiv preprint arXiv:2502.16982. Cited by: §1. Nair (2025) P. Nair Softmax is 1/21/2-lipschitz: a tight bound across all ℓp _p norms. arXiv preprint arXiv:2510.23012. External Links: Link Cited by: Appendix A, §1.2. Pang et al. (2026) T. Pang, Y. Fang, Z. Liu, S. Deng, L. Hsiung, S. Yu, and Y. Yang Htmuon: improving muon via heavy-tailed spectral correction. In Findings of the Association for Computational Linguistics: ACL 2026, p. 36504–36535. Cited by: §1. 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. In International Conference on Machine Learning, p. 49069–49104. Cited by: Appendix A, §1.2, §1, §1. Popoviciu (1935) T. Popoviciu Sur les équations algébriques ayant toutes leurs racines réelles. Mathematica 9 (129-145), p. 20. Cited by: §3.2.1. Press and Wolf (2017) O. Press and L. Wolf Using the output embedding to improve language models. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers, p. 157–163. Cited by: §1.1, §3.3. Raffel et al. (2020) C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research 21 (140), p. 1–67. Cited by: §4.1. Shah (2026) K. Shah Token geometry. arXiv preprint arXiv:2607.01455. External Links: Link Cited by: Appendix A, §1.2, §4. Si et al. (2025) C. Si, D. Zhang, and W. Shen Adamuon: adaptive muon optimizer. arXiv preprint arXiv:2507.11005. Cited by: §1. Team et al. (2025) K. Team, Y. Bai, Y. Bao, Y. Charles, C. Chen, G. Chen, H. Chen, H. Chen, J. Chen, N. Chen, et al. Kimi k2: open agentic intelligence. arXiv preprint arXiv:2507.20534. Cited by: §1. Touvron et al. (2023) H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample LLaMA: open and efficient foundation language models. External Links: 2302.13971, Link Cited by: §4.1. Yang et al. (2023) G. Yang, J. B. Simon, and J. Bernstein A spectral condition for feature learning. arXiv preprint arXiv:2310.17813. Cited by: §1. Zhao et al. (2024) J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian GaLore: memory-efficient llm training by gradient low-rank projection. In International Conference on Machine Learning, p. 61121–61143. Cited by: §4.1. Appendix A Related Work The original Muon (Jordan et al., 2024) optimizer uses approximate polar factor updates for hidden linear layers and assigns embeddings and language model heads to AdamW. Bernstein and Newhouse (Bernstein and Newhouse, 2025) already derive the embedding geometry used by MuonIO: their 1→RMS1 norm formulation, equivalent to 1→21→ 2 up to a dimension-dependent factor, gives an update that normalizes each column, though this choice is not empirically evaluated. Scion’s (Pethick et al., 2025) operator norm LMO framework includes column normalization for one-hot inputs and row normalization as an RMS→∞RMS→∞ output option, which is motivated by control of the largest output coordinate rather than softmax stability. Separately, Nair (2025) establishes softmax Lipschitz bounds under matched ℓp _p norms. Our analysis instead uses the mixed ℓ∞→ℓ1 _∞→ _1 relation to motivate the 2→∞2→∞ norm for a language model head through total variation stability. Several recent optimizers use related normalization directions or otherwise give special treatment to token-facing matrices. SCALE applies column normalization to every layer and uses momentum only for the final layer (Glentis et al., 2026). After transposition, its final layer normalization resembles one component of MuonIO, but SCALE is otherwise a distinct algorithm motivated by observed gradient variance and token frequency imbalance. By contrast, Lau and Su (2026) use a different principle. They require updates to be equivariant to vocabulary reordering and orthogonal transformations of hidden space. The resulting family includes the token normalization used here. For untied heads, their rule also centers momentum, removes common logit shifts, and chooses scale separately. Ember (Shah, 2026) also treats embeddings and language model heads specially, but follows a different formulation and algorithm using factorized second moments rather than a Muon-style update. Appendix B Proofs Proof of Proposition 3.3. First suppose V is even. Take i=1Vp_i= 1V for all i (i.e., =z=0). Set i=1u_i=1 when i is even and i=−1u_i=-1 when i is odd. Starting from ‖J()‖1=∑i=1Vi|i−m|\|J(z)u\|_1= _i=1^Vp_i|u_i-m|, we have m=0,|i−m|=1∀i,m=0,|u_i-m|=1\ ∀ i, hence ‖J()‖1=∑i=1Vi|i−m|=∑i=1V1V=1.\|J(z)u\|_1= _i=1^Vp_i|u_i-m|= _i=1^V 1V=1. When V is odd, a uniform p cannot place mass exactly 12 12 on the coordinates with i=+1u_i=+1, so we double the mass of one coordinate instead: take =(0,…,0,ln2)z=(0,…,0, 2), so that =(1V+1,…,1V+1,2V+1).p= ( 1V+1,…, 1V+1, 2V+1 ). Choose ∈±1Vu∈\± 1\^V with V=+1u_V=+1 and exactly V−32 V-32 of the first V−1V-1 coordinates equal to +1+1 (an integer since V is odd). The mass on the set i:i=+1\i:u_i=+1\ is then 2V+1+(V−3)/2V+1=4+V−32(V+1)=12, 2V+1+ (V-3)/2V+1= 4+V-32(V+1)= 12, so again m=0m=0 and |i−m|=1|u_i-m|=1 for all i, whence ‖J()‖1=∑i=1Vi=1\|J(z)u\|_1= _i=1^Vp_i=1. (Conceptually, this is the even case with V+1V+1 uniform coordinates, two of which — sharing the sign +1+1 — are merged into a single coordinate of double mass.) ∎ Proof of Proposition 3.5. Step 1: the supremum equals 11. Assume first that V is even and let ∈±1Vu∈\± 1\^V alternate in sign as in Proposition 3.3. For ε>0 >0 set =−ε2z=- 2u and ′=ε2z = 2u, so that ‖′−‖∞=ε\|z -z\|_∞= . Along the line (t)=tz(t)=tu, the softmax takes only two distinct values: pi(t)=2V⋅etet+e−tp_i(t)= 2V· e^te^t+e^-t on the V/2V/2 coordinates with ui=+1u_i=+1, and pi(t)=2V⋅e−tet+e−tp_i(t)= 2V· e^-te^t+e^-t on the rest. Summing the coordinatewise differences between t=ε/2t= /2 and t=−ε/2t=- /2 gives ‖softmax(′)−softmax()‖1=2eε/2−e−ε/2eε/2+e−ε/2=2tanh(ε/2),\|softmax(z )-softmax(z)\|_1=2\, e /2-e^- /2e /2+e^- /2=2 ( /2), (23) so the ratio in (15) equals 2tanh(ε/2)/ε→12 ( /2)/ → 1 as ε↓0 0. For V odd, apply the same construction to the first V−1V-1 coordinates and set the last coordinate of both z and ′z to −M-M; letting M→∞M→∞ recovers the same limit. Combined with the upper bound of Lemma 3.4, the supremum equals 11. Step 2: the supremum is not attained. Suppose equality held in Lemma 3.4 for some finite ′≠z . Write =′−u=z -z, let γ(t)=+tγ(t)=z+tu for t∈[0,1]t∈[0,1] be the segment from the proof of Lemma 3.4, and let (t):=softmax(γ(t))∈ΔVp(t):=softmax(γ(t))∈ ^V denote the softmax along this segment. Since t↦‖J(γ(t))‖1t \|J(γ(t))u\|_1 is continuous and bounded above by ‖∞\|u\|_∞, equality of the integral bound in the proof of Lemma 3.4 forces ‖J(γ(t))‖1=‖∞\|J(γ(t))u\|_1=\|u\|_∞ for every t∈[0,1]t∈[0,1]. Let XtX_t take value uiu_i with probability pi(t)p_i(t), as in the proof of Lemma 3.2. Equality in Jensen’s inequality requires |Xt−[Xt]||X_t-E[X_t]| to be almost surely constant, and equality in Popoviciu’s inequality requires XtX_t to be supported on −‖∞,+‖∞\-\|u\|_∞,+\|u\|_∞\ with mass 12 12 each. Together these give m(t):=(t)⊤=[Xt]=0m(t):=p(t) u=E[X_t]=0 and Var(Xt)=‖∞2Var(X_t)=\|u\|_∞^2 for every t∈[0,1]t∈[0,1]. Hence m≡0m≡ 0 and so m′≡0m ≡ 0. But differentiating along the segment gives m′(t)=⊤J(γ(t))=Var(Xt)=‖∞2>0m (t)=u J(γ(t))u=Var(X_t)=\|u\|_∞^2>0, which is a contradiction. Hence the inequality in Lemma 3.4 is strict for every finite pair, consistent with Step 1, where 2tanh(ε/2)<ε2 ( /2)< strictly for every ε>0 >0. ∎ Proof of Proposition 3.7. We show that every logit pair in (15) is realized by an lm_head triple with the same ratio. Fix any nonzero ∈ℝdy ^d. For ,′∈ℝVz,z ^V with ′≠z , define the rank-one matrices =⊤‖22,Δ=(′−)⊤‖22,L= z\,y \|y\|_2^2, = (z -z)\,y \|y\|_2^2, (24) so that =Ly=z and (+Δ)=′(L+ )y=z . The t-th row of Δ is (′−)t/‖22(z -z)_t\,y/\|y\|_2^2, whose ℓ2 _2 norm is |(′−)t|/‖2|(z -z)_t|/\|y\|_2. Proposition 2.2 then gives ‖Δ‖2→∞‖2=‖′−‖∞.\| \|_2→∞\,\|y\|_2=\|z -z\|_∞. (25) Consequently, the ratio in (18) evaluated at (,Δ,)(L, ,y) equals the ratio in (15) evaluated at (,′)(z,z ). Taking the supremum over ′≠z and invoking Proposition 3.5, the supremum in (18) is at least 11; by Lemma 3.6 it is at most 11. For non-attainment, suppose some triple (,Δ,)(L, ,y) with Δ≠ 0, ≠y 0 attained equality in (18). If (Δ)=( )y=0, the left hand side of Lemma 3.6 vanishes while the right hand side is positive, so equality fails. Otherwise, equality forces in particular ‖softmax(′)−softmax()‖1=‖′−‖∞\|softmax(z )-softmax(z)\|_1=\|z -z\|_∞ for the finite pair =≠′=(+Δ)z=Ly =(L+ )y, contradicting Step 2 of the proof of Proposition 3.5. ∎ Proof of Remark 3.8. Using the dual norm characterization ‖q=sup‖q∗≤1⟨,⟩\|y\|_q= _\|v\|_q^*≤ 1 ,y and bilinearity of the inner product, ‖p→q \|A\|_p→ q =sup‖p≤1‖q = _\|x\|_p≤ 1\|Ax\|_q =sup‖p≤1sup‖q∗≤1⟨,⟩ = _\|x\|_p≤ 1\; _\|v\|_q^*≤ 1 ,Ax =sup‖q∗≤1sup‖p≤1⟨⊤,⟩ = _\|v\|_q^*≤ 1\; _\|x\|_p≤ 1 v,x =sup‖q∗≤1‖⊤‖p∗ = _\|v\|_q^*≤ 1\|A v\|_p^* =‖⊤‖q∗→p∗. =\|A \|_q^*→ p^*. Specializing to p=2p=2, q=∞q=∞ (so p∗=2p^*=2, q∗=1q^*=1) gives (20). ∎ Appendix C Exact FLOP and Optimizer-State Accounting This appendix derives the counts reported in Table 3. Let one input or output parameter matrix have shape V×dV× d (or its transpose) and let N:=VdN:=Vd be its number of entries. We count FLOPs in the following way. One scalar addition, subtraction, multiplication, division, comparison or square root counts as one FLOP; consequently, a fused multiply–add counts as two FLOPs. Operations that have in total O(1)O(1) FLOPS throughout one optimizer call are ignored. AdamW. Denote the learning rate, weight decay and the betas with γ,λ,β1,β2γ,λ, _1, _2. Let kM_k and kV_k denote the first and second moment buffers. Absorb the learning rate, bias corrections, and decoupled weight decay into scalars that are computed once per tensor. A bias-corrected AdamW update can then be organized into the four stages in Table 4. In the last two rows, kD_k denotes the corrected square root denominator. Table 4: AdamW tensor-arithmetic count for one matrix and one optimizer update. The elementwise product and division are denoted by ⊙ and ⊘ , respectively. Stage Tensor operation FLOPs First moment k←β11−β1tk−1+1−β11−β1tkM_k← _11- _1^tM_k-1+ 1- _11- _1^tG_k 3N3N Second moment k←β21−β2tk−1+1−β21−β2tk⊙kV_k← _21- _2^tV_k-1+ 1- _21- _2^tG_k _k 4N4N Corrected denominator k←k+εD_k← V_k+ 2N2N Decay and adaptive step k+1←(1−γλ)k−γ(k⊘k)W_k+1←(1-γλ)W_k-γ(M_k _k) 4N4N Total 13N=13Vd13N=13Vd The first moment update uses two multiplications and one addition per entry. The second moment update additionally squares the gradient, giving three multiplications and one addition per entry. The corrected denominator uses one square root and one addition per entry. Finally, the parameter update uses two multiplications, one division, and one subtraction per entry. Thus the exact FLOP count is 13N13N. AdamW keeps both kM_k and kV_k, for 2N2N optimizer state entries. MuonIO. We count one iteration of Algorithm 1 on the row-oriented matrix k∈ℝV×dW_k ^V× d from (21). Write rj:=∥(k+1)j,:∥2r_j:=\|(M_k+1)_j,:\|_2 for the norm of row j. Each rjr_j is computed once and reused in both the normalization and the sum sk+1s_k+1. We then form a single per-row scale cj=ak+1/maxrj,εc_j=a_k+1/ \r_j, \ with ak+1=λdsk+1a_k+1=λ d\,s_k+1. This fuses the normalization and scaling into the parameter update, so k+1U_k+1 is never materialized. Table 5 gives the resulting count. Table 5: MuonIO tensor-arithmetic count for one matrix and one iteration of Algorithm 1. The constant λdλ d is computed once, before training. Stage Operation FLOPs Momentum k+1←βk+(1−β)~kM_k+1← _k+(1-β) G_k 3N3N Row norms rj←∥(k+1)j,:∥2r_j←\|(M_k+1)_j,:\|_2, j∈[V]j∈[V] V[d+(d−1)+1]=2NV[d+(d-1)+1]=2N Sum norms sk+1←∑j=1Vrjs_k+1← _j=1^Vr_j V−1V-1 Form scalar scale ak+1←(λd)sk+1a_k+1←(λ d)\,s_k+1 11 Form row scales cj←ak+1/maxrj,εc_j← a_k+1/ \r_j, \, j∈[V]j∈[V] 2V2V Fused parameter update (k+1)j,:←(k)j,:−cj(k+1)j,:(W_k+1)_j,:←(W_k)_j,:-c_j(M_k+1)_j,:, j∈[V]j∈[V] 2N2N Total 7N+3V=7Vd+3V7N+3V=7Vd+3V For each row of length d, the norm takes d squarings, d−1d-1 additions, and one square root, or exactly 2d2d FLOPs. Summing over the V rows gives 2Vd=2N2Vd=2N. Forming cjc_j costs one division and one comparison per row. The fused update, which equals line 6 of Algorithm 1, then costs one multiplication and one subtraction per matrix entry. The only persistent optimizer tensor is the momentum k+1M_k+1, so the optimizer state contains N entries. Appendix D Experiment Settings and Hyperparameters We tune optimizer hyperparameters on the C4 validation set, minimizing validation loss or perplexity. For the four-group and IO variants of Polar Express, we use Bayesian optimization with a maximum of 100 trials per sweep and logarithmic learning rate search ranges of [0.005,0.08][0.005,0.08] for hidden weights, [0.001,0.3][0.001,0.3] for embeddings, [0.0003,0.2][0.0003,0.2] for the output head, and [0.0003,0.03][0.0003,0.03] for gains and biases. We search the Polar Express iteration count over 5,6,8\5,6,8\. For Muon, we perform a grid search over hidden weight learning rates 0.005,0.01,0.02\0.005,0.01,0.02\ and auxiliary learning rates 0.0003,0.0005,0.001\0.0003,0.0005,0.001\, since it has the same learning rates for hidden_gains_biases and nonhidden_params. For Muon Tuned, MuonIO and Ember, we independently search discrete embedding and output head learning rate grids spanning [0.001,0.3][0.001,0.3] and [0.0003,0.5][0.0003,0.5], respectively, fixing the hidden weight learning rate to 0.020.02 and 0.040.04, respectively. Table 6: Optimizer hyperparameters for the 18 selected LLaMA pre-training runs on C4. Values use three decimal places; nonzero values below 10−310^-3 use scientific notation. μh _h is hidden weight momentum; (β1A,β2A)( _1^A, _2^A) are Adam/AdamW coefficients; βI/O _I/O is the separate coefficient for non-Adam I/O updates. PE denotes Polar Express. Method Size ηh _h ηi _i ηo _o ηa _a μh _h β1A _1^A β2A _2^A βI/O _I/O Plain Muon 60M 0.010 3×10−43× 10^-4 3×10−43× 10^-4 0.001 0.950 0.900 0.950 – 130M 0.010 3×10−43× 10^-4 3×10−43× 10^-4 0.001 0.950 0.900 0.950 – 1B 0.019 3×10−43× 10^-4 3×10−43× 10^-4 0.023 0.950 0.950 0.950 – Muon Tuned 60M 0.020 0.040 0.020 0.020 0.900 0.900 0.950 – 130M 0.020 0.040 0.020 0.020 0.900 0.900 0.950 – 1B 0.016 0.197 0.004 0.015 0.900 0.900 0.950 – Ember 60M 0.040 0.100 0.020 0.020 0.900 0.900 0.950 0.900 130M 0.040 0.100 0.010 0.020 0.900 0.900 0.950 0.900 1B 0.033 0.273 0.002 0.017 0.900 0.900 0.950 0.900 MuonIO 60M 0.040 0.300 0.005 0.020 0.900 0.900 0.950 0.900 130M 0.040 0.300 0.005 0.020 0.900 0.900 0.950 0.900 1B 0.031 0.185 0.001 0.020 0.900 0.900 0.950 0.900 PE Tuned 60M 0.030 0.047 0.004 0.012 0.900 0.900 0.950 – 130M 0.014 0.034 0.013 0.025 0.950 0.900 0.950 – 1B 0.016 0.104 0.005 0.020 0.900 0.900 0.950 – PE IO 60M 0.036 0.040 9.086×10−49.086× 10^-4 0.026 0.900 0.900 0.950 0.900 130M 0.028 0.161 3.199×10−43.199× 10^-4 0.014 0.900 0.900 0.950 0.900 1B 0.026 0.056 0.002 0.017 0.900 0.900 0.950 0.900 Parameter groups. Table 6 lists the selected optimizer hyperparameters. The learning rates ηh _h, ηi _i, ηo _o, and ηa _a correspond to hidden weights, embeddings, the language model head, and gains/biases, respectively. Muon Tuned denotes the Muon baseline with separately tuned I/O learning rates. PE Tuned and PE IO denote the corresponding Polar Express configurations. Momentum and Adam coefficients. The coefficients (β1A,β2A)( _1^A, _2^A) apply to gains/biases for every method and additionally to embeddings and the language model head for Plain Muon, Muon Tuned, and PE Tuned. For MuonIO and PE IO, βI/O _I/O is the I/O momentum coefficient; their I/O updates use no second moment coefficient. For Ember, βI/O _I/O is the I/O second moment decay coefficient; its I/O updates use no first moment coefficient. A dash in the table indicates that the I/O layers use the displayed Adam coefficient pair. Hidden weight momentum μh _h is specified separately. Shared training settings. All runs use C4, a global batch size of 512 sequences, a sequence length of 256, and bfloat16 precision. The learning rate follows a cosine decay schedule to 0.1 of its peak value, with a warmup fraction of 0.1. The embedding and language model head weights are untied. Batching. We describe each batching configuration by a tuple (microbatch size per GPU, gradient accumulation steps, number of GPUs), whose product is the global batch size of 512. At the 60M scale, runs use (128,4,1)(128,4,1), except Muon Tuned and MuonIO, which use (64,8,1)(64,8,1). At the 130M scale, runs use (64,8,1)(64,8,1), except Plain Muon, which uses (128,4,1)(128,4,1). At the 1B scale, runs use (64,8,1)(64,8,1), except MuonIO, which uses (64,2,4)(64,2,4). Model configuration. For the 60M, 130M, and 1B models, respectively, the hidden widths are (512,768,2048)(512,768,2048), the numbers of transformer layers are (8,12,24)(8,12,24), the numbers of attention heads are (8,12,32)(8,12,32), and the feed-forward widths are (1376,2048,5461)(1376,2048,5461). All models use an RMSNorm epsilon of 10−610^-6 and a configured maximum position count of 2048, while the training sequence length is 256.