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

TL;DR

This work analyzes how spectral gradient methods mitigate anisotropy-driven misalignment in nonlinear phase retrieval by reducing training dynamics to a 3D invariant manifold capturing signal, spike, and bulk components. It proves that SpecGD’s sign-based, matrix-adaptive updates prevent spike amplification and yield dimension-independent transition times, accelerating alignment and noise contraction relative to standard gradient descent. The results are supported by a two-stage dynamics framework and rigorous proofs, with numerical experiments showing robustness to broader covariances and finite-sample regimes. The findings advance understanding of spectral gradient methods and offer practical insights for training under anisotropic data, with potential impact on deep learning optimization and feature learning under structured inputs.

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.

Figures (8)

• Figure 1: Qualitative summary of the learning dynamics of SpecGD (top) and GD (bottom). The parameter matrix is projected to the teacher signal, spurious spike, and isotropic bulk components, and the respective magnitudes (or coefficients) are denoted by $a_k$, $b_k$, and $c_k$, at time $k$, formally defined in Section \ref{['sec:ODE']}. $\eta$ and $\Tilde{\eta}$ are learning rates for SpecGD and GD, and $d$ is data dimension. In our theoretical analysis, we track how these three coefficients behave dynamically. (i) Two stages: all coefficients grow in Stage I and the signal coefficient approaches to $\Theta(1)$ (alignment) in Stage II; (ii) Spike dominance: spike $b_k$ dominates the others without the scale invariance of SpecGD; (iii) Spike saturation: after $b_k$ stops growing, the dynamics enter Stage II and $a_k$ begins to dominate (alignment); (iv) Transition time: SpecGD has no dimensional dependency in its transition time $T_1$ and hence exits both stages faster thanks to the synchronous growth in Stage I.
• Figure 2: Final Frobenius alignment for GD (left) and SpecGD (right). Parameters: $d=300$, $m=100$, $T=1000$, and $\rho_0=10^{-2}$.
• Figure 3: Log-scale plots of the dynamics ($d=m=300$, $\eta=10^{-3}$, $\rho_0=10^{-2}$, $\lambda=10$). Left: GD. Right: SpecGD.
• Figure 4: Population training dynamics. Left: Population loss $\mathcal{L}(W^{(t)})$. Right: Alignment $\mathrm{Align}(t)$.
• Figure 5: Loss evolution.

Theorems & Definitions (67)

• Remark 2.1
• Remark 2.2
• Remark 2.3
• Lemma 3.0
• Theorem 4.1: Stage I: uniform growth
• Theorem 4.2: Stage II: alignment
• Theorem 4.3
• Proposition 5.0
• Remark 5.1
• Proposition 5.1