← BlogResearch
June 10, 20268 min

Same Dictionary, Different Geometry

Recent work shows Muon representations are more robust and transferable. We ask what that means anatomically: at 124M scale, firing-pattern matching says the sparse features are the same; Muon recovers the AdamW feature set up to seed noise, packaged into a hotter, sparser, more crowded, and more seed-reproducible residual-stream geometry.

TL;DR. Muon's better representations are not a different sparse dictionary. They look like the same dictionary, packed into a different geometry. We trained matched GPT-2-class models with Muon and AdamW on MC-1, held validation loss fixed, and compared what formed inside: a Muon model's features match an AdamW model's as well as two AdamW seeds match each other, while the packaging shifts everywhere we looked, with weights spread across 2–10× more directions, a residual stream running 4–7× hotter, fewer and rarer live SAE features packed closer together. And Muon finds the shared feature set more consistently across random seeds than AdamW does.

Optimizers are switching under the field's feet. Muon-family pretraining has moved from speedrun curiosity to production frontier runs in under two years, on one argument: it reaches the same loss faster. Loss, though, is a one-number summary of a model. Two models can match it exactly and still differ in what they computed to get there, which is precisely the kind of difference that matters if you build on top of the model's internals, as interpretability work does. So before asking whether Muon is a better optimizer, we asked a quieter question: does the optimizer change what ends up inside?

Muon is steepest descent under the spectral norm [1, 2]. With momentum Mt=βMt1+(1β)GtM_t = \beta M_{t-1} + (1-\beta)G_t and SVD Mt=UtΣtVtM_t = U_t \Sigma_t V_t^\top, the update is

Wt+1  =  WtηUtVtW_{t+1} \;=\; W_t - \eta\, U_t V_t^\top

which is MtM_t with Σt\Sigma_t replaced by the identity: every direction in the update gets equal strength, however loudly or quietly it appeared in the gradient. AdamW instead rescales each coordinate by a running second-moment estimate and never sees the matrix. An update rule that flat should leave a fingerprint in the weights, and recent experiments find it. Muon-trained matrices carry more uniform spectra than Adam-trained ones, in pretraining [3] and in fine-tuned adapters [4], while standard training drifts the other way, toward rank minimization [5].

Recent work shows Muon representations are more robust and transferable: Ruan et al. [8], two days before this post, demonstrate both on corrupted inputs and downstream tasks, with higher effective rank in the hidden states. We ask what that means anatomically: are the sparse features different? At 124M scale, firing-pattern matching says no. Muon recovers essentially the AdamW feature set up to seed noise, but packages it into a hotter, sparser, more crowded, and more seed-reproducible residual-stream geometry. The rest of this post is the evidence for that sentence, and the expectation it overturns: if Muon spreads weight mass across hundreds of directions where Adam concentrates it into a few dozen, you would predict genuinely different features. The prediction fails, and the way it fails is the interesting part.

The bench

We wanted the cleanest possible A/B test, so we borrowed one. The modded-nanogpt speedrun history [6] contains a 124M-parameter GPT-2-class trainer (RoPE, QK-norm, ReLU², plain softmax attention) that Keller Jordan used for his own optimizer comparison in October 2024, with hyperparameters he tuned per optimizer over dozens of attempts. We took both arms verbatim.

  • Model: 124M GPT-2 class, 12 layers, d=768
  • Data: FineWeb, identical order for every run
  • Arms: AdamW (lr 0.0018) vs Muon (lr 0.02), embed/head on AdamW in both
  • Seeds: 3 per arm, varying init only
  • Comparison point: matched validation loss, not matched steps
  • Hardware: 8×B200 on MC-1, one full run ≈ 8 minutes

Two of these choices carry most of the weight. Comparing at matched validation loss removes the lazy confound: a worse model has different features because it is worse, and Muon reaches any given loss sooner than AdamW. We save checkpoints the first time each run crosses fixed loss waypoints and compare there. Final losses at 5,100 steps: Muon 3.292 ± 0.001, AdamW 3.350 ± 0.002 across seeds. The gap reproduces Jordan's published comparison.

fineweb val loss against training step3.183.694.194.695.193.83.63.53.43.351k2k3k4k5ktraining stepfineweb val lossadamwmuon
Validation loss, six runs (three seeds per optimizer). Dashed lines are the matched-loss waypoints: every comparison below happens at checkpoints where both arms sit at the same loss.

Second, every measurement below uses identical instruments. The sparse autoencoders that follow share one init, one data stream, one set of hyperparameters across all six models. The model under the microscope is the only variable.

The fingerprint is real

First we confirmed the known result on our bench. We summarize a spectrum by its stable rank,

srank(W)  =  WF2W22  =  iσi2σ12\mathrm{srank}(W) \;=\; \frac{\lVert W\rVert_F^2}{\lVert W\rVert_2^2} \;=\; \frac{\sum_i \sigma_i^2}{\sigma_1^2}

which counts the directions carrying non-negligible mass: 1 for a rank-one matrix, min(m,n) for a perfectly flat spectrum. At matched validation loss, the stable rank of Muon-trained weight matrices sits far above AdamW's, in every weight class of the network:

  • MLP in · Muonmuon
    115srank
  • MLP in · AdamWadamw
    21srank
  • attention Q · Muonmuon
    67srank
  • attention Q · AdamWadamw
    27srank
  • attention out · Muonmuon
    136srank
  • attention out · AdamWadamw
    61srank

peak 160 stable rank

Stable rank at matched validation loss (3.40 waypoint), mean over seeds and layers. Attention K/V and MLP out follow the same pattern, at 1.7–2.6×.

The gap is widest early in training (10× for MLP input matrices at val loss 3.8) and narrows as AdamW slowly accumulates rank. Raw weight spectra inherit structure from the init, so we recomputed everything on WWinitW - W_{\mathrm{init}}, the pure accumulated update. The numbers barely move (MLP in: 5.6×, attention Q: 2.5×); the fingerprint comes from training, and the init contributes close to nothing.

So far, so consistent with the literature. The model stores its computation across many more weight directions under Muon. The question that motivated all of this is whether any of it reaches the features.

Same features

Our microscope is a top-k sparse autoencoder. It encodes a residual-stream activation xR768x \in \mathbb{R}^{768} as

z=TopKk(Wenc(xbdec)+benc),x^=bdec+Wdeczz = \mathrm{TopK}_k\big(W_{\mathrm{enc}}(x - b_{\mathrm{dec}}) + b_{\mathrm{enc}}\big), \qquad \hat{x} = b_{\mathrm{dec}} + W_{\mathrm{dec}}\, z

with k = 32 features active out of 12,288. Each row of WdecW_{\mathrm{dec}} is a feature direction, and ziz_i on a given token says how strongly feature i fired there.

Comparing features across two independently trained models has a trap in it. Each model picks its own basis for the residual stream, so you cannot compare SAE decoder directions across models. We checked: even two seeds of the same optimizer share no matching decoder vectors (mean best cosine 0.14), and pairs that share an init match slightly better than pairs that do not, which tells you the coordinate frame comes from the init. Direction matching across models measures basis rotation, not feature content.

The basis-free version: run two models over the same million tokens, record when each SAE feature fires, and match features by the correlation of their firing patterns. For feature i of model A, the match score is

mi  =  maxjρ(ziA,  zjB)m_i \;=\; \max_j\, \rho\big(z_i^{A},\; z_j^{B}\big)

the best Pearson correlation against any of model B's features over the shared tokens. A feature is what it responds to, not where it points.

We did this for all 15 pairs of our six models. Within-optimizer pairs set the ceiling, since two seeds of the same recipe can only agree up to SAE training noise:

mean best feature match0.5830.5940.6050.6170.628mean best feature matchmuon <-> muonmuon <-> adamwadamw <-> adamw
All 15 model pairs at layer 6: per-pair mean best match with 95% bootstrap CIs over features. Every within-Muon pair sits above every other pair; cross-optimizer pairs interleave with within-AdamW. Share of features with a strong (>0.7) twin: 43% / 40% / 39%.

The middle group is the result. Crossing the optimizer boundary costs nothing: a Muon model's features match an AdamW model's features as well as two AdamW seeds match each other (0.599 vs 0.594, nine cross pairs vs three within pairs, consistent across every pairing). Changing the optimizer perturbs feature identity about as much as changing the random seed. The picture replicates at layer 9, where the packaging differences described below are most extreme: within-Muon 0.610, cross 0.586, within-AdamW 0.585. Same ordering, same conclusion, at the depth most hostile to it.

There is a second result in that chart. Muon agrees with Muon (0.618) more than AdamW agrees with AdamW (0.594). Train the same recipe twice with different seeds and you get a more similar pair of models under Muon. The feature set is shared property; Muon retrieves it more reproducibly. We did not predict this and have not seen it reported.

Different packaging

Feature identity survives the optimizer swap; feature demographics shift. Two more definitions: reconstruction error is the fraction of variance unexplained, FVU=xx^2/xxˉ2\mathrm{FVU} = \sum\lVert x - \hat{x}\rVert^2 / \sum\lVert x - \bar{x}\rVert^2, and decoder crowding is the mean over features of the cosine to the nearest other feature direction. Pointing the identical SAE at each model's layer-6 residual stream, at matched loss: reconstruction quality ties (FVU 0.119–0.140 vs 0.121–0.147), and everything else separates with no overlap across seeds.

  • Muonof 12,288 features
    69dead per 1k
  • AdamWof 12,288 features
    11dead per 1k

peak 90 dead per 1k

Dead features per 1,000 (never fire on held-out data), mean over seeds. Muon's live features also fire more rarely (median rate 7.0×10⁻⁴ vs 8.4×10⁻⁴) and pack closer together (mean nearest-neighbor cosine 0.310 vs 0.289).

The same dictionary-finder, given a Muon model, comes back with a smaller live vocabulary of rarer, more specialized features packed closer together in direction space.

We tried to kill this three ways. Muon runs its residual stream 4 to 7 times hotter than AdamW, so we renormalized activations and retrained: the gaps survive. Muon's matched-loss checkpoints see fewer tokens, so we recompared at matched tokens instead: numbers unchanged. We repeated the measurement at other depths: layer 9 separates harder (26× on dead features), layer 3 shows almost no separation. The packaging effect grows with depth.

  • layer 3 · Muonmuon
    2,800mean ‖x‖
  • layer 3 · AdamWadamw
    550mean ‖x‖
  • layer 6 · Muonmuon
    4,200mean ‖x‖
  • layer 6 · AdamWadamw
    700mean ‖x‖
  • layer 9 · Muonmuon
    7,915mean ‖x‖
  • layer 9 · AdamWadamw
    1,072mean ‖x‖

peak 8,500 mean activation norm

Residual-stream activation norms by depth. Muon runs the stream 4–7× hotter at every layer, and the gap grows with depth.

A single-seed spot check at 350M (24 layers, d=1024) reproduces the loss gap (3.13 vs 3.19), the norm inflation (3.6×), the crowding gap, and the rarer firing. The dead-feature fraction does not replicate: both arms jump above 20% dead there and the ordering flips, which tells us that metric is the most sensitive to SAE configuration (our 16×, k=32, 50M-token recipe is undertrained for a d=1024 stream). The three metrics that survive the scale jump are the ones we lean on.

Side by side, the two results invert the prediction we started with. Flat weight spectra suggested features smeared across more directions. The features came back fewer, rarer, and more tightly packed, while being the same features. Weight-space geometry and feature-space geometry moved in opposite directions. Ruan et al. [8] sharpen the puzzle from the other side: Muon's hidden states span a higher-effective-rank subspace, yet our SAEs find a smaller live dictionary inside that subspace. We do not have a mechanism story we trust yet; the residual-norm inflation and the depth dependence look like clues.

Why you might care

The interpretability toolchain grew up on Adam-family models. SAE recipes, dead-feature lore, sparsity budgets, activation-scale assumptions: all calibrated against one optimizer's idea of a residual stream. Frontier labs now ship Muon-family pretraining (Kimi K2 most visibly [7]). Our result cuts both ways for that transition. The features you are hunting are the same ones, so circuit-level findings should transfer. The population statistics around them shift enough to matter for tooling, starting with dead-feature rates and activation scales.

The reproducibility result is the one we want to chase. If Muon reliably converges to a more canonical version of the model, that is a statement about the loss landscape and about how much of "what a model learns" is recipe versus accident. It would also make Muon-trained models the better substrate for interpretability research, where seed-to-seed feature drift is a tax on every experiment.

Limitations

Mostly one scale: the 350M spot check covers the packaging metrics with a single seed per arm, and feature-identity matching has not been run at 350M (it needs multiple seeds per arm for the within-optimizer ceiling). One architecture, one data distribution. Feature matching at two layers of the 124M models. SAEs are a proxy for features, k=32 top-k SAEs at one expansion factor are a specific proxy, and firing-pattern correlation over 1M tokens is noisy for the rarest features, though the noise applies to all three comparison groups equally. The within-optimizer ceiling (0.62 mean match) is itself well below 1.0, so all claims about identity live relative to that ceiling, not to perfect feature recovery.

What's next

Three threads: multi-seed 350M runs for feature matching at scale, a mechanism for the norm inflation, and the experiment this result sets up directly, which is whether Muon's noise-amplification at small singular values can be shaped without losing its reproducibility advantage.


References

[1] Jordan, "Muon: An optimizer for hidden layers in neural networks", 2024.

[2] Chen et al., "Muon Optimizes Under Spectral Norm Constraints", arXiv 2025.

[3] "To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters", arXiv 2026.

[4] "Can Muon Fine-tune Adam-Pretrained Models?", arXiv 2026.

[5] Yunis et al., "Approaching Deep Learning through the Spectral Dynamics of Weights", arXiv 2024.

[6] Jordan et al., modded-nanogpt and the October 2024 optimizer comparison record, GitHub 2024.

[7] Moonshot AI, "Kimi K2 Technical Report", 2025.

[8] Ruan, Zhang, Wang, Zhang, "Muon Learns More Robust and Transferable Features than Adam", arXiv 2026.

Experiments ran on MC-1, Maincode's 8×B200 node. Code, reproducible training logs, and all pairwise matching results: github.com/maincodeMax/muon-feature-study.