cs.LGMay 15, 2026

Grokking as Structural Inference: Transformers Need Bayesian Lottery Tickets

Authors: Kai HidajatSolden StollJoseph An

Organizations: Department of Computer Science University of Washington Seattle, WA 98195

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.

Explore similar work

CardsList
  1. Incremental Learning of Sparse Attention Patterns in Transformers

    Feb 22, 2026Oğuz Kaan Yüksel, Rodrigo Alvarez Lucendo, Nicolas FlammarionTransformer AttentionTransformer Architectures

  2. TPR-Attention for Combinatorial Generalization

    Aug 31, 2026Melisa Civelekoğlu, Isabeau Prémont-SchwarzNeural ArchitecturesInductive Bias