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

Abstract: 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.

• The paper demonstrates how spectral gradient descent eliminates anisotropy-driven variance amplification, ensuring robust signal alignment.
• It presents a rigorous analysis showing that SpecGD achieves dimension-free, accelerated convergence and maintains stability with larger learning rates.
• Empirical results confirm that SpecGD overcomes misalignment issues in nonlinear phase retrieval, leading to improved feature learning in high-dimensional data.

Spectral Gradient Descent Mitigates Anisotropy-Driven Misalignment: Analysis in Phase Retrieval

Introduction

The paper "Spectral Gradient Descent Mitigates Anisotropy-Driven Misalignment: A Case Study in Phase Retrieval" (2601.22652) presents a comprehensive theoretical and empirical analysis of Spectral Gradient Descent (SpecGD) methods—of which Muon is a prime example—within the context of high-dimensional nonlinear phase retrieval with anisotropic covariates. The analysis directly addresses foundational questions regarding how gradient-based optimization interacts with substantial data anisotropy, especially in regimes where dominant variance directions are misaligned from the signal of interest. The work places particular emphasis on the dynamics induced by SpecGD relative to standard Gradient Descent (GD), elucidating both qualitative and quantitative divergences.

Problem Setting and Motivations

The nonlinear phase retrieval model investigated is $y = (x^\top w_\star)^2 + \nu$, with $x \sim \mathcal{N}(0, Q)$ and $Q$ taking a spiked covariance form, $Q = I_d + \lambda vv^\top$, with $v^\top w_\star = 0$. This setup ensures that the largest variance direction in the data is uninformative for signal recovery, creating a stringent test for feature learning by first-order methods. The network model considered can be equivalently viewed as a one-hidden-layer quadratic neural network, $f_W(x) = \sum_{j=1}^m (w_j^\top x)^2$, or, in matrix notation, as learning a rank-$m$ positive semidefinite matrix $M = WW^\top$.

The work is motivated by recent successes of the Muon optimizer and related spectral methods, which, via polar decomposition, update weights in a basis-adaptive, scale-invariant direction. Empirical gains observed in deep learning pipelines (including LLM training (Liu et al., 24 Feb 2025, AI et al., 4 May 2025)) raise direct questions: What structural properties of these optimizers explain the observed robustness and accelerated signal alignment in highly anisotropic settings? Is the effect a straightforward consequence of regularized step sizes, or do spectral updates fundamentally alter feature-learning dynamics?

Dynamics on Invariant Manifolds and Phase Structure

A pivotal aspect of the analysis is the identification of a three-dimensional invariant manifold on which both GD and SpecGD dynamics confine. The population loss is invariant under rotation preserving $w_\star$ and $v$, allowing for an exact reduction of the matrix-valued updates $(M_k)$ to scalar coefficient dynamics $(a_k, b_k, c_k)$ associated with projections onto the signal, spike, and isotropic bulk directions, respectively.

GD Dynamics: Under strong anisotropy (large $\lambda$), GD disproportionately amplifies the spike direction ($b_k$) during early iterations, leading to a regime where the parameter norm becomes severely dominated by an uninformative direction. As a result, alignment with the true signal ($a_k$) is suppressed, and instability emerges if learning rates are not tuned to accordance with the data anisotropy. Only after a delayed transition does signal growth take over—see the characteristic two-stage behavior, with protracted Phase I (spike-dominated) growth and delayed Phase II (signal amplification).

SpecGD Dynamics: SpecGD, by contrast, replaces the raw gradient with a sign-based, basis-adaptive update via the polar factor. This alters the evolution such that all coefficients grow synchronously during early training, without multiplicative advantage for spike directions. As a result, spike amplification is suppressed ab initio, and the signal coefficient $a_k$ begins its ascent to optimal alignment earlier and more robustly, with transition times demonstrably independent of dimension.

This behavior is succinctly visualized in the following learning dynamics summary (Figure 1):

Figure 1: Summary of the learning dynamics of SpecGD (top) and GD (bottom) on the coefficients $a_k$ (signal), $b_k$ (spike), and $c_k$ (bulk).

Theoretical Analysis and Main Results

A sequence of theorems provides a precise characterization of the two-stage dynamics for both GD and SpecGD. For SpecGD, Stage I (uniform growth) is completed in $O(1)$ iterations and is free of dimensional dependence, while Stage II (signal alignment) follows with rapid convergence to the correct signal direction. For GD, Stage I is protracted—$O((\eta\lambda)^{-1}\log d)$ when $\lambda$ is large—and both signal alignment and loss contraction are retarded by strong anisotropy.

The core claims validated both analytically and empirically are:

• SpecGD eliminates anisotropy-induced variance amplification: Unlike GD, which amplifies uninformative directions under large $\lambda$, SpecGD ensures all manifold components grow comparably in the initial phase, with no instability from large learning rates in signal-orthogonal high-variance directions.
• Phase transitions occur at dimension-free timescales for SpecGD: The onset of signal learning in SpecGD is unaffected by dimensionality, in stark contrast to GD, as shown in heatmaps of final alignment across $(\lambda, \eta)$ grids (Figure 2).

  Figure 2: Final Frobenius alignment for GD (left) and SpecGD (right) as a function of learning rate and spike intensity. SpecGD remains robust to large $\lambda$ and large $\eta$.

• Signal alignment for SpecGD is immediate and monotonic post-transition: Once the growth phase ends, $a_k$ increases monotonically, $b_k$ and $c_k$ contract or saturate at order $(1+\lambda)^{-1}$ and $d^{-1}$, respectively, and overall alignment rapidly approaches $1$.

Empirical Validation and Robustness

Extensive numerical experiments confirm the theoretical predictions and extend them. Empirical evidence—on both population and minibatch settings—demonstrates the persistent advantage of SpecGD over GD even in scenarios with general (power-law) input spectra far beyond spiked covariance. Notably, in regimes where the teacher signal is aligned with the tail of the covariance matrix (a particularly adverse scenario for feature learning), SpecGD achieves alignment while GD remains orthogonal for extensive iterations (Figure 3).

Figure 4: Under power-law covariance, SpecGD achieves signal alignment while GD remains misaligned, despite similar minimization of empirical loss.

Further, the robustness of SpecGD to learning rate aggresiveness is substantially higher—GD diverges with moderately large $\eta$ in strong anisotropy, while SpecGD yields controlled training and loss evolution (Figure 5).

Figure 6: Loss evolution for GD and SpecGD under various learning rates; GD is unstable unless $\eta$ is tuned with respect to $\lambda$, SpecGD is robust.

Implications and Future Directions

This work makes a bold claim with theoretical and numerical support: The use of spectral (polar-based) gradient updates fundamentally alters the geometry and coupling of learning trajectories under data anisotropy. By nullifying spurious amplification, SpecGD enables more robust feature learning and signal alignment, with substantial implications for the design of optimizers in high-dimensional, heavy-tailed, or imbalanced real-world data regimes.

Practical implications include:

• Justification for deploying matrix- or layer-aware spectral optimizers in settings with adverse covariance spectra or adverse signal-to-noise alignments.
• Quantitative learning rate guidance: SpecGD tolerates much larger step sizes without instability, obviating the need for fine-grained LR tuning under anisotropy.
• Alignment and loss contraction are decoupled from dimension for SpecGD, implying accelerated convergence in large-scale models.

Theoretically, the analysis opens new directions:

• Derivation of sharp scaling laws for representation learning under general covariance structures, extending to power-law and heavy-tailed regimes.
• Extension of spectral update analysis to richer nonlinear models (multi-index, deep networks) and online/noisy regimes.
• The detailed phase-wise decomposition here provides a canonical template for studying dynamics under other forms of normalization and non-Euclidean geometric constraints.

Conclusion

This work rigorously demonstrates that spectral gradient descent mechanisms—via basis-adaptive, scale-invariant updates—are highly effective in mitigating anisotropy-driven misalignment intrinsic to GD and its variants. The result is more rapid, robust feature learning and direct alignment with the signal, regardless of high-dimensional ambient geometry or adverse data variance. These insights have concrete relevance to both the practice and theory of optimization in contemporary high-dimensional inference and deep learning.