PaperlayerView exact revision on arXiv

Source-linked AI summary

Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

Yongzhong Xu

arXiv:2602.16746v3cs.LGcs.AI

TL;DR

Grokking remains difficult to explain because memorization can persist long after training accuracy is perfect. This paper combines PCA trajectory analysis, commutator-defect curvature measurements, and causal interventions to study that transition. It finds low-dimensional confinement with transverse curvature accumulation, while orthogonal gradient flow is necessary but not sufficient for grokking.

  • Problem

    Grokking is a delayed transition from memorization to generalization, and whether its low-dimensional trajectory manifold is dynamically invariant and predictive remains unresolved.

  • Method

    The paper applies PCA to attention-weight trajectories, projects commutator defects onto the resulting execution manifold, and tests causal effects of orthogonal gradient flow and curvature.

  • Results

    Curvature grows orthogonally and precedes generalization by 600–1600 steps, while suppressing orthogonal gradient flow prevents grokking and boosting curvature defects has no effect.

  • Takeaways & Limitations

    The findings support a geometric account in which grokking involves escape from a metastable regime of low-dimensional confinement and transverse curvature accumulation.

  • Takeaways & Limitations

    The experiments use small Transformer models and synthetic modular-arithmetic tasks, so generalization to large language models and real-world datasets remains unclear.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies grokking as a geometric transition in optimization dynamics, combining low-dimensional trajectory analysis with curvature measurements and causal interventions. It finds low-dimensional confinement, transverse curvature accumulation, and an asymmetric causal relationship between orthogonal motion and grokking.

  • A single principal component captures 68–83% of weight-space trajectory variance, indicating a rank-1 execution submanifold.The trajectory predominantly evolves within a low-dimensional subspace during grokking.
  • ρ ≈1.000 across 36 conditions, with an exec/random projection ratio of ≈2–3×, shows that commutator defects are predominantly orthogonal to the execution subspace.The curvature therefore does not empirically deflect the trajectory out of its learned subspace.
  • 10–1000× curvature growth precedes generalization by 600–1600 steps, while non-grokking controls show only 30–50× growth without generalization.Curvature onset is necessary but not sufficient for grokking.
  • The defect-to-generalization lead time follows a power law, ∆t ∝ t_grok^α, with α = 1.27 ± 0.03 and R2 = 0.97 across a 300× learning-rate sweep.The predictive window improves from 24% at η = 3×10−3 to 95% at η = 3×10−5.
  • Suppressing orthogonal gradient flow prevents grokking, whereas boosting curvature defects has no effect, establishing an asymmetric causal relationship.Orthogonal motion is necessary but artificially increased curvature is not sufficient.
  • The findings replicate across the learning-rate range, a 200× timescale difference between regimes, four operations, and three random seeds.

2 Experimental Setup

The experiments use small Transformers trained on modular arithmetic under canonical grokking conditions, with comparisons across operations, learning regimes, and recorded attention-weight trajectories. Grokking is operationalized using sustained test accuracy thresholds and a 90% test-accuracy grok step.

  • Model and data: The model predicts f(a, b) mod p from two integer tokens with p = 97.
  • Model and data: The architecture is a 2-layer, 4-head Transformer encoder with dmodel = 128, dff = 256, no dropout, and approximately 290k parameters.
  • Training: Training uses AdamW at learning rate 10−3, weight decay λ = 1.0 or 0.0 for controls, batch size 512, and a 50/50 train/test split.
  • Training: A model is grokked when test accuracy reaches 98% for three consecutive evaluations; the grok step is the first step reaching 90% test accuracy.No non-grokking operation exceeds 78% test accuracy in any condition.
  • Evaluation: Four of six modular-arithmetic operations grok under the main hyperparameters, while two do not.
  • Regime comparison: The slow regime uses lr = 5 × 10−5, λ = 0.1, three Transformer layers, and reaches grokking at approximately 570k steps versus approximately 3k in the fast regime.

3 Methods

The methods define an execution manifold from PCA of attention-weight trajectories and use commutator defects to probe curvature relative to that manifold. Projection, random-baseline comparisons, and trajectory-alignment measures test whether optimization remains confined while curvature develops transversely.

  • PCA eigenanalysis: PCA applies SVD to centered, flattened weight-change trajectories and measures each component’s explained variance.PC1 measures the fraction of trajectory variance captured by a single direction.
  • PCA eigenanalysis: The execution manifold is M = {W0 + VKα : α ∈ R^K}, spanned by the top-K PCA directions through the initial weights.When PC1% exceeds 70%, it is effectively rank-1.
  • Commutator defect: The commutator defect measures how much successive gradient updates differ when mini-batches A and B are applied in opposite orders.Flat regions yield commuting steps, whereas curved regions produce order dependence without requiring Hessian computation.
  • Commutator defect: The scale-normalized defect is estimated from K = 9 independent samples at each measurement point, using the median for robustness.
  • Manifold projection: The invariance measure ρ quantifies the residual commutator fraction orthogonal to the PCA subspace; ρ ≈1 indicates transverse curvature and empirical invariance.The analysis treats this as finite-window, discrete-time empirical invariance rather than strict continuous-flow invariance.
  • Controls and alignment: Projection onto five random K-dimensional bases supplies a baseline, with exec/random ratios above 1 indicating excess commutator energy captured by the PCA subspace.Trajectory–curvature alignment is also compared with random-vector expectations.

4 Results

The results identify a low-dimensional, empirically invariant execution manifold in grokking: attention-weight trajectories concentrate in one principal direction, while curvature is predominantly transverse to that subspace.

  • 4.1 Weight Evolution is Rank-1: 68–83% of trajectory variance is captured by the first principal component across grokking conditions, indicating essentially one-dimensional weight evolution.The first eigenvalue dominates across operations, and concentration increases as grokking progresses.
  • 4.1 Weight Evolution is Rank-1: 5–20 standard deviations above the random-walk baseline, grokking PC1% values show that trajectory concentration is not explained by smoothness alone.No-weight-decay controls also have moderately high PC1%, but the null comparison distinguishes grokking concentration from a smooth-trajectory artifact.
  • 4.2 The Execution Manifold Exhibits Empirical Invariance: ρ ≈1.000 across 36 conditions, with commutator defects predominantly orthogonal to the PCA subspace, demonstrates empirical invariance of the execution manifold.The residual fraction is defined as ρ = ∥δ⊥∥/∥δ∥, and the near-unity result indicates that curvature is confined to the normal bundle.
  • 4.2 The Execution Manifold Exhibits Empirical Invariance: Curvature does not deflect the optimization trajectory out of its learned subspace despite substantial curvature in the ambient parameter space.The dominant curvature component lies outside the directions used by the model for learning.
  • 4.2 The Execution Manifold Exhibits Empirical Invariance: 1.8–2.9× more commutator energy is captured by the execution basis than by an equal-dimensional random subspace, supporting geometrically structured alignment.The comparison uses random K-dimensional subspaces with K = 24 across four grokking operations.

4.3 Curvature Explodes Orthogonally During Grokking

During grokking, curvature increases sharply in directions orthogonal to the execution manifold while the trajectory remains empirically confined to it; defect onset consistently precedes generalization.

  • 10–1000× higher commutator defect occurs in grokking operations than in non-grokking controls, concentrated predominantly outside the PCA manifold.The execution manifold remains empirically invariant while orthogonal curvature changes substantially.
  • The PCA execution basis captures 1.8–2.9× more commutator energy than an equal-dimension random subspace, confirming geometrically structured alignment.The random-subspace control uses K = 24 dimensions.
  • The trajectory does not align with curvature directions: mean absolute cosine similarity is indistinguishable from a random-vector baseline, indicating avoidance of high-curvature directions.
  • Defect onset is defined by exceeding 10× the early-training baseline and an absolute threshold of 20; it begins before test accuracy rises in every grokking run.
  • 600–1600 steps separate defect onset from 90% test accuracy across 12 grokking runs, with mean lead time 1117 steps and p = 2−12 ≈2.4×10−4.
  • 30–50× defect growth in non-grokking operations does not produce generalization, so curvature onset is necessary but not sufficient.Grokking operations reach 500–2000× baseline, with zero overlap in total growth magnitude between groups.
  • Suppressing orthogonal gradient flow prevents grokking, establishing the mechanistic necessity of curvature growth.

4.5 Regime Invariance

The geometric pattern persists across substantially different learning rates and hyperparameter regimes, although alignment dynamics and timescales vary quantitatively with the learning rate.

  • Regime Invariance: Despite a 200× difference in grokking timescale, 10× difference in weight decay, and different layer counts, the execution manifold remains invariant with ρ ≈1.000.Defect onset still precedes grokking by hundreds of thousands of steps in the slow regime.
  • Learning Rate Sweep: Across six learning rates, the same four operations grok while two complex operations never grok, leaving the grok/no-grok boundary invariant to learning rate.
  • Learning Rate Sweep: Grokking speed scales roughly linearly with learning rate, from ∼136k steps at η = 3×10−5 to ∼200 at η = 10−2.
  • Learning Rate Sweep: Maximum defect reaches 10^4 at low learning rate but only ∼20–60 at high learning rate, suggesting greater orthogonal curvature accumulation during slower optimization.
  • Learning Rate Sweep: The predictive lead time decreases monotonically with learning rate, while a log-log fit gives α = 1.27 ± 0.03 and R2 = 0.97.At η ≥3×10−3, grokking occurs in fewer than 1k steps and defect onset is concurrent rather than predictive.
  • Alignment Dynamics: Alignment is below baseline at η = 10−4, above baseline at η = 10−2, and transitions between these behaviors at η = 10−3.These patterns replicate across both evaluated operations.
  • Dynamical Regimes: Low learning rates produce overdamped trajectories with prolonged confinement and defect buildup, whereas high rates produce underdamped trajectories with stronger curvature-direction alignment.
  • Phase Portrait: Across addition and multiplication, grokking occurs when trajectories exit a metastable region characterized by high curvature defect and suppressed mobility.

4.7 Causal Interventions on Learning Dynamics

Interventions distinguish the necessary direction of optimization motion from curvature accumulation itself: suppressing execution-manifold motion delays or prevents grokking, whereas injected curvature does not accelerate it.

  • Experimental Design: The experiments test necessity, sufficiency, and specificity while preserving the underlying architecture and dataset.
  • Necessity: Suppressing motion along PCA execution-manifold directions delays grokking monotonically, with full suppression preventing generalization in 0/12 seeds.Partial suppression uses s = 0.25–0.75; full projection uses s = 1.0.
  • Specificity: Random projections have little effect at intermediate strengths, supporting specificity to PCA directions rather than generic low-dimensional constraints.At full strength, both PCA and random projections prevent grokking because confinement to any 16-dimensional subspace is too restrictive.
  • Sufficiency: Neither commutator-aligned nor random orthogonal kicks accelerate grokking across amplitudes α ∈{50, 100, 200, 500}.All 27/27 runs generalized at statistically indistinguishable times of ∼3200 steps.
  • Necessity: Across four operations, s = 0.75 delays grokking by 600–800 steps, while complete suppression prevents grokking in all tested seeds.The dose–response relationship is consistent across modular addition, subtraction, multiplication, and quadratic addition.
  • Combined Interpretation: The intervention results confirm necessity of execution-manifold motion, reject sufficiency of artificially increased defects, and support PCA-direction specificity.

5 Spectral Mechanism Underlying the Commutator Transition

Weight-matrix spectra reveal a staged symmetry-breaking mechanism for the commutator transition: near-degeneracy destabilizes representations, then dominant-mode alignment coincides with grokking.

  • Spectral timeline: g23 gradually shrinks before the leading modes become nearly degenerate, with g12 reaching a minimum of 0.002–0.02.This sequence indicates compression of the sub-leading spectrum followed by top-mode near-degeneracy.
  • Three dynamical regimes: During competition, nearly degenerate top modes destabilize the representation basis; during instability, symmetry breaking drives maximal non-commutativity; during alignment, WQ and WK share a basis as grokking occurs.The three regimes are competition, instability, and alignment.
  • Mechanism: Near-degenerate singular modes can rotate under perturbations, producing large commutator defects, whereas σ1 ≫σ2 locks the representation orientation and reduces ordering ambiguity.The proposed mechanism connects singular-vector instability to the rise and subsequent decline of commutator defects.
  • Interpretation: The spectral symmetry-breaking transition provides a mechanistic account of grokking, with transient representation instability resolved as one attention mode becomes dominant.Weight-SVD complements trajectory PCA by measuring the eigenstructure of learned operators rather than update statistics.
  • Empirical support: The g23 decline precedes grokking by ∼1000 steps and appears in 12 of 12 grokking runs, compared with 1 of 12 matched controls.The study identifies this weight-matrix gap closing as a reproducible diagnostic associated with the transition.

6 Discussion and Theoretical Connections

The paper frames grokking as escape from a low-dimensional metastable regime in which transverse curvature accumulates before generalization. Its discussion connects this geometry to dynamical-systems phases, implicit regularization, scaling laws, circuit stabilization, and unresolved limits of prediction and generality.

  • Geometric account: Grokking is interpreted as escape from a metastable regime formed by low-dimensional confinement and accumulating transverse curvature.The proposed sequence is confinement, barrier growth, and escape.
  • Dynamical-systems interpretation: PC1 captures 68–83% of trajectory variance as optimization compresses onto a rank-1 execution manifold.The trajectory remains confined while curvature accumulates in the manifold’s normal bundle.
  • Dynamical-systems interpretation: 10–1000× defect growth accompanies transverse instability while the trajectory remains empirically invariant, followed by reorganization into a lower-anisotropy regime.The invariance measure is approximately 1.000, and the post-grokking trajectory has reduced curvature anisotropy.
  • Dynamical-systems interpretation: The leading singular values of WQ and WK become nearly degenerate before one mode dominates and the attention operators align into a shared eigenbasis.This spectral symmetry breaking is presented as the mechanism sharpening the transverse-instability picture.
  • Implicit regularization: The intervention results indicate that motion along task-specific learned directions is necessary, whereas generic low-dimensional constraints and curvature boosting are insufficient.This supports geometric structure as an implicit-regularization mechanism rather than norm or margin control alone.
  • Scaling and phase diagram: α = 1.27 ± 0.03 (R2 = 0.97) across 43 runs shows that defect-onset lead time grows super-linearly with grokking timescale.The paper reports distinct overdamped, critically damped, and underdamped learning-rate regimes.
  • Limitations and outlook: The evidence is limited to 2–3-layer, ∼290k-parameter Transformers on modular arithmetic, and several diagnostics require expensive computation that is difficult to scale.The paper also states that a complete theoretical characterization of the phase transitions remains open.
  • Limitations and outlook: The commutator defect detects landscape reorganization early but does not predict which runs will generalize without regularization-sensitive signals.Non-grokking operations and unregularized training can show curvature growth without generalization.

7 Related Work

The related-work discussion situates the paper at the intersection of grokking studies, loss-landscape geometry, and intrinsic-dimensionality research. Its contribution is to extend low-dimensional trajectory findings with invariance tests, curvature dynamics, and causal interventions.

  • Grokking: Prior grokking work examined representation learning, weight decay, circuit formation, phase transitions, and adaptive-optimizer dynamics.These strands provide complementary accounts of delayed generalization in algorithmic tasks.
  • Loss-landscape geometry: The commutator defect connects this study to loss-landscape geometry, Lie brackets of gradient fields, natural gradients, and Fisher-information analyses.It measures non-commutativity in optimization flow as a geometric signal.
  • Intrinsic dimensionality: Building on intrinsic-dimensionality results, the paper establishes empirical invariance of the execution manifold and curvature confinement to its normal bundle.It adds random-baseline controls, temporal prediction analysis, and causal tests of orthogonal gradient flow.

8 Conclusion

The conclusion identifies grokking with a rank-1 execution manifold whose normal-bundle curvature grows before generalization, while emphasizing that curvature onset alone is not sufficient. Across broad training conditions, causal interventions show that orthogonal gradient flow is necessary for the transition.

  • Core conclusion: The weight trajectory lies on a rank-1 empirically invariant execution manifold, with loss-landscape curvature confined to its normal bundle.This summarizes the paper’s central geometric finding.
  • Temporal ordering: Curvature growth precedes generalization by 600–1600 steps, but moderate curvature growth also occurs without generalization.Thus, onset is a necessary precondition rather than a sufficient predictor.
  • Robustness: The findings hold across six modular arithmetic operations, three seeds, two weight-decay settings, a 100× learning-rate sweep, and two qualitatively different regimes.The paper reports replication across substantially different training timescales.
  • Causal intervention: Suppressing orthogonal gradient flow prevents grokking with a monotonic dose–response, whereas artificially boosting curvature defects has no effect.The interventions distinguish mechanistic necessity from sufficiency.
  • Reproducibility: The code and figures are publicly available, with full reproduction requiring approximately 9 hours on one Apple M-series machine.This is the paper’s stated reproducibility resource and compute requirement.

A Additional Figures

The additional figures reinforce the paper’s temporal, geometric, phase-diagram, and causal claims across learning rates, operations, regimes, seeds, and matrix-level analyses. They show that curvature timing and intervention specificity are robust, while alignment and spectral trajectories vary with regime.

  • Regime dependence: The slow regime replicates qualitative findings over 200× longer training, although trajectory–curvature alignment differs quantitatively.Alignment is below baseline at η = 10−4 and above baseline at η = 10−2.
  • Temporal ordering: Defect onset precedes grokking by ∼30k steps at η = 10−4, whereas onset is concurrent at η ≥3×10−3.The figure contrasts slow and fast learning-rate regimes.
  • Learning-rate scaling: α = 1.27 ± 0.03 (R2 = 0.97, n = 43) describes lead-time scaling with grokking step.Lead fraction rises from 24% at η = 3×10−3 to 95% at η = 3×10−5.
  • Causal interventions: At full suppression, grokking fails in 0/12 seeds across four operations, with monotonic delay as suppression increases.This dose–response is shown for gradient projection interventions.
  • Causal interventions: PCA projection delays grokking at s = 0.25–0.75 while random projection has no effect; at s = 1.0, both constraints eliminate grokking.The intermediate-strength separation supports specificity to the PCA manifold.
  • Causal interventions: Kicks up to 500× the gradient-step norm along commutator or random orthogonal directions do not accelerate grokking beyond baseline variability.The result argues that orthogonal defects are not sufficient to induce the transition.
  • Phase portraits: Phase portraits trace grokking through competition, instability, and alignment, whereas memorizing trajectories are diffuse random walks without directed escape.The spectral portrait links instability to near-degenerate leading modes and later operator alignment.
  • Spectral dynamics: The spectral gap minimum precedes the SGD defect spike, which precedes the matrix commutator peak, with test accuracy rising last.This ordering is shown on a common [0, 1] scale.
Loading 2602.16746v3…