Structure-Specific Representational Priors Causally Control the Grokking Delay

Gunner Levi Howe

cs.LGJul 5, 2026

Structure-Specific Representational Priors Causally Control the Grokking Delay

Authors: Gunner Levi Howe

Abstract

Grokking -- generalization long after training-set interpolation -- has been accelerated by structure-agnostic interventions (gradient filtering, weight-norm clamping, geometric penalties). Whether the delay specifically measures the time to form task-structured representations has remained observational. We test it causally by injecting representational priors of varying content into a one-layer transformer learning modular addition, via a supervised-contrastive loss whose positives encode (i) the task's true structure ( $(a+b) \bmod p$ ), (ii) a coherent-but-wrong sibling ( $(a-b) \bmod p$ ), or (iii) a random partition -- all with identical loss form, strength, class sizes, and geometry. Whether generalization occurs follows a clean gradation: true 22/30 runs, sibling (same periodic features, wrong combination) 14/15, random (only memorizable) 0/20 (Fisher $p=1.3\times10^{-7}$ ). A weight-norm-matched control replaying the norm trajectory onto plain cross-entropy generalizes 0/15, ruling out the norm as mediator. Probes show structure formation precedes and predicts generalization in all runs. Only the true structure also accelerates grokking (up to $2.75\times$ ), but this is dose-dependent and bimodal. We then confirm the mechanism by prediction: because the acceleration is gated by a weight-norm side-effect, clamping the norm during training yields a reliable, standalone accelerator with a median $8.6\times$ speedup (up to $22\times$ on the fastest seeds, under 1000 epochs), growing monotonically as the norm is held lower; the residual stalls also vanish, though significant only pooled over the two mitigations run at both strengths ( $0/40$ vs $6/20$ , $p=7.7\times10^{-4}$ ), not per method. The grokking delay is, causally, the time to form the right representational structure -- decided at the level of features, not labels.

Explore similar work

Jul 14, 2026cs.LG

What Makes a Representational Prior Work? Feature Families, Label-Free Invariances, and Critical Windows in Grokking

Companion work showed the grokking delay is causally the time to form task-structured representations, injectable via a contrastive prior. Here we characterize what makes such a prior work, across four axes, in 188 new runs. Content: a coherent, learnable prior built from the wrong feature family (magnitude bands) blocks generalization like a random partition (1/15 vs 0/20 grok;

p=0.43

between them), confirming the companion's prediction that priors act at the level of the circuit's features. Supervision: a fully label-free invariance prior -- positives are commuted pairs

(a,b)\sim(b,a)

only -- generalizes in 15/15 runs at a median

2.7\times

speedup, more reliably than the label-supervised prior itself (

p=0.038

), and combined with a weight-norm clamp yields the strongest method we test (median

17\times

, 5/5) -- strongest meaning reliably fast: plain cross-entropy with a clamp matches this speed only at the exact critical norm, while the prior keeps it fast across the entire clamp range. Timing: the prior is only needed early -- applied solely during the first 2000 epochs (4% of budget) it generalizes 10/10 at

2.7\times

, beating continuous application (8/10,

1.25\times

) and a duration-matched later window (

2.1\times

). Setting: the dissociation replicates on modular multiplication and across depths and normalization variants, and a clamp sweep quantifies the companion's central claim: structure injection flattens the weight-norm delay-law exponent about 17-fold (plain cross-entropy slows

31\times

per +10 norm units, a lower bound as higher cells are censored, versus

1.22\times

with the prior). Honest boundary: tasks that generalize before memorizing have no delay to control. Feature-family alignment decides whether a prior permits generalization; invariance content suffices for acceleration without labels; a brief early window captures nearly all of the benefit.

Gunner Levi Howe

Prior Knowledge Structural Generalization Grokking Model Pretraining Invariance Delay+1

Ask PDF

May 15, 2026cs.LG

Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets

Why does a Transformer that has memorized its training set wait thousands of steps before it generalizes? Existing accounts locate this delay in norm minimization, feature emergence, or the late discovery of sparse subnetworks. These explanations capture important parts of the transition, but ignore a constraint unique to attention-based models: if attention discards an informative token, no bounded downstream computation can recover it. We formalize attention as an implicit Bayesian posterior over the task dependency graph and prove that generalization requires two separable conditions: a familiar Goldilocks bound on MLP capacity, coinciding with norm-based theories of grokking, and a novel Bayesian structural condition requiring attention to place sufficient mass on every informative token. This decoupling explains delayed generalization as delayed structural inference. Early in training, the MLP memorizes through unaligned features, drives the cross-entropy loss near zero, and thereby starves attention of structural gradient. Weight decay must then erode memorization before the missing graph becomes learnable, yielding the known inverse-weight-decay delay, which we derive as a structural waiting time. We then prove that this explaining-away delay can be bypassed by a KL-based structural intervention, yielding an inverse-intervention-strength scaling law for the grokking time. Experiments on algorithmic sequence tasks isolate structure from capacity and show that this Bayesian ticket matches or outperforms lottery-ticket transfer.

Kai Hidajat, Solden Stoll, Joseph An

Attention-Based Models Transformer Architectures Structural Inference Grokking Memorization Delay+2

Ask PDF

Jul 13, 2026cs.LG

How to Tame Grokking: Representation Geometry as a Control Signal

Grokking is a phenomenon in which neural networks initially memorize training data and only later exhibit strong generalization after prolonged optimization. Despite extensive recent study, the factors influencing the emergence and timing of grokking remain incompletely understood. We investigate the relationship between representation geometry and delayed generalization. We find that dimensionality collapse consistently precedes the onset of grokking in all evaluated settings. Motivated by these observations, we introduce Geometric Dimensionality Regularization (GeomDR), a simple spectral regularizer that modifies the effective dimensionality of hidden representations during training. Across modular addition, modular division, and permutation composition tasks, GeomDR consistently alters grokking dynamics and can substantially accelerate the onset of generalization depending on the intervention schedule and target dimensionality. In several settings, grokking is accelerated by up to 52 times relative to standard AdamW training. Similar qualitative effects are observed in both multilayer perceptrons and transformers. Together, these results suggest that representation geometry can serve as an effective control signal for grokking and provide evidence that geometric interventions offer a practical approach for studying and influencing delayed generalization in neural networks.

Maksim A Kazanskii

Grokking Multilayer Perceptrons Representation Space Dimensionality Reduction

Ask PDF