Early-Warning Signals of Grokking via Loss-Landscape Geometry

Yongzhong Xu

cs.LGFeb 19, 2026

Early-Warning Signals of Grokking via Loss-Landscape Geometry

Authors: Yongzhong Xu

Abstract

Grokking -- the abrupt transition from memorization to generalization after prolonged training -- has been linked to confinement on low-dimensional execution manifolds in modular arithmetic. Whether this mechanism extends beyond arithmetic remains open. We study two sequence-learning benchmarks: SCAN compositional generalization and Dyck-1 depth prediction. Across both tasks and a wide range of learning rates, the commutator defect -- a curvature measure derived from non-commuting gradient updates -- rises well before generalization, with lead times following a superlinear power law (alpha approximately 1.18 for SCAN, approximately 1.13 for Dyck), consistent with prior results on modular arithmetic. Weight-space PCA reveals that spectral concentration is not a universal precursor; the commutator defect is. Causal interventions demonstrate a mechanistic role: amplifying non-commutativity accelerates grokking (roughly 32% on SCAN, roughly 50% on Dyck), while suppressing orthogonal gradient flow delays or prevents it. The three task families form a spectrum of causal sensitivity -- modular arithmetic is rigid, Dyck is responsive, SCAN is intermediate -- yet suppression delays or prevents grokking in all cases, establishing necessity as a universal finding. These results identify the commutator defect as a robust, architecture-agnostic, causally implicated early-warning signal for delayed generalization in transformers.

Figures & tables

	Modular Arithmetic	Dyck	SCAN
Architecture	Encoder-only	Causal decoder	Enc–Dec
$d_{\text{model}}$	128	128	256
Layers	2	2	3+3
Heads	4	4	4
$d_{\text{ff}}$	256	256	512
Parameters	$\sim$ 290k	$\sim$ 150k	$\sim$ 1.5M

Table 1: Architecture comparison across three task families. The Dyck and modular arithmetic architectures share the same encoder structure; SCAN uses a full encoder–decoder.

Learning Rate	Seed	Grok Step	Onset Step	Lead Time	Lead Fraction	Notes
$10^{-5}$	42	115,000	3,000	112,000	97.4%
$5{\times}10^{-5}$	42	14,500	1,000	13,500	93.1%
	137	27,000	2,000	25,000	92.6%
	2024	14,000	2,000	12,000	85.7%
$10^{-4}$	42	6,500	700	5,800	89.2%
	137	7,000	1,700	5,300	75.7%

Table 2: SCAN: Per-run defect onset and grokking analysis. Of 13 runs, 12 grok and 11 show both grokking and a detectable defect onset with positive lead time. Grokking is defined as test sequence accuracy

\geq 0.98

sustained for 3 consecutive evaluations.

Learning Rate	$N_{\text{seeds}}$	Mean Grok Step	Mean Lead Time	Mean Lead Fraction
$10^{-5}$	1	115,000	112,000	97.4%
$5{\times}10^{-5}$	3	18,500	16,833	90.5%
$10^{-4}$	3	6,900	5,767	83.7%
$5{\times}10^{-4}$	3	7,267	5,933	65.9%
$10^{-3}$	1	1,700	900	52.9%

Table 3: SCAN: Learning-rate-averaged scaling statistics. Five learning rates spanning two orders of magnitude.

Learning Rate	Seed	Grok Step	Onset Step	Lead Time	Lead Fraction	Notes
$3{\times}10^{-5}$	42	68,000	3,000	65,000	95.6%
$10^{-4}$	42	29,000	13,500	15,500	53.4%
	137	44,000	1,000	43,000	97.7%
	2024	32,000	9,000	23,000	71.9%
$5{\times}10^{-4}$	42	5,800	600	5,200	89.7%
	137	6,800	400	6,400	94.1%

Table 4: Dyck: Per-run defect onset and grokking analysis. Of 16 runs, 15 grok and 14 show both grokking and a detectable defect onset with positive lead time.

Learning Rate	$N_{\text{seeds}}$	Mean Grok Step	Mean Lead Time	Mean Lead Fraction
$3{\times}10^{-5}$	1	68,000	65,000	95.6%
$10^{-4}$	3	35,000	27,200	74.4%
$5{\times}10^{-4}$	3	7,800	7,100	91.5%
$10^{-3}$	3	4,900	4,700	95.3%
$3{\times}10^{-3}$	2	1,650	900	46.4%
$10^{-2}$	3	6,200	5,700	92.2%

Table 5: Dyck: Learning-rate-averaged scaling statistics. Six learning rates spanning over two orders of magnitude.

Condition	Modular Arithmetic [ Xu, 2026b ]	SCAN	Dyck
Baseline	$\checkmark$	$\checkmark$	$\checkmark$
1A-kick (aggressive boost)	No effect	$\times$ (destabilized)	$\checkmark$ (50% faster)
1A-noise (mild boost)	No effect	$\checkmark$ (32% faster)	$\checkmark$ (40% faster)
1B-project (hard suppress)	$\times$ (prevented)	$\checkmark$ (2.3 $\times$ slower)	$\checkmark$ (50% slower)
1B-penalty (soft suppress)	$\times$ (prevented)	$\times$ (prevented)	$\checkmark$ (delayed)

Table 6: Comparison of causal intervention results across three task families.

\checkmark

= grokked (faster or slower),

\times

= did not grok, — = not tested. All three tasks confirm that suppression delays or prevents grokking (necessity). The tasks differ on boosting: modular arithmetic shows no effect, Dyck shows acceleration, and SCAN shows an intermediate pattern (mild boosting helps, aggressive boosting destabilizes).

Dataset	$\alpha$	$R^{2}$	$N_{\text{runs}}$	(Slowest LR)	Architecture	Grokking?
				Lead Frac.		PC1 Precedes
Modular Arithmetic [ Xu, 2026b ]	$1.27\pm 0.03$	0.97	43	95% ( $\eta{=}3{\times}10^{-5}$ )	Encoder-only	Yes
SCAN	1.18	0.990	11	97% ( $\eta{=}10^{-5}$ )	Enc–Dec	No
Dyck	1.13	0.908	14	96% ( $\eta{=}3{\times}10^{-5}$ )	Causal LM	Yes

Table 7: Cross-dataset comparison of defect onset scaling laws. All three tasks show positive lead times at every tested learning rate and super-linear scaling (

\alpha>1

). The modular arithmetic results are from Xu [2026b] .

Dataset	Seed	SGD Spike	Comm Peak	Grok Step	SGD Lead
Dyck	42	200	600	600	400
Dyck	137	200	600	1,400	1,200
Dyck	2024	200	600	1,000	800
SCAN	42	700	1,500	3,000	2,300
SCAN	137	1,700	1,500	4,000	2,300
SCAN	2024	1,000	1,500	2,500	1,500

Table 8: Temporal ordering of spectral events relative to grokking. SGD spike = commutator defect onset (native 100-step resolution); Comm peak = matrix commutator inflection (

\lVert[W_{Q},W_{K}]\rVert_{F}

); Grok = first step where test accuracy

\geq 0.95

Dataset	Run	Grok	$g_{23,e}$	$g_{23,g}$	Decline	Mono
Dyck	$\lambda{=}1$ , s42 (dense)	600	3.88	0.081	48.1 $\times$	N
Dyck	$\lambda{=}1$ , s137	1,400	11.3	0.240	47.2 $\times$	Y
Dyck	$\lambda{=}1$ , s2024	1,000	10.1	2.23	4.5 $\times$	Y
SCAN	$\lambda{=}1$ , s42	3,000	27.7	0.555	49.9 $\times$	Y
SCAN	$\lambda{=}1$ , s137	4,000	26.6	0.858	31.0 $\times$	Y
SCAN	$\lambda{=}1$ , s2024 (dense)	3,000	8.21	0.166	49.5 $\times$	N

Table 9: Gram matrix eigenvalue gap replication.

g_{23,e}

: peak

g_{23}

before grokking;

g_{23,g}

g_{23}

at grok step; Decline:

g_{23,e}/g_{23,g}

. All grokking runs show

k^{*}=1

Figure 1: SCAN: Defect predicts grokking across learning rates. Commutator defect (solid lines, left axis, log scale) and test sequence accuracy (dashed lines, right axis) for seed 42 at five learning rates. Dotted vertical lines mark defect onset; dash-dot lines mark grokking step. At every learning rate, the defect spike precedes the accuracy transition. Compare with Figure 10 in Xu [2026b] for the analogous modular arithmetic result.

Figure 2: Dyck: Defect predicts grokking across learning rates. Same format as Figure 1 but for the Dyck depth prediction task at six learning rates ( $3{\times}10^{-5}$ through $10^{-2}$ ). The defect onset precedes grokking in all cases. At $\eta=10^{-2}$ , the non-monotonicity (longer grokking time than $\eta=10^{-3}$ ) is visible, driven by accuracy oscillation.

Figure 3: SCAN: Lead time scaling law. (A) Lead time vs. learning rate. (B) Lead fraction vs. learning rate. (C) Defect onset step vs. grok step with power-law fit ( $\alpha\approx 1.18$ , $R^{2}=0.990$ , $n=11$ ). Colored points show individual seeds; black diamonds show LR means. Compare with Figure 11 in Xu [2026b] ( $\alpha=1.27$ , $R^{2}=0.97$ , $n=43$ ).

Figure 4: Dyck: Lead time scaling law. Same format as Figure 3 but for Dyck ( $\alpha\approx 1.13$ , $R^{2}=0.908$ , $n=14$ ). Fifteen runs across 6 learning rates and 3 seeds show positive lead time. All three task families (this work + modular arithmetic) exhibit super-linear scaling ( $\alpha>1$ ).

Panel of Figure 5 — Figure 5: PC1 trajectory dissociation. (a) On SCAN, PC1 variance continues increasing through the grokking step (red dashed line), with no turnover before grokking. (b) On Dyck, PC1 turnover (green dotted line) occurs well before grokking, with a lead of 2,300 steps. This contrasts with modular arithmetic [ Xu, 2026b ] , where PC1 turnover precedes grokking (matching Dyck). The commutator defect precedes grokking on all three tasks, making it a more universal signal.

Panel of Figure 6 — Figure 6: Spectral concentration is specific to grokking. Top-5 principal components for each weight matrix in the deepest layer. Solid bars: grokking condition ( $\lambda=1.0$ ); transparent bars: no weight decay ( $\lambda=0$ ). Grokking models exhibit strong PC1 dominance (70–90%), while controls show diffuse spectra. This matches Figure 1 in Xu [2026b] .

Panel of Figure 7 — Figure 7: Attention layers dominate spectral concentration. PC1 explained variance decomposed by weight type. Attention projections ( $W_{Q},W_{K},W_{V},W_{O}$ ) show 70–85% concentration vs. 40–60% for MLP weights ( $W_{\text{up}},W_{\text{down}}$ ), indicating the grokking transition is primarily mediated by the attention mechanism. This is consistent with the attention-focused analysis of Xu [2026b] .

Panel of Figure 8 — Figure 8: Integrability breakdown at the pre-grok phase. (a) Exec/random projection ratios for three bases (Weight SVD, $\Delta W$ SVD, Gradient SVD) across four training phases. Ratios spike to 2–3 $\times$ at the pre-grok phase, indicating structured non-commutativity in the learning subspace. Compare with the exec/random ratio of 1.8–2.9 $\times$ reported in Xu [2026b] for modular arithmetic. (b) No-weight-decay controls show no integrability breakdown.

Panel of Figure 9 — Figure 9: SCAN: Causal interventions on the commutator defect. (a) Defect (solid, log scale) and test accuracy (dashed) under five conditions. Mild boosting (1A-noise) accelerates grokking by ${\sim}32\%$ , but aggressive boosting (1A-kick) destabilizes the encoder–decoder dynamics and prevents grokking entirely. (b) Grokking step by condition. SCAN occupies an intermediate position between modular arithmetic (where boosting has no effect) and Dyck (where even aggressive boosting helps).

Figure 10: SCAN: Dose–response relationship. Grokking speed as a function of intervention strength. Suppression interventions show a monotonic dose–response, confirming graded causality. Compare with the Dyck dose–response ( Figure 12 ) and Figure 14 in Xu [2026b] for modular arithmetic.

Panel of Figure 11 — Figure 11: Dyck: Causal interventions on the commutator defect. (a) Defect (solid, log scale) and test accuracy (dashed) under five conditions. Both boosting interventions (1A-kick, 1A-noise) accelerate grokking; both suppression interventions (1B-project, 1B-penalty) delay it. Unlike SCAN, even aggressive boosting (1A-kick) accelerates rather than destabilizes. (b) Grokking step by condition. Dyck is the most responsive task: ${\sim}50\%$ acceleration from 1A-kick, ${\sim}50\%$ delay from 1B-project.

Figure 12: Dyck: Dose–response relationship. Grokking speed as a function of intervention strength. The monotonic relationship confirms graded causality for both boosting and suppression. Compare with the SCAN dose–response ( Figure 10 ) and Figure 14 in Xu [2026b] for modular arithmetic.

Figure 13: Instability at $\eta=10^{-2}$ on Dyck. Raw test accuracy (gray), smoothed accuracy (blue: 200-step window; red: 500-step window), and commutator defect (orange, log scale) for all three seeds. The 0.98 threshold is shown in green. At this learning rate, accuracy oscillates violently (80+ threshold crossings). Seed 42 achieves sustained grokking only at step 11,900 despite first exceeding 0.98 at step 900. This instability regime was not observed in modular arithmetic [ Xu, 2026b ] and motivates the sustained-accuracy grokking criterion.

Panel of Figure 14 — Figure 14: Spectral narrative test. All quantities normalized to $[0,1]$ and overlaid on a single time axis for grokking runs ( $\lambda=1.0$ ). Plotted quantities: $\sigma_{1}-\sigma_{2}$ (spectral gap, red), $\sigma_{2}-\sigma_{3}$ (sub-leading gap, orange), $\lVert[W_{Q},W_{K}]\rVert_{F}$ (matrix commutator, dashed purple), $\log_{10}(\mathcal{D})$ (SGD defect, dash-dot green), test accuracy (dotted cyan). Vertical green line marks grokking step. In both datasets, the SGD defect and matrix commutator peak before grokking, while spectral gap opening continues long after.

Panel of Figure 15 — Figure 15: Phase portraits: spectral gap vs. non-commutativity. Training trajectories in the $(\sigma_{1}-\sigma_{2},\lVert[W_{Q},W_{K}]\rVert_{F})$ state space, colored by training step (left panels) and test accuracy (right panels). Grokking trajectories sweep from the top-left (high commutator, low gap: modes competing) to the bottom-right (low commutator, large gap: aligned). Event markers: gray circle = init, green triangle = SGD spike, red diamond = commutator peak, orange star = grokking. The trajectory bends rightward at the commutator inflection point, coinciding with grokking onset.

Panel of Figure 16 — Figure 16: Grokking vs. memorizing in phase space. Grokking runs (warm colors, 3 seeds) traverse the full phase space from competition (top-left) to alignment (bottom-right), while no-weight-decay controls (cool colors, 3 seeds) remain confined to a small region near initialization. Orange stars mark grokking events at the trajectory elbow. This topological distinction—sweep vs. stagnation—holds across all seeds and both datasets, with zero false positives.

Explore similar work

Feb 18, 2026

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Grokking -- the delayed transition from memorization to generalization in small algorithmic tasks -- remains poorly understood. We present a geometric analysis of optimization dynamics in transformers trained on modular arithmetic. PCA of attention weight trajectories reveals that training evolves predominantly within a low-dimensional execution subspace, with a single principal component capturing 68-83% of trajectory variance. To probe loss-landscape geometry, we measure commutator defects -- the non-commutativity of successive gradient steps -- and project them onto this learned subspace. We find that curvature grows sharply in directions orthogonal to the execution subspace while the trajectory remains largely confined to it. Importantly, curvature growth consistently precedes generalization across learning rates and hyperparameter regimes, with the lead time obeying a power law in the grokking timescale. Causal intervention experiments show that motion along the learned subspace is necessary for grokking, while artificially increasing curvature is insufficient. Together, these results support a geometric account in which grokking reflects escape from a metastable regime characterized by low-dimensional confinement and transverse curvature accumulation. All findings replicate across this learning-rate range, a qualitatively different slow regime (lr=5e-5, wd=0.1, 3 layers), and three random seeds, though alignment dynamics differ quantitatively between regimes. Causal intervention experiments establish that orthogonal gradient flow is necessary but not sufficient for grokking: suppressing it prevents generalization with a monotonic dose-response across four operations, while artificially boosting curvature defects has no effect.

Yongzhong Xu

Grokking Neural Network Training Dynamics Gradient Descent Dynamics Loss Landscape Geometry

Ask PDF

Jun 11, 2026

Learned Structure Shows Up Earlier in the Task's Own Basis: A Controlled Look at a Synchronisation Precursor in Grokking

Analyses of grokking usually read hidden activations in one basis, most often the Fourier basis modulo

p

. We ask how much the answer to "has the network organised yet?" depends on that choice. The measure FSD

_{\rm gen}

averages each unit's activation over inputs with the same answer and counts how many units share the same strongest irreducible representation. For modular multiplication we read one answer-averaged profile in three orders: discrete-logarithm order, in which the characters of the multiplicative group become Fourier modes; natural order; and random orders. Across 60 networks at

p\in\{59,97,113\}

, the discrete-log reading reached a fixed threshold before both other readings in 60 of 60 networks (59 of 60 when it had to hold 500 steps). Because untrained networks already score higher in discrete-log order, we also required an equal rise above each reading's untrained level; the discrete-log reading was still first in at least 16 of 20 networks at every prime. At

p=59

and

p=97

this happened during the memorisation plateau in 35 of 39 networks; at

p=113

in 12 of 20. Memorisation preceded grokking in all 179 runs that grokked, by more than any synchronisation signal on average, so a lead over grokking is not by itself evidence of a detected circuit. Synchronisation time is only weakly related to grokking time (pooled

ρ=0.18

). On six finite groups, read with the same measure in every basis, the irreducible-representation reading came before random element orders on all six; with equal-rise thresholds this held in every case on four groups, in 7 of 10 on

\mathbb{Z}_5\times\mathbb{Z}_5

and in none on

D_{16}

, where about half the profile power lies in one-dimensional representations the measure leaves out. We also correct a conjugate-pair artefact that roughly halves the measure on abelian groups. This version withdraws the causal and predictive claims of v1.

Achyuthan Sivasankar

New York University

Weight Decay Frequency-Domain Feature Learning Neural Network Training Dynamics Grokking

Ask PDF

Feb 19, 2026

The Geometry of Multi-Task Grokking: Transverse Instability, Superposition, and Weight Decay Phase Structure

Grokking -- the abrupt transition from memorization to generalization long after near-zero training loss -- has been studied mainly in single-task settings. We extend geometric analysis to multi-task modular arithmetic, training shared-trunk Transformers on dual-task (mod-add + mod-mul) and tri-task (mod-add + mod-mul + mod-sq) objectives across a systematic weight decay sweep. Five consistent phenomena emerge. (1) Staggered grokking order: multiplication generalizes first, followed by squaring, then addition, with consistent delays across seeds. (2) Universal integrability: optimization trajectories remain confined to an empirically invariant low-dimensional execution manifold; commutator defects orthogonal to this manifold reliably precede generalization. (3) Weight decay phase structure: grokking timescale, curvature depth, reconstruction threshold, and defect lead covary systematically with weight decay, revealing distinct dynamical regimes and a sharp no-decay failure mode. (4) Holographic incompressibility: final solutions occupy only 4--8 principal trajectory directions yet are distributed across full-rank weights and destroyed by minimal perturbations; SVD truncation, magnitude pruning, and uniform scaling all fail to preserve performance. (5) Transverse fragility and redundancy: removing less than 10% of orthogonal gradient components eliminates grokking, yet dual-task models exhibit partial recovery under extreme deletion, suggesting redundant center manifolds enabled by overparameterization. Together, these results support a dynamical picture in which multi-task grokking constructs a compact superposition subspace in parameter space, with weight decay acting as compression pressure and excess parameters supplying geometric redundancy in optimization pathways.

Yongzhong Xu

Grokking Neural Network Training Dynamics Weight Decay Multi-Task Learning

Ask PDF