Grokking Explained: Delayed Structural Inference in Attention—What Happens During Those 10,000 Steps?

Forum topic · ✨步子哥 · 2026-05-19

Summary

Grokking is a striking machine learning phenomenon where a Transformer first memorizes training data, then—after thousands of apparently stagnant training steps—suddenly generalizes, e.g., jumping from ~30% to 95%+ validation accuracy on modular arithmetic. A new paper by Hidajat et al. attributes the delay to attention's role as an information bottleneck that implicitly performs Bayesian inference over the task's dependency structure. The authors decompose generalization into two separable conditions: a Goldilocks constraint on MLP capacity (too much capacity enables memorization) and a Bayesian structural condition requiring attention to place non-zero weight on every informative token. The long plateau arises from "explaining away": once the MLP memorizes the data, gradient signals vanish, so attention receives no feedback about task structure until weight decay erodes the memorization. Adding a KL-divergence prior on attention distributions dramatically shortens the wait, with a clear scaling law—doubling intervention strength roughly halves grokking time. Open questions remain, including the relationship to the Lottery Ticket Hypothesis and generalization to noisier tasks like natural language.

Grokking: From Memorization to Sudden Understanding—What Happens During Those 10,000 Steps?

*English translation of a zhichai.net forum post discussing recent research on the grokking phenomenon.*

What Is Grokking, and Why Is It So Puzzling?

Grokking is a striking phenomenon in machine learning: a model first memorizes its training data, then—after a long wait—suddenly understands and generalizes. Concretely, when a Transformer is trained on a small dataset such as modular arithmetic, it quickly memorizes every training example, reaching 100% training accuracy, while performing near random (~30%) on unseen combinations of the same operation. Training continues for thousands, even tens of thousands of steps with no visible progress. Then, around some point (e.g., roughly step 12,000), validation accuracy spikes—from 30% to over 95% within a few dozen steps, as if the model suddenly "got it." The term comes from Robert Heinlein's coined Martian word *grok*, meaning to understand completely and intuitively.

The puzzle is not that the model eventually generalizes, but why it waits so long. Intuitively, once the model has perfectly memorized the training data, it should also "pick up" the underlying rule along the way. Instead, it lingers between memorization and understanding for over ten thousand steps with no progress, then transitions instantly. What is happening inside the model during this long silent period?

Previous Explanations: Norm Minimization, Feature Emergence, and the Lottery Ticket Hypothesis

Researchers have proposed several accounts of the grokking plateau, though each tends to overlook the specifics of the attention mechanism:

Norm minimization: During memorization, the model learns large-norm weights that fit the training set but do not generalize. Regularization such as weight decay gradually shrinks these weights; once the norm becomes small enough, the model is forced toward generalizing solutions. This aligns with the Goldilocks capacity constraint: too much capacity invites memorization via large weights, while moderate capacity encourages generalizable features.
Feature emergence: Early on, the model relies on low-level, simple features sufficient for memorization. Over time, more abstract, generalizable features emerge, shifting the model from memorization to understanding.
Lottery Ticket Hypothesis: Randomly initialized networks contain sparse, trainable "winning ticket" subnetworks. In the grokking setting, the model has not yet found the generalizing parameter subset early in training; when it eventually does, generalization appears suddenly.

Each captures part of the mechanism, but none accounts for a constraint unique to attention: attention is an information bottleneck. If attention drops a piece of key information at one layer, no subsequent computation can recover it.

Attention as a Gambler

A new paper by Hidajat et al. starts from the attention mechanism itself. Attention, they argue, is essentially a gambler placing bets: it decides which tokens matter and which can be ignored. This "betting" can be described in Bayesian terms—attention is implicitly inferring the task's dependency structure, assigning each token a weight akin to a posterior probability of its relevance.

The key point: if attention drops a highly informative token, no amount of downstream finite computation can recover it. Unlike convolutional or fully connected networks, where later layers can potentially reconstitute lost signals, attention offers no second chance.

Therefore, generalization requires attention to place sufficient mass on every informative token—weights large enough that the token is not erased by Softmax normalization. Only then can the model reason over the complete task structure.

Two Separable Conditions: Structural Inference and Bayesian Lottery Tickets

The authors decompose Transformer generalization into two independent requirements:

1. The Goldilocks constraint on MLP capacity: The feedforward layers must have *moderate* capacity. Too little capacity means the model cannot fit the data at all; too much invites memorization. Only the middle ground forces generalizable representations. This matches prior weight-norm-based accounts. 2. The Bayesian structural condition on attention (the new contribution): Attention must correctly infer the dependency graph, treating its weights as a posterior over the task's dependencies. Generalization requires non-zero posterior mass on every informative node. If a key token is ignored, the model cannot learn the correct structure and falls back on memorization.

These conditions are separable—satisfying one does not imply the other. This decomposition reveals the root of the delay: delayed generalization is, at heart, delayed structural inference.

Explaining Away: Why Must Understanding Wait for Memory to Fade?

Why is structural inference delayed? The answer lies in the "explaining away" effect. Early in training, the MLP has enough capacity to memorize the data via shortcut features unrelated to the task's essence. Once cross-entropy loss approaches zero, gradient signals become vanishingly weak, so attention receives almost no useful feedback about task structure. Worse, because the MLP already "explains" all outputs through memorization, attention has no incentive to identify which tokens matter—a classic case of explaining away in causal reasoning: one variable (MLP memorization) explains the outcome, eliminating the influence of another (attention's structural inference).

Only when weight decay gradually erodes the MLP's memorization—raising the loss and restoring gradient signal—can attention begin learning the neglected dependencies. Grokking waits because it is waiting for memory to fade. This explains the observed inverse relation between grokking time and weight decay strength: stronger decay erodes memory faster, so attention gets feedback sooner. The delay is a structural waiting time, arising from memorization suppressing structural inference.

Breaking the Wait: A Scaling Law for Structural Intervention

If the delay stems from attention lacking structural gradients, can we give attention direct structural guidance? The authors add a KL divergence term to the objective, pulling attention distributions toward a prior (e.g., encouraging more uniform attention so no token is prematurely ignored).

The results are striking: KL intervention dramatically shortens grokking time, and a clean scaling law emerges—doubling the intervention strength roughly halves the waiting time. By giving attention a "hint," the memorization-to-understanding transition can be nearly linearly accelerated, bypassing the explaining-away bottleneck even while the MLP continues to memorize.

Open Questions and Limitations

Relation to the Lottery Ticket Hypothesis: The paper's title invokes "Bayesian Lottery Tickets," engaging directly with the lottery ticket view. Lottery tickets ask *which parameters* generalize; this work asks *which information* the model uses. Whether these accounts are complementary or competing—and whether they can be unified—remains unclear.
Task scope: Experiments focus on algorithmic tasks (modular arithmetic, parity) with clean, discrete dependency structures. For messier tasks like natural language syntax, dependency graphs are noisy and ambiguous; the Bayesian-inference behavior of attention—and the mechanism of delayed generalization—may differ substantially.
Other architectures: Grokking is not unique to Transformers; it has been observed in RNNs and MLPs, which lack explicit attention. Whether delayed structural inference explains grokking there remains an open question.

Summary: Structure and Lottery Tickets in Conversation

Hidajat et al. decompose grokking into two separable conditions—a capacity-based Goldilocks constraint and a Bayesian structural-inference condition on attention—showing that delayed generalization is delayed structural inference, caused by memorization suppressing inference via explaining away. Through structural intervention they shorten the wait and uncover a scaling law between intervention strength and grokking time.

The work answers "what happens during those 10,000 steps" and offers a practical way to accelerate understanding. But as the authors acknowledge, the relationship between lottery tickets and structural inference, and the theory's applicability to more complex tasks, remain open. The dialogue between structural inference and the Lottery Ticket Hypothesis has only just begun.