When do spectral gradient updates help in deep learning?

TL;DR

This paper introduces a layerwise condition predicting when spectral gradient updates outperform Euclidean ones in deep learning, supported by theoretical analysis and experiments on neural networks and language models.

Contribution

It proposes a simple criterion based on matrix ratios to identify regimes where spectral methods are advantageous, with theoretical proofs and empirical validation.

Findings

Spectral updates outperform Euclidean updates when the nuclear-to-Frobenius ratio is high.

Post-activation matrices exhibit low stable rank at initialization in various models.

In trained models, activation stable rank remains low, favoring spectral methods.

Abstract

Spectral gradient methods, such as the recently popularized Muon optimizer, are a promising alternative to standard Euclidean gradient descent for training deep neural networks and transformers, but it is still unclear in which regimes they are expected to perform better. We propose a simple layerwise condition that predicts when a spectral update yields a larger decrease in the loss than a Euclidean gradient step. This condition compares, for each parameter block, the squared nuclear-to-Frobenius ratio of the gradient to the stable rank of the incoming activations. To understand when this condition may be satisfied, we first prove that post-activation matrices have low stable rank at Gaussian initialization in random feature regression, feedforward networks, and transformer blocks. In spiked random feature models we then show that, after a short burn-in, the Euclidean gradient's nuclear-to-Frobenius ratio grows with the data dimension while the stable rank of the activations remains bounded, so the predicted advantage of spectral updates scales with dimension. We validate these predictions in synthetic regression experiments and in NanoGPT-scale language model training, where we find that intermediate activations have low-stable-rank throughout training and the corresponding gradients maintain large nuclear-to-Frobenius ratios. Together, these results identify conditions for spectral gradient methods, such as Muon, to be effective in training deep networks and transformers.