Where Muon’s Advantage Lives: Small-Scale Causal Localization of an Optimizer’s Effect on Learned Features
Abstract
Does an optimizer change what a language model learns, or mainly how it stores and reaches it? We study Muon versus AdamW on a 124M GPT-2-class model, matching sparse-autoencoder features across independently trained models by firing-pattern correlation and calibrating against within-optimizer seed baselines. Under this fixed top-k SAE protocol, cross-optimizer feature agreement reaches the AdamW seed ceiling, while Muon changes spectra and yields a more seed-reproducible representation. We then apply Muon to individual attention sub-circuits. Muon on the value circuit WV WO is largely sufficient to recover the loss and reproducibility gains, whereas Muon on the query/key scoring circuit WQ WK^T raises stable rank to near-full-Muon levels yet recovers neither. A rank-promoting AdamW control rules out endpoint W_OV stable rank as a sufficient knob; it does not by itself prove update geometry is the unique mechanism. An attention-architecture control moves the same feature-matching metric far below its seed ceiling, showing the metric can detect feature changes. All main models are below 3B parameters; FLOP budgets and artifacts are reported in Appendix G.