Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets

TL;DR

This paper models attention in Transformers as a Bayesian posterior over task dependencies, explaining delayed generalization (grokking) as a structural inference problem and proposing interventions to accelerate learning.

Contribution

It introduces a Bayesian structural perspective on attention, decouples capacity and structural conditions for generalization, and proposes a KL-based intervention to speed up grokking.

Findings

Attention as Bayesian posterior explains grokking delays.

Structural interventions can bypass the delay and accelerate generalization.

Bayesian tickets outperform lottery tickets in transfer tasks.

Abstract

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.