Low-Dimensional and Transversely Curved Optimization Dynamics in Grokking

TL;DR

This paper investigates the geometric properties of optimization dynamics in transformers during grokking, revealing low-dimensional confinement and transverse curvature growth as key factors in the transition to generalization.

Contribution

It provides a geometric analysis showing that grokking involves low-dimensional dynamics and curvature growth orthogonal to the main trajectory, with causal experiments confirming their roles.

Findings

Training evolves mainly within a low-dimensional subspace.

Curvature grows sharply in directions orthogonal to this subspace before grokking.

Orthogonal gradient flow is necessary for grokking, but increasing curvature alone is insufficient.

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.