Structure-Specific Representational Priors Causally Control the Grokking Delay

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<sup>{-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 across methods ($0/40$ vs $6/20$, $p=7.7\times10<sup>{-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.

• The paper demonstrates that injecting task-aligned representational priors via supervised contrastive losses achieves up to 2.75× faster generalization in modular addition tasks.
• It rigorously isolates structural effects from weight-norm influences, showing that class-like representational formation precedes the generalization jump.
• Norm-control techniques, especially norm-clamping, when combined with structural priors, yield robust acceleration with median improvements up to 8.6×.

Causal Control of Grokking Delay by Structure-Specific Representational Priors

Introduction and Motivation

The phenomenon of grokking in neural networks—test set generalization arising long after the model has interpolated the training data—is a focal point for understanding implicit learning dynamics under rich supervision, especially in algorithmic tasks such as modular addition (Power et al., 2022). Traditionally, acceleration or modulation of the grokking delay has only been achieved using structure-agnostic methods, such as gradient filtering, weight-norm clamping, or geometric constraints on representations (Lee et al., 2024, Khanh et al., 11 Jun 2026, Tiwari et al., 30 Jun 2026). However, a key mechanistic hypothesis has remained untested: whether the grokking delay causally reflects the time to develop task-relevant representational structure, or is instead an ineluctable property of optimization pathologies.

This work provides the first rigorous causal test by directly injecting structure-specific representational priors via supervised contrastive losses. Specifically, the hidden representations of a one-layer transformer trained on modular addition are regularized using three structurally distinct positive-set definitions: (1) the true modular addition equivalence classes, (2) coherent-but-incorrect modular subtraction classes (requiring the same periodic features, but a different combination), and (3) a random partition (uninformative with respect to the task). All architectural and optimization details are held fixed, enabling unambiguous attribution of effects to the structure of the prior.

Experimental Setup and Methodology

Training is performed on modular addition modulo 97 with a one-layer decoder-only transformer, purposefully omitting LayerNorm to maintain interpretability of weight norm interventions (Nanda et al., 2023). The core experimental manipulation involves the addition of an auxiliary supervised contrastive (SupCon) loss (Khosla et al., 2020) atop the cross-entropy objective, parameterized by strength $\lambda$. The essential experimental variable is the definition of the positive set (the structure):

• True structure: Examples sharing $(a+b)\bmod p$.
• Wrong-but-coherent: Examples sharing $(a-b)\bmod p$.
• Shuffled: Fixed random permutation preserving class-size distribution.

All variants share class size, representational geometry, and computational overhead, isolating structural content as the only variable. Weight-norm-matched controls and Grokfast (Lee et al., 2024) acceleration baselines are included for mechanistic clarity.

Strong Empirical Results: Structural Content Dictates Generalization

The key result is a sharp, structure-specific control of generalization (Figure 1). Injecting the task's true structure as the contrastive prior produces the earliest generalization events observed—up to $2.75\times$ faster than baseline, with 22/30 seeds achieving test accuracy within a strict budget limit.

Figure 1: Survival curves: fraction of seeds not yet reaching 95% test accuracy, per auxiliary-loss strength $\lambda$; the true-structure prior yields rapid generalization and earliest transitions in the dataset.

Conversely, the random partition (shuffled) condition produces zero generalization under otherwise identical pressure, and the wrong-but-coherent structure yields intermediate results: high grokking reliability (14/15) but no median acceleration. This dissociation clearly shows that the coherence and compatibility of the injected structure with the task-level feature family—not mere label-level correctness or optimization pressure—is the critical determinant of both the likelihood and the timescale of generalization.

Figure 3: Dose–response of epochs to generalization versus auxiliary-loss strength $\lambda$; the true-structure prior produces a range from acceleration to trapping, with shuffled never generalizing regardless of strength.

Dissecting Mechanisms: Not Norm-Mediated, but Representation-Limited

Both the true and shuffled contrastive conditions are found to inflate weight norms substantially, yet only the true structure yields generalization. When the weight-norm trajectory of each intervention is replayed as a strict control (norm-matched), the generalization benefit evaporates entirely—all controls fail to generalize (0/15) and collapse into logit-scale saturation, confirming that weight-norm effects are not sufficient or necessary for the representational benefit (Figure 4).

Figure 5: Survival curves for delay-mitigating interventions; norm-clamp (holding low weight norm) produces stall-free, rapid generalization while default Grokfast stalls all seeds in composition.

Crucially, representation-timing probes corroborate a mechanistic link: building class structure in held-out representations consistently precedes the generalization jump and is absent in random structure and norm-matched controls, matching observational accounts of "representation-first" grokking dynamics.

Figure 4: Representation-timing probes—embedding Fourier concentration, class cosine gap, and CKA—all peak before accuracy jumps under the true structure, but not in the null controls.

Revealing a Bidirectional, Dose-Dependent Lever

The use of strong, structure-specific priors is bidirectional: the same intervention that produces the fastest transitions also stalls a non-negligible fraction of seeds. This is interpreted as a stochastic "race" between representational structure seeding and weight-norm-driven saturation; when the prior seeds class clusters before saturation, rapid generalization occurs; missing this window leads to trapping in a slow regime. Notably, random partitions at high strength can even delay memorization, demonstrating causal control over both the direction and magnitude of the grokking delay.

Norm Control as Reliable Acceleration: Mechanism-Informed Mitigation

Building on the insight that norm inflation competes with structure seeding, three independent strategies are shown to reliably retain the speedup: (1) annealing the auxiliary loss ($\lambda$-anneal), (2) stepwise norm-replay, and (3) direct norm-clamping. The norm-clamp method yields a median $8.6\times$ acceleration, with all seeds grokking and delays collapsed to the shortest timescales observed.

Figure 7: Median epochs to generalization and delay by intervention at $\lambda=1.0$; norm-clamp achieves the greatest reliability and acceleration.

This monotonic relationship—lower held norm corresponds to greater acceleration—demonstrates that norm-control mechanisms can be combined with structure-specific priors to reliably achieve rapid generalization, as predicted by the mechanistic "race" account.

Implications and Future Directions

The findings bear several important theoretical and practical implications:

• Causal confirmation of representation-limited grokking: This is strong evidence that generalization delay is not fixed by optimization but set by representational structure formation, validating representation-timing theories (Chou et al., 26 May 2026, Sivasankar, 11 Jun 2026, Liu et al., 2022) and placing empirical boundaries on the weight-norm delay law (Khanh et al., 11 Jun 2026).
• Feature-level specificity: The critical determinant is not task-label correctness but expressivity within the feature family required by the generalizing circuit; wrong-but-coherent priors (e.g., subtraction) permit—but do not accelerate—grokking, while unaligned priors prevent it entirely.
• Limitations of geometric regularization: Purely geometric or anti-saturation effects of contrastive losses cannot substitute for structure-specific information.
• Non-compositionality of structure and optimization accelerators: Adding structure-specific and optimization-driven accelerators (Grokfast) fails to produce super-additive effects and may destabilize learning unless carefully tuned.
• Transfer limitations and future needs: Demonstrated control is for clean algorithmic tasks; scaling to data with unknown or ambiguous structure will require self-supervised or invariance-driven structure priors.

Conclusion

This paper establishes, through tightly controlled interventional designs, that the grokking delay in neural sequence models is causally dictated by the formation of task-suitable representational structure, and this process is directly manipulable by injecting precise structural priors via contrastive objectives. The work not only resolves a key mechanistic debate regarding the nature of delayed generalization, but also introduces norm-aware mitigations capable of converting potentially unreliable acceleration into robust, rapid generalization—a foundation for future work on efficient structural learning in broader domains.