On the Convergence Behavior of Preconditioned Gradient Descent Toward the Rich Learning Regime
Authors: Shuai Jiang
Eric C Cyr
Ben S. Southworth
Alexey Voronin
Venue: ICLR 2026 Poster (2026)
Keywords: spectral bias
preconditioned gradient descent
grokking
optimization dynamics
neural tangent kernel
higher-order methods
Primary area: learning theory
TL;DR: Preconditioning conjugate gradient reduces spectral bias and delayed generalization by exploring the NTK space faster.

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

Published as a conference paper at ICLR 2026
ON THE CONVERGENCE BEHAVIOR OF PRECONDI- TIONED GRADIENT DESCENT TOWARD THE RICH LEARNING REGIME
Shuai Jiang
Alexey Voronin abvoron@sandia.gom
Eric C. Cyr eccyr@sandia gov
sjiang@sardia.ggom
Ben Southworth southworth@lanle
ABSTRACT
Spectral bias, the tendency of neural networks to learn low frequencies first, can be both blessing and curse. While it enhances the generalization capa- bilities by suppressing high-frequency noise, it can be a limitation in scientific tasks that require capturing fine-scale structures. The delayed generalization phe- nomenon known as grokking is another barrier to rapid training of neural networks. Grokking has been hypothesized arise as learning transitions from the NTK the feature-rich regime. This paper explores the impact of preconditioned gradient descent (PGD), such 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 evi- dence that grokking represents transitional behavior between the lazy regime characterized by the NTK and the rich regime. These findings deepen our under- standing of the interplay between optimization dynamics, spectral bias, and the phases of neural network learning.
INTRODUCTION
Neural networks (NNs) are cornerstones of modern machine learning, demonstrating remarkable generalization performance using highly over-parameen parameterred networks across wide range of tasks, including image classification, natural language processing, and scientific applications. Though already mature, major challenges still plague the training of NNs. One difficulty is the tendency of NNs to learn low-frequency components first before slowly converging to high-frequencies, the so-called "spectral bias" or "F-Principle" (Rahaman al
2019: XuLetal 2019).
One explanation of spectral bias is through the lens of the neural tangent kerelss (NTKs), where (laz all 2018) argues that the mode-dependent convergence results from disparities in the eigen- values of kernel. In some sense, spectral bias serves as buttress against overfitting, allowing NNs learn general patterns rather than overfitting on noise, and without relying on other techniques like regularization or dropout. However, not all applications using neural networks desire slow con- vergence to high frequency data. As a result, there has been debate over the usefulness of using higher-order optimization methods (Wadia all 2021, Amar et all 2IZ1; Buttelh et all (X12)) Throughout this document we use the phrase "higher-order to indicate optimizers beyond gradient descent and its relatives, most notably Adam. Specifically explore the of Tauss-Newton and Levenberg-Marquare Marquardt that use curvature information from the Hessian or its approximate.
The discussion of generalization of higher-order methods extends to the concept of grokking, phe- nomenon where generalization on test data occurs long after the model has memorized the training data. First in algorithmic datasets (Power et all, 2022), many theories have been suggested

---

Published as a conference paper at ICLR 2026
which attempt to explain the delay including regularization with weights, dynamics of adaptive op- timizers, and circuit efficiencies. We focus on two interpretations (Kumar et all, 2024, Zhou et all 20240, which argue that grokking arises result of transitioning from "lazy" "TK-inmmdtted regime to a "rich", feature-learning regime. In light of these two interpretations, we argue that higher-order gradient descent methods allow for faster exploration of the NTK subspace dia- grammed in Figure A thereby allowing training to enter rich regime faster. We show theoretical and numerical evidence that higher-order methods can explore the NTK regime rapidly, but numerical evidence indicates they can struggle with generalization.
In this paper, we investigate the interplay between spectral bias, grokking, and higher- methods, focusing specifically on preconditioned gradient descent (PGD) in the forms of Heuuss-Newtnn and Levenherg- -eeeberr--aaruurdd.. Specifically, we:
Show that using PGD allows one to greatly accelerate convergence to all frequency modes in the lazy, or NTK, regime (Figure R)
Provide evidence that grokking is a transitional behavior by using PGD to explore the lazy regime uniformly, without spectral bias, thereby eliminating the characteristic delay in generalization (Figure BD
Propose that the lack of generalization observed in (Wadia et all 2021: et 2024) stems from higher-order methods remaining close to the lazy regime. However, perhaps counterintuitively, we find that generalization can be achieved by transitioning to first-order methods after the lazy regime is exhausted (Figure
We aim to provide an in-depth analysis of how these factors influence the training dynamics and generalization capabilities of neural networks.
BACKGROUND
Spectral Bias and the Neural Tangent Kernel
Spectral bias, the tendency of neural networks to learn lower-frequency components faster than higher-frequeng frequency ones, has emerged as a crucial aspect of understanding implicit regularization and generalization in deep learning (Kabaman et all, 2019: Xu 2019; Murray et al, 2022). This inherent inductive bias arises from the interplay between architecture, initialization, and gradient dynamics (lacot et all, 2018; Allen-Zhu et all, 2019 Huang and Yau, 2020, Roberts all, 2022).
While spectral bias can act as form of implicit regularization in some tasks, it becomes hindrance solutions is essen-
in scientific and engineering applications where convergence to tial (Xu et all 2025; Wang all 2022). Theoretical analyses have further elucidated the mechanisms behind this spectral learning, linking it to the eigenstructure of the Neural Tangent Kernel (NTK) and the evolution of the Fourier spectrum of the learned function during training (Bowman, (TZE)) Consequently, understanding and potentially mitigating this spectral bias has become a significant area of research, particularly for tasks that require accurate reconstruction of fine-scale structure. The proposed approaches to combat this issue range from architectural adjustments (Jagtap and Karniadaks 2020, Trask al 2022: Sulzmann all, 2020; Hong et al,, 2022: |Wag and Laj, 2024 to optimization techniques that reshape the loss landscape or rescale the gradient flow (Tancik ct all 2020: Amari et all 2020; Chen all 2024).
Beyond Gradient Descent
Stochastic gradient descent takes steps single direction scaled by a single learning rate. Progress is limited by the slowest NTK eigenmode. corresponding to the flattest direction in the parameter space, so training lingers the "lazy" regime before transitioning into the feature-rich regime. The Adam optimizer speeds up the convergence by assigning each parameter its own learning rate, so small gradients receive higher effective learning rates (Kingma and Ba, 20TE). This effectively shortens the lazy regime, but still ignores pross-parammttr interactions, so that convergence in the ill- conditioned directions still remains relatively slow. Curvature-aue awar methods replace a fixed/static learning rate with smarter rescaling of the gradients, e.g. (Wadia et all 2021: Amar et all 2020 Buifell et all 2024). incorporating cross poos--arameter interaction to improve convergence. Given an operator M, that approximates local curvature of the loss function I., the optimization parameter
2

---

Published as a conference paper at ICLR 2026
update becomes
i Pali =0, -
For symmetric positive definite (SPD) Mr several equivalent interpretations of eq. (I) exist:
view it gradient computed in the M,-induced product:
i11 as the "natural gradient" arising in a Riemannian space with non Non-Euclidean distance metric, which adjusts the descent path to reflect local curvature (Amari, 1998);
(iii) finally change-of- basis, where parameters are transformed by VM, gradients computed in the transformed space, and the result is mapped back to the original space (Cum B all 1998: Desjardins et al. 2015).
Choosing an effective Mj central challenge. Substituting the Hessian matrix for M, gives the classical Newton step, which has quadratic convergence near minima, but comes with a high computational price. Natural-gradient descent, stemming from the Fisher Information Matrix, also encodes second-order curvature information and is positive-definite by construction (Amari, 1998: Amari all 2019), making it preferred operator to second-order methods. For mean square error (MSE) loss, the Fisher information matrix (FIM) coincides with the Gauss-Newton method (Martens. 2021: Schraudolph, 2002)) which only requires the Jacobian. Levenberg- Margragrd (LM) modifies the Gauss-Newton approach by adding diagonal damping
guaranteeing numerical stability and preventing overly aggressive parameter updates (Benzi, 2002: Morc, 2006). Kronecker-Factten Approximation Curvature (K-FAC) approximates the FIM using block -ock--iagonal matrix based on the products of layer-wise statistics (Martens and Grosss,, 2013; Botev et all, (0077).
Ultimately, the goal choose a computable M, that conditions the optimization landscape so that the error is reduced uniformly across all modes.
Grokking
Grokking is a phenomenon of delayed generalization first observed in algorithmic tasks (Power et all 2022). During training, a model will first overfit the training data, showing poor test performance. Then after a prolonged period with minimal further reduction in training loss, the model suddenly be- gins to generalize, leading to an increase in test accuracy. Several hypotheses have been proposed to explain the delayed generalization. For instance, some theories point to the dynamics of adaptive op- timizers (Thlla al ((22)), or architectural bias transformers that favors simpler, lew-sensitivity functions that generalize better (Hhattamishra et all 2022).. Another line of research suggests that grokking happens because, over time, training gradually pushes the model away from memorization patterns towards simpler and more general representations that explain the data better (Barak et all 2022: Narma et all, 2023; Liu et all, 20226).
Two recent papers provide complementary views on grokking from spectral bias perspective. The first, by (Kumar et all 2024), hypothesizes that grokking occurs as result of inefficient training which initially stays confined to the NTK subspace and spectral bias limits it to learning only the wowest-freqenne features. Only later does the model escape this regime and move towards the generalization manifolds, resulting in a sudden improvement in accuracy. The second, by (Zhou et all (022)), argues that grokking mainly arises from spectral mismatch in the training and test data. Due to spectral bias, the model first learns low-frequency modes that are dominant in the train set but may not be predictive for the test set. The generalization emerges once the model begins to learn higher-frequency components that coincide with the test data.
EXPLORING THE NTK REGIME
3.1 MITIGATING SPECTRAL BIAS WITH PRECONDITIONING
Here we discuss how spectral bias can be tempered by the use of PGD. Let f(x,0) : R + R be standard MLP where x is the input and 0 E I2P the network parameters. Specifically, is an MLP of depth and constant width W with the form
(
Lou--
(...
0 (1
36(1) ) 36(1-11 + (1()
f(z, 0) = =
3

---

Published as a conference paper at ICLR 2026
We assume the same initialization and scaling of the weight matrices those discussed in (accol al 2018). @ Define f(0) as the vectorized -ehorized-shorthandd of the quantities {f(z,, )) For sim-
plicity, consider the least-squares regression problem ming (()) with ((0) = f y
where y is the N dimensional labels. We consider the continuous gradient flow underlying gradient descent by taking step size +
00 ",, (I(O,) - y) = δt - (of(0())) - y),
pn = , Vof((o))
with (t) signifying the continuous flow of the weights. For sake of notation, let J; = Vof(O(t)
be the N x p Jacobian matrix at time t. Thus, in function space, we have the usual dynamics
aS((()) (s(((1)) (0(1) = JJJf (f(o(t)) - y).
(2)
Ot
00 Ōt
Now define the error e(t) f(((t)) and the (time-dependent) NTK matrix K, == J,J Note that K, is symmetric positive semi-definite (assuming sufficiently overparametrized), K, is also strictly positive definite with high probability (Bowman 2022% Telgarsky 2021),, with an orthogonal basis of eigenvectors. Then from (2) we can write out an error evolution equation with respect the eigendecompossiiin of K,
Lemma 3.1. For 1 S i S n, let A = diag(A,) 2 be the eigenvalues of K,, and é; the error constant associated with the ith eigenvector. Then continuous gradient flow of é; takes the form
ot é = -x(e)e
because the matrix K, and the corresponding eigenvalue and basis of eigenmodes are evolving nonlinearly with e. However, as the width W of f tends towards infinity, we have that K, K xo where K is the (constant in time) symmetric positive definite neural tangent kernel (NTK), In this regime, we arrive at a linear decoupled evolution in error modes,
Here have that A depends
i -x
Öt
Equation (L) precisely describes spectral bias, because the convergence of each mode is defined by the corresponding eigenvalue of K and global error convergence is defined by the condition number of the NTK matrix. In particular, the learning rate must be small enough for stable evolution of the largest eigenvalue of K but this means error modes associated with small eigenvalues converge like 1 - \ς/AN < 1 for k < N. Empirically, only O(1) eigenvalues are "large", so most modes converge slowly (Murray et al, 2022)
As in standard numerical linear algebra though, we can apply preconditioning to normalize contours towards a more isotropic landscape for convergence (Benzi, 2XX2).. Let μ >> be a regularization parameter, and consider the
(LM) algorithm (Moré, 2006), which evolves the weights according to 0n-1 1 = O, Tul + J J.) -f - y). In the context of least squares, this is akin ridge regression (Hoerl and Kennard 1970). We note that the inversion of the matrix dell-defined due to the inclusion of the μ factor. Performing similar gradient flow manipulations above, we arrive at the following continuous dynamics for the LM preconditioned error evolution a 3 JJ + J J, e, where J, implicitly depends on e. The following shows that the conditioning of the dynamics is greatly improved, meaning that spectral bias during training is reduced.
Lemma 3.2. Let A and (1) be as before. Then the LM I.M-preconditioned continuous gradient flow of é; takes the form
>x(e) 0 4 = μ ++A(e)
24
8t
As before, in the NTK/infinite width regime, may drop the dependence of X; e. Then, we that the mapping from X, +
_ - ,サン greatly improves the conditioning compared with Equation (3).
While the theory highly depends the NTK initialization for strict analysis, the experiments more standard initializations unless otherwise indicated.
4

---

Published as a conference paper at ICLR 2026
Assuming A O, for gradient descent have the condition number ROD
Ax X
whereas for LM,
we have KLM
Av A
AutH AN+H
< R. general. In partieular, if y = X1, then KLLM RV
As F O. LM converges Gauss-Newton (GN) iteration, which takes the form 0 0n. J(), (f() -)) where j denotes the matrix pseudoinverse due to the fact that the Jacobian may be extremely ill-conditioned singular. We assume that the pseudo-inverse is calculated with a cutoff of e. Note that the LM dynamics can be interpreted as a trust region variation the GN optimization steps.
Lemma 3.3. Let A and ei(t) be as before, then if Gauss-Newton precoddtioeedd gradient descent is used, then
Ōt é = - (e)
This means that the GN dynamics result in essentially all modes converging at uniform rate (up to conditioning tolerance E), at cost of the smallest eigenvalues (and typically geometrically highest frequencies) X < e not converging due to numerical necessity of the pseudo-inverse.
Of course, in practice for GN or LM, the additional computation is not trivial. Inverting the large matrix (μl + J Jt) can be computationally expensive. To address this, the resulting linear system typically solved with iterative methods such as the conjugate gradient algorithm (Gargıamı et all 2020: Caretal 20T9),, by applying identities like the Deceman-Morrisonn-- Wooiooorrrooouryyy formula (Ren and Golddarl 2019). These solvers are often paired with a line search to determine an appropriate step size (Milller and Zeinhofer, 2023: Inin et all 2024);; detailed computational discussions deferred to the appendix.
While the above analysis is straightforward, the results primarily meaningful in the NTK regime where K, is linear or nearly linear in its evolution (e.g., either via initialization, scaling or large width (Chizal et all 2019; Jacol all, () This the case even with more sophisticated con- vergence analyses, e.g., (Cayci, 2024), with convergence in the rich regime a largely open question. However, this theory demonstrates how early preconditioning can accelerate training and exploration of the NTK generally linear regime.
3.2 GROKKING BEYOND THE NTK REGIME
Kumar et all (2024) first showed that neither weight norms nor adaptive optimizers necessary for grokking. In Figure II, we see the classical grokking behavior on MNIST (Deng. 2012) as described in eil (20226) using two-layer MLP with AdamW, Adam and PGD with LM dynamics. The exact values of the accuracy notwithstanding note that the weight norms can increase, decrease or stay the same depending on the optimizer as one trains and the model generalizes. The same grokking behavior appears for non-adaptive gradient descent, meaning grokking can occur in spite of adaptivity or weight decay (Thilak all. (022).
Zhou ct al (2114) and Kumar et al (2024)) both argue that grokking occurs due to a mismatch between "ideal" training dynamics and reality. By "ideal", we refer an optimization scenario where feature learning occurs immediately and uniformly across all relevant modes, without delays caused by suboptimal model initialization and optimizer dynamics. In reality, the practical behavior of neural networks is quite different; convergence is often biased toward certain modes or delayed due to suboptimal initiation and optimization dynamics. The former argues that grokking occurs due to rrquuency-dependen convergence, where spectral bias causes the model to fit low-frequency components first, delaying generalization of higher frequency modes present in the test data. The latter argues that neural networks tend stay in the lazy training regime, characterized by the subspace defined by
f(x, ~ f(x 0n) + (0
(5)
where Op are the initial weights. The authors hypothesized that feature learning can only occur after escaping the lazy regime. Both views argue that grokking stems from spectral bias combined with prolonged confinement to the lazy training regime.
TThe test accuracy of the PGD can be worse compared to first-order methods as discussed in (Buftelli ct 2004 Wadia all (02T)..
5

---

Published as a conference paper at ICLR 2026
08
Weeppe Nerm
80
0.4
110000
Accuracy (Adamil) onlly)
Accuracy [LM only)
1.0
AlcaMs
1.0
Agnre1
... A.color 7aat
Apns-2
0.8
0.8
0.6
deiby
10 Optimization Steps
1.0
10
Optiminst Steps
Figure 1: MNIST grokking induced by multiplying the initialization by Top left: Train (solid) and test (dotted) accuracy for different optimizers for a = LM is able to shorten generalization delay, but cannot obtain as high 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 pronounced delay between train and test accuracy. Bottom Right: LM (f fixed) compresses the test-delay across a, but attains lower final test accuracy than first-order methods.
Building upon these theories, we present additional evidence through the usage of PGD as detailed in the diagram in Figure As discussed theoretically in Section B.ll preconditioning accelerates convergence in the NTK regime, particularly in the attenuation of spectral bias. If grokking stems from frequency mismatch and prolonged lazy-regime confinement, then PGD should reduce the time-to- to-eeerrllizationn by accelerating NTK exploration. This is observed in Figures E and &. While PGD compresses the delay, final generalization can be lower; switching to first-order methods after the lazy regime restores accuracy opposite to typical PDE practice where one often finishes with second-order.
4 EXAMPLES
To evaluate the theoretical predictions and discussions from Section EL we consider convergence and grokking experiments. These experiments show how different optimization approaches behave in the NTKKlazy regime and highlight transitions into the feature-rich regime, where our assumptions begin fail. The results from Section I are:
Higher-order PGD methods such as LM and GN accelerate convergence of all frequency modes relative to SGD/Adam performance, this confirms predictions outlined in Sec- B.I. GN achieves uniform exponential decay across all frequencies (Lemma B.3) while LM interpolates between SGD and GN behavior depending in the damping parameter jr (Lemma B.2).
Hidher-order methods' efficacy is limited to the NTK/layy regime. The methods have di- minished performance when transitioning into the rich regime, where non-linear feature learning dominates and the linear curvature approximation Equation (5) does not apply.

---

Published as a conference paper at ICLR 2026
Spectral bias
No spectral bios
Proced
100
With SGD, spectral bias reflects ill-conditioned NTK curvature, resulting in the trajectory from Wy to wxIx that bends along level sets, so progress differs across directions. Top-middle: Preconditioning (LM, μ > 0) uses curvature/Hessian information (Gauss-Newton) to rescale directions, producing more direct path. Top-right: As μ - 0 (GN), updates nearly equal- ize progress across directions on the NTK manifold, effectively removing spectral bias. Bottom: Optimization first approaches the NTK solution wXTx on the lazy subspace (plane): the LM/GN end- point w can under-generalle relative to the true target w". Switching to first-order method moves off-subspacc and recovers final generalization.
Figure 2: Top-left:
To further understand the connection between optimizers and generalization, we investigate grokking behavior in the context of training and test loss across multiple tasks. The main obser- vations from Section H are:
Preconditioning greatly compresses the delay between memorization and generalization seen in grokking across range of tasks, by providing uniform convergence through each mode in the NTK subspace.
This supports recent theories proposed by (Kumar et all 2024: Zhou et all 2024), that overfitting or adaptivity are not the main reason for grokking. We propose that spectral bias plays an important role.
4.1 CONVERGENCE RESULTS
17 and 1 consider the regression problem of fitting To numerically realize the effects of Lemmas
a MLP to u(z) = Σk= k sin ((2k + 1)xx k) uniform grid of 100 points drawn from [o, 1]. We employ MSE as our Ioss function for NN consisting of two layers and hidden dimension of 80. The network is initialized using Glorot normal, and trained using SGD or preconditioned variant with constant learning rate 1e-2. Of particular interest the error in the frequency space
ec()) FFT; (u(x) ((0()),))
The first ten modes are plotted in Figure t In particular, the slope of the error is in accordance with Lemmas 3.7 and 3.. with GN's errors converging uniform, exponential rate for all frequen- cies. and with LM converging to GN as μ - 0. There also appears a clear demarcation between the NTK convergence and rich regime, with preconditioning' effectiveness ending the rich regime.
We next solve the Poisson equation on [0, 1]2 with homogeneous Dirichlet boundary condition using PINNs (Raissi forcing functions corresponding to differing frequency solutions. We choose forcing functions such that the solutions are u(x y) = sin(mner) sin(Tmy) with (n,, m) E ((1, 1). (2, 2), (3, 3)). In Fig-
all all
2019) with a shallow network consisting width 256 dimensions with various
Adam and other Adam-like optimizers can be interpreted as preconditioned gradient descent. Using them with GN or LM preconditioning can cause unexpected results the application of two preconditioners is not well understood.
7

---

Published as a conference paper at ICLR 2026
Figure 3: Mode-wise FFT error (first 10 frequencies) under SGD, LM (4 E {0.5, 0.1)), and GN. Higher-order PGD attenuates spectral bias: GN yields near-uniform decay across modes; LM inter- polates between SGD and GN.
ure H. we show the loss arising from using SGD, Adam and the LM optimizers with learning rates of 1e-3, le-2 and 1e-1 respectively.
SGD and Adam initially show fast loss decay, likely due to fast elimination of dominant low- frequency emror modes. In contrast, the LM training takes more tempered path. The LM opti- mizer performs noticeably better compared to Adam as frequency increases, in particular, note that the slope with which the decreases seems be uniform with respect to different frequencies, which suggests superior handling of higher frequency components as derived in Lemma
SGD
Avaar
UM p-m..)
15'4
Than8a Epcch
Trassg Epods
Trakkng Zpac)
Figure 4: PINNs training loss SGD, Adam and LM dynamics with μ = 0.1 for low (blue), medium (orange) and high (green) frequency forcing functions.
4.2 GROKKING
We present additional evidence that grokking is due to the need to explore the NTK regime efficiently before generalizing, and that PGD reduces grokking. We repeat several grokking experiments the literature, such as those introduced in Kumar all (2024)); etll (2(22))) Zhou et all (2024)).
Test, a=
-=1
Accuracy
80
0.8
Test, a=
0.6
Train. 4.55
Tast, 1.5
Train, a-2
Test, a=
100
107
10a
10"
307
107
10'
30a
10"
Training Epochs
Training Epochs
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 a becomes larger which is considerably alleviated by applying PGD.
8

---

Published as a conference paper at ICLR 2026
5GD
U8 p- 3/2
50D
UM p/.
Norma LLargestFFo
Normalzeed
0.8
9.6
0.2
0.3
10
10a
15
10
10
10'
10
10'
10
10'
15
10
15
Trannin Spach
Training Spach
Training Spoch
Traniing Spoch
Figure 6: Polynomial regression grokking induced by output scaling Left two: Train (solid) vs. test (dotted) loss. As a -> oo, SGD shows an increasing delay between train and test (grokking). LM explores the lazy NTK regime faster, reducing the delay across a. Right two: Error in the largest FFT frequency along ID subspace. The solid lines indicate the subspace {cz; 1 1} where is in the train set, and the dashed lines indicate the same subspace where z; = (1, (I,..., 1). SGD decreases the training and "testing" errors at different times while PGD greatly attenuates the disparity between train and test error.
We first consider modular addition task, consisting of fitting shallow MLP to the values for addi- tion under the ring. Grokking is induced through the use of scaling the model output by a², which has larger NTK regime as a + oc (Chizat et all 2019). In Figure HL we show the train/test accu- racy and losses obtained using SGD and LM. The boxed area indicates the the curves corresponding to the LM method. While it is not surprising that train loss decreases much faster when using higher- order methods, the fact that testing dynamics are similar with respect to scaling a suggest that the ability to explore the NTK uniform rate is highly valuable to accelerate generalization.
This can be more clearly seen through the nigh-dimmnsional polynomial regression task as presented in (Kumar all 2024 $5). Grokking is again induced via scaling the output of the shallow MLP. The left two plots of Figure B again show that the simple usage of preconditioning greatly reduces the delay when a - xo. The gap between training loss and generalization is independent of as the NTK regime is explored at a more uniform rate.
In the right plots of Figure K we show the error in the largest FFT frequency along two different ID subspaces: the first subspace spanned by the one piece of training data, and the second along the vector (L.....J 1) "test" data. In the case of SGD, the error in the training subspace quickly con- verges while the testing subspace has error dependent on the scaling a. This suggests that spectral bias causing slower convergence of those unseen modes. However, in the case of PGD, the train- ing/testing subspaces all converge at roughly the same time as all modes converge at roughly the same rate.
We replicate the classical grokking experiment as presented in Power et all (2(222) which demon- strated grokking on the modular arithmetic task using two-layer decoder transformer with vanilla Adam and cross-entropy loss." _ In Figure E. we see that the introduction of PGD reduces when gen- cralization happens, but again, it seems full generalization is not achievable using just strict PGD. In fact, it can be seen the loss values the right that the validation loss seems to be static (and even slightly increasing) as more PGD iterations are used, further suggesting that using Hessian information can reduce the time which generalization happens, but has a tendency to greatly overfit. One can recover full generalization by switching to Adam again, which we show in Appendix I
Finally, let examine the grokking induced on MNIST by scaling initial weights by introduced in aa 40228); here a corresponds to the "size" of the NTK regime with a -> ac correspond- ing to larger lazy regime (Chizal eL all MT)).. We see on the middle and right plots of Figure that the delay is again uniformly reduced by using PGD, however the final generalization is clearly far weaker. This is observed for the full MNIST dataset and other architectures in (Wadia all 2021: Buffelli et all 2024). Fortunately, we are able to recover the exact (or better) testing data performance by using first-order methods after the higher-order methods, which seen in Figure M where we use 2000 iterations of LM iterations before 20000 Adam W iterations. This again suggests
"Cross-entropy loss necessitates the of generalized Gauss--Neton.

---

Published as a conference paper at ICLR 2026
1.0
Alssbm-- 7pt
A1coMa
10
0.6
Lo52
10
6.2
Iterations
Iterations
Figure 7: Liveebber--Marquare reduces grokking but doesn't 't generalize alone. Left: Accuracy on MNIST using LM for the first 2000 iterations before switching to AdamW, demarcated by the verti- cal black bar. For reference, the shaded dotted lines denote the continuation of LM without AdamW. Higher-order methods effectively reduces the delay but tend to remain near the lazy regime, which results in a final generalizability gap. However, applying AdamW after LM recovers the generaliz- ability, leveraging the benefits of first-order methods in the final stages. Right: Corresponding loss.
CCN
DCM
Mam
Jdb
517000 055
0.4
127
107 32
107
Optrrccttoon Steps
Optmiaatinn Stuza
Figure 8: Classical modular addition example using a transformer. Solid lines indicate train and dashed indicate validation. Left: Using generalized Gauss-Newton results in a faster rise in vali- dation but final validation stagnates severely (45% for GGN compared to 100% for Adam). Right: Loss values for GGN seems to be more tempered for the validation set, however, increasing iter- ation count of GGN doesn't seem encourage generalization. No weight decay is used for both optimizer.
that GN tend to stick near the NTK subspace rather than explore the rich regime, reinforcing those observations of (Wadia et all 2021: Buttelli et all 2024).
LIMITATIONS AND CONCLUSIONS
We showed that PGD, in theory and empirically, reduces the effect of spectral bias in the NTK/lazy regime. Using this lens, we reinforced the theories suggested by Kumar all (2022)); Zhou all (2024) that grokking arises as result of the NNs' tendency to slowly explore the NTK first, by show- ing that PGD uniformly reduces the delay to generalization. Our approach does not address another likely contributor to grokking: train/test dataset sizes. This is also not a wholesale endorsement of the use of high-order methods: while they accelerate entry into the rich regime, they often struggle to achieve high, final generalization. We recover strong generalization by switching to a first-order method (e.g., Adam) after using high-order methods, thus suggesting training procedures that begin with PGD and then transition to first-order methods once the lazy/linear regime is exhausted. Further work, especially the study of convergence results in the rich regime (Woodworth et all 2020), is of particular interest.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We have made every effort to ensure that the results in this paper are fully reproducible. All algorith- mic details are provided in the main text, with additional derivations, proofs, and implementation specifics included in the appendix. Experimental settings, hyperparameters, and data processing steps described in detail to enable faithful replication.
Upon completion the required internal information review, we will publicly release a complete codebase containing all scripts necessary reproduce every experiment and figure in this paper https://q1thui com/ sandialabs, The repository will include instructions for reproduc- ing results, generating figures, and replicating all reported metrics.
ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
This paper describes objective technical results and analysis. Any subjective views opinions that might be expressed in the paper do not necessarily represent the views of the U.S. Department of Energy or the United States Government.
This article has been authored by an employee of National Technology & Engineering Solutions of Sandia, LLC under Contract No. DE-NA0003525 with the U.S. Department of Energy (DOE). The employee owns all right, title and interest in and to the article and is solely responsible for its contents. The United States Government retains and the publisher, by accepting the article for publication, acknowledges that the United States Government retains a non-exclusive, paid-up, ir- revocable, world-wide license to publish reproduce the published form of this article or allow others to do so, for United States Government purposes. The DOE will provide public access these results of federally sponsored research in accordance with the DOE Public Access Plan https: /twp://wwwwwwm www. www.weerrr....mmm wwweery..ovovoe.... /own Loads/ doe publppublii
pccess-plan The work per- formed at Sandia National Laboratories was supported by the U.S. Department of Energy. Office of Science, Office of Advanced Scientific Computing Research, DyGenAI project, and the SEA- CROGS project the MMICCs program. Additional support was received from Interlab Laboratory Directed Research and Development program at Sandia.
This work was funded in part by the National Nuclear Security Administration Interlab Laboratory Directed Research and Development program under project number 20250861ER. This paper de- scribes objective technical results and analysis. Any subjective views or opinions that might ex- pressed in the paper do not necessarily represent the views of the U.S. Department of Energy or the United States Government. Sandia National Laboratories is a multimission laboratory managed and operated by National Technology and Engincering Solutions of Sandia, LLC., wholly owned subsidiary of Honeywell International, Inc., for the U.S. Department of Energys National Nuclear Security Administration under contract DE-NA-0003525. SAND2026- 1906C.. The research was performed under the auspices of the National Nuclear Security Administration of the U.S. Depart- ment of Energy at Los Alamos National Laboratory, managed by Triad National Security, LLC under contract 89233218CNA00000, Los Alamos National Laboratory Report LA--U-25-30551.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances neural information processing systems, 32, 2019.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276. 1998.
Shun-ichi Amari, Ryo Karakida, and Masafumi Oizumi. Fisher information and natural gradient learning in random deep networks. In The 22nd International Conference on Artificial Intelligence and Statisties, pages 694-7022 PMLR, 2019.
Shun-ichi Amari. Jimmy Ba, Roger Grosse, Xuechen Li, Atsushi Nitanda, Taiji Suzuki, Denny Wu, and Ji Xu. When does preconditioning help or hurt generalization? arXiv preprint arXiv:2006. 10732, 2020.
11

---

Published as a conference paper at ICLR 2026
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hid- den progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-21764, 2022.
Stefania Bellavia, Greta Malaspina, and Benedetta Morini. Inexact gauss-newton methods with matrix approximation by sampling for nonlinear least-squares and systems. Mathematics of Com- putation, 2025.
Michele Benzi. Preconditioning techniques for large linear systems: a survey. Journal of computa- fional Physics, 182(2):418 477, 2002.
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions, arXiv preprint arXiv:2211.12316. 2022.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical gauss-newton optimisation for deep learning. In International Conference on Machine Learning, pages 557-565. PMLR, 2017.
Benjamin Bowman. On the Spectral Bias of Neural Networks in the Neural Tangent Kernel Regime. University of California, Los Angeles, 2023.
Davide Buffelli, Jamie McGowan, Wangkun Xu, Alexandru Cioba, Da-shan Shiu, Guillaume Hen- nequin, and Alberto Bernacchia. Exact, tractable gauss-newton optimization in deep reversible architectures reveal poor generalization. arXiv preprint arXiv:2411.0799, 2024.
Tianle Cai, Ruiqi Gao, Jikai Hou, Siyu Chen, Dong Wang, Di He, Zhihua Zhang, and Liwei Wang. tram-gauss-neweee method: Learning overparameterized neural networks for regression prob- lems. arXiv preprint arXiv: 11675, 2019.
Coralia Cartis, Jaroslav Fowkes, and Zhen Shao. A randomised subspace gauss- method for nonlinear ceast-squares. arXiv preprint arXiv:2211.05727,20000 2022.
Semih Cayci. Gauss-newton dynamics for neural networks: A riemannian optimization perspective. arXiv preprint arXiv:2412.140. 2024.
Wenqian Chen, Amanda A Howard, and Panos Stinis. Self-adaptive weights based on balanced residual decay rate for physics-informed neural networks and deep operator networks. arXiv preprint arXiv:2407.01613, 2024.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019.
Le Cun, L Bottou, Orr, and K Muller. Efficient backprop, neural networks: tricks of the trade. Lecture notes in computer sciences, 1524(5-50):23, 1998.
Li Deng. The mnist database of handwritten digit images for machine learning research [best of the web]. IEEE signal processing magazine, 41-,224 2012. 2012.
Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, et al. Natural neural networks. Advances in neural information processing systems, 28, 2015.
Matilde Gargiani, Andrea Zanelli, Moritz Diehl, and Frank Hutter. On the promise of the stochastic generalized gauss-newton method for training dnns arXiv preprint arXiv:2006.02400 02409, 2020.
Arthur E Hoerl and Robert W Kennard. Ridge regression: Biased estimation for nonorthogonal
problems. Technometrics, 12(1):55-67, 1970.
Qingguo Hong, Jonathan W Siegel, Qinyang Tan, and Jinchao Xu. On the activation function dependence of the spectral bias of neural networks. arXiv preprint arXiv:2208 0rXi4: 2022.
Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hierar- chy. In International conference on machine learning, pages 4542-4551. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gener- alization in neural networks. Advances in neural information processing systems, 31, 2018.
12

---

Published as a conference paper at ICLR 2026
Ameya D Jagtap and George Em Karniadakis. Extended thlysics-informed neural networks (xpinns): A generalized space-time domain decomposition based deep learning framework for nonlinear partial differential equations. Communications in Computational Physics, 28(5), 2020.
Anas Jnini, Flavio Vella, and Marius Zeinhofer. Gauss-newton natural gradient descent for physics- informed computational fluid dynamics. arXiv preprint arXiv:2402. 2024.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Mikalai Korbit, Adeyemi D Adeoye, Alberto Bemporad, and Mario Zanon. Exact gauss-newre optimization for training deep neural networks. arXiv preprint arXiv:2405.14402 2024.
Tanishq Kumar, Blake Bordelon, Samuel J. Gershman, and Cengiz Pehlevan. Grokking as the transition from lazy to rich training dynamics, 2024.
06110
Fei-Fei Li, Justin Johnson, Ehsan Adeli, and Zane Durante. Cs23 In: Convolutional neural networks for visual recognition. cttps://cs230 github. 1o/] 2025. URL attps://cs231n
github.1c; Stanford University, Course Notes. Accessed: 2025-06-17.
Ziming Liu, Ouail Kitouni, Niklas $ Nolte, Eric Michaud, Max Tegmark, and Mike Williams. To- wards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651-3444
2022a.
Ziming Liu, Eric J Michaud, and Max Tegmark. Omnigrok: Grokking beyond algorithmic data.
arXiv preprint arXiv:2210.01 2022b.
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21 21(146):1-76, 2020.
James Martens and Roger Grosse Optimizing neural networks with Kronecker-factored approx-
imate curvature. In International conference on machine learning, pages 2408-2417. PMLR. 2015.
Jorge J Moré. The Gevenberg-Maqquare algorithm: implementation and theory. In Numerical analysis: proceedings of the biennial Conference held at Dundee, June 28-July 1, 1977, pages 105-116. Springer, 2006.
Johannes Müller and Marius Zeinhofer. Achieving high accuracy with pinns via energy natural gradient descent. In International Conference Machine Learning, pages 25471-25485. PMLR. 2023.
Michael Murray, Hui Jin, Benjamin Bowman, and Guido Montufar. Characterizing the spectrum of the NTK via power series expansion. arXiv preprint (.XXx:2211.078:00 2022.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting small algorithmic datasets. arXiv preprint arXiv:2201.0217. 2022.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville, On the spectral bias of neural networks. In International conference machine learning, pages 5301-5310. PMLR, 2019.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Yi Ren and Donald Goldfarb. Efficient subsampled gauss newton and natural gradient methods for training neural networks. arXiv preprint 123333,906.0233, 2019.
Daniel A Roberts, Sho Yaida, and Boris Hamin. The principles of deep learning theory, volume 46. Cambridge University Press Cambridge, MA, USA, 2022.
13

---

Published as a conference paper at ICLR 2026
Nicol N Schraudolph. Fast curvature matrix- vector products for recond-order gradient descent. Neu- ral computation, 14(7):1723-1733) 2002.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in neural information processing systems, 33:7462-7473, 2020.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains, Advances in neural information processing systems, 33:7537-7547, 2020.
Matus Telgarsky. Deep learning theory lecture notes. https: hopps://mi0: htt.///mmmccmm
no eddu
ditzl 2021. Version: 2021-10-27 v0.0-e7150f2d (alpha).
Vimal Thilak, Etai Littwin, Shuangfei Zhai, Omid Saremi, Roni Paiss, and Joshua Susskind. The slingshot mechanism: An empirical study of adaptive optimizers and the grokking phenomenon. arXiv preprint arXiv:2206. 220.::2206008877, 2022.
Nathaniel Trask, Amelia Henriksen, Carianne Martinez, and Eric Cyr. Hierarchical partition of unity networks: fast multilevel training, In Mathematical and Scientific Machine Learning, pages 271-286. PMLR, 2022.
Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pages 1747-1756. PMLR, 2016.
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint 230.022300002390, 2023.
Neha Wadia, Daniel Duckworth, Samuel S Schoenholz, Ethan Dyer, and Jascha Sohl-Dickstein. Whitening and second order optimization both make information in the dataset unusable dur- ing training, and reduce prevent generalization. In International Conference Machine Learning, pages 10617-10629. PMLR, 2021.
Sifan Wang, Xinling Yu, and Paris Perdikaris. When and why pinns fail to train: A neural tangent kernel perspective. Journal of Computational Physics, 449:110768, 2022.
Yongji Wang and Ching Yao Lai. Multi-stage neural networks: Function approximator of machine precision. Journal of Computational Physics, 504:1 12865, 2024.
Blake Woodworth, Suriya Gunasckar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference Learning Theory, pages 3635-3673. PMLR, 2020.
Zhi-Qin John Xu, Yaoyu Zhang, and Yanyang Xiao. Training behavior of deep neural network in frequency domain. In Neural Information Processing: 26th International Conference, ICONIP 2019, Sydney, NSW, Australia, December 12-15, 2019, Proceedings, Part I 26, pages 264-274. Springer, 2019.
Zhi-Qin John Xu, Lulu Zhang, and Wei Cai. On understanding and overcoming spectral biases of deep neural network learning methods for solving pdes. arXiv preprint wwXiv:2501.09987 2025,
Zhangchen Zhou, Yaoyu Zhang, and Zhi-Qin John Xu. A rationale from frequency perspective for grokking in training neural network. arXiv preprint arXiv:2405. 24..22205..77777 17479, 2024.
14

---

Published as a conference paper at ICLR 2026
A PROOFS IN SECTION 1.1
We first present the proofs for the results in Section 6.]
Proof of Lemma 1 Consider the SVD J. U,A, / where for notational purposes, repre-
sent the singular values their square root form A Then by direct calculation
Ju(uI J.J U,A V (WVVV + VAAV 1/2 Uy
= =UAA (MI + A.) A.)- UU
-UAAUU
where A, = Addag
X.(e) (+)((e))
Proof of Lemma B.1
Using the same SVD as the above proof, let JT) "Vā be the pseudo-inverse where the diagonal matrix is defined entree-wise
( <S
A
X((e X(e)
E
otherwise
X.(e)
with some user-chosen truncation parameter. Thus,
JJ =
where now the singular values are either Is or Os depending on the magnitude relative to E.
B COMPUTATIONAL DETAILS
Since in the case of the pure Gauss-Newton steps, the pscudo-inverse is hard to compute and we only perform it for the simplest models and batch sizes. In such cases, the Jacobian matrices constructed and the SVD is taken as described above. There are recent work which focuses on Gauss-Newton using inexact randomized methods (Cartis al, 2022: Bellavia et all 2025), but we choose to focus more on the Levenberg-Marn Marquardt formulations as we found it, not only to be easier to compute, but also more stable in its training dynamics.
Computing the inverses for for even moderately-size batch sizes and models. There are several ways skirt around this issue, including using conjugate gradient (Gargiam et all, (020), the Sherman-Morriss soati--ooossoooooobuee (SMW) identity (Ren and Goldarb, 2019), the Duncan Guttman identity (Korbit at 2024)), or, departing Duncan-Guttman
(μI + J,) naively memory intensive
from the PGDs discussed earlier, techniques like Kronecker products in K-FAC (Martens aneeeanddiinnn Grosse 2015)
In general, we opted to use SMW approach, meaning that one needs to compute
(μI J. J.) = 1 J
I
J.
(6)
p
/
In particular, note that JTJ is X n; in general, the batch size n are far smaller than the number of parameters p meaning the matrix inverse scales far better.
Furthermore, modern MI. frameworks allow one reuse the computational graphs constructed after doing backpropagation. In particular, the products against J. and J in Equation (6) can done in matrix-free manner; for example using vip, jvp from Jax. For simplicity, we opted to construct the inner matrix-inverss explicitly, but one in theory could apply matrix-free conjugate gradient to avoid that too.
In the case of cross- cross-entropy loss, the so called generalized Gauss-Newtor (GGN) need to be used. In particular, the preconditioner is now J H,J, where H is the second derivative of the loss with the usual modifications for the Dwvenberg-arquuaae refer the reader to Botey et all (2017) for more details regarding GGN.
stabilization of adding scaled identity. We
15

---

Published as a conference paper at ICLR 2026
122
10-1
10-1
10"1
10"1
20
63
00
0
20
60
Thaning Soach
Saining Spoch
Figure 9: Plot of the residuals in the first ten frequencies of the FFT of residuals for SGD, and preconditioning with GN on a problem with slowly decaying Fourier modes for fitting Equation (α)).
B.1 2D REGRESSION AND REGRESSION HYPERPARAMETERS
The results of Section I holds true even in functions with slowly decaying modes, such as a 2D function with discontinuities. Using the same network in the ID case, fit it to the following discontinuous function on [0, 1|2
1.0
if + < 0.5
2.0
if 1.5
u(x,z) =
0.75
if 0.5 and 2 0.5
0.0
otherwise
total of 1600 points were sampled. Again, to verify the lemmas from above, we show the residual in the function's 's frequency space, which plot in Figure g only for the case of SGD and GN. As predicted, the errors decay in a uniform manner for the Gauss-Newton method, implying an exponential convergence for all modes to the "end" of the lazy regime.
The details for the 1D and 2D regression are the same, and is detailed in Table II.
Table 1: Hyperparameters for regression problems
Parameter
Value
Number of Layers
Hidden Dimension
80
Kernel Initialization
Kaiming Uniform
Bias Initialization
Zeros
Activation Function
tanh
Output Dimension
Learning Rate
10 10-2
Batch Size (ID)
100 (full batch)
Batch Size (2D)
400
B.2 PINNS HYPERPARAMETERS
The usage of PGD in PINNs is not new, and one can obtain great performance using GN/LM meth- ods methods compared Adam and BFGS PINNs problem by using dynamic line search for the learning rate (Innn et all 2024: Miller and Zeinholar (023). We chose LM with fixed learn- ing rate (rather than dynamic line-search) to better emulate continuous time dynamics. We found experimentally that Gauss-Newton, without regularization can be unstable for a fixed learning rate.
The details for the PINNs example is shown in Table M
16

---

Published as a conference paper at ICLR 2026
Table Hyperparameters for PINNs
Parameter
Value
Number of Hidden Layers
Hidden Dimension
256
Kernel Initialization
Kaiming Uniform
Bias Initialization
Zeros
Activation Function
tanh
Learning Rate
1 3, 1,1x 2000 10
Total interior points
640
Total boundary points
40
Batch Size
200
B.3 GROKKING HYPERPARAMETERS
B.3.1 MODULAR ARITHMETIC
We follow the experiment as used in (Chizat all 2019); in particular, the shallow MLP uses quadratic activation, and the outputs are scaled with the scale and divided by the input dimension and hidden dimension. The remaining details are shown in Table
Table Hyperparameters for Modular Arithmetic Model
Hyperparameter
Value
Modulo parameter
23
Input dimension
46
Train Data Fraction
0.9
Hidden Dimension
100
Lpocchs
1000
Scale
[0.5. 1.0, 1.5,
Learning Rate
10 10-2/8²
Batch size
Full batch
LM Regularization Parameter
00.07, 0.0125, 0.005, 0.0025]
We also considered the modular addition task and model exactly as described in Power et all (2022) and implemented with standard PyTorch modules with further changes. However, had to implement three additional changes compared to the other examples due to the model size and complexity of transformers. The first, discussed in the main text, the change from vanilla Gauss- Newton to Generalized Gauss-Newton due to the cross-entropy loss. The slightly larger model size also
of transformers (4 10 parameters)
had to conjugate gradient to solve
(μI J. H.J.)h=
rather than the SMW approach from above. Finally, also Armijo line search as found without the line search, the losses and accuracies tended to oscillate more. As such, we use large learning rate. The damping schedule is smooth decay schedule over 200 iterations 10 starting at 1:
10-1
if 2 200 otherwise
A(1
exp (In(109) + 200 (In(10-1) - In(10)))
The remaining hyperparameters are detailed in Table
B.3.2 POLYNOMIAL REGRESSION
The main implementations based (Kumar all 2024, $5) which we refer the reader to for model definition, and exact data generation formulation. The remaining hyperparameters are shown in Table EL
17

---

Published as a conference paper at ICLR 2026
Table 4: Hyperparameters for Transformer Modular Addition
Hidden Dimensions
128
Layers
2
Head
Modulo Parameter
97
Training Percentage of Dataset
50%
Learning Rate (Adam)
10
Learning Rate (LM)
Weight Decay (All)
0
Batch Size
512
Conjugate Gradient
max iters 150, residual threshold 10
Line Search
10
0.5, max iters 10
Table Polynomial Regression Parameters
Input Dimension
100
Hidden Layer Neurons
500
Scaling é
0.25
Training Data Points
450
Test Data Points
TO00
Learning Rate
0.5 x
Training Iterations
60000
Batch size
Full batch
LM Regularization Parameter
0.1
B.3.3 MNIST CLASSIFICATION
When considering lack of generalization when using PGD the MNIST task (Figure vs Figure M). natural question is what is the loss value? In Figure III, show the loss plots from pure PGD run versus an AdamW run. Interestingly, the loss values are generally lower for PGD and clearly decreasing the draw backs of lg loss regression a classification problem (Van Den Oord el all 2016 Li Ci al
the classification eπror fails to noticeably change. This is indicative of
2025).
Our implementation is based on (Liu et all 20226). For the regularization parameter LM, we opted to dynamically scaling. We found that the best performances is achieved when the regularization is small, however, if using a constant learning rate, rather than an adaptive line search logarithmic interpolation to ease the parameter
for example, this be unstable. We hence
down over 500 iterations using the following
10-4
if 2
A(1)
exp (In(10-))
500
-In110-))In(10-1 (110-)))) 2))
otherwise
which ensures a smooth decay. When we perform the switching from preconditioned training to AdamW, use the same learning rate as if training using AdamW for the entire time. The remain- ing details are shown Table
C MNIST WITH
Typically, MNIST classification is performed using cross-entropy loss, however our experiment above uses MSE loss. The motivation for our choice of MSE is to replicate the experiments presented in Li et all (2)2245)). where grokking on the MNIST dataset was first introduced. In Appendix of the latter paper Lu et all (20226), the authors discussed how the use of cross entropy loss also allows for grokking, albeit a more limited fashion. However, we found that the grokking induced using cross entropy the MNIST dataset required small dataset (only 200 total samples used)
18

---

Published as a conference paper at ICLR 2026
Table 6: Hyperparameters for the MNIST grokking experiment.
Hidden Layer Width
250
Layers
2
Activation Function
Relu
Kernel Initializer
Glorot
Training Points
1000
Batch Size
200
Learning Rate (SGD, LM)
10 4
Learning Rate (Adam/AdamW)
10-3
Weight Decay (for AdamW)
0.1
Table 7: Hyperparameters for the MNIST grokking experiment with cross-entropy. The changes in
width, layers, weight scaling, and training points are all obtained from Lt (202260.
Hidden Layer Width
200
Layers
3
Activation Function
Relu
Kernel Initializer
Glorot
Initial weight scaling a
100
Training Points
200
Batch Size
200
Learning Rate (SGD/Adam)
10
Learning Rate (LM)
10
Damping parameter for LM (A())
1
Weight Decay (for all)
0.01
yielding training loss is equal to zero. Meaning that generalization is strictly achieved through weight norm decrease. B
In Figure M. show data from the MNIST grokking using cross-entropy discussed in Lu (20225); full hyperparameters are detailed Table D. The accuracies in Figure DT all platcaus around 55% before subsequent jump to 75%, which the authors of Liu all (2022b) argue is still grokking. The mechanism of the generalization, as seen on in Figures DDR and Id is clearly due to the weights decaying zero, forcing the network to generalize as the losses are machine zero. In such case, the use of different optimizers will have similar performance as the only changes in gradients come from the weight decay implementation. We observe that SGD and GGN behave
*The fact that grokking is entirely due to weight norm changes doesn't 't refute our above observation that weight norm alone is not necessary in all cases.
Ap
Mpt5-1/
10/
14:
1850
1004
32
109
301
107
301
Figure 10: Plot of loss resulting from MNIST data; AdamW on the left and PGD on the right. The loss values attained by PGD is lower than Adam W. however the classification error doesn't 't seem to reflect this.
19

---

Published as a conference paper at ICLR 2026
10
101
10-0
300
200 baet
12-14
Adam tut
6:0 mann
G6h
12
Shene
Triiiing 0og
(a) Accuracy
(b) Loss
Adbm
5GD
1000
400
200
(c) Weight norms
Figure 11: Grokking behavior on MNIST using cross-entropy under different optimizers. Grokking can only be induced artificially by using very small training sets, leading to training losses of zero, meaning an optimizer will not impact performance for the grokking portion. Top: accuracy and loss trajectories. Bottom: weight norm decay driving generalization.
essentially like Adam, with generalization arising solely from weight norm decay. We were not able replicate this behavior increase the dataset size. This suggests that, least for the MNIST dataset, MSE does "reveal" more grokking. Nevertheless, we show in the modular arithmetic exam- ple, that Gauss-Newton seems reduce the generalization point even for transformers.
D CONTINUATION FOR TRANSFORMER MODULAR ARITHMETIC
We observed in Figure R that PGD shortens the gap between training accuracy reaching 100% and validation accuracy rising. However, GaussNewton optimization substantially reduces the achiev- able validation accuracy (roughly 45%, compared to nearly 100% for Adam). In the MNIST setting (Figure Z),, we were able continue training from the PGD-found minimum and recover strong generalization simply by switching to AdamW.
In the transformer modular- arithmiicii task, however, the minimum identified by PGD appears to lie "far" from generalizing solution. As shown in Figure I12 performing the same continuation strategy-switching from PGD to Adam (without weight decay, to match the original setup) requires a comparable amount iterations before achieving full generalization usual Adam, and the validation loss spikes dramatically in the process. This suggests that, at least this modular example, the GaussNewton minimum is not easily continued into generalizing one. We emphasize that results are not intended as an endorsement for fully adopting pseudosecond-ordee methods such as GGN, but rather an investigation into their optimization dynamics.
20

---

Published as a conference paper at ICLR 2026
1.2 Adam Cantinue
Condrur tall loo)
Adem (val (o4s)
0.5
Coas-trt-pyyon
3.0
10.0
7.5
5.0
25
Figure 12: Accuracy and loss from the modular addition using transformer task where "Continue" indicates using GGN for the first 3000 gradient steps, and switching Adam. For reference, the pure Adam is also plotted. The black bar indicates where the switching takes place.
Aghan2
0.
0.8
0.4
0.4
3.0
1110
10:4
10-4
307
100
30
100
Figure 13: Plots Figures and Ø with additional seeds. Top plot corresponds to Figure II and bottom plot corresponds to Figure M Shaded area indicate range, and lines indicate median.
E ADDITIONAL SEEDS
Figure we show the same plots from Figures and D replicated with more seeds. The areas between the min/max are shaded, and the median lines are plotted for each quantity. Note that variation minimal, indicating robustness across this and other experiments.
21


Reviews:
[{"review_id": "gfvcqqNmf2", "reviewer": "Reviewer_NxWh", "summary": "The paper connects between grokking, spectral bias, and higher-order GD optimization (*preconditioned* gradient decent). Throughout the paper, the authors argue that grokking can be accelerated/alleviated using higher-order gradient descent algorithms because these methods allow for faster exploration of the NTK subspace and therefore enter into so-called *rich-regime* faster. The authors have justified their claims by showing that well-preconditioned gradient descents such as Gauss-Newton or Levenberg-Marquardt that mitigate the spectral bias also reduce grokking artefacts through experiments.", "questions": "- The main idea of tuning optimizer to compress the grokking interval is also addressed in previous works such as [Grokfast](https://arxiv.org/abs/2405.20233), which seems to be tackling on the momentum and gradient filtering. How is the idea of the authors differ from previous works?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper is generally well written and easy to read.\n- Claims are supported by rich experimental results. The authors have provided various types of optimization algorithms and tasks to demonstrate their arguments. Furthermore, all theoretical results are linked to the experimental justification. All experimental details are provided in the appendices.\n- The discussion is honest, with statement of limitations.", "weaknesses": "Some issues prevent me from giving a higher score; I will raise my score if the issues are resolved.\n\n- The theoretical part shows that specific preconditioned gradient descents suffer less from spectral bias *within the NTK domain*, and the linked toy experiments prove that empirically, which is well addressed. However, linking this with the main claim: “the preconditioned gradient descent accelerates grokking since it reduces spectral bias within the NTK domain” only involves empirical experiments as in Figure 5. I believe additional theoretical bridge between this gap will strengthen the claim.\n- The *preconditioned gradient descent* is rather a general category, and there should be certain criteria that determine “good precondition” (likewise for “good high-order GD”), which seems to be not explicitly stated and discussed throughout the paper. Rather, exemplar cases such as Adam, GN and LM are individually discussed. I believe providing unifying lens will further clarify the argument.\n- Figures 1, 3, and 7 are quite messy. Perhaps scaling up some plot components including labels and plot areas will help visualize better.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NfeoYDohXa", "reviewer": "Reviewer_oNGp", "summary": "The paper aims to understand the impact of preconditioned gradient descent on grokking.  They show that preconditioned methods converge uniformly along all spectral modes; they claim that preconditioned methods have diminished performance in the 'rich regime'; they argue that this supports recent theories that explain grokking using spectral bias rather than weight norm or adaptivity; and they argue hat preconditioned methods tend to remain in the lazy regime, which causes worse generalization.", "questions": "* The paper asserts, based on Figure 3, that Gauss-Newton / Levenberg Marquardt preconditioning are helpful only in the lazy/NTK regime, and that their effectiveness goes away in the rich/feature-learning regime.  I am not convinced by the current evidence.  The loss curves do show a slowdown at a particular point, but how do I know that this is when we leave the NTK regime?  And, is it actually true that the preconditioning is not helping at all past this point?  (If we stopped preconditioning, would we still train as fast?).  This is a pretty sweeping claim that is being made based on a single experimental setting.\n\n* In Figure 5, it looks like the LM runs across different $\\alpha$'s do not grok, and get similar test accuracy to one another, whereas in Figure 6 it looks like they do grok, and some get bad test accuracies.  Is my understanding correct?  If so, do you know why there is a discrepancy?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper studies an interesting question (the intersection between grokking, rich/lazy training dynamics, and preconditioning).", "weaknesses": "* The paper's first main claim, on lines 62-63, that Gauss-Newton preconditioning ameliorates the spectral bias in optimization, is almost immediate from prior works.\n* The paper's second main claim, on lines 64-66, that \"higher order gradient descent methods allow for faster exploration of the NTK subspace, thereby allowing training to enter rich regime faster,\" does not seem to me to be a well-defined statement.  What does it mean to \"explore the NTK subspace faster\"?  It's not like GD and PGD take the same path, with PGD moving faster on the first part of the path -- GD and PGD take different paths.  So I don't know what it means to \"explore the NTK subspace faster.\"  How is this different from just saying that PGD trains faster?\n* The paper's third main claim, on lines 67-71, is that higher-order methods generalize worse because they stay close to the lazy regime.  I don't know understand why this would be the case, or what evidence is provided for this proposition.  Moreover, it seems to contradict the second main claim (that higher-order methods \"explore the NTK subspace faster\").\n* The paper frames itself as building upon two papers, Kumar et al '24 and Zhou et al '24, aimed at understanding grokking.  Yet, upon reading these two papers, I feel that they are giving quite separate explanations for grokking (and the latter paper doesn't really seem to be studying grokking, since they seem to rely on a mismatch between the train and test set). Thus, I'm not sure how a paper can built on both of these prior works simultaneously, if they are disagreeing with one another.\n\nFurther, see the questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ehKizzxgeu", "reviewer": "Reviewer_Gxam", "summary": "This paper explores how Preconditioned Gradient Descent (PGD) affects two major phenomena in neural network training: Spectral Bias and \"Grokking.\"\nNeural networks (NNs) tend to learn low-frequency information first (spectral bias). While this can aid generalization, it also slows down learning for tasks that require high-frequency details. Concurrently, the \"Grokking\" phenomenon (where test accuracy suddenly improves long after the model has overfit the training data) hinders rapid training.\nGrokking is a delayed behavior arising from the learning dynamics transitioning from the \"lazy\" NTK (Neural Tangent Kernel) regime to the \"rich\" feature-learning regime, and this delay is closely linked to spectral bias.\n PGD overcomes spectral bias and accelerates grokking. The authors propose a hybrid training strategy: first use PGD to rapidly exhaust the Lazy regime, then switch to a first-order method (like AdamW). Experiments demonstrate that this method leverages both PGD's fast convergence and the first-order method's strong generalization capability in the Rich regime, ultimately achieving the best of both worlds.", "questions": "Are the new parameters difficult to tune or not?\nA experiment on Transformer will be very strong to show your contribution", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper clearly links three complex concepts: spectral bias, grokking, and the NTK/rich learning regimes. It is a potential problem for optimizer study.\nThis \"higher-order first, then first-order\" strategy is counter-intuitive to traditional fine-tuning but is well-supported by experiments (Fig 7), offering significant practical guidance.", "weaknesses": "Computational Cost and Scalability remain issues. The computational cost of PGD is extremely high because of it need inverse related to the Jacobian, it is a problem.\nRegarding the rich regime, the paper primarily focuses on reaching this state faster, but it doesn't deeply analyze the complex, non-linear feature-learning dynamics that actually occur within it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4mXCeLfsmw", "reviewer": "Reviewer_mDac", "summary": "The paper analyzes how preconditioned gradient methods (PGD) shape spectral bias and grokking dynamics. In the neural tangent kernel (NTK) regime, the authors theoretically show that Levenberg–Marquardt (LM) reduces spectral bias by re-scaling eigenmodes, and that as LM’s damping vanishes, it approaches Gauss–Newton (GN), which further equalizes per-mode convergence rates. Empirically, they evaluate on equation fitting and PDE solving with PINNs and find that PGD accelerates early error decay and compresses the train–test generalization delay associated with grokking, which aligns with the previous analysis. On MNIST classification, PGD speeds early loss reduction and shortens the grokking delay. However, it can the generalization can be weaker. The author suggests to use first-order methods after the higher-order methods, which solves the problem.", "questions": "1. Do the observed effects still hold true on mid-scale or large-scale datasets or models? It would be better if such experiments could be added to the paper for completion.\n2. On MNIST, the authors use MSE loss. Can you explain why you use that instead of cross-entropy loss, since the authors also acknowledge the issue of using MSE loss as stated in the supplementary material? Will the conclusions and findings remain the same with cross-entropy loss?\n3. Is there any way to show “all modes converging at a uniform rate” in the MNIST task?\n4. The paper claims to first use first-order methods after the higher-order methods. However, how can we know when the training leaves the NTK regime?\n5. How did you choose the learning rate for LM/GN, and how does it interact with the damping schedule? How sensitive is the method to the learning rate, since the authors say GN can be unstable at fixed LR?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The theory analysis looks good to me, showing how LM narrows conditioning and how GN limit equalizes per-mode convergence, linking preconditioning to reduced spectral bias.\n\nExperiments are conducted to examine and support the theory analysis.\n\nThe structure of the paper is well-structured with clear notations and helpful visualizations that make the paper easy to follow.", "weaknesses": "The novelty is limited. While the theory is well constructed, the contribution is mainly a clarified perspective on known conditioning effects of LM and the GN limit in the NTK regime.\n\n\nThe conducted experiments are small, and it is unclear whether these findings transfer to real and large-scale settings.\n\n\nThere are some questionable experiment design choices, see questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)