Induction Heads Interpolate N-Grams

Technical Summary: Induction Heads Interpolate N-Grams

Problem Statement
Induction heads, identified as the primary circuit for in-context learning (ICL) in transformers, are traditionally characterized in the "hard-attention" limit. In this regime, the circuit performs exact pattern matching and maximum-likelihood (ML) counting. While mechanistically appealing, this view is statistically incomplete for high-order dependencies: exact $k$-gram matches are often too rare in finite samples to support reliable prediction, and the ML estimator assigns zero probability to unseen continuations. Classical language models address this via smoothing and back-off mechanisms (e.g., Jelinek–Mercer, Katz back-off), which redistribute probability mass to unseen events and interpolate across context lengths. The central question addressed by this work is whether induction heads implement a richer class of estimators beyond strict counting, specifically capable of performing such smoothing.

Methodology
The authors analyze transformers trained on order-$k$ Markov chains, a setting that balances analytical tractability with the finite-sample trade-offs inherent in context-based prediction.

1. Theoretical Framework: The study utilizes a "disentangled" transformer architecture (removing MLPs and using concatenation for residual streams) to facilitate exact analysis. The authors construct a specific two-layer transformer with Relative Positional Encoding (RPE) and prove it implements a soft context-matching estimator.
2. Mechanism Construction:
   • Layer 1 (Copying Heads): Uses $k$ attention heads to copy tokens from specific relative lags ($t-1, \dots, t-k$), constructing a representation of the current context.
   • Layer 2 (Context Matching): Uses a single attention head to compare the query context against all candidate contexts in the history. The attention scores are designed to compute a shift-aligned inner product, weighting matches based on the number of overlapping tokens (Hamming overlap) rather than requiring exact matches.
   • BOS Token: The authors incorporate a Beginning-of-Sequence (BOS) token as a sequence-independent attention target to introduce additive pseudo-counts.
3. Statistical Analysis: The authors derive the closed-form estimator implemented by this circuit and compare it against Bayes-optimal predictors under two task priors:
   • Independent Dirichlet Prior: Where contexts are uncorrelated.
   • Hierarchical Dirichlet Prior: Where higher-order transitions are centered on their lower-order suffixes, creating correlations between contexts sharing a suffix.

Key Contributions
The paper establishes that induction heads implement two complementary smoothing mechanisms when operating at finite attention-weight scales:

1. Soft Context Matching (Interpolation):

   • The circuit aggregates contributions from both exact and partial context matches.
   • Matches are weighted exponentially by their overlap size ($|M|$).
   • This induces a data-dependent interpolation across context orders, analogous to Jelinek–Mercer smoothing. Unlike classical JM smoothing, which uses fixed global weights, the transformer's weights depend on the realized counts in the specific sequence.
   • The estimator unifies mixture models (like JM) and back-off strategies (like Katz): when exact matches are rare, the model automatically "backs off" to lower-order partial matches without explicit rules.

2. BOS as Additive Pseudo-Counts:

   • The BOS token acts as a fixed, sequence-independent sink for attention.
   • This mechanism allows the model to add constant pseudo-counts to context-specific transition counts.
   • Theoretically, this recovers add-$\alpha$ smoothing (e.g., Laplace smoothing) in the limit of large attention weights, providing a circuit-level explanation for how architectural conventions implement prior-like regularization.

3. Empirical Validation:

   • The authors train both disentangled and standard (fully parameterized) transformers on Markov chain tasks.
   • Results show that trained models recover the predicted attention patterns:
     • Without BOS, they rely on soft context matching to interpolate across orders.
     • With BOS, they utilize pseudo-counts to achieve near-Bayes-optimal performance under independent priors.
     • Under hierarchical priors, where lower-order contexts provide structured evidence, the models outperform fixed add-$\alpha$ baselines by leveraging the interpolation mechanism.
   • Mechanistic analysis confirms that trained standard transformers learn the specific weight structures (e.g., block-shifted matrices) predicted by the theory.

Results

• KL Divergence: In settings where add-$\alpha$ smoothing is Bayes-optimal (Independent Dirichlet), trained transformers with BOS match the Bayes-optimal predictor. In hierarchical settings, where structured evidence exists, the interpolation mechanism allows transformers to outperform fixed smoothing baselines.
• Attention Patterns: Visualizations of trained models confirm the two-stage circuit: Layer 1 performs lag-specific copying, and Layer 2 performs soft matching, distributing attention mass across full and partial matches based on overlap.
• Asymmetric Lag Weights: In hierarchical tasks, trained models learn to assign higher weights to matches on the most recent token (lag 1) compared to older tokens (lag 2), consistent with the prior's structure where recent tokens are better predictors of the next state.

Significance
The paper bridges the gap between mechanistic interpretability of induction heads and classical statistical smoothing. It demonstrates that transformers do not merely perform unregularized counting; rather, they learn to regularize in-context estimation. By dynamically interpolating across context orders and utilizing architectural tokens (BOS) for additive smoothing, induction heads implement sophisticated statistical procedures that adapt to the data distribution, explaining their robustness in in-context learning tasks.