Structure-Specific Representational Priors Causally Control the Grokking Delay

Grokking — generalization arriving long after training-set interpolation — has recently beenshown to be controllable by structure-agnostic interventions: gradient filtering, weight-normclamping, and geometric penalties on hidden representations. Whether the delay specificallymeasures the time to form task-structured representations has remained an observational claim.We test it causally. Training a one-layer transformer on modular addition, we add a supervisedcontrastive auxiliary loss on hidden representations whose positive sets encode the task’s trueequivalence structure (c = (a + b) mod p), and compare against an exactly-matched controlwhose positive sets encode a shuffled structure of identical class-size distribution, loss form,strength, and geometry. The true-structure prior generalizes in 22/30 runs — up to 2.75×faster than baseline, the fastest transitions we observe, though the acceleration is dose-dependent(median per-seed speedup arises only at the strongest dose we test; weaker doses slow the mediantransition) — while the shuffled prior generalizes in 0/20 (Fisher exact p = 1.3×10−7). Becausethe auxiliary loss inflates the weight norm, and the norm alone is known to set the grokkingtimescale, we replay each intervention run’s exact norm trajectory onto plain cross-entropytraining: 0/15 of these norm-matched controls generalize, collapsing instead into logit-scalesaturation — so the acceleration is not norm-mediated. Representation probes (embeddingFourier concentration, class-cluster formation, CKA) show the true-structure prior builds classclustered representations during the apparent plateau, ahead of and predictive of each run’sgeneralization time. The grokking delay is, causally, the time to form the right representationalstructure: injecting that structure can collapse the delay, and injecting the wrong structure atequal strength prevents generalization entirely