Targeted muon method improves training for language and vision AI
Spectrally Targeted Muon
Machine Learning
Summary
Training AI models can be tricky because some directions that help learning are very small and often ignored. The paper introduces a way to focus on certain parts of these learning steps by only adjusting the big or small parts of updates, instead of all parts. This new method helps the model learn better and faster, especially for language models where many update parts are very small at first. The authors also compare this method to others like AdamW and show it builds useful learning directions faster.
What this means in practice
- •For deep learning engineers: Improve training speed and effectiveness of large language and vision models by selectively orthogonalizing update directions during optimization.
- •For machine learning infrastructure teams: Integrate targeted orthogonalization methods to optimize resource use in training AI models where gradient update quality affects convergence time.
Authors
Vishrut Goyal, Rohan Ramkumar
Abstract
The Muon optimizer orthogonalizes each update matrix, setting all of its singular values to one, and has proven highly effective for training large language models. It remains unclear, however, whether this success comes from amplifying small singular directions that gradient descent neglects or from suppressing large, degenerate directions that disrupt training. We introduce Spectrally Targeted Muon, which orthogonalizes only the singular values above or below a threshold $τ$, so that varying $τ$ interpolates between normalized SGD and Muon. It isolates the relevant singular subspaces with projections computed by Newton-Schulz iteration on a shifted Gram matrix, so no SVD is needed. We evaluate these variants on the CIFAR-10 and NanoGPT speedruns, tracking the effective rank of gradient, update, and weight matrices and a new metric, the alignment of updates with the tangent space of the weight matrix's isospectral manifold. We find three things. First, the small singular values of the momentum are not noise. Orthogonalizing everything except the few largest singular values of each matrix nearly matches Muon while touching only a small fraction of the momentum, whereas orthogonalizing only the top falls well short even though it holds almost all of it. On language models every momentum singular value is far below one, so targeted orthogonalization can only amplify, and Muon wins by making directions that are too small to train on at their raw scale trainable. Second, shrinking the largest singular values is what keeps the parameter spectrum flat; this is cheap, and it is not what drives the loss. Third, AdamW differs from Muon mainly in how slowly it builds structure, which explains its slower start and why warmup helps AdamW but only hurts Muon.