The mechanistic basis of data dependence and abrupt learning in an in-context classification task

TL;DR

This work investigates what data distributions and simple architectures promote in-context learning (ICL) in transformer models, using a controlled setting with parameters $K$, $L$, $N$, $B$, $\varepsilon$, and $\alpha$. A minimal two-layer attention-only network trained on a deliberately structured dataset demonstrates an abrupt ICL transition preceding the induction-head formation. It presents a two-parameter induction-head model and a three-nested-logits phenomenological loss landscape that reproduce the full-network dynamics and explain the slow-to-fast learning transition. The results imply that induction heads and intrinsic curricula may underlie sudden zero-shot abilities in large language models and offer directions for mechanistic interpretability and curriculum-based training.

Abstract

Transformer models exhibit in-context learning: the ability to accurately predict the response to a novel query based on illustrative examples in the input sequence. In-context learning contrasts with traditional in-weights learning of query-output relationships. What aspects of the training data distribution and architecture favor in-context vs in-weights learning? Recent work has shown that specific distributional properties inherent in language, such as burstiness, large dictionaries and skewed rank-frequency distributions, control the trade-off or simultaneous appearance of these two forms of learning. We first show that these results are recapitulated in a minimal attention-only network trained on a simplified dataset. In-context learning (ICL) is driven by the abrupt emergence of an induction head, which subsequently competes with in-weights learning. By identifying progress measures that precede in-context learning and targeted experiments, we construct a two-parameter model of an induction head which emulates the full data distributional dependencies displayed by the attention-based network. A phenomenological model of induction head formation traces its abrupt emergence to the sequential learning of three nested logits enabled by an intrinsic curriculum. We propose that the sharp transitions in attention-based networks arise due to a specific chain of multi-layer operations necessary to achieve ICL, which is implemented by nested nonlinearities sequentially learned during training.

Figures (12)

• Figure 1: (a) Input sequences consist of $N$ item-label pairs followed by a target. Items are drawn from $K$ classes assigned to $L \leq K$ labels. At least one item belongs to the same class as the target. The network is tasked to predict the label of the target. The number of classes ($K$), their rank-frequency distribution ($\alpha$), within-class variability ($\varepsilon$) and the number of items from a single class in an input sequence ($B$) parameterize the data distribution. (b) IWL is measured using input sequences where the items' and target's classes are randomly sampled. ICL is measured using items and targets from novel classes and by swapping the label of an existing class in the context. (c) Network architecture. (d) Loss and accuracy curves for six seeds (dark lines show averages over the seeds). Here, $B=2, K = 512$.
• Figure 2: In-weights (top row) and in-context accuracy (bottom row) against the number of classes ($K$), burstiness ($B$), within-class variability ($\varepsilon$) and the exponent of the rank-frequency distribution ($\alpha$). Here $K = 1024, \alpha = 0, B = 1, \varepsilon = 0.1$ except when that parameter is varied.
• Figure 3: (a) IC accuracy curve ($p_C = 0.8, B = 1, K = 256$) shows a slow learning phase followed by the abrupt transition to zero loss. (b) The layer 1 and 2 attention maps $p^{(1)}$ (top matrices) and $p^{(2)}_{q.}$ (bottom vectors) before and after the abrupt transition (marked in the IC curve in panel (a)).
• Figure 4: Progress measures for six seeds aligned based on when the IC accuracy crosses 50%. The color-progress measure pairings are orange: (ILA1), green: (TILA2), blue: (CLA), red: (TLA2), black: IC accuracy. See text for more details.
• Figure 5: An illustration of the four operations performed by an induction head.