Reproduce “What needs to go right for an induction head? A mechanistic study of in-context learning circuits and their formation” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Aaditya K. Singh, Ted Moskovitz, Felix Hill, Stephanie C.Y. Chan, Andrew M. Saxe

The authors study how induction-like in-context learning forms in a small, two-layer attention-only transformer trained to label a query handwritten character from two labeled examples. They clamp activations during training to separate previous-token processing, query-key matching, and label copying. These components interact to produce a sharp loss drop and 99.99% training accuracy, but the account is limited to synthetic Omniglot experiments.

Key insights

  • The abrupt emergence of in-context learning capabilities in transformers stems from the collective, non-linear interaction of multiple smoothly learning subcircuits.
  • Induction heads consist of three causally distinct subcircuits: Layer 1 Previous Token (PT) operations, the Layer 2 QK match, and the Layer 2 value copy mechanism.
  • Task-specific data characteristics, such as the number of classes or labels, differentially influence the learning rates and formation timing of these individual subcircuits.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.