🤖 AI Summary
This work elucidates the statistical estimation mechanism underlying inductive heads in Transformers during in-context learning, specifically addressing challenges of limited context matching and data sparsity. By training Transformers on k-order Markov chains and leveraging controlled synthetic data, attention analysis, and disentangled architectures, the study establishes the first theoretical connection between inductive head behavior and classical n-gram smoothing techniques—namely Jelinek-Mercer and Dirichlet smoothing. The findings reveal that inductive heads implement adaptive interpolation smoothing through soft context matching and pseudo-counts derived from start tokens, effectively learning regularized contextual estimates. Under conditions where pseudo-counts are near-optimal or lower-order contexts provide structured signals, this mechanism enables Transformers to match or even surpass traditional count-based n-gram models in performance.
📝 Abstract
Induction heads are attention circuits believed to underlie in-context learning in transformers, yet a precise characterization of the estimators they implement remains elusive. We study transformers trained on order-$k$ Markov chains and identify two complementary smoothing mechanisms. First, at finite attention-weight scale, the circuit implements a soft context-matching estimator: it aggregates contributions from exact and partial context matches, weighted exponentially by their overlap, and induces a data-dependent interpolation across context orders analogous to Jelinek-Mercer smoothing. Second, a beginning-of-sequence (BOS) token induces additive pseudo-counts, recovering Dirichlet-style smoothing. We construct a disentangled transformer implementing both mechanisms and show that trained transformers recover the predicted attention patterns. Across settings where pseudo-count smoothing is optimal or lower-order contexts provide structured evidence, trained transformers match or outperform classical count-based baselines. Our results bridge mechanistic interpretability of induction heads with classical statistical smoothing, revealing that transformers learn to regularize in-context estimation rather than simply count.