On the Convergence Behavior of Preconditioned Gradient Descent Toward the Rich Learning Regime

TL;DR

This work investigates how preconditioned gradient descent (PGD) methods, notably Gauss-Newton (GN) and Levenberg–Marquardt (LM), influence spectral bias and grokking in neural networks. By analyzing continuous-time dynamics, it shows that GN achieves near-uniform, mode-wise convergence in the NTK (lazy) regime, while LM tunably accelerates convergence across modes through the damping parameter $\mu$, thereby mitigating spectral bias. The results suggest grokking is a transitional phenomenon between the NTK-dominated lazy regime and a feature-learning rich regime, and that PGD can substantially reduce the delay to generalization, though final generalization may benefit from transitioning to first-order methods after the lazy phase. Overall, the study clarifies the interplay between optimization dynamics, spectral bias, and learning phase transitions, and proposes practical training strategies that start with PGD and switch to first-order optimization to achieve robust generalization.

Abstract

Spectral bias, the tendency of neural networks to learn low frequencies first, can be both a blessing and a curse. While it enhances the generalization capabilities by suppressing high-frequency noise, it can be a limitation in scientific tasks that require capturing fine-scale structures. The delayed generalization phenomenon known as grokking is another barrier to rapid training of neural networks. Grokking has been hypothesized to arise as learning transitions from the NTK to the feature-rich regime. This paper explores the impact of preconditioned gradient descent (PGD), such as Gauss-Newton, on spectral bias and grokking phenomena. We demonstrate through theoretical and empirical results how PGD can mitigate issues associated with spectral bias. Additionally, building on the rich learning regime grokking hypothesis, we study how PGD can be used to reduce delays associated with grokking. Our conjecture is that PGD, without the impediment of spectral bias, enables uniform exploration of the parameter space in the NTK regime. Our experimental results confirm this prediction, providing strong evidence that grokking represents a transitional behavior between the lazy regime characterized by the NTK and the rich regime. These findings deepen our understanding of the interplay between optimization dynamics, spectral bias, and the phases of neural network learning.

Figures (13)

• Figure 1: MNIST grokking induced by multiplying the initialization by $\alpha$. Top left: Train (solid) and test (dotted) accuracy for different optimizers for $\alpha=8$. LM is able to shorten generalization delay, but cannot obtain as high a generalization accuracy. Top right: Weight norms corresponding to top left; grokking occurs regardless of whether norms grow, decay, or remain stable. Bottom left: AdamW exhibits a pronounced delay between train and test accuracy. Bottom Right: LM ($\mu$ fixed) compresses the test-delay across $\alpha$, but attains lower final test accuracy than first-order methods.
• Figure 2: Top-left: With SGD, spectral bias reflects ill-conditioned NTK curvature, resulting in the trajectory from $w_0$ to $w^*_{\text{\tiny NTK}}$ that bends along level sets, so progress differs across directions. Top-middle: Preconditioning (LM, $\mu>0$) uses curvature/Hessian information (Gauss-Newton) to rescale directions, producing a more direct path. Top-right: As $\mu\to 0$ (GN), updates nearly equalize progress across directions on the NTK manifold, effectively removing spectral bias. Bottom: Optimization first approaches the NTK solution $w^*_{\text{\tiny NTK}}$ on the lazy subspace (plane); the LM/GN endpoint $w^*_\mu$ can under-generalize relative to the true target $w^*$. Switching to a first-order method moves off-subspace and recovers final generalization.
• Figure 3: Mode-wise FFT error (first 10 frequencies) under SGD, LM ($\mu\in\{0.5,0.1\}$), and GN. Higher-order PGD attenuates spectral bias: GN yields near-uniform decay across modes; LM interpolates between SGD and GN.
• Figure 4: PINNs training loss SGD, Adam and LM dynamics with $\mu = 0.1$ for low (blue), medium (orange) and high (green) frequency forcing functions.
• Figure 5: Accuracy of the modulo task trained using SGD and LM with similar initialization. The LM dynamics is highlighted in the box. Without preconditioning, grokking is observed as $\alpha$ becomes larger which is considerably alleviated by applying PGD.

Theorems & Definitions (5)

• Lemma 3.1
• Lemma 3.2
• Lemma 3.3
• proof : Proof of Lemma \ref{['lem:lm-error']}
• proof : Proof of Lemma \ref{['lem:gn-error']}