Spectral optimizers such as Muon have recently shown strong empirical performance in large-scale language model training, but the source and extent of their advantage remain poorly understood. We study this question through the linear associative memory problem, a tractable model for factual recall in transformer-based models. In particular, we go beyond orthogonal embeddings and consider Gaussian inputs and outputs, which allows the number of stored associations to greatly exceed the embedding dimension. Our main result sharply characterizes the recovery rates of one step of Muon, SGD, and Newton's method on the logistic regression loss under a power law frequency distribution. We show that the storage capacity of Muon significantly exceeds that of SGD, and even matches Newton's method while only using first-order information. Moreover, Muon saturates at a larger critical batch size. We further analyze the multi-step dynamics under a thresholded gradient approximation and show that Muon achieves a substantially faster initial recovery rate than SGD, while both methods eventually converge to the information-theoretic limit at comparable speeds. Experiments on synthetic tasks validate the predicted scaling laws. Our analysis provides a quantitative understanding of the signal amplification of spectral preconditioners and lays the groundwork for establishing scaling laws across more practical language modeling tasks and optimizers.
Figures & tables
Figure 1: (a) Capacity achieved by one Muon and GD step on the population objective; Muon improves the storage capacity when frequency is power-law distributed with exponent α>1 . (b) Critical batch size for the first Muon and SGD step ( α=1.5 ); the Muon capacity saturates at a much larger batch size than SGD.
Figure 2Figure 3
Figure 2: Capacity scaling after one population Muon and GD step. We set N=100,000 and vary d,α . Each experiment is repeated 16 times. For each α , we fit the dimension exponents of the mean capacity dCα (dashed lines), and then find the best fit of exponents Cα in the form of Cα=c1+αc2 (solid lines). Observe that Muon achieves much higher storage than GD, and the exponents are consistent with Theorems 4.1 , 4.3 .
Figure 3: Capacity scaling after one Muon and SGD step on empirical loss. We set N=100,000,α=1.5 , and vary the minibatch size B . Each experiment is repeated 16 times. The dashed red line indicates the information-theoretic rate, and the horizontal dashed lines in Figure 3(b) correspond to the d1+2α1 ceiling; the predicted critical batch sizes are given by their intersections. Observe that Muon offers capacity gain over SGD only at sufficiently large B , and the empirical critical batch sizes match well with our predictions.
Figure 4: Capacity after T Muon steps on the population cross-entropy loss. We set N=250,000 , η=2d . Figures 4(a) , 4(b) , 4(c) report the capacity at T=2,3,4 , respectively (see Figure 2(b) for T=1 ); Figure 4(d) presents the capacity at large T : we run Muon for up to 500 steps and early stop when the capacity improvement over 10 steps drops below 0.5% . Figure 4(e) compares the fitted dimension exponents against predictions of Theorem 5.4 ; observe that the exponents agree except at small α and large T .
Figure 5: Capacity scaling of multi-step Muon and GD. We set N=100,000 , α=1.5 . (a) Population update: for GD we implement an increasing learning rate schedule (see Theorem 5.5 ) with η1=0.01d ; for Muon we use a fixed step size η=d . Observe that the benefit of Muon is most visible in the “early phase” of training (the initial plateau of GD in the first 3 steps is due to small η1 chosen for numerical stability). (b) Capacity of minibatch Muon vs. total sample size B×T ; for each batch size B , we run minibatch Muon for T=20 steps with η=d . Dashed red line indicates the information-theoretic rate (BT)1/α .
Figure 6: Capacity scaling after one (population) Muon and Newton step in the anisotropic setting: we choose ui∼N(0,d1Id) , vi∼N(0,Ξv) , where Ξv is a trace-normalized diagonal matrix with λi(Ξv)≍i−κ , κ≥0 . We set N=100,000 and vary d,α . For Newton’s method we add a ridge regularization λ=10−8 for numerical stability when the preconditioner is rank-deficient. Observe that when κ=0 (isotropic, Figure 6(a) ), Muon and Newton both achieve d1+2α1 capacity, but as κ increases (Figures 6(b) , 6(c) ), the performance of Muon worsens while the Newton update remains invariant.
Figure 7: ID (left two) and OOD (right two) accuracy on the in-context recall task as a function of model dimension, for Muon, AdamW, and SGD, with batch size 256 at iterations 128 and 1024 . For each (dim,optimizer) pair, the learning rate and batch size are chosen to maximize accuracy.
Figure 8: OOD and memory recall accuracy as a function of batch size B , for Muon, AdamW, and SGD (columns left to right), with different curves per model dimension, at iteration 1024 . For Figure 8(b) , we use a two-layer transformer with no feed-forward layers to avoid redundancies between the value matrix and the subsequent MLP layer. For each (B,dim) pair, the learning rate is chosen to maximize accuracy.
Figure 9: OOD accuracy as a function of model dimension for Muon, AdamW, and SGD (columns left to right), with batch size 256 at iteration 512 . Each curve corresponds to a different power-law exponent for (a) the output distribution α0 ; (b) the trigger distribution αt , with α=0 being the uniform distribution and larger α concentrating probability mass on fewer tokens (where we expect adaptive optimizers to be beneficial). For each (dim,α) pair, the learning rate is chosen to maximize OOD accuracy.