The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains (2402.11004v1)

Abstract: LLMs 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$.

• The paper demonstrates that transformers exhibit a phased learning process, initially predicting uniformly before mastering unigram then bigram statistics via induction heads.
• Methodology involves training models on synthetic sequences generated from Dirichlet-based Markov chains, using one-hot embeddings to capture transition probabilities.
• Empirical results highlight the impact of simplicity bias and suggest curriculum-based strategies could accelerate the transition to complex n-gram solutions.

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

Introduction

The paper examines the emergence of in-context learning (ICL) capabilities in LLMs via a synthetic sequence modeling task grounded in Markov chains. By training transformers on sequences from randomly generated Markov chains, the paper investigates the formation of statistical induction heads, which facilitate accurate next-token predictions using bigram statistics of the context. The research demonstrates that these models go through distinct phases during training, beginning with uniform predictions and progressing through in-context unigram and bigram solutions.

Figure 1: Left: Training process for ICL-MC using transition matrices; Right: Output distribution comparison through training phases.

In-Context Learning of Markov Chains

Task Setup

The task involves training transformers on sequences produced by Markov chains, each defined by a transition matrix sampled from a Dirichlet distribution. The objective is to predict the next token given a sequence, focusing initially on bigram models, an $n$-gram case where $n=2$.

Phases of Learning and Induction Heads

Transformers trained on this task display a multi-phase learning process. Initially, the model's output aligns with a uniform distribution. As training progresses, the model captures unigram statistics before transitioning to accurately predict using bigram statistics.

The emergence of induction heads, a critical component in transformers detecting and leveraging context patterns, is validated. These heads function by correlating current tokens with their historical co-occurrences, boosting prediction probabilities accordingly.

Figure 2: Attention evolution in training, revealing induction-like behavior where tokens attend to preceding token patterns.

Theoretical Insights and Empirical Validation

Model Construction

The research proposes a theoretical framework illustrating how a two-layer attention-only transformer can encapsulate bigram statistics. Utilizing one-hot embeddings, the model incrementally adjusts weights to reflect transition probabilities, enabling the comprehension of ICL dynamics.

Empirical Observations

Empirical results confirm that transformers adopt a hierarchical learning approach, initially stabilizing at unigram predictions before transitioning abruptly to bigram solutions. This shift is characterized as a 'phase transition', aligning with prior findings on induction head formation [elhage2021mathematical].

Implications of Simplicity Bias

Impact on Learning Dynamics

The presence of simpler unigram solutions appears to delay the adoption of complex bigram solutions due to inherent model biases favoring simplicity. Experimentation with varied data distributions demonstrates that reducing unigram utility accelerates convergence to the bigram phase.

Insights from a Minimal Model

A simplified model helps elucidate this phenomenon, revealing that while initial gradient steps emphasize simpler correlations, subsequent adjustments are required for accurate bigram learning. Such insights accentuate the potential for curriculum-based learning strategies to expedite model training.

Figure 3: Training loss and strategy alignment over time, marking clear phase distinctions in adopting n-gram solutions.

Generalization to Higher n-Grams

The research extends to $n$-grams beyond bigrams, particularly examining trigrams. When trained on trigrams, transformers likewise exhibit sequential learning phases, progressing from unigrams through higher-order $n$-gram models, with performance gradually aligning with Bayes-optimal predictions.

Figure 3: Hierarchical convergence in training with trigrams, underscoring multi-stage strategy adoption consistent with increasing algorithmic complexity.

Conclusion

The paper's analysis of ICL through the lens of Markov chains sheds light on the sequential learning capabilities of transformers and the pivotal role of induction heads. It emphasizes the importance of phase transitions and their dependence on biases toward simpler solutions, providing a valuable framework for understanding LLM behavior. Future work could extend these findings to natural language contexts, exploring how these principles manifest in more complex linguistic environments.