🤖 AI Summary
This work investigates whether the delayed onset of grokking is determined by the emergence of task-specific representational structure. By introducing a structured supervised contrastive loss in a single-layer Transformer, the study provides the first causal evidence that grokking depends on the formation of correct feature-level structure—not merely label-level alignment. Experiments demonstrate that authentic structural priors accelerate grokking by up to 2.75×; even incorrect but consistent structures permit grokking, whereas random structures entirely suppress generalization. Across 95 trials, the emergence of representational structure consistently precedes and reliably predicts the onset of generalization, underscoring the critical role of structural priors in enabling grokking.
📝 Abstract
Grokking -- generalization arriving long after training-set interpolation -- can be accelerated by structure-agnostic interventions: gradient filtering, weight-norm clamping, geometric penalties on hidden representations. Whether the delay specifically measures the time to form task-structured representations has remained an observational claim. We test it causally by injecting representational priors of varying structural content into a one-layer transformer learning modular addition: a supervised-contrastive auxiliary loss whose positives encode (i) the task's true equivalence structure ($(a+b) \bmod p$), (ii) a coherent-but-wrong sibling structure ($(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 structure 22/30 runs; sibling structure, which needs the same periodic features but the wrong combination, 14/15; random partition, satisfiable only by memorization, 0/20 (Fisher exact $p = 1.3 \times 10^{-7}$). A weight-norm-matched control replaying each intervention's norm trajectory onto plain cross-entropy generalizes in 0/15, collapsing into logit-scale saturation, ruling out the norm as mediator. Representation probes show structure formation precedes and predicts generalization in all 95 runs. Only the true structure also accelerates grokking, up to $2.75\times$ faster than baseline, but the acceleration is dose-dependent, bimodal across seeds, and a net wall-clock win only in its strongest cases given the contrastive term's overhead. The grokking delay is, causally, the time to form the right representational structure, where "right" is decided at the level of features rather than labels: coherent-but-wrong structure leaves grokking intact, random structure abolishes it, and only the true structure hastens it.