Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory

Abstract: Spectral optimizers such as Muon have recently shown strong empirical performance in large-scale LLM 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 and SGD 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 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 Muon and lays the groundwork for establishing scaling laws across more practical language modeling tasks and optimizers.

• The paper demonstrates that Muon accelerates item recovery and boosts capacity beyond traditional SGD limits through spectral orthogonalization.
• It employs rigorous theoretical analysis and empirical experiments to quantify sharp capacity scaling laws and critical batch-size thresholds in associative memory tasks.
• The findings offer practical insights for enhancing transformer memory recall and guiding efficient training protocols in large-scale language models.

Sharp Capacity Scaling of Spectral Optimizers in Associative Memory Learning

Overview

The paper "Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory" (2603.26554) addresses the quantitative advantage of spectral matrix-based optimizers, specifically Muon, over conventional stochastic gradient descent (SGD) in the context of associative memory tasks. Associative memory models formalize factual recall in transformer architectures, where weight matrices implement input-output mappings between atomic facts, capturing the superposition principle: embedding many associations in a low-dimensional space.

The analysis sharply characterizes the scaling laws of item recovery and critical batch size for Muon and SGD, using a linear associative memory with Gaussian embeddings and power-law frequency distributions—a regime relevant to realistic language modeling. Strong empirical and theoretical results are provided, showing both a superior storage efficiency and accelerated early learning dynamics for Muon, especially under large batch sizes and heavily-tailed distributions.

Theoretical Analysis: One-Step Recovery and Capacity Scaling

The authors investigate the storage efficiency of Muon versus SGD via both population and minibatch training dynamics, focusing on the item recovery rate after a single optimization step in the associative memory objective.

Muon leverages spectral orthogonalization (polar factor) of the gradient, amplifying weak spectral signals and enabling superposition beyond the orthogonality-limited bound of $N \leq d$. The theoretical results show:

• Muon one-step capacity: Recovers up to $$\widetilde{\Theta}(\min\{d^{1+1/2\alpha}, B^{1/\alpha}\})$$ items, for a power-law frequency exponent $\alpha>1$. This exceeds the conventional $d$-item limit and approaches the information-theoretic optimal $d^2$ scaling via superposition.
• SGD one-step capacity: Limited to $$\widetilde{\Theta}(\min\{d^{1/2\alpha}, B^{1/\alpha}\})$$, saturating at a much smaller critical batch size.

Muon thus provides a batch-size-dependent amplification, saturating at $$B^\star_\text{Muon} = \widetilde{\Theta}(d^{\alpha + 1/2})$$, compared to $B^\star_\text{SGD} = \widetilde{\Theta}(d^{1/2})$, as formally and empirically validated.

Figure 1: Capacity scaling with embedding dimension $d$; Muon demonstrates higher recovery rates, particularly as $d$ increases under power-law frequency distributions.

This scaling is further substantiated in the population regime with Gaussian embeddings, where the critical batch size and dimension exponents match theoretical predictions.

Multi-Step Dynamics and Acceleration

The theoretical framework is extended to multi-step gradient trajectories via thresholded (deflation-based) approximations:

• Muon multi-step recovery: After $t$ steps, achieves $\widetilde{\Theta}(d^{2-(1-1/2\alpha)^t})$ capacity, converging exponentially toward the optimal $d^2$ limit.
• SGD multi-step recovery: Recovers $d_t$ via a recursion that remains strictly suboptimal for initial steps, requiring at least $\lceil 2\alpha \rceil$ steps to reach $d_t \gtrsim d$.

Muon thus achieves substantial acceleration in the early phase when gradients are anisotropic; eventually, both optimizers converge at comparable rates as the gradient becomes isotropic.

Figure 2: Gradient descent (population); demonstrates the slower scaling of SGD relative to Muon at initial steps.

Figure 3: Gradient descent (minibatch); illustrates batch-size-dependent saturation for SGD.

Figure 4: Capacity at $T=2$; confirms Muon's accelerated scaling in the early learning phase.

Figure 5: Multi-step GD vs.~Muon (population); Muon's advantage is prominent in the early steps.

Empirical Validation: Synthetic and Transformer-Based Tasks

Extensive experiments substantiate the theoretical scaling laws:

• Synthetic linear memory: Muon outperforms SGD both in item recall and final loss, in both population and minibatch settings, as shown in fitted power-law exponents and batch-size plateaus.
• Transformer in-context recall: Muon exhibits superior out-of-distribution (OOD) memory recall across embedding dimensions, batch sizes, and power-law exponents for trigger/output distributions.

  Figure 6: ID (left two) and OOD (right two) accuracy on the in-context recall task as a function of model dimension; Muon consistently outperforms SGD and AdamW across metrics.

  

Figure 7: % Memory recall accuracy $R(\bW_V^2)$ against batch size $B$; Muon achieves higher recall accuracy more rapidly with increasing batch size.

Relaxing the embedding orthogonality reveals that Muon reliably stores $d^{1+1/2\alpha}$ items, leveraging spectral amplification for empirical factual recall.

Figure 8: OOD accuracy as a function of model dimension for varying power-law exponents; Muon retains robustness as the tail becomes heavier, unlike AdamW and SGD.

Bold and Contradictory Claims

The authors demonstrate (both theoretically and empirically) that Muon directly enables superposition-based capacity scaling far beyond orthogonality-limited bounds, recovering more than $d$ items after a single step. The claim that Muon's critical batch size is dramatically larger than SGD is quantitatively precise and empirically validated.

Additionally, the one-step recovery rate of spectral estimators is proven to be optimal under broad invariant conditions, showing no further improvement is possible even with more sophisticated spectral maps.

Implications and Future Directions

Practical

This analysis provides rigorous foundation for the empirical success of spectral optimizers (like Muon) in large-scale LLM training, emphasizing their advantages in high-capacity factual recall, especially at large batch sizes and under power-law item distributions. Critical batch size characterization informs efficient training protocol design, particularly when deploying Muon or similar matrix-adaptive optimizers for memory-intensive transformer workloads.

Theoretical

The characterization of capacity scaling laws opens avenues for rigorous understanding of spectral preconditioning and signal amplification in high-dimensional learning. The proof techniques using resolvent representations, perturbative expansions, and invariance principles are highly applicable to more complex, compositional or multi-hop reasoning tasks. Open questions remain (as conjectured) about end-to-end guarantees outside the deflation heuristic and in more anisotropic settings, as well as for momentum-based spectral methods.

Future Developments

Speculative directions include: extending scaling laws to anisotropic embeddings and non-linear architectures; formalizing momentum-based spectral update limits; benchmarking spectral-memory scaling in practical LLMs under realistic data distributions; investigating compositional and multi-hop reasoning scaling with spectral optimizers; and integrating spectral updates into adaptive methods for robust generalization under heavy-tailed distributions.

Conclusion

This work rigorously establishes sharp capacity scaling laws for spectral optimizers—most notably, Muon—in high-dimensional associative memory learning. Through both analytic and empirical demonstrations, it is evidenced that Muon enables dramatically higher factual recall and faster early training than vanilla SGD, especially under large batches and heavy-tailed item frequencies. These results provide quantitative foundation for spectral methods in LLMs and highlight compelling directions for both practical optimizer development and theoretical scaling law research in AI.