Spectral Gradient Descent Mitigates Anisotropy-Driven Misalignment: A Case Study in Phase Retrieval

Spectral gradient methods, such as the Muon optimizer, modify gradient updates by preserving directional information while discarding scale, and have shown strong empirical performance in deep learning. We investigate the mechanisms underlying these gains through a dynamical analysis of a nonlinear phase retrieval model with anisotropic Gaussian inputs, equivalent to training a two-layer neural network with the quadratic activation and fixed second-layer weights. Focusing on a spiked covariance setting where the dominant variance direction is orthogonal to the signal, we show that gradient descent (GD) suffers from a variance-induced misalignment: during the early escaping stage, the high-variance but uninformative spike direction is multiplicatively amplified, degrading alignment with the true signal under strong anisotropy. In contrast, spectral gradient descent (SpecGD) removes this spike amplification effect, leading to stable alignment and accelerated noise contraction. Numerical experiments confirm the theory and show that these phenomena persist under broader anisotropic covariances.

💡 Research Summary

The paper investigates why spectral gradient methods, exemplified by the Muon optimizer, outperform conventional first‑order methods such as SGD or Adam in deep learning. To obtain a tractable yet representative setting, the authors study a nonlinear phase‑retrieval problem that is mathematically equivalent to training a two‑layer neural network with a quadratic activation and fixed second‑layer weights. The input data are drawn from an anisotropic Gaussian distribution whose covariance matrix Σ has a “spiked” structure: a dominant eigenvalue λ₁ with eigenvector v₁ that is orthogonal to the true signal w*. This creates a high‑variance direction that carries no information about the target.

The authors first analyze standard gradient descent (GD). Because the gradient of the loss is pre‑conditioned by Σ, the component of the parameter vector along v₁ (the “spike coefficient” bₖ) is amplified by a factor proportional to λ₁/λ₂ (λ₂ being the bulk eigenvalue). During the early escaping stage, bₖ grows much faster than the signal coefficient aₖ, causing a variance‑induced misalignment: the model’s parameters become dominated by an uninformative direction, and alignment with w* (measured by the cosine similarity) is delayed. This phenomenon worsens as the anisotropy (ρ = λ₁/λ₂) increases, and forces GD to use a smaller learning rate for stability.

Spectral gradient descent (SpecGD) modifies the update rule by taking the singular value decomposition of the gradient matrix G = U Σ Vᵀ, discarding the singular values, and reconstructing an update Δθ = U sign(Σ) Vᵀ. In effect, SpecGD preserves only the directional information of the gradient while normalizing its magnitude to one (or to its sign). This “scale‑invariant” update is performed in an adaptive basis defined by the current gradient rather than the ambient Euclidean basis.

A key technical contribution is the reduction of both GD and SpecGD dynamics to a three‑dimensional invariant manifold spanned by (aₖ, bₖ, cₖ), where cₖ captures the isotropic bulk component. On this manifold, SpecGD’s update becomes a sign‑based gradient step that treats the three coefficients symmetrically. Consequently, during the first stage (growth stage) all three coefficients increase at comparable rates; the spike does not dominate. In the second stage (alignment stage) the bulk and spike coefficients saturate while the signal coefficient continues to grow, leading to rapid alignment with w*.

The authors derive closed‑form expressions for the transition time T₁ between the two stages. For SpecGD, T₁ = Θ(log d) and is essentially independent of the anisotropy level, whereas for GD, T₁ = Θ((λ₁/λ₂)·log d), showing a strong dependence on the spiked variance. Moreover, SpecGD tolerates larger learning rates (η = O(1)) compared with GD, which must satisfy η = O(λ₂/λ₁) for stability.

Extensive simulations validate the theory. Experiments with single‑spike, multi‑spike, and power‑law spectra, across dimensions d = 100–2000 and both population‑level and finite‑sample regimes, demonstrate that SpecGD consistently achieves faster loss decay, higher final cosine similarity (>0.95), and robustness to extreme anisotropy where GD either diverges or converges extremely slowly. The empirical results also confirm that the three‑dimensional reduction accurately predicts the observed trajectories of aₖ, bₖ, and cₖ.

In summary, the paper shows that the advantage of spectral gradient methods stems from a scale‑invariant, direction‑preserving update that neutralizes variance‑driven amplification of uninformative directions. This mechanism yields balanced learning of all principal components, eliminates the spike‑dominance pathology of GD, and leads to earlier, more stable alignment. The findings provide a rigorous explanation for the empirical success of Muon‑type optimizers and suggest that similar spectral updates could benefit training of large‑scale, highly anisotropic models such as Transformers.