Paper deep dive
An Isotropy-Preserving Spectral Cap for Muon: Theory and Three Case Studies
Jiachun Li
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 90%
Last extracted: 7/23/2026, 2:37:45 AM
Summary
This preliminary report proposes a 'spectral cap' for the Muon optimizer to preserve isotropy in large language model training. The authors argue that Muon's matrix-sign step removes the natural 1/||W||_F brake present in SGD, causing Frobenius and spectral norms to drift outward faster (t^{1/2} vs t^{1/4}). A lightweight spectral cap, which projects out the first-order growth of the top singular direction, controls output covariance without freezing training. The method is validated on three case studies: nanoGPT feed-forward projections, a 64-expert MoE router, and FlashAttention query/key projections, showing improved isotropy and prevention of failure modes like expert collapse or attention head divergence.
Entities (10)
Relation Signals (6)
Spectral Cap → controls → Spectral Norm
confidence 95% · a lightweight 'spectral cap' ... can control the output covariance ... without freezing training
Muon → causesfasternormgrowth → Frobenius Norm
confidence 92% · Muon's matrix-sign step removes that brake, so both the Frobenius and spectral norms drift outward faster (t^{1/2} versus t^{1/4})
Spectral Cap → increases → Isotropy
confidence 90% · In each case the cap increases isotropy
Muon → removesbrake → SGD
confidence 90% · plain SGD carries a built-in 1/||W|| brake on its update size, whereas Muon's matrix-sign step removes that brake
Spectral Cap → prevents → MoE Router
confidence 85% · prevents a concrete failure, while leaving validation loss essentially unchanged ... a router collapsing to a single expert
Scale Invariance → underlies → Muon
confidence 80% · This preliminary report proposes a unified framework built on a single idealizing assumption -- exact scale invariance ... Under this assumption ... Muon's matrix-sign step removes that brake
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Muon and related matrix-sign optimizers are increasingly used to pre-train large language models, but their effect on the internal geometry of individual weight matrices is not well understood. This preliminary report proposes a unified framework built on a single idealizing assumption -- exact scale invariance of the loss under weight rescaling, which holds approximately in normalization-heavy networks. Under this assumption, plain SGD carries a built-in 1/||W|| brake on its update size, whereas Muon's matrix-sign step removes that brake, so both the Frobenius and spectral norms drift outward faster (t^{1/2} versus t^{1/4}). We further observe that the spectral-norm perturbation has a non-negative second-order term. This implies that a lightweight "spectral cap" -- which projects out only the first-order growth of the single top singular direction from each update -- can control the output covariance W K_X W^T without freezing training: the weight keeps learning through non-top directions, top-direction rotation, and top switching. We relate this cap to the min-entropy (H-infinity) of the singular-value spectrum. We then study three systems trained with Muon: a nanoGPT feed-forward projection, a 64-expert mixture-of-experts router, and the query/key projections of a bf16 FlashAttention block. In each case the cap increases isotropy and, at the margins -- a router collapsing to a single expert, and the near-divergence of one attention head -- prevents a concrete failure, while leaving validation loss essentially unchanged. We emphasize that the scale-invariance assumption is strong and that these small-scale results are preliminary; comments are welcome.
Tags
Links
- Source: https://arxiv.org/abs/2607.19771v1
- Canonical: https://arxiv.org/abs/2607.19771v1
Trouble viewing inline? Open PDF directly →
Full Text
64,804 characters extracted from source content.
Expand or collapse full text
An Isotropy-Preserving Spectral Cap for Muon: Theory and Three Case Studies (Preliminary report — comments welcome) Jiachun Li (2026-06 (draft)) Abstract This report has four parts. Part 1 (theory) proposes a single explanatory framework. In a scale-invariant network the loss is insensitive to the scale ‖W‖F\|W\|_F of a weight matrix. Under this idealisation, plain SGD carries a built-in 1/‖W‖F1/\|W\|_F “brake” on its update size, whereas the matrix-sign step used by Muon removes that brake, so both the Frobenius norm and the spectral norm drift outward faster. We then note that the perturbation of the spectral norm has a non-negative second-order term (a convexity effect): even after the first-order top-direction component of an update is projected away, the weight can still learn through three mechanisms — growth of non-top singular directions, rotation of the top direction, and top-direction switching. This is the core reason a spectral cap can control the spectrum without freezing training. Parts 2, 3 and 4 test the idea on three concrete systems: a nanoGPT feed-forward projection, a 64-expert MoE router, and the query/key projections of a FlashAttention block run in bf16. For each system we state precisely what the weight W, the input activation X, and the covariance KXK_X are, together with the algorithm and cost of finding the top singular pair (u1,v1)(u_1,v_1). The three case studies at a glance: Part System W (matrix the cap controls) X (input activation) Main finding 2 nanoGPT FFN mlp.c_proj ∈ℝ384×1536 ^384× 1536 GELU hidden, X∈ℝM×1536X ^M× 1536 data-cap lowers the top-share of the c_proj output covariance KYK_Y from 0.740.74 to 0.300.30 with no change in val loss 3 64-expert MoE router router ∈ℝ64×384 ^64× 384 post-RMSNorm token, X∈ℝM×384X ^M× 384 selection cap pushes the top-share of Kz+bK_z+b to 0.120.12–0.520.52 across all six layers; without it, layer 0 collapses to a single dominant expert (top-share 0.9950.995) 4 FlashAttention Q/K WQ(h),WK(h)∈ℝ64×768W_Q^(h),W_K^(h) ^64× 768 (per head) token embedding, X∈ℝM×768X ^M× 768 at iteration 86K a single head (L1 H2) reaches λmax=133 _ =133K and is caught in real time by the cap; the uncapped baseline crashes in the same window caveatbox Status of this report. This is a preliminary technical report. The theory in Section 1 rests on a fairly strong idealisation — exact scale invariance of the loss with respect to weight rescaling (Assumption 1). Real normalisation-heavy networks are only approximately scale invariant, so several of the theoretical claims should be read as clean limiting statements rather than established facts. We give a first empirical sanity check in Section 1.4.4, but confirming (or refuting) the assumption and its consequences will require substantially more experiments. We share the report at this stage precisely to invite feedback; corrections and counter-examples are very welcome. Contents 1 Theory: a spectral cap and isotropy under Muon 1.1 Setup and goal 1.2 The scale-invariance assumption 1.2.1 Corollary B (weak form): the gradient is orthogonal to W 1.2.2 Corollary A (scale form): the gradient scales like 1/‖W‖F1/\|W\|_F 1.3 SGD vs Muon: who keeps the 1/‖W‖F1/\|W\|_F brake? 1.3.1 SGD update scale 1.3.2 Muon update scale: msign annihilates the 1/‖W‖F1/\|W\|_F factor 1.4 Finding 1: Muon grows the Frobenius / spectral norm faster than SGD 1.4.1 Norm increment 1.4.2 SGD: ‖Wt‖F∼t1/4\|W_t\|_F t^1/4 1.4.3 Muon: ‖Wt‖F∼t1/2\|W_t\|_F t^1/2 1.4.4 Empirical check of scale invariance via cPc_P and cGc_G 1.4.5 Spectral norm: Muon moves ‖W‖2\|W\|_2 more than SGD when ‖W‖F\|W\|_F is large 1.5 Finding 2: after a cap, the spectral-norm second order is ≥0≥ 0 1.5.1 Symmetric dilation 1.5.2 Why the second order being ≥0≥ 0 is exactly what we want 1.6 Extension: the cap targets H∞H_∞ entropy; other entropies give other cap directions 1.7 Remark: we do not actually need KX=IK_X=I 2 Application 1: nanoGPT FFN (proof of concept) 2.1 What KXK_X is here 2.1.1 Estimating KYK_Y and finding the top eigenpair (q,λ1)(q, _1) 2.2 Experimental setup 2.3 Result: the three methods have equal val loss 2.4 What the cap changes: spectral norm and isotropy side by side 2.5 Participation-rank figure 2.6 Part 2 summary 3 Application 2: 64-expert MoE router 3.1 What KXK_X is here 3.1.1 Algorithm: estimating Kz+bK_z+b and finding the top eigenpair 3.2 What isotropy means in the MoE 3.3 Core result: Kz+bK_z+b isotropy across the three methods 3.4 L0 / L5 load balancing: a phase transition and LFB failure 3.4.1 The L0 story: G0 collapses to rank-1 3.4.2 The L5 story: G3 sel cap wins 3.5 Val loss and a late-decay reversal 3.6 Load-balance visualisation 3.7 G0’s spectral norm is large and very anisotropic, yet most layers balance — why? 3.8 Part 3 summary 4 Application 3: FlashAttention bf16 + L1 H2 rescue 4.1 What KXK_X is here 4.1.1 Algorithm and cost 4.2 The bf16 FA failure mechanism (from the reference paper) 4.3 F3 experimental setup 4.4 The full L1 H2 crisis and rescue 4.4.1 Precise definitions of the diagnostic metrics (read before the table) 4.4.2 The main L1 H2 diagnostic table (key iterations) 4.4.3 The physical events, stage by stage (L1 H2 only) 4.4.4 Counterfactual: the uncapped baseline over the same window 4.5 The physical meaning of excess_mean (the cap’s counterfactual metric) 4.6 Why our cap fixes what the paper’s patch does not 4.7 Figures 4.8 Part 4 summary 5 Conclusions across the three case studies A Figure index References 1 Theory: a spectral cap and isotropy under Muon 1.1 Setup and goal Consider a matrix weight W∈ℝm×nW ^m× n that maps an input activation X∈ℝN×nX ^N× n linearly to the output Y=XW⊤∈ℝN×mY=XW ^N× m. We care about the second moments at the input and output: KX:=1NX⊤X∈ℝn×n,KY:=1NY⊤Y=WKXW⊤∈ℝm×m.K_X:= 1NX X ^n× n, K_Y:= 1NY Y=WK_XW ^m× m. KXK_X describes how the input energy is distributed over directions in the n-dimensional feature space, and KYK_Y is the output energy actually seen by the downstream module (attention, softmax, or another FFN block). For the derivations in this section we take the simplifying case KX≈I,K_X≈ I, and defer the general KX≠IK_X≠ I case to Section 1.7. Under KX=IK_X=I, KY=WW⊤,λ1(KY)=σ1(W)2,trKY=‖W‖F2,K_Y=W , _1(K_Y)= _1(W)^2, _Y=\|W\|_F^2, so the natural “top-share” quantity is q:=λ1(KY)trKY=σ1(W)2‖W‖F2=1stable rank(W).q:= _1(K_Y)trK_Y= _1(W)^2\|W\|_F^2= 1stable rank(W). The failure mode we want to avoid is KYK_Y degenerating towards (approximate) rank one: a single top eigenvalue λ1(KY) _1(K_Y) much larger than the rest, i.e. q→1q→ 1. This shows up as a concrete physical failure in each of the three case studies: • nanoGPT FFN: KYK_Y is the covariance the FFN writes into the residual stream; a high top-share means the FFN writes almost entirely along one direction. • MoE router: KYK_Y is the covariance of the selection score over tokens; top-share →1→ 1 means all tokens route to the same expert (rank-1 collapse). • FlashAttention Q/K: λmax(KQ(h)) _ (K_Q^(h)) directly controls the magnitude of the attention scores; when it is too large, the bf16 mantissa can no longer separate multiple near-maximal entries in a row, softmax rounds in a biased way, and training destabilises. So we use λ1(KY) _1(K_Y) as the failure indicator for isotropy and try to stop it from growing too fast under Muon updates. That is what the spectral cap does. The two subsections that follow establish the two structural claims: 1. Muon and SGD affect ‖W‖F\|W\|_F and ‖W‖2\|W\|_2 differently (Sections 1.2–1.4); 2. a spectral cap on ‖W‖2\|W\|_2 can control anisotropy at a low order without freezing the weight (Sections 1.5–1.6). 1.2 The scale-invariance assumption Many normalisation-heavy networks (pre-LN / RMSNorm transformers, and other dense layers whose output is normalised) are almost insensitive to the overall scale of a weight matrix: RMSNorm/LN divides out ‖W‖F\|W\|_F, so L(cW)≈L(W)L(cW)≈ L(W) for c>0c>0. We idealise this into an exact statement. Assumption 1 (Exact scale invariance). For the weight matrices we study, the loss is invariant under positive rescaling: L(cW)=L(W)for all c>0.L(cW)=L(W) all c>0. caveatbox Assumption 1 is a strong idealisation and is the main load-bearing assumption of the theory. In practice normalisation makes networks only approximately scale invariant (biases, weight decay, residual connections, and the last unnormalised layer all break it). Every consequence below — the 1/‖W‖F1/\|W\|_F brake, the t1/4t^1/4 vs t1/2t^1/2 norm-growth rates, and the orthogonality ⟨W,G⟩F≈0 W,G _F≈ 0 — inherits this caveat. We give a first empirical check of ⟨W,G⟩F≈0 W,G _F≈ 0 in Section 1.4.4, but whether the assumption holds tightly enough for the quantitative predictions to survive is an open question that needs more measurement across architectures and training stages. 1.2.1 Corollary B (weak form): the gradient is orthogonal to W Differentiating L(cW)=L(W)L(cW)=L(W) in c: dcL(cW)=⟨W,∇WL(cW)⟩F=0for all c>0. ddcL(cW)= W, _WL(cW) _F=0 all c>0. Evaluating at c=1c=1 gives ⟨W,∇WL(W)⟩F=0. \; W, _WL(W) _F=0.\; Geometrically: W (the radial direction, which is exactly the scale-invariant direction) receives no raw-gradient component. The raw gradient ∇WL _WL lives in the subspace orthogonal to ‖W‖F\|W\|_F. 1.2.2 Corollary A (scale form): the gradient scales like 1/‖W‖F1/\|W\|_F Scale invariance means L depends only on the direction Θ:=W/‖W‖F :=W/\|W\|_F, so we can write L(W)=ℓ(Θ)L(W)= ( ). With r:=‖W‖Fr:=\|W\|_F, the differential of r along a perturbation H is dr[H]=⟨W‖W‖F,H⟩F=⟨Θ,H⟩F,dr[H]= W\|W\|_F,H _F= ,H _F, using r=⟨W,W⟩r= W,W . Hence, for Θ=W/r =W/r, dΘ[H]=Hr−Wr2dr[H]=1r(H−Θ⟨Θ,H⟩F)=1rPΘ⟂H,d [H]= Hr- Wr^2dr[H]= 1r (H- ,H _F )= 1r\,P_ H, where PΘ⟂H:=H−Θ⟨Θ,H⟩FP_ H:=H- ,H _F is the projection onto the Frobenius-orthogonal complement of Θ . By the chain rule on L(W)=ℓ(Θ(W))L(W)= ( (W)), ⟨∇WL,H⟩F=⟨∇Θℓ,dΘ[H]⟩F=1r⟨∇Θℓ,PΘ⟂H⟩F=1r⟨PΘ⟂∇Θℓ,H⟩F _WL,H _F= _ ,d [H] _F= 1r _ ,P_ H _F= 1r P_ _ ,H _F (using that PΘ⟂P_ is a self-adjoint projection). Therefore ∇WL(W)=1rH(Θ),H(Θ):=PΘ⟂∇Θℓ, \; _WL(W)= 1r\,H( ), H( ):=P_ _ ,\; and taking Frobenius norms, ‖∇WL‖F=‖H(Θ)‖F‖W‖F∼1‖W‖F.\| _WL\|_F= \|H( )\|_F\|W\|_F 1\|W\|_F. The direction-only quantity ‖H(Θ)‖F\|H( )\|_F is O(1)O(1) in ‖W‖F\|W\|_F, so the raw gradient’s Frobenius norm shrinks like 1/‖W‖F1/\|W\|_F: the larger W grows, the smaller the raw gradient — a natural brake. 1.3 SGD vs Muon: who keeps the 1/‖W‖F1/\|W\|_F brake? 1.3.1 SGD update scale SGD uses the raw gradient: DSGD:=Wt+1−Wt=−ηGt=−ηrtH(Θt),D_ SGD:=W_t+1-W_t=-η\,G_t=- ηr_tH( _t), so ‖DSGD‖F=η‖H(Θt)‖F‖Wt‖F=O(η‖Wt‖F).\|D_ SGD\|_F= η\,\|H( _t)\|_F\|W_t\|_F=O\! ( η\|W_t\|_F ). As ‖W‖F\|W\|_F grows, the SGD update shrinks like 1/‖W‖F1/\|W\|_F. 1.3.2 Muon update scale: msign annihilates the 1/‖W‖F1/\|W\|_F factor The RMS-matched Muon update (dropping momentum for clarity) is DMuon:=Wt+1−Wt=−ηsPt,Pt=msign(Gt),s=0.2max(m,n).D_ Muon:=W_t+1-W_t=-η s\,P_t, P_t=msign(G_t), s=0.2 (m,n). Here msignmsign is defined through the SVD: for M=UΣV⊤M=U V , msign(M):=UV⊤msign(M):=UV . It is scale-invariant, since αM=U(αΣ)V⊤α M=U(α )V for α>0α>0 has the same singular vectors: msign(αM)=msign(M),α>0.msign(α M)=msign(M), α>0. Because Gt=H(Θt)/rtG_t=H( _t)/r_t with 1/rt>01/r_t>0 a positive scalar, Pt=msign(Gt)=msign(H(Θt)rt)=msign(H(Θt)).P_t=msign(G_t)=msign\! ( H( _t)r_t )=msign (H( _t) ). The 1/‖Wt‖F1/\|W_t\|_F factor is annihilated by msignmsign. So PtP_t is O(1)O(1) in ‖Wt‖F\|W_t\|_F (Pt⊤Pt=VV⊤P_t P_t=V is a rank-k projection, ‖Pt‖F2=k\|P_t\|_F^2=k with k=rank(G)k=rank(G)), and ‖DMuon‖F=ηs‖Pt‖F=ηsk=O(η).\|D_ Muon\|_F=η s\,\|P_t\|_F=η s k=O(η). Comparing the two: ‖DSGD‖F=O(η‖W‖F),‖DMuon‖F=O(η).\|D_ SGD\|_F=O\! ( η\|W\|_F ), \|D_ Muon\|_F=O(η). Muon has lost the 1/‖W‖F1/\|W\|_F brake. This is the mechanism behind the claim “Muon drifts the norm faster than SGD”, made precise in Section 1.4. 1.4 Finding 1: Muon grows the Frobenius / spectral norm faster than SGD 1.4.1 Norm increment For any update D (SGD or Muon), the change in squared Frobenius norm is ‖W+D‖F2−‖W‖F2=2⟨W,D⟩F+‖D‖F2.\|W+D\|_F^2-\|W\|_F^2=2 W,D _F+\|D\|_F^2. Writing Rt:=‖Wt‖F2R_t:=\|W_t\|_F^2, ΔRt:=Rt+1−Rt=2⟨Wt,Dt⟩F+‖Dt‖F2, R_t:=R_t+1-R_t=2 W_t,D_t _F+\|D_t\|_F^2, with two contributions: • radial term 2⟨Wt,Dt⟩F2 W_t,D_t _F: how much the update pushes W along itself, i.e. grows ‖W‖F\|W\|_F; • update-energy term ‖Dt‖F2≥0\|D_t\|_F^2≥ 0: always non-negative. 1.4.2 SGD: ‖Wt‖F∼t1/4\|W_t\|_F t^1/4 Because ⟨W,G⟩=0 W,G =0 (Corollary B), the SGD radial term vanishes: 2⟨Wt,DSGD⟩F=−2η⟨Wt,Gt⟩F=0,2 W_t,D_ SGD _F=-2η W_t,G_t _F=0, so ΔRtSGD=‖DSGD‖F2=η2‖Gt‖F2. R_t SGD=\|D_ SGD\|_F^2=η^2\|G_t\|_F^2. By Corollary A, ‖Gt‖F2=‖H(Θt)‖F2/Rt=C2/Rt\|G_t\|_F^2=\|H( _t)\|_F^2/R_t=C^2/R_t (with C=‖H(Θt)‖FC=\|H( _t)\|_F an O(1)O(1) direction-only quantity), so ΔRtSGD∼η2C2Rt. R_t SGD η^2C^2R_t. As an ODE, R˙=a/R R=a/R with a:=η2C2a:=η^2C^2, i.e. 12(R2)˙=a 12 (R^2)=a, giving R(t)2=R(0)2+2at⟹R(t)∼2at(t large),R(t)^2=R(0)^2+2at R(t) 2at (t large), hence ‖Wt‖FSGD=R(t)∼t1/4\|W_t\|_F SGD= R(t) t^1/4 (precisely ηCt1/4 η C\,t^1/4). 1.4.3 Muon: ‖Wt‖F∼t1/2\|W_t\|_F t^1/2 For Muon, Dt=−ηsPtD_t=-η sP_t with Pt=msign(Gt)P_t=msign(G_t). As a baseline, ⟨Wt,Dt⟩F=−ηs⟨Wt,Pt⟩F≈0 W_t,D_t _F=-η s W_t,P_t _F≈ 0 (the sign matrix stays roughly orthogonal to W; see the empirical check below). Then ΔRtMuon≈‖Dt‖F2=η2s2‖Pt‖F2=η2s2kt, R_t Muon≈\|D_t\|_F^2=η^2s^2\|P_t\|_F^2=η^2s^2k_t, where kt=rank(Gt)≤min(m,n)k_t=rank(G_t)≤ (m,n) is roughly a constant k, so ΔRtMuon∼b R_t Muon b with b:=η2s2kb:=η^2s^2k. This is independent of RtR_t, so R(t)=R(0)+bt∼t(t large),‖Wt‖FMuon∼t1/2R(t)=R(0)+bt t (t large), \|W_t\|_F Muon t^1/2 (precisely ηskt1/2η s k\,t^1/2). So under the same schedule Muon grows the Frobenius norm as t1/2t^1/2 against SGD’s t1/4t^1/4. 1.4.4 Empirical check of scale invariance via cPc_P and cGc_G Define the (cosine) alignments cG,t:=⟨Wt,Gt⟩F‖Wt‖F‖Gt‖F,cP,t:=⟨Wt,Pt⟩F‖Wt‖F‖Pt‖F.c_G,t:= W_t,G_t _F\|W_t\|_F\|G_t\|_F, c_P,t:= W_t,P_t _F\|W_t\|_F\|P_t\|_F. The theory predicts that if scale invariance holds then cG,t≈0c_G,t≈ 0 (this is Corollary B). We measured cGc_G and cPc_P on a nanoGPT step-48K checkpoint over a fixed diagnostic batch, averaged over 24 two-dimensional weights (6 layers × c_attn, c_proj, mlp.c_fc, mlp.c_proj): Run @ step cGc_G |cG||c_G| cPc_P |cP||c_P| B0 (plain Muon) +0.0000+0.0000 0.00520.0052 −0.0007-0.0007 0.01040.0104 M4plain (proj cap) −0.0003-0.0003 0.00330.0033 +0.0024+0.0024 0.00490.0049 The alignments are within a few times 10−310^-3 of zero, consistent with approximate scale invariance at this checkpoint. We stress that this is a single sanity check, not a validation of the assumption across the whole trajectory. 1.4.5 Spectral norm: Muon moves ‖W‖2\|W\|_2 more than SGD when ‖W‖F\|W\|_F is large Assume σ1(W) _1(W) is simple, with top singular vectors (u1,v1)(u_1,v_1). First order, σ1(W+D)=σ1(W)+u1⊤Dv1+O(‖D‖F2σ1−σ2). _1(W+D)= _1(W)+u_1 Dv_1+O\! ( \|D\|_F^2 _1- _2 ). For SGD, ‖DSGD‖F=O(η/‖W‖F)\|D_ SGD\|_F=O(η/\|W\|_F), so |u1⊤DSGDv1|=O(η/‖W‖F)|u_1 D_ SGDv_1|=O(η/\|W\|_F) and (using σ1=‖W‖Fθ1 _1=\|W\|_F _1 with θ1=O(1) _1=O(1)) |Δσ1SGD|σ1=O(η‖W‖F2). | _1 SGD| _1=O\! ( η\|W\|_F^2 ). For Muon, ‖DMuon‖F=O(η)\|D_ Muon\|_F=O(η), so |u1⊤DMuonv1|=O(η)|u_1 D_ Muonv_1|=O(η) and |Δσ1Muon|σ1=O(η‖W‖F). | _1 Muon| _1=O\! ( η\|W\|_F ). Taking the ratio, |Δσ1Muon|/σ1|Δσ1SGD|/σ1∼‖W‖F. | _1 Muon|/ _1| _1 SGD|/ _1 \|W\|_F. The larger ‖W‖F\|W\|_F, the more Muon disturbs the spectral norm relative to SGD. 1.5 Finding 2: after a cap, the spectral-norm second order is ≥0≥ 0 A spectral cap (at tolerance ρ=0ρ=0) removes the top component of an update D0D_0: D=D0−αu1v1⊤,α=u1⊤D0v1,D=D_0-α\,u_1v_1 , α=u_1 D_0v_1, so that u1⊤Dv1=0u_1 Dv_1=0. Then to first order σ1(W+D)=σ1(W)+0⏟cap kills 1st order+(2nd order)+⋯ _1(W+D)= _1(W)+ 0_cap kills 1st order+(2nd order)+·s Is the second-order term a problem (does it still push σ1 _1 up)? We show it is non-negative, so the cap does not freeze the weight. 1.5.1 Symmetric dilation For W∈ℝm×nW ^m× n, form the symmetric dilation ℋ(W)=(0W⊤0)∈ℝ(m+n)×(m+n),H(W)= pmatrix0&W\\ W &0 pmatrix ^(m+n)×(m+n), whose eigenvalues are ±σ1,±σ2,…± _1,± _2,… (padded with zeros), so σ1(W)=λmax(ℋ(W)) _1(W)= _ (H(W)). The standard second-order eigenvalue perturbation of a symmetric matrix A+EA+E, λ1(A+E)=λ1+z1⊤Ez1+∑k≠1(zk⊤Ez1)2λ1−λk+⋯, _1(A+E)= _1+z_1 Ez_1+ _k≠ 1 (z_k Ez_1)^2 _1- _k+·s, with zkz_k the eigenvectors of A, applied to the dilation with aj:=uj⊤Dv1a_j:=u_j Dv_1, bj:=u1⊤Dvjb_j:=u_1 Dv_j for j>1j>1, gives σ1(W+D)=σ1+u1⊤Dv1+∑j>1[(aj+bj)24(σ1−σj)+(aj−bj)24(σ1+σj)]+O(‖D‖F3). _1(W+D)= _1+u_1 Dv_1+ _j>1\! [ (a_j+b_j)^24( _1- _j)+ (a_j-b_j)^24( _1+ _j) ]+O(\|D\|_F^3). Every second-order term is non-negative (σ1−σj>0 _1- _j>0 is the spectral gap; σ1+σj>0 _1+ _j>0; numerators are squares). 1.5.2 Why the second order being ≥0≥ 0 is exactly what we want With the cap enforcing u1⊤Dv1=0u_1 Dv_1=0, the weight W can still evolve — and KYK_Y can still accumulate energy in non-top directions, becoming more isotropic — through three mechanisms: 1. Non-top eigenvalue growth: components of D along ujvj⊤u_jv_j (j>1j>1) directly raise σj _j, spreading the energy; 2. Top-vector rotation: aj=uj⊤Dv1≠0a_j=u_j Dv_1≠ 0 or bj=u1⊤Dvj≠0b_j=u_1 Dv_j≠ 0 rotate the top direction toward a new one through the second-order term; 3. Top switching: if some σj _j grows past σ1 _1, the next step’s cap acts on the new top direction. So a spectral cap is not about freezing W: it converts “keep pouring energy into the same top direction” into “explore new directions in the orthogonal complement”. Since σ12/‖W‖F2=1/stable rank _1^2/\|W\|_F^2=1/stable rank, capping the top while other directions grow raises the stable rank, and KY=WW⊤K_Y=W becomes more isotropic. 1.6 Extension: the cap targets H∞H_∞ entropy; other entropies give other cap directions Write the squared-singular-value distribution of W as pi=σi2∑jσj2,∑ipi=1,p_i= _i^2 _j _j^2, _ip_i=1, and the usual Rényi entropies H∞=−logpmax,H2=−log∑ipi2,H1=−∑ipilogpi.H_∞=- p_ , H_2=- _ip_i^2, H_1=- _ip_i p_i. The spectral cap == H∞H_∞. We have eH∞=1pmax=‖W‖F2σ12=stable rank(W),e^H_∞= 1p_ = \|W\|_F^2 _1^2=stable rank(W), and −dH∞=d(logpmax)=2dlogσ1−2dlog‖W‖F,-dH_∞=d( p_ )=2\,d _1-2\,d \|W\|_F, so, keeping the dominant piece, dH∞<0dH_∞<0 iff u1⊤Dv1/σ1>⟨W,D⟩/‖W‖F2u_1 Dv_1/ _1> W,D /\|W\|_F^2. Enforcing u1⊤Dv1≤0u_1 Dv_1≤ 0 (our cap) protects H∞H_∞. Participation rank == H2H_2. With PR(W)=eH2=1∑ipi2=(∑iσi2)2∑iσi4=(trKY)2‖KY‖F2,PR(W)=e^H_2= 1 _ip_i^2= ( _i _i^2)^2 _i _i^4= (trK_Y)^2\|K_Y\|_F^2, a more “democratic” measure of the number of effective directions, and −dH2∝∑ipi2dlogσi2,-dH_2 _ip_i^2\,d _i^2, the corresponding cap projection direction is ∑ipi2uivi⊤ _ip_i^2\,u_iv_i — it touches all directions uivi⊤u_iv_i , weighted by pi2p_i^2 (a multi-modal cap). Shannon H1H_1. Similarly −dH1=∑ilogpidpi-dH_1= _i p_i\,dp_i, with direction ∑ilogpiuivi⊤ _i p_i\,u_iv_i . We use the H∞H_∞ (spectral) cap because it is the cheapest: 1. the top pair (u1,v1)(u_1,v_1) is available from a few steps of power iteration; 2. it controls σ1 _1 directly (attention score =σ1= _1; MoE rank-1 collapse =σ1= _1); 3. H∞H_∞ collapse fires exactly when σ1 _1 dominates. Because the second order is ≥0≥ 0 (Section 1.5), capping only σ1 _1 lets the remaining update pour energy into non-top directions, so trKYtrK_Y grows while σ1 _1 is held — p1p_1 falls, H∞H_∞ rises indirectly, and H2H_2 and H1H_1 improve for free. All three case studies use the H∞H_∞ cap. 1.7 Remark: we do not actually need KX=IK_X=I If KX≠IK_X≠ I, with KY=WKXW⊤K_Y=WK_XW , the “spectral to control” is not that of W but of A:=WKX1/2,KY=AA⊤,σi(A)2=λi(KY).A:=WK_X^1/2, K_Y=A , _i(A)^2= _i(K_Y). So the object to cap is the spectral of A, per application: • nanoGPT FFN: KXK_X is the GELU-hidden second moment, anisotropic, so we cap KYK_Y (a data-cap, not W); • MoE router: KXK_X is the post-RMSNorm token covariance, relatively isotropic (RMS removes the global scale); • FlashAttention: KXK_X is the layer-input embedding covariance, anisotropic, so F3 caps KQ(h)=WQ(h)⊤KXWQ(h)K_Q^(h)=W_Q^(h) K_XW_Q^(h), not WQW_Q. KX=IK_X=I is only a pedagogical simplification (Muon still removes the brake, the cap still applies); in practice the cap is on the data-coupled KYK_Y, not on the raw W. 2 Application 1: nanoGPT FFN (proof of concept) 2.1 What KXK_X is here Module: the mlp.c_proj layer, which projects the GELU-activated hidden (dhidden=1536d_ hidden=1536) back into the residual stream (dembd=384d_ embd=384): W∈ℝdembd×dhidden=ℝ384×1536,X∈ℝM×dhidden.W ^d_ embd× d_ hidden=R^384× 1536, X ^M× d_ hidden. X is the GELU output, with M≈32×1024=32768M≈ 32× 1024=32768 tokens per batch. KXK_X: the second moment of the GELU-hidden activation, KX=1MX⊤X∈ℝ1536×1536K_X= 1MX X ^1536× 1536, strongly anisotropic (after GELU only a few channels are typically active). KYK_Y: the covariance the FFN writes into the residual stream, KY=WKXW⊤∈ℝ384×384K_Y=WK_XW ^384× 384. This is the geometric object the layer actually contributes downstream, and the one we want isotropic. 2.1.1 Estimating KYK_Y and finding the top eigenpair (q,λ1)(q, _1) At each Muon step: 1. take X∈ℝM×1536X ^M× 1536 from the forward-pass cache (sub-sampling M=4096M=4096 tokens to avoid the memory cost of the full batch); 2. compute Y=XW⊤∈ℝM×384Y=XW ^M× 384, cost O(M⋅dembd⋅dhidden)O(M· d_ embd· d_ hidden); 3. do not form KYK_Y explicitly (downstream only needs λ1 _1 and its q); 4. find the top eigenpair by power iteration: from a random q0∈ℝ384q_0 ^384, iterate q←KYq/‖KYq‖Fq← K_Yq/\|K_Yq\|_F, with each KYq=WKXW⊤qK_Yq=WK_XW q done by matrix–vector products (never forming KYK_Y). Cost per step: each power iteration is O(Mdhidden+Mdembd)O(Md_ hidden+Md_ embd); with npower=4n_ power=4, total ≈4×4096×1920≈31≈ 4× 4096× 1920≈ 31 MFLOPs. Against the ∼ forward+backward this is about 0.003%0.003\%; the cap overhead is negligible. Cap first-order projection, given λ1 _1 and q: B=∂λ1∂W=2qg⊤,g=1M∑m(q⊤ym)xm,H←H−max(0,⟨B,H⟩−ρλ1)B‖B‖F2.B= ∂ _1∂ W=2qg , g= 1M _m(q y_m)\,x_m, H← H- \! (0, B,H -ρ _1 ) B\|B\|_F^2. Geometry of qg⊤qg (important): the rank-1 matrix B=2qg⊤B=2qg is the direction in weight space, at the current W, along which λ1(KY)=λ1(WKXW⊤) _1(K_Y)= _1(WK_XW ) grows fastest. Here q is the most sensitive output-side direction (top eigenvector of KYK_Y), g is the input-side direction most aligned with q (in the KXK_X metric), and the outer product qg⊤qg couples them. Adding B to W raises λ1 _1 the fastest — the direction we least want W to move along. The cap projects out the positive part ⟨B,H⟩−ρλ1 B,H -ρ _1 of the candidate update; the remainder lies in the orthogonal complement of B, so W still updates but no longer along the λ1 _1-growth direction. ρ=0ρ=0 is a hard cap (λ1 _1 does not grow to first order); ρ>0ρ>0 leaves a buffer. 2.2 Experimental setup Setting Value Model nanoGPT 6L / 6H / nembd=384n_ embd=384 / block 1024 Optimiser RMS-matched Muon (2D) + AdamW (norms / embed / lm_head) Precision bf16 attention, fp32 weights LR ×10× 10 cosine decay 10−2→10−310^-2→ 10^-3 Iterations 50K WD 0.00.0 (removed, to test whether the cap works on its own) Data OpenWebText Hardware 2×2× RTX 4090 Three setups: • D0: no cap (baseline); • D1: W-cap on W=W= mlp.c_proj (directly controls σ1(W) _1(W); corresponds to the KX=IK_X=I simplification); • D1d: data-cap on KYK_Y (the actual KX≠IK_X≠ I case). 2.3 Result: the three methods have equal val loss Figure 1: nanoGPT 50K val-loss curves. D0/D1/D1d essentially overlap over the whole run, with a final difference |Δval|<0.01| |<0.01 nats. This is the proof of concept: the spectral cap (projecting out first-order top-mode growth) does not hurt val loss, so the model keeps learning through the “update learns new directions in the orthogonal complement” mechanism of Section 1.5. Final val loss @ 50K: D0 =3.8680=3.8680, D1 =3.8763=3.8763, D1d =3.8742=3.8742. The differences are <0.01<0.01 nats, well below the seed-to-seed noise (∼0.05 0.05). 2.4 What the cap changes: spectral norm and isotropy side by side Table A: per-layer σ1(W) _1(W) @ 50K L D0 D1 (W-cap) D1d (data-cap) L0 98.9 85.0 100.6 L1 163.3 116.6 165.0 L2 137.4 117.2 139.6 L3 156.4 128.5 157.8 L4 219.4 117.4 164.5 L5 211.2 145.2 191.6 D1 lowers σ1 _1 by 14–46%; D1d barely moves it. Table B: per-layer KYK_Y top-share @ 50K L D0 D1 (W-cap) D1d (data-cap) L0 0.992 0.993 0.704 L1 0.151 0.109 0.039 L2 0.060 0.052 0.053 L3 0.071 0.071 0.066 L4 0.742 0.603 0.299 L5 0.499 0.288 0.286 D1d improves L1/L4 isotropy by 60–74%. The two caps control different objects: D1 presses σ1(W) _1(W) directly (parameter space); D1d presses λ1(KY)=σ1(WKX1/2)2 _1(K_Y)= _1(WK_X^1/2)^2 directly (output space). Because KX≠IK_X≠ I the two are not equivalent: D1 lowers the spectrum of W, but KXK_X redistributes the energy onto the strong data directions and the KYK_Y top-share barely moves; D1d controls KYK_Y directly while the spectrum of W can stay high, because KX1/2K_X^1/2 orthogonalises it away. This confirms the design choice at the end of Section 1.7: cap the data-coupled KYK_Y (no need to assume KX=IK_X=I). Capping only the spectrum of W lets the strong directions of KXK_X keep the KYK_Y top-share almost unchanged, and the cap fails to do its job. 2.5 Participation-rank figure Figure 2: Participation rank tr(KY)2/∥KY∥F2tr(K_Y)^2/\|K_Y\|_F^2 (i.e. eH2e^H_2 of Section 1.6, max =384=384) of the FFN output KYK_Y across the six mlp.c_proj layers. D1d (green) raises the L1 PR from D0’s (blue) 34 to 96 (nearly 3×3×); D1 (red) is in between. L0 is the hardest — all three are near 1 (i.e. KYK_Y is essentially rank-1 at L0). Reading: the rise in PR (i.e. higher H2H_2 entropy) is the core benefit of the D1d cap. The cap stops the continued pouring into the H∞H_∞ direction (σ1 _1); the remaining update lets other singular values grow, the energy distribution flattens, and H2H_2 improves automatically. This is exactly the “H∞H_∞ cap also improves H2H_2” point of Section 1.6. 2.6 Part 2 summary • the cap does not hurt val loss ⇒ empirical support for the “second order still learns new directions” mechanism of Section 1.5; • the data-cap (D1d) is more on-target than the W-cap (D1) ⇒ support for the “cap KYK_Y, do not assume KX=IK_X=I” design of Section 1.7; • PR rises from 34 to 96 ⇒ support for “H∞H_∞ cap improves H2H_2”. nanoGPT’s own val loss does not move, though; to see a real gain from the cap we turn to the MoE and FA systems below. 3 Application 2: 64-expert MoE router 3.1 What KXK_X is here Module: the router of a MoE layer — a linear layer mapping a token embedding to the logits of E=64E=64 experts: Wr∈ℝE×d=ℝ64×384,X∈ℝM×d.W_r ^E× d=R^64× 384, X ^M× d. X is the post-RMSNorm token embedding, M the number of tokens sampled per step (the cap uses M=4096M=4096). KXK_X: the covariance of the post-RMSNorm token embedding, KX=1MX⊤X∈ℝ384×384K_X= 1MX X ^384× 384, relatively isotropic (RMS unifies the overall scale). KzK_z = expert-centred covariance of the raw router logits: Z=XWr⊤∈ℝM×E,Zc=ZC,C:=IE−1E⊤,Kz=1MZc⊤Zc∈ℝ64×64.Z=XW_r ^M× E, Z_c=ZC, C:=I_E- 1E11 , K_z= 1MZ_c Z_c ^64× 64. Why the factor C: softmax and top-k routing are insensitive to a per-token global shift (adding the same constant ctc_t to all experts does not change the ordering). C removes this gauge freedom: Z+c⊤→(Z+c⊤)C=ZCZ+c1 →(Z+c1 )C=ZC. So KzK_z is the geometric quantity that actually affects routing. Kz+bK_z+b = selection-score covariance: once the loss-free bias (LFB) adds a per-expert bias b∈ℝEb ^E, top-k is taken on the selection score s=z+bs=z+b, not on raw z. So the geometry that truly decides routing is Kz+b=1M((Z+b⊤)C)⊤((Z+b⊤)C).K_z+b= 1M ((Z+1b )C ) ((Z+1b )C ). C removes only the per-token shift; b is a per-expert shift, and b→b+cb→ b+c is also removed by C (CbCb has mean 0). So the direction the LFB actually learns, CbCb, is retained in Kz+bK_z+b. 3.1.1 Algorithm: estimating Kz+bK_z+b and finding the top eigenpair At each Muon step: 1. take the cached X∈ℝM×384X ^M× 384 (M=4096M=4096 tokens); 2. compute Z=XWr⊤∈ℝM×64Z=XW_r ^M× 64, add b to get Z+bZ+b, multiply by C to get Zc∈ℝM×64Z_c ^M× 64; 3. form K=Zc⊤Zc/M∈ℝ64×64K=Z_c Z_c/M ^64× 64 directly (E=64E=64 is small, so O(ME2)O(ME^2) is affordable); 4. power iteration on K for (q,λ1)(q, _1), four iterations at O(E2)O(E^2); 5. cap projection so that Δλ1≤0 _1≤ 0 to first order. Cost: O(ME2+MEd)≈100O(ME^2+MEd)≈ 100 MFLOPs, again negligible against forward+backward. Three setups: • G0: LFB only (no cap); • G1: LFB + cap on the raw KzK_z (without b); • G3: LFB + cap on Kz+bK_z+b (the selection-score geometry). 3.2 What isotropy means in the MoE The top-share of Kz+bK_z+b, top_share=λ1/trKz+btop\_share= _1/trK_z+b, measures how much of the variance of the selection score across tokens is concentrated in one direction: top_share→1E−1≈0.016:energy spread equally over expert directions;top\_share→ 1E-1≈ 0.016:energy spread equally over expert directions; top_share→1:every token’s s=z+b varies along a single fixed direction.top\_share→ 1:every token's s=z+b varies along a single fixed direction. The physics of top-share →1→ 1 (rank-1): each token t’s selection score is st≈ct⋅dmax∈ℝE,s_t≈ c_t· d_ ^E, with dmaxd_ a fixed direction and ctc_t a token-dependent scalar. Then top-k(st)top-k(s_t) is the same as top-k(dmax)top-k(d_ ) — all tokens route to the same k experts. This is the “one-expert dominates” rank-1 collapse. Why the LFB cannot fix rank-1: the LFB has only E per-expert scalars, adjusting each expert’s threshold along the direction dmaxd_ . If the geometry is one-dimensional, 63 of the 64 LFB degrees of freedom are wasted — it can move an expert up or down but cannot make different tokens pick different experts. 3.3 Core result: Kz+bK_z+b isotropy across the three methods The three tables below are the final Kz+bK_z+b geometry (seed=42, fast MoE; G0v2/G1v2 are freshly run baselines, G3 is the previously finished run at the same seed). The three metrics are complementary: • λmax _ : absolute scale — variance energy in the strongest direction; • top-share =λmax/trKz+b= _ /trK_z+b: relative concentration (0 isotropic, 1 rank-1); • participation rank =(trKz+b)2/‖Kz+b‖F2=(trK_z+b)^2/\|K_z+b\|_F^2: effective number of directions (upper bound E−1=63E-1=63). Kz+bK_z+b λmax _ (smaller = more isotropic) L0 L1 L2 L3 L4 L5 G0 (no cap) 82,517 5,179 1,824 729 453 8,062 G1 (raw cap on KzK_z) 2,052 2,173 66 71 121 7,142 G3 (sel cap on Kz+bK_z+b) 521 435 102 37 13 209 Huge λmax _ gaps: G0 L0 =82,517=82,517 is 158×158× G3 L0 =521=521; G0 L5 =8,062=8,062 is 38×38× G3 L5 =209=209. The G1 raw cap leaves a clearly higher residual at L0/L5 than G3. Kz+bK_z+b top-share L0 L1 L2 L3 L4 L5 G0 (no cap) 0.995 0.667 0.653 0.422 0.356 0.864 G1 (raw cap on KzK_z) 0.479 0.645 0.341 0.211 0.418 0.803 G3 (sel cap on Kz+bK_z+b) 0.272 0.521 0.305 0.119 0.173 0.267 Kz+bK_z+b participation rank (max =E−1=63=E-1=63) L0 L1 L2 L3 L4 L5 G0 (no cap) 1.0 2.1 2.3 4.8 5.1 1.3 G1 (raw cap) 3.2 2.2 6.6 11.2 4.8 1.5 G3 (sel cap) 8.8 3.3 5.6 20.1 11.7 8.3 Reading: 1. G0’s L0 is a full rank-1 collapse: top-share =0.995=0.995, PR =1.0=1.0 (exactly one effective expert direction), λmax=82,517 _ =82,517 far above every other layer; 2. G1 (raw cap on KzK_z) improves Kz+bK_z+b indirectly: the raw cap controls KzK_z (without b), but once KzK_z is held, b grows in other directions to restore balance, so Kz+bK_z+b is also better — just not as good as G3; 3. G3 (sel cap on Kz+bK_z+b) wins on all three metrics: across 66 layers ×3× 3 metrics =18=18 cells, G3 takes 16; the other two (PR at L2 / L1) go marginally to G1; 4. G3’s L5 improvement is the largest: from G0’s top-share 0.860.86 / λmax _ 80628062 down to 0.270.27 / 209209 (38×38×). L5 has the weakest LFB feedback (last layer), so directly protecting the geometry matters most there. 3.4 L0 / L5 load balancing: a phase transition and LFB failure Normalisation of the hard-load table: write F as the “fraction of tokens that include expert e in their top-k” — summed over experts this is k=4k=4, with ideal uniform value k/E=4/64=0.0625k/E=4/64=0.0625. The table below has F multiplied by k=4k=4 (i.e. sum-to-k instead of sum-to-1), so its maximum is 1.01.0 (one expert appears in the top-4 of every token). Effective experts per layer (max =64=64) L0 L1 L2 L3 L4 L5 G0 (no cap) 31.3 64.0 64.0 64.0 63.9 31.0 G1 (raw cap) 62.2 51.7 63.9 63.9 63.9 14.3 G3 (sel cap) 46.3 63.9 63.9 64.0 63.9 30.1 Hard-load max F (uniform =0.0625=0.0625; rank-1 →1.0→ 1.0) L0 L1 L2 L3 L4 L5 G0 (no cap) 0.999 0.066 0.069 0.069 0.072 0.820 G1 (raw cap) 0.071 0.139 0.071 0.070 0.074 0.894 G3 (sel cap) 0.219 0.075 0.072 0.067 0.075 0.236 3.4.1 The L0 story: G0 collapses to rank-1 G0 L0 hard-max =0.999=0.999 — literally the same expert appears in the top-4 of every token (the hard-count signature of a first-order rank-1 collapse). effE =31.3=31.3 reflects that the other 3 top-4 slots are still spread over many experts, but the top-1 expert is fixed. G1 raw cap: hard-max =0.071≈=0.071≈ uniform 0.0625, a clean rescue of the L0 rank-1 collapse. effE =62.2=62.2 is near the ceiling of 64. G3 sel cap: hard-max =0.219=0.219, in between — not as clean as G1 (partial recovery), but far better than G0’s 0.999. effE =46.3=46.3. 3.4.2 The L5 story: G3 sel cap wins L5 hard-max: • G0: 0.820 (the top expert takes 82% of the tokens’ top-4); • G1: 0.894 (worse — the G1 raw cap has a side effect at L5); • G3: 0.236 (33–4×4× better than G0/G1, close to uniform 0.0625). This matches the Kz+bK_z+b top-share ordering of Section 3.3 exactly: G3 breaks the L5 geometry, and the hard-load follows. 3.5 Val loss and a late-decay reversal Run @ step 50K (seed=42 fast) val loss G0 (no cap) 3.7351 G1 (raw cap) 3.6660 G3 (sel cap) 3.6920 The three final val losses are within 0.07 nats. Figure 3: seed=42 val-loss trajectory, three methods overlaid. Key observation: during step 10K–30K, G0’s val is slightly below G1/G3 (the cap constrains the router’s directional freedom); during the lr decay at step 40K–50K (warmdown start marked by the dashed line), the capped methods G1/G3 overtake G0. This is the “slower early, more stable late” signature of the cap: it prevents G0’s fast early loss drop bought with over-anisotropic geometry, but G0’s later spectral collapse worsens, and the capped methods are more stable through the warmdown. 3.6 Load-balance visualisation Figure 4: seed=42 final per-expert load histogram (3 methods × 6 layers). In each subplot the horizontal axis is the 64 experts (sorted by load), the vertical axis the load fraction, with the ideal 1/E=0.01561/E=0.0156 as a dashed line. L0 (top row): G0 (blue) has its first expert take 25% (the rank-1 collapse, visualised); G1 and G3 both flatten to near-ideal. L1–L4: the three methods are essentially the same. L5 (bottom row): G1’s top 4 experts take 65% (a hard collapse); G3 repairs the G1 L5 side effect. Figure 5: Per-layer load CV (E⋅std(F)/mean(F)E·std(F)/mean(F), lower is more balanced, ideal =0=0). L0: G0 =1.96=1.96 (very imbalanced), G1 =0.30=0.30 (near perfect), G3 =0.92=0.92 (in between). L1 is the reverse: G1 =1.07=1.07 (side effect), G3 =0.07=0.07 (repaired). L5: G3 =1.50<=1.50< G0 =1.88<=1.88< G1 =2.86=2.86. 3.7 G0’s spectral norm is large and very anisotropic, yet most layers balance — why? Back to G0’s Kz+bK_z+b table: L1–L4 have top-share 0.350.35–0.670.67 and λmax _ from hundreds to thousands — still quite anisotropic (far from the ideal top-share ≈0.016≈ 0.016); yet their hard-max is all 0.0660.066–0.0720.072 (near uniform 0.0625). Explanation (with Section 3.2): 1. Muon injects a fixed-energy update into WrW_r each step; without a cap, λ1(Kz+b) _1(K_z+b) keeps rising, but the LFB uses b to “pull the direction back to balance”; 2. in the “moderately anisotropic” regime top-share ∈[0.35,0.7]∈[0.35,0.7], the LFB’s 64 degrees of freedom are enough to level the hard count — it shifts each expert’s selection threshold so that different tokens pick different experts; 3. but the phase transition to top-share →1→ 1 (rank-1 collapse) breaks the LFB: at rank-1, all tokens’ scores are scalar multiples of a single direction, and a threshold shift cannot make different tokens pick different experts; 4. G0 L0 / L5 sit right at this transition (top-share 0.860.86–0.990.99), so their hard-max explodes; 5. G1 / G3 push all layers below the transition through the cap, leaving room for the LFB. This also explains why most of G0’s layers look fine: the cap is not a necessary condition for balance, but a preventative against the rank-1 phase transition. Its value is in the worst case (L0 / L5), not the typical case (L1–L4). 3.8 Part 3 summary • the selection cap (G3) wins almost across the board on Kz+bK_z+b isotropy over all six layers; • G0’s L0 enters the rank-1 phase transition at top-share =0.995=0.995 ⇒ the LFB fails ⇒ hard-max →1.0→ 1.0; the G1/G3 caps repair it (G1 cleanly, =0.071=0.071; G3 partially, =0.219=0.219); • during lr decay the capped methods overtake the baseline, possibly because a healthier geometry finds good directions more easily at small LR; • at L5 the cap is not always perfect on the hard count, but it does improve the geometry. 4 Application 3: FlashAttention bf16 + L1 H2 rescue 4.1 What KXK_X is here Module: the Q and K projections of each attention block. Each layer has 12 heads, each projecting the dembd=768d_ embd=768 input embedding into a dh=64d_h=64 Q/K subspace: WQ(h),WK(h)∈ℝdh×dembd=ℝ64×768.W_Q^(h),\,W_K^(h) ^d_h× d_ embd=R^64× 768. KXK_X: the input-embedding covariance at layer ℓ , KX(ℓ)=1MXℓ⊤Xℓ∈ℝ768×768K_X^( )= 1MX_ X_ ^768× 768, where XℓX_ is the residual-stream token embedding at the entry of layer ℓ . Strongly anisotropic (a typical transformer input space has a few dominant directions). Cap target == per-(layer, head) data covariance: KQ(ℓ,h)=WQ(ℓ,h)KX(ℓ)WQ(ℓ,h)⊤=1M(Q(ℓ,h))⊤Q(ℓ,h)∈ℝ64×64,K_Q^( ,h)=W_Q^( ,h)K_X^( )W_Q^( ,h) = 1M(Q^( ,h)) Q^( ,h) ^64× 64, with Q(ℓ,h)=XℓWQ(ℓ,h)⊤Q^( ,h)=X_ W_Q^( ,h) the query matrix of that head on the current batch; the K side is symmetric. Why per-(layer, head): attention failure is a local loss of control at one head in one layer. A layer-aggregate cap reacts too slowly to a single head’s phase change, so we run the cap separately on each (layer, head). 4.1.1 Algorithm and cost At each Muon step: 1. take Q(ℓ,h)∈ℝM×64Q^( ,h) ^M× 64 per head from the forward-pass cache (M=4096M=4096 tokens per rank); 2. form KQ(ℓ,h)=(Q(ℓ,h))⊤Q(ℓ,h)/M∈ℝ64×64K_Q^( ,h)=(Q^( ,h)) Q^( ,h)/M ^64× 64, cost O(Mdh2)≈17O(Md_h^2)≈ 17 MFLOPs per head; 3. power iteration on KQ(ℓ,h)K_Q^( ,h) for (q,λ1)(q, _1), four iterations at O(dh2)O(d_h^2); 4. cap projection so that Δλ1≤0 _1≤ 0; repeat on the K side. Total cost (12 layers × 12 heads × 2 (Q+K) =288=288 caps per step): ∼5 5 GFLOPs, about 5% of the ∼100 100 GFLOP forward+backward — still acceptable. 4.2 The bf16 FA failure mechanism (from the reference paper) The reference paper on low-precision FlashAttention failures [1] identifies the mechanism: FlashAttention v2 uses block-wise online softmax. For each query row the score matrix is split into blocks, each block’s max and log-sum-exp are computed, and the blocks are merged. With a bf16 mantissa of 7 bits, at a score scale of ∼104 10^4 the resolution is only ∼80 80. Failure chain: 1. late in training, some layer’s WQ(ℓ,h)W_Q^( ,h) and WK(ℓ,h)W_K^( ,h) accumulate energy along the strong directions of KX(ℓ)K_X^( ); 2. λmax(KQ(ℓ,h))=σ1(WQ(ℓ,h)KX(ℓ)1/2)2 _ (K_Q^( ,h))= _1(W_Q^( ,h)K_X^( )1/2)^2 blows up; 3. the attention score sij=qi⋅kj/dhs_ij=q_i· k_j/ d_h, bounded by |sij|≤λQλK/dh|s_ij|≤ _Q _K/ d_h, grows with it; 4. when a row has many entries near its row max (large near-max count), the cross-block max-subtraction loses precision in bf16; 5. the softmax accumulates error into the attention output; 6. loss creep, then crash. Timeline of the uncapped baseline: • iter 70K: onset, loss =3.33=3.33 (creep begins); • iter 100K: late creep, loss =3.76=3.76; • iter 110K: crash, loss =4.31=4.31; • iter 120K: blown, loss =5.19=5.19. The paper fixes this with a patched (STABLE) softmax that adds a numerical safeguard inside attention. Our F3 leaves softmax untouched and instead uses the cap to stop λmax(KQ(ℓ,h)) _ (K_Q^( ,h)) from reaching the failure threshold. Why the cap should work: it guarantees Δλmax≤0 _ ≤ 0 per step to first order. The crash path needs λ to grow from ∼103 10^3 to ∼105 10^5 over 20K steps, i.e. an average Δλ≈5 λ≈ 5 per step. Additionally, because the largest direction is capped each step, the matrix explores other directions rather than pouring energy into the same direction — which also avoids producing many equal scores. 4.3 F3 experimental setup Setting Value Model GPT-2 small (12L / 12H / nembd=768n_ embd=768 / block 1024) Optimiser RMS-matched Muon (2D) + AdamW Precision bf16 attention (STABLE=0, the paper’s original unstable softmax, no patch) LR 5e-4, WD =0=0, warmup 2000 Batch 327,680 tok/iter Cap CAP_MODE=qk_data, per-(layer, head) ρ=0ρ=0 hard cap Iterations 120K Monitoring diag_attn_layers=1,11 (L1 and L11 record fine-grained score / near-max; all 12 layers are capped and log λmax _ ) Comparison: the same configuration without a cap — the uncapped baseline run. 4.4 The full L1 H2 crisis and rescue Indexing note: the paper uses 1-indexed layers, so the paper’s “L2” is our L1 (the second layer). Among the 24 monitored (layer, head) pairs, L1 H2 (the third head of the paper’s L2) has the largest score.max peak, 27,264 at iter 86K. It corresponds to the bf16 FA failure entry the paper observes at its L2. This subsection tracks this one head throughout. 4.4.1 Precise definitions of the diagnostic metrics (read before the table) The diagnostics are computed on one captured forward pass (diag_capture=True). The attention score tensor is S∈ℝB×H×Tq×Tk,Sb,h,i,j=qb,h,i⋅kb,h,jdh,S ^B× H× T_q× T_k, S_b,h,i,j= q_b,h,i· k_b,h,j d_h, with B the diagnostic batch size, H=12H=12 heads, Tq=Tk=1024T_q=T_k=1024, dh=64d_h=64. The metrics below are all defined for a fixed head h (here h=2h=2, i.e. L1 H2). The causal mask sets j>ij>i to −∞-∞, and only finite entries are counted in the sums/maxima. (1) score.max (per head, the single largest entry): score.max(h):=maxb,i,jSb,h,i,j.score.max^(h):= _b,i,j\;S_b,h,i,j. Code: S.amax(dim=(0,2,3)). This is the single largest entry over all (batch, query, key) triples, not the average of the row maxima. Physically: the worst-case bf16 stress (the entry most squeezed by the bf16 mantissa). (2) Row-tie indicator + near-max count: first the per-query row max rowmaxb,i(h):=maxjSb,h,i,jrowmax^(h)_b,i:= _jS_b,h,i,j, then the per-query row-tie count (how many entries in a row are within a 10−310^-3 tolerance of the row max), nb,i(h):=|j:Sb,h,i,j≥rowmaxb,i(h)−10−3|.n^(h)_b,i:= |\j:S_b,h,i,j ^(h)_b,i-10^-3\ |. (3) near.mean (per head, the mean row-tie count): near.mean(h):=1BTq∑b,inb,i(h).near.mean^(h):= 1BT_q _b,i\;n^(h)_b,i. Healthy attention (one dominant softmax peak) gives ≈1.0≈ 1.0. A large value means most rows have several ties, so the paper’s biased-rounding mechanism fires frequently. (4) near.max (per head, the worst row’s tie count): near.max(h):=maxb,inb,i(h).near.max^(h):= _b,i\;n^(h)_b,i. Value K means at least one query has K entries all within row max ±10−3± 10^-3 — the worst-case trigger strength. (5) λmax(KQ(h)) _ (K_Q^(h)) and q_excess: see Section 4.1 and Section 4.5. q_excess is the cap’s excess_mean on the Q side at that step (already including LR; units == first-order change of λ1 _1 per Muon step). (6) top1 prob =⟨maxjPb,h,i,j⟩b,i= _jP_b,h,i,j _b,i with P=softmax(S)P=softmax(S): the mean top probability per row. High == peaked attention. 4.4.2 The main L1 H2 diagnostic table (key iterations) All numbers are specifically for L1 H2, not layer-aggregate. iter λmax(KQL1,H2) _ (K_Q^L1,H2) λmax(KKL1,H2) _ (K_K^L1,H2) q_excess k_excess score.max near.mean near.max top1 prob 200 3.2 3.2 0.000 0.000 2 1.02 4 0.012 5K 13.2 8.0 0.021 0.011 7 1.02 3 0.057 30K 189.9 156.9 0.098 0.087 64 1.53 11 0.114 60K 502.4 872.8 0.133 0.130 209 1.91 12 0.246 70K 4,527 11,435 5.000 4.330 1,144 6.54 32 0.334 75K 57,088 97,092 49.31 42.39 13,568 2.60 15 0.632 80K 61,633 124,017 27.89 28.46 24,064 1.60 9 0.781 86K 101,698 133,208 peak 22.39 15.95 27,264 peak 1.39 7 0.845 90K 93,315 129,284 17.47 17.71 25,600 1.38 8 0.845 100K 31,657 76,344 12.49 9.75 11,136 1.23 6 0.901 120K 1,008 31,741 0.42 0.17 2,272 1.16 5 0.921 Column definitions are in Section 4.4.1. In one line: score.max is the single global max entry (not an average), near.mean the mean per-row tie count, near.max the worst-row tie count, and q_excess already includes LR and is on the same scale as λ1 _1. 4.4.3 The physical events, stage by stage (L1 H2 only) Stage A (0–60K): a healthy plateau. λ grows slowly from 3 to 873; score.max from 2 to 209; near.mean =1.91=1.91 (normal attention is ≈1≈ 1). The cap is already trimming but the excess is small (0.130.13). Stage B (60K–70K): entering the danger zone. In 10K steps λ(K)=873→11,435λ(K_K)=873→ 11,435 (13×13×), score =209→1144=209→ 1144 (5×5×), and near.mean rises from 1.91 to 6.54 (3.4×3.4×). This is exactly the entry into the uncapped baseline’s creep. Stage C (70K–86K): the cap enters crisis mode. The cap excess escalates: 60K:0.13→ 70K:5.0→ 75K:49(peak crisis)→ 86K:22.60K:0.13\;→\;70K:5.0\;→\;75K:49\;(peak crisis)\;→\;86K:22. The meaning of q_excess =49=49 (see Section 4.5): at that step the cap trims a first-order Δλ λ-growth direction of magnitude 49 nats out of the Muon update, i.e. if uncapped, the next step’s λmax _ would rise by 49. But the cap is a first-order projection and the second order can still grow λ: from 70K to 86K, λ actually rises from 11K to 133K (the cap presses each step’s Δλ λ from an uncapped 50+ down to ∼10 10, so λ grows about 5×5× more slowly). Stage D (86K, peak): the crisis peak. λK=133,208 _K=133,208, score =27,264=27,264 (about half the bf16 max of 65504). But near.mean =1.39=1.39, near.max =7=7 (both down sharply). Why near-max falls: attention becomes peaked. At 70K the attention top1 prob is 0.330.33 (spread); at 86K it is 0.850.85 (very peaked). Spread attention with several near-equal entries per row gives a high near-max; peaked attention with one dominant entry gives a low near-max. This is not designed into the cap — it is a by-product of the transformer’s own learning dynamics; but the cap itself avoids an exponential blow-up while encouraging continued learning in new directions, so training settles into a stable direction (see the figures) in which most tokens’ attention is highly concentrated. Stage E (86K–120K): monotone cool-down. The cap keeps redirecting W away from the high-λ basin. λK=133K→76K→32K _K=133K→ 76K→ 32K, score =27K→11K→2K=27K→ 11K→ 2K. The system fully recovers; by 120K the cap excess is =0.17=0.17, below the level at step 200. 4.4.4 Counterfactual: the uncapped baseline over the same window iter F3 (with cap) uncapped baseline 60K loss =3.13=3.13, λKL1H2=873 _K^L1H2=873 loss =3.03=3.03, healthy 70K loss =3.11=3.11, λKL1H2=11.4 _K^L1H2=11.4K cap crisis enters loss =3.33=3.33, onset 80K loss =3.04=3.04, near.max =9=9 (bf16 critical) loss =3.50=3.50, creep 86K loss =3.00=3.00, λ peak 133K, cap trims 22 nats/step loss ∼3.55 3.55 100K loss =3.11=3.11, λ already down to 76K loss =3.76=3.76, late creep 110K loss =3.07=3.07 (still stable) loss =4.31=4.31, crash 120K loss =2.94=2.94, λ=32λ=32K, fully recovered loss =5.18=5.18 150K (not run) loss =7.63=7.63, blown F3’s loss stays stable through the entire crash window (70K–110K) while the uncapped baseline crashes. This is the strongest direct evidence of the cap’s value. 4.5 The physical meaning of excess_mean (the cap’s counterfactual metric) Definition: excess_meant(ℓ,h,Q):=max(0,⟨Bt(ℓ,h),Ht(ℓ,h)⟩F−ρλ1(ℓ,h)),excess\_mean_t^( ,h,Q):= \! (0,\; B_t^( ,h),H_t^( ,h) _F-ρ _1^( ,h) ), where Bt(ℓ,h)=∂λ1(KQ(ℓ,h))/∂WQ(ℓ,h)=2q1g1⊤B_t^( ,h)=∂ _1(K_Q^( ,h))/∂ W_Q^( ,h)=2q_1g_1 is the WQW_Q direction that grows λ1 _1 fastest; ρ=0ρ=0 is a hard cap. On the scale of H (important): Ht(ℓ,h)H_t^( ,h) is the actual candidate update of that Muon step, H=−ηMuonadj⋅P,P=msign(M),ηMuonadj=0.2max(m,n)⋅ηtH=- _ Muon adj· P, P=msign(M), _ Muon adj=0.2 (m,n)· _t H already carries the LR and the RMS-matched scale — it is not the raw msignmsign direction nor an unscaled momentum. So excess_mean is in units of the change of λ1 _1 per actual Muon step, not “per unit LR”. Geometric reading: first-order Taylor, λ1(WQ+H)≈λ1(WQ)+⟨B,H⟩F _1(W_Q+H)≈ _1(W_Q)+ B,H _F (no separate η factor, since H already includes it). Hence excess_mean=X⇔“if the cap did not act this step, the next step’s λ1 would rise by X (LR included)”. \;excess\_mean=X\; \;``if the cap did not act this step, the next step's $ _1$ would rise by $X$ (LR included)′.\; The units match λ1 _1 itself (our λmax(KQ) _ (K_Q) ranges from ∼3 3 early to ∼105 10^5 at 86K; the excess is a scalar on the same scale). Verification on L1 H2: at iter 86K, excess_mean =22=22. From 86K to 100K (14K steps), an uncapped expected rise of ∼10 10–2222 per step (mean ∼15 15) would accumulate a first-order expected λ increase of ∼200,000+ 200,000+. Instead λ falls by 56K (133K→76K133K→ 76K). 4.6 Why our cap fixes what the paper’s patch does not The paper fixes the problem with a patched (STABLE) softmax — a numerical safeguard inside attention. We use the cap to stop λmax _ at the source, before it reaches the failure threshold. The two act at different places: • the paper changes attention-internal numerics ⇒ even if λ is large, softmax does not blow up; • we change the Muon optimiser ⇒ λ never gets large enough to trouble softmax. Advantages of our approach: 1. the cap does not touch attention internals, so it is compatible with any attention implementation (FlashAttention versions, xFormers, custom kernels); 2. the cap yields excess_mean, a real-time counterfactual diagnostic — the patched softmax is an “after-the-fact fix” and cannot tell you “how this step would have blown up”; 3. the cap keeps λ in a safe range while continuing to learn new directions, so training settles into a stable state (a peaked attention). 4.7 Figures Figure 6: Loss trajectory: F3 vs uncapped vs patched (horizontal axis up to 120K, the steps we actually ran F3 for). Blue == F3 (ours, with the QK cap and no softmax patch); red == uncapped (no cap, no patch, known to crash); grey == patched (no cap, with the softmax patch). The red shading == the 70K–110K crash window. F3 (blue) traverses the whole crash window with loss stable at 3.03.0–3.13.1; the uncapped run (red) begins creeping at 80K, crashes at 104K (loss 4.28→4.494.28→ 4.49), and reaches 5.18 by 120K. F3 vs uncapped is a strictly controlled comparison: same 327K tok/iter, same Muon, same GPT-2S, same data, same lr schedule — the only variable is whether the QK covariance cap is on. Figure 7: Global grad-norm trajectory (log-y axis, up to 120K). The global grad norm ‖g‖2=∑p‖∇pL‖F2\|g\|_2= _p\| _pL\|_F^2 (over all parameters p, before clipping), one value per train step. Blue == F3 (continuous trace); red dots == the uncapped baseline (one summary point per 1000 iters). Dashed black == grad-clip max-norm =1.0=1.0. The two overlap before 56K (grad norm ∼0.4 0.4–0.50.5); from ∼56 56K the uncapped run explodes — 64K==52, 72K==609, 80K==886, 88K==8K, 96K==29K, 104K=6.7×105=6.7× 10^5, 120K=8.1×108=8.1× 10^8 — eight orders of magnitude. F3 stays <0.5<0.5 throughout (except an iter-86K spike to 1.41 that falls back within seconds, the moment the L1 H2 cap acts). A log y axis is needed to see the full uncapped trajectory. Figure 8: The 86K crisis at L1 H2 (this figure is specifically for L1 H2). Left axis: λmax(KQL1,H2) _ (K_Q^L1,H2) (blue) and λmax(KKL1,H2) _ (K_K^L1,H2) (red) over the run; right axis: the corresponding q_excess (green) and k_excess (orange). Red shading == the 80K–95K crisis window. λK _K rises from 873 at 60K to a peak of 133K at 86K (150×150×), and the cap excess rises in step from 0.13 to 22/step (170×170×). The cap intercepts hardest exactly at the λ peak. Figure 9: Two bf16-safety quantities of the L1 H2 attention scores. Left: score.max (the single largest score entry over (B,Tq,Tk)(B,T_q,T_k) for this head, not a row-max average, log scale); right: near_max.mean (the number of “near row max” entries per query token, averaged over batch and query). The bf16 max, 65504, is marked in red. Both spike at iter 86K then recover: score.max 209→27,264→2,272209→ 27,264→ 2,272; near_max.mean 1.9→6.51.9→ 6.5 (70K) →1.4→ 1.4 (86K, attention now peaked) →1.2→ 1.2 (120K). Neither breaks the bf16 mantissa threshold, so the softmax error does not accumulate. 4.8 Part 4 summary • F3 fully covers the uncapped baseline’s crash window (70K–110K) with a stable loss (Fig. 6); • L1 H2 undergoes a real near-explosion at iter 86K (λK=133 _K=133K, score =27=27K, near half the bf16 ceiling), and the cap intervenes in real time with excess =22=22 nats/step; • the cap’s counterfactual diagnostic excess_mean quantifies, in real time, “what the next step would have done” — information the patched softmax cannot provide; • the rescue is a combination of three mechanisms: (i) the cap trims the first-order growth rate; (i) the training dynamics’ own attention-peaking lowers near-max; (i) the cap keeps redirecting W away from the high-λ basin. 5 Conclusions across the three case studies Part W (capped) X K=WKXW⊤K=WK_XW Failure / goal 2 nano FFN mlp.c_proj 384×1536384× 1536 GELU hidden ℝ1536R^1536 FFN out covariance ℝ384×384R^384× 384 isotropy, PR ↑ 3 MoE router router 64×38464× 384 RMSNorm token ℝ384R^384 sel-score cov Kz+b∈ℝ64×64K_z+b ^64× 64 rank-1 collapse as top-share →1→ 1 4 FA Q/K per-head 64×76864× 768 layer input ℝ768R^768 per-head KQ(ℓ,h)∈ℝ64×64K_Q^( ,h) ^64× 64 λmax _ blow-up ⇒ bf16 softmax error The unified framework (Section 1): 1. in the scale-invariant regime Muon removes the 1/‖W‖F1/\|W\|_F brake ⇒ the Frobenius/spectral norm drifts more strongly than under SGD; 2. the spectral cap removes the first-order top mode but keeps a non-negative second order ⇒ the weight still learns in new directions; 3. the spectral cap == an H∞H_∞ entropy cap; controlling H2H_2/H1H_1 would need other projection directions. The unified experimental observations: 1. KXK_X has a different dimension and physical meaning in each application, but the cap algorithm is always “sample-estimate KYK_Y + power iteration + first-order projection”, at negligible cost O(M⋅small)O(M·small); 2. a data-coupled cap (directly on KYK_Y) is more on-target than a pure W-cap (nanoGPT); 3. the cap is decisive at the margins / in crises (MoE L0 rank-1, FA L1 H2 at 86K) and, in the typical case, leaves the result unchanged while protecting the geometry. caveatbox Limitations and open questions. (i) The whole theoretical story rests on Assumption 1 (exact scale invariance); the quantitative predictions — the 1/‖W‖F1/\|W\|_F brake, the t1/4t^1/4 vs t1/2t^1/2 norm growth, and ⟨W,G⟩F≈0 W,G _F≈ 0 — should be tested much more broadly than the single checkpoint of Section 1.4.4 before being relied upon. (i) The experiments are at small scale (nanoGPT / GPT-2 small, tens of thousands of steps) and single-seed in places; the effects at larger scale are untested. (i) The FlashAttention result compares against one uncapped log rather than a matched ensemble. We regard these as a coherent first set of observations rather than a settled account, and we welcome comments, corrections, and counter-examples. Appendix A Figure index Path Content nano_D_val_loss.pdf nanoGPT D0/D1/D1d val loss (Fig. 1) nano_D_isotropy_pr.pdf nanoGPT FFN KYK_Y participation rank (Fig. 2) g3_val_loss_3way.pdf MoE G0/G1/G3 val loss (Fig. 3) g3_per_layer_cv_3way_seed42.pdf MoE seed=42 per-layer load CV (Fig. 5) g3_per_expert_load_3way_seed42.pdf MoE seed=42 per-expert load histogram (Fig. 4) fa_F3_loss_vs_paper.pdf F3 vs baseline loss trajectory (Fig. 6) fa_F3_grad_norm.pdf F3 grad-norm trajectory (Fig. 7) fa_F3_L2_crisis.pdf L1 H2 86K λmax _ + excess crisis (Fig. 8) fa_F3_score_nearmax.pdf L1 H2 score.max + near_max (Fig. 9) References [1] H. Qiu and Q. Yao. Why low-precision transformer training fails: an analysis on Flash Attention. In International Conference on Learning Representations (ICLR), 2026. arXiv preprint arXiv:2510.04212. [2] J. Liu et al. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982, 2025. [3] K. Jordan. Muon: an optimizer for hidden layers in neural networks. https://kellerjordan.github.io/posts/muon/, 2024. [4] A. Karpathy. nanoGPT. https://github.com/karpathy/nanoGPT, 2023.