Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets
For researchers studying grokking and generalization in Transformers, this provides a theoretical framework that unifies norm-based and structural explanations, with practical implications for accelerating generalization.
The paper explains why Transformers exhibit delayed generalization (grokking) by decoupling it into a capacity condition (Goldilocks bound on MLP) and a structural condition (attention must place sufficient mass on informative tokens). It proves that weight decay erodes memorization before structural inference occurs, and introduces a KL-based intervention that bypasses this delay, yielding an inverse-intervention-strength scaling law for grokking time.
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.