Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Yongzhong Xu

cs.LGFeb 18, 2026

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Authors: Yongzhong Xu

Abstract

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.

Figures & tables

Operation	Formula	Groks?	Grok step
add	$(a+b)\bmod 97$	Yes	$\sim$ 2900
sub	$(a-b)\bmod 97$	Yes	$\sim$ 3400
mul	$(a\times b)\bmod 97$	Yes	$\sim$ 2600
x2_y2	$(a^{2}+b^{2})\bmod 97$	Yes	$\sim$ 1900
x2_xy_y2	$(a^{2}+ab+b^{2})\bmod 97$	No	—
x3_xy	$(a^{3}+ab)\bmod 97$	No	—

Table 1: Operations tested. Grok step is the mean step at which test accuracy reaches 90%, averaged over 3 seeds.

Panel of Figure 1 — Figure 1: Weight trajectories during grokking are rank-1. (a) PC1% across operations: grokking runs (wd=1.0) show 68–83% variance in a single component. (b) Eigenspectrum showing dominant first eigenvalue.

Panel of Figure 2 — Figure 2: PCA concentration is genuine and increases over training. (a) Z-scores above random-walk null. (b) Expanding-window PC1% over training.

Panel of Figure 3 — Figure 3: The execution manifold exhibits empirical invariance under the optimization dynamics. Commutator defect vectors are predominantly orthogonal to the PCA subspace, with curvature confined to the normal bundle.

Panel of Figure 4 — Figure 4: Random subspace control confirms that the PCA projection is geometrically structured, not a dimensionality artifact. Exec/random ratio $\approx 1.8$ – $2.9\times$ across operations.

Panel of Figure 5 — Figure 5: Curvature explodes during grokking but remains orthogonal to the learned subspace.

Figure 6: Temporal ordering of curvature growth and generalization. Top four panels: grokking operations (3 seeds each), showing defect (solid) rising before test accuracy (dashed). Bottom two panels: non-grokking controls show moderate defect growth ( $30$ – $50\times$ baseline) but no generalization. Dotted vertical lines mark defect onset; green regions mark grokking.

Panel of Figure 7 — Figure 7: Temporal ordering: curvature growth precedes generalization.

Metric	Fast regime	Slow regime
Learning rate	$10^{-3}$	$5\times 10^{-5}$
Weight decay	1.0	0.1
Layers	2	3
Adam $\beta_{2}$	0.98	0.999
Grok step (add, mean)	$\sim$ 2,900	$\sim$ 570,000
Invariance ( $\rho$ )	$\approx 1.000$	$\approx 1.000$

Table 2: Hyperparameter regimes and key metrics.

Panel of Figure 8 — Figure 8: Regime invariance: qualitative findings replicate in the slow regime ( $200\times$ longer training), though alignment dynamics differ quantitatively.

Figure 9: Phase diagram of grokking dynamics across learning rates ( $\eta\in\{3\!\times\!10^{-5}\text{--}10^{-2}\}$ , $\lambda=1.0$ , 3 seeds per cell, 108 runs). (A) Grok fraction: the phase boundary between grokking and non-grokking operations is invariant to learning rate. (B) Mean grok step (log scale): grokking speed scales with $\eta$ . (C) Mean max defect (log scale): curvature explosion is largest at low $\eta$ . (D) Mean lead time (onset step $-$ grok step): defect onset is most predictive at low $\eta$ .

Figure 10: Defect and test accuracy trajectories for addition across six learning rates. At low $\eta$ (e.g., $10^{-4}$ ), defect onset precedes grokking by $\sim$ 30k steps (92% of training); at high $\eta$ ( $\geq 3\!\times\!10^{-3}$ ), grokking is nearly instantaneous and onset is concurrent.

Figure 11: Lead time scaling across learning rates. (A) Power-law fit: lead time $\propto$ grok step α with $\alpha=1.27\pm 0.03$ ( $R^{2}=0.97$ , $n=43$ ). Each point is one grokking run; the dashed line shows the best fit, the dotted line shows linear scaling ( $\alpha=1$ ). (B) Lead fraction (lead / grok step) increases monotonically as $\eta$ decreases: from 24% at $\eta=3\!\times\!10^{-3}$ to 95% at $\eta=3\!\times\!10^{-5}$ . (C) Onset step vs. grok step for each operation.

Figure 12: Trajectory–curvature alignment at three training phases across learning rates. At $\eta=10^{-4}$ (blue), alignment stays below the random baseline (dotted), consistent with overdamped dynamics. At $\eta=10^{-2}$ (green), alignment exceeds the baseline, consistent with underdamped exploration. $\eta=10^{-3}$ (orange) shows intermediate behavior. Both operations exhibit the same pattern.

Figure 13: Phase portrait of grokking dynamics in curvature–trajectory space. We plot the trajectory–curvature alignment ratio (mean $|\!\cos(\Delta\theta,\delta)\!|$ normalized by random baseline) against the commutator defect magnitude for three learning rates ( $\eta=10^{-4},10^{-3},10^{-2}$ ). Each polyline traces the evolution from memorization ( mem ) through defect onset ( spike ) to the post-grokking regime ( post ), shown for addition (circles) and multiplication (squares); arrows indicate the direction of training. The horizontal dashed line indicates isotropic alignment. Training at high learning rates remains in an underdamped regime characterized by strong alignment and low defect, while low learning rates produce overdamped dynamics with large defect accumulation and weak alignment. Intermediate learning rates interpolate between these regimes. Stars mark checkpoints where test accuracy exceeds $90\%$ . Grokking corresponds to escape from a metastable region of high curvature defect and reduced mobility, with regime-dependent relaxation dynamics.

Figure 14: Dose–response curve for gradient projection across all four grokking operations. Each panel shows mean grok step (3 seeds) vs. suppression strength $s$ . At $s=1.0$ , grokking fails universally (0/12 seeds). Dashed line: baseline (no intervention). The monotonic delay and complete suppression at full strength replicate across all operations.

Figure 15: PCA-specific suppression control. Left: grok step vs. suppression strength for PCA projection (blue) and random projection (green); dashed line is baseline. Right: grok success rate. At intermediate strengths ( $s=0.25$ – $0.75$ ), PCA projection monotonically delays grokking while random projection has no effect. At $s=1.0$ , both kill grokking (any 16-dim constraint is too restrictive). The dose–response separation confirms the geometric specificity of the PCA manifold.

Figure 16: Sustained directional kicks along the commutator (red) vs. random orthogonal (gray) directions, with kick magnitudes up to $500\times$ the gradient step norm applied every 50 steps. Left: mean grok step (3 seeds); right: grok success rate. Neither direction accelerates grokking beyond baseline variability (dashed line), confirming that the orthogonal defect is not sufficient to induce the phase transition.

Figure 17: Phase portrait of grokking in the $(\sigma_{1}-\sigma_{2},\;\left\|[W_{Q},W_{K}]\right\|_{F})$ plane for modular addition (seed 42, layer 0). Training follows a consistent trajectory through three regimes: a competition phase where the top singular modes are nearly degenerate, an instability phase where the commutator peaks as one mode begins to dominate, and an alignment phase where the operators converge to a shared eigenbasis and grokking occurs. Left: colored by training step. Right: colored by test accuracy—the trajectory remains at chance (red) throughout competition and instability, turning green precisely as it enters the alignment region.

Figure 18: Phase portraits for grokking (top, $\mathrm{wd}=1$ ) vs. memorizing (bottom, $\mathrm{wd}=0$ ) across four operations (seed 42). Grokking trajectories exhibit a characteristic loop with directional flow through competition $\to$ instability $\to$ alignment. Memorizing trajectories are diffuse random walks in the same phase space, with much higher commutator values and no directed escape.

Figure 19: All quantities on a common $[0,1]$ scale (seed 42, layer 0) for the four grokking operations. The spectral gap $\sigma_{1}-\sigma_{2}$ (red) reaches its minimum before the SGD defect spike (green, dashed-dotted), which in turn precedes the matrix commutator peak (purple, dashed). Test accuracy (cyan, dotted) rises last. Vertical dashed line: grokking step.

Operation	$\boldsymbol{\omega}$	Grok step	$g_{23}^{\mathrm{early}}$	$g_{23}^{\mathrm{grok}}$	Decline	$R_{\mathrm{early}}$	$k^{*}_{\mathrm{term}}$	Decl.
$a{+}b$	1.0	2500–3100	$14.6\pm 0.2$	$0.56\pm 0.4$	$50\times$	$1.40\pm 0.05$	1	✓
	0.0	—	$33.0\pm 9.3$	—	—	$2.72\pm 0.22$	—	—
$a{\times}b$	1.0	2600–3000	$14.7\pm 0.1$	$0.39\pm 0.2$	$49\times$	$1.41\pm 0.05$	1	✓
	0.0	—	$23.6\pm 10.7$	—	—	$2.85\pm 0.11$	—	—
$a{-}b$	1.0	3300–3700	$14.4\pm 0.1$	$0.48\pm 0.2$	$35\times$	$1.31\pm 0.02$	1	✓
	0.0	—	$33.2\pm 2.8$	—	—	$2.62\pm 0.23$	—	—

Table 3: Rolling-window Gram matrix spectral analysis (

W=10

, mean

\pm

std over 3 seeds).

g_{23}^{\mathrm{early}}

: peak

g_{23}

before grokking;

g_{23}^{\mathrm{grok}}

g_{23}

at the grokking step; Decline

=g_{23}^{\mathrm{early}}/g_{23}^{\mathrm{grok}}

;

R

: gap ratio at signal rank

k^{*}

Figure 20: Local integrability: proj/full ratio $\rho_{\mathrm{local}}$ over training for grokking operations (solid) and non-grokking controls (dashed). Grokking operations show $\rho_{\mathrm{local}}\approx 1.0$ with a dip during memorization and recovery post-grokking; non-grokking controls remain below $0.5$ . Random baseline (gray band) at $\sqrt{K/P}\approx 0.005$ .

Figure 21: Basis-independent sign flip: exec/random ratio for three independent bases (weight SVD, displacement SVD, gradient SVD) across training. All three show ratio $>1$ during memorization (commutator aligns with weight structure) and $<1$ post-grokking (commutator exits the subspace), confirming the structural nature of the transition.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 22: PC1% heatmap by operation and weight matrix (last layer, wd=1.0). All weight matrices show high PC1% for grokking operations.

Figure 23: Per-weight-matrix PC1% comparison across operations (grok vs. no-wd).

Panel of Figure 24 — Figure 24: Regime comparison for PCA concentration.

Panel of Figure 25 — Figure 25: Commutator analysis details.

Panel of Figure 26 — Figure 26: Converse analysis: the weight trajectory avoids high-curvature directions.

Figure 27: Temporal traces for $(a+b)\bmod 97$ with 3 seed overlays, showing consistency of the invariance and defect patterns across seeds.

Explore similar work

Feb 19, 2026

Early-Warning Signals of Grokking via Loss-Landscape Geometry

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.

Yongzhong Xu

Grokking Neural Network Generalization Loss Landscape Geometry

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

Sep 30, 2026

Misalignment of Low-Loss Regions Causes Grokking

Grokking refers to the delayed emergence of validation-set generalization after a model has already overfit the training set. Although first observed in small algorithmic tasks trained with transformers, its underlying mechanism remains unsettled. In this work, we develop an analysis framework based on mode connectivity and the geometry of low-loss regions. The framework predicts that the standard modular-arithmetic setting does not always produce grokking: under a symmetry-preserving train/validation split, we observe a stable anti-grokking case in which validation performance does not recover. This counterexample challenges several existing correlational explanations of grokking. More broadly, our analysis framework and results further suggest that grokking arises when the low-loss regions induced by the training and validation partitions are misaligned. Once these regions become well aligned, training hyperparameters alone cannot produce grokking and the observed dynamics collapse to either trainable or non-trainable behavior.

Yongding Tian, Zaid Al-Ars, Maksim Kitsak +1

Computer Engineering Lab Delft University of Technology Delft, NL · HDL TypeTech Delft, NL · Network and Architecture Service Delft University of Technology Delft, NL +1

Loss Landscape Geometry Grokking

Ask PDF