Distinct mechanisms underlying in-context learning in transformers

TL;DR

This paper provides a detailed mechanistic analysis of in-context learning in transformers trained on Markov chains, revealing four algorithmic phases and distinct subcircuits responsible for adaptive computation.

Contribution

It characterizes the four phases of in-context learning in transformers and identifies the underlying subcircuits and conditions influencing their operation.

Findings

Transformers exhibit four distinct algorithmic phases during in-context learning.

Two key boundaries, $K_1^*$ and $K_2^*$, depend on data diversity and influence memorization and generalization.

Theoretical analysis explains the transition from 1-point to 2-point generalization and the loss landscape features.

Abstract

Modern distributed networks, notably transformers, acquire a remarkable ability (termed `in-context learning') to adapt their computation to input statistics, such that a fixed network can be applied to data from a broad range of systems. Here, we provide a complete mechanistic characterization of this behavior in transformers trained on a finite set $S$ of discrete Markov chains. The transformer displays four algorithmic phases, characterized by whether the network memorizes and generalizes, and whether it uses 1-point or 2-point statistics. We show that the four phases are implemented by multi-layer subcircuits that exemplify two qualitatively distinct mechanisms for implementing context-adaptive computations. Minimal models isolate the key features of both motifs. Memorization and generalization phases are delineated by two boundaries that depend on data diversity, $K = |S|$. The first ($K_1^\ast$) is set by a kinetic competition between subcircuits and the second ($K_2^\ast$) is set by a representational bottleneck. A symmetry-constrained theory of a transformer's training dynamics explains the sharp transition from 1-point to 2-point generalization and identifies key features of the loss landscape that allow the network to generalize. Put together, we show that transformers develop distinct subcircuits to implement in-context learning and identify conditions that favor certain mechanisms over others.