PaperlayerView exact revision on arXiv

Source-linked AI summary

The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains

Benjamin L. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, Nikolaos Tsilivis

arXiv:2402.11004v1cs.LG

TL;DR

The mechanisms behind in-context learning remain incompletely understood. This paper studies the question with Markov-chain sequence prediction, finding that transformers form statistical induction heads and pass through simpler strategies before learning bigram statistics.

  • Problem

    Although in-context learning is useful, its underlying mechanisms are not fully understood.

    Source: §1 · p. 1

  • Method

    The paper trains transformers to predict the next token from sequences sampled from different Markov chains, focusing mainly on bigram models.

    Source: §1 · p. 2

  • Results

    Transformers develop statistical induction heads approaching Bayes-optimal prediction, progress from uniform to unigram to bigram strategies, and show evidence that unigram solutions delay optimal learning.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The controlled task exposes discrete incomplete solutions on the way to solving ICL-MC, providing a setting to investigate how simpler solutions relate to more accurate circuits.

    Source: §4 · p. 11

    The study focuses on uniformly random transition probabilities, stationary initial states, and chains with 2 or 3 states.

    Source: §2.1 · p. 5

Abstract

from arXiv · show

Large language models have the ability to generate text that mimics patterns in their inputs. We introduce a simple Markov Chain sequence modeling task in order to study how this in-context learning (ICL) capability emerges. In our setting, each example is sampled from a Markov chain drawn from a prior distribution over Markov chains. Transformers trained on this task form \emph{statistical induction heads} which compute accurate next-token probabilities given the bigram statistics of the context. During the course of training, models pass through multiple phases: after an initial stage in which predictions are uniform, they learn to sub-optimally predict using in-context single-token statistics (unigrams); then, there is a rapid phase transition to the correct in-context bigram solution. We conduct an empirical and theoretical investigation of this multi-phase process, showing how successful learning results from the interaction between the transformer's layers, and uncovering evidence that the presence of the simpler unigram solution may delay formation of the final bigram solution. We examine how learning is affected by varying the prior distribution over Markov chains, and consider the generalization of our in-context learning of Markov chains (ICL-MC) task to $n$-grams for $n > 2$.

1 Introduction

The paper introduces in-context learning of Markov chains as a controlled setting for studying how transformers learn from patterns in their context. It finds that models develop statistical induction heads through staged learning, with evidence that simpler unigram solutions may delay the bigram solution.

1.1 Related Work

  • Earlier controlled ICL studies span supervised few-shot tasks and discrete stochastic processes, while this work focuses on learning context statistics without cleanly separated inputs and outputs.Source: §1 · p. 1

1 Introduction

  • Transformers learn statistical induction heads that compute next-token probabilities from prior occurrences of the current token, achieving performance near the Bayes-optimal predictor.These heads use bigram statistics from the context.Source: §1 · p. 2
  • Training passes through uniform, unigram, and bigram prediction phases, with rapid loss drops between phases as models learn increasingly complex strategies.The study also considers generalizing the task to n-grams.Source: §1 · p. 2, Figure 1 · p. 2
  • Evidence suggests that simpler in-context unigram solutions delay learning the optimal solution; removing their usefulness from training examples speeds convergence on the original distribution.Source: §1 · p. 3
  • The transition from the inadequate unigram solution to the correct bigram solution is tied to alignment between model layers.The first layer’s learning signal depends on how close the second layer is to its correct weights.Source: §1 · p. 3
  • With relative position embeddings, the first layer develops a bias toward odd-numbered lookbacks, although only looking back one token is clearly useful.The simplified model’s theoretical analysis and real-transformer experiments both show this pattern.Source: §1 · p. 3

2 Setup

The setup defines in-context learning of randomly sampled Markov chains and contrasts unigram and bigram prediction strategies. It also introduces a simplified transformer that can represent the bigram solution.

2.1 ICL-MC Task

  • The task is to infer transition probabilities from a sampled Markov-chain sequence and use them to predict the next state.Transition matrices are drawn from a prior; the primary setup samples each row from a Dirichlet distribution.Source: §2.1 · p. 5

2.2 Potential Strategies to Solve ICL-MC

  • The unigram strategy treats tokens as independent and estimates state frequencies, whereas the bigram strategy accounts for dependencies between adjacent tokens.The unigram strategy forms a posterior belief about the stationary distribution from state counts.Source: §2.2 · p. 5
  • The bigram strategy's assumption of a uniformly sampled first token differs from the stationary-distribution initialization used in experiments and theory, but the difference becomes negligible as context grows.The bigram statistics approach the Bayes-optimal solution under either initialization as context length grows.Source: §2.2 · p. 6

2.3 Architectures: Transformers and Simplifications

  • A single-head, two-layer attention-only transformer can find bigram statistics in the in-context Markov-chain task.The simplified model is expressive enough to represent the optimal bigram solution.Source: §2.3 · p. 6
  • The simplified transformer's first layer looks back one token, while its second layer attends to matching prior tokens to aggregate bigram counts.The resulting output approaches probabilities given by empirical bigram statistics.Source: §2.3 · p. 7

3 Empirical Findings and Theoretical Validation

Transformers and a minimal model learn Markov-chain predictions in stages, moving from a unigram solution toward the Bayes-optimal bigram solution. The minimal-model analysis links this progression to interactions between layers, and experiments extend hierarchical learning to trigrams.

3.1 Transformers in-context learn Markov Chains hierarchically

  • Models reach near-Bayesian performance through a unigram plateau followed by a rapid transition to the bigram strategy; one-layer transformers fail to make this transition.The model’s similarity shifts from the unigram strategy toward the bigram strategy around the second loss drop.Source: §3.1 · p. 7, Figure 3 · p. 8

3.2 Theoretical Insights from the Minimal Model

  • When the data distribution removes unigram signal, transformers learn faster without a long intermediate plateau, including when evaluated on the original distribution.The minimal-model analysis likewise finds no heavy initial preference for the first coordinate in this special case.Source: §3.1 · p. 8, §3.2 · pp. 10–11
  • The minimal model also converges to the bigram solution, spending significantly less time, if any, at the representable unigram solution.The experiments use deterministic constant initialization and minimize cross-entropy with SGD.Source: §3.2 · p. 9
  • In the minimal model, the second layer first moves toward the correct solution, while the first layer’s positional embeddings initially favor the current token and develop the useful one-back coordinate after another step.The analysis describes the learning signal for the first layer as tied to how closely the second layer approaches its correct weights.Source: §3.2 · p. 9
  • The minimal model’s positional embeddings amplify even coordinates more than odd ones, including even coordinates unnecessary for the optimal solution; transformers show a similar pattern before the induction head forms.The transformer’s positional embeddings show the pattern shortly before the second loss drop.Source: §3.2 · pp. 9–10

3.3 Beyond Bigrams: n-gram Statistics

  • With three attention heads and context length 200, the trigram model converges close to Bayes optimal through four learning stages, while single-head models do not outperform bigrams.The stages correspond to different algorithms, and the strategies include unigrams, bigrams, and trigrams.Source: Figure 5 · p. 10, §3.3 · p. 11

3.2 Theoretical Insights from the Minimal Model

  • In a simplified minimal-model setting with curriculum learning, two gradient-descent steps suffice to reach the bigram solution when learning rates and second-layer weight decay are appropriately chosen.The first step can align the second layer; the second uses hard examples to identify the relevant relative embedding.Source: §3.2 · p. 11

4 Conclusion and Discussion

The paper presents ICL-MC as a controlled setting for studying in-context learning and induction heads. It points to natural-language training and the role of simpler incomplete solutions as questions for future investigation.

4 Conclusion and Discussion

  • ICL-MC provides a controlled setting for studying in-context learning and the emergence of statistical induction heads.The authors identify the extension of these findings to natural-language training data as a direction to explore.Source: §4 · p. 11
  • The authors propose investigating how far their findings extend to natural-language training data, given that induction heads are common in LLMs.This is presented as a natural direction for further work.Source: §4 · p. 11
  • The paper raises whether simple, incomplete solutions in language modeling and other rich settings speed or slow the formation of more accurate, complex circuits.The conclusion frames this as a question for future investigation.Source: §4 · p. 11

A Proofs

These proofs analyze how staged learning in a transformer leads to the bigram solution and connect the two-stage parameter updates to observed training dynamics. They cover the initial diagonal bias, the later bigram-relevant update, and the consistency of dynamics across seeds despite variation in phase-transition positions.

Setup and notation; Model

  • The minimal model assumes a two-state Markov chain and represents bigram counts with a one-step lag, while unigram counts use the current token.The transition-matrix rows are drawn from a Dirichlet prior; unless otherwise stated, its concentration parameter is uniform over the simplex.Source: A Proofs · p. 15

Training

  • The first update makes diagonal entries of Wk grow about 4 times faster than off-diagonal entries, and the second makes v_2 dominate the other coordinates.The second-step dominance is established in the large-sequence-length regime.Source: A Proofs · pp. 15, 16, 25
  • With a = b, the leading gradient terms toward the unigram solution disappear, leaving less unigram-learning signal.The analysis notes that the cubic O(t^3) terms in the Wk gradient disappear, leaving leading terms of order O(t^2).Source: A Proofs · pp. 25–26
  • The paper’s main technical result demonstrates that two steps of stochastic gradient descent lead the model to the bigrams solution; prior calculations note less “signal” toward the unigrams solution when a = b.Source: A Proofs · p. 26

A Proofs

  • Training dynamics are consistent across ten two-layer attention-only transformers with different random seeds, although the exact phase-transition positions vary.The seeds vary initialization and training data.Source: A Proofs · p. 28, Figure 8 · p. 28
  • The proof’s two-step analysis yields a limiting prediction as t→∞, with an O(1/t^2) first-step rate for learning the second layer and a larger O(1/t) second-step rate that cleans up that layer and learns the first layer.The authors connect the immediate first loss drop and later positional-embedding “grokking” plateau in experiments to these two stages.Source: A Proofs · p. 28

B Experimental Details

The experiments train two-layer attention-only transformers on ICL-MC and compare them with a minimal model, using sequences of 100 tokens. Additional experiments vary the number of symbols and the distributions used to test learning.

B Experimental Details

  • Transformers use AdamW at learning rate 3e−5, batch size 64, and hidden dimension 16; the minimal model uses SGD at learning rate 2e−3 and batch size 64.For 3-grams, the minimal model's learning rate is 3e−2.Source: B Experimental Details · p. 28
  • The experiments evaluate model outputs at the last token of each example.Source: B Experimental Details · p. 28
  • For k = 2 Figure 4a distributions, p in (0,1) sets µ = x + p(1 − 2x) from uniform x; y is sampled uniformly within 0.2 of µ, clipped to [0,1], and with probability 1/2, a = x and b = y.Source: B Experimental Details · p. 28
  • For more than two or three symbols, sequence length 200 makes the unigram–bigram difference large enough to reveal a unigram phase before the final loss drop.The passage notes a plateau before the final drop in either case.Source: Figure 9 · p. 29

C Additional Experiments

Additional experiments test single-layer models, altered test distributions, and comparisons with a minimal model. They also inspect layer attention and positional patterns during training.

C Additional Experiments

  • Figure 10 displays part of attention matrix A across time steps, with first-layer attention above second-layer attention.Source: Figure 10 · p. 29
  • Figure 11 compares two-layer transformers with the minimal model using KL-divergence test loss on in-distribution, iid-token, and another described distribution.The caption is truncated before specifying the final distribution.Source: Figure 11 · p. 30
  • On uniformly random doubly stochastic transition matrices, both models' in-distribution test loss quickly reaches the unigram algorithm's level.In this setting, unigram-based guesses are as good as uniform-probability guesses.Source: C Additional Experiments · p. 30
  • For k = 2, the transformer and minimal model comparison shows an alternating positional-encoding pattern at 64,000 examples seen.Source: Figure 12 · p. 30
  • Single-layer transformers show no notable gain beyond the unigram strategy, even after eight times longer training, across the tested MLP and positional-encoding variants.The caption states this holds with or without MLPs and with absolute or relative positional encodings.Source: Figure 13 · p. 31
Loading 2402.11004v1…