Structure-Specific Representational Priors Causally Control the Grokking Delay

Engineering notes

Engineering notes will be added by the aipentium editorial team.

Chinese explanation / 中文解读

中文解读待补充：本站会优先为大语言模型、生成式AI、ChatGPT相关技术、计算机视觉、深度学习等高价值论文补充中文说明。

Original 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.

Links and sources

• Official / arXiv page
• PDF from original source

Comments