Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory

TL;DR

This paper analyzes the capacity and efficiency of spectral optimizers like Muon in learning associative memory, revealing their superior recovery rates and scaling properties compared to SGD and Newton's method.

Contribution

It provides a theoretical characterization of spectral optimizer performance in a tractable associative memory model, highlighting their advantages over traditional methods.

Findings

Muon exceeds SGD in storage capacity and matches Newton's method using only first-order info.

Muon saturates at a larger critical batch size, enabling better scaling.

Experimental results validate the theoretical scaling laws and recovery rates.

Abstract

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.