A reproduction of the induction-heads mechanism (Olsson et al., 2022) in the smallest possible setting: a 2-layer attention-only transformer trained on synthetic sequences where a randomly-positioned first half is followed by its exact repetition. The 2-layer model develops the canonical two-head induction circuit and solves the task; a 1-layer baseline cannot, because the circuit requires two-layer composition.
This is the project where mechanistic interpretability meets the current frontier of transformer research. Induction heads are the most-discussed mechanistic finding of the past several years; the toy version reproduced here is the simplest setting in which the circuit can be cleanly isolated, identified, and ablated.
-
Phase 1 — Capability emergence. A 2-layer attention-only transformer trained on random-boundary repeat sequences develops in-context learning, with final second-half cross-entropy loss of
$0.05$ nats versus chance$\log V = 4.16$ nats. A 1-layer baseline reaches only$1.34$ nats on the same task, a factor-of-$27$ gap. The 2-layer loss curve exhibits a sharp transition around training step$\sim$ $800$; the 1-layer curve descends gradually without a clear transition. -
Phase 2 — Mechanistic identification. The induction circuit is identified by visualization and ablation. Layer-1 head 2 is the dominant previous-token head (off-diagonal-1 attention strength
$0.60$ , ablation impact$+3.48$ nats). The induction-reading role in layer 2 is distributed across three heads rather than concentrated in one (ablation impacts$+0.18, +0.76, +0.83$ nats). The pair-ablation result confirms an upstream-downstream circuit structure: ablating both the previous-token head and the most-impactful layer-2 head drives loss to$3.80$ nats, higher than the 1-layer baseline. -
Phase 3 — Circuit formation timeline. Training-time tracking reveals that the previous-token attention pattern (layer-1 head 2 strength crossing 0.3) and the in-context-learning capability (second-half loss dropping below half-initial) emerge simultaneously, within a narrow window around training step
$\sim$ $700$–$900$. Post-saturation refinement: the previous-token head continues to specialize even after loss has saturated. Transient representational competition: layer-1 head 0 briefly attempts the previous-token role during the transition window before retracting.
git clone https://github.com/QuantumPouya19/induction-heads.git
cd induction-heads
python -m venv .venv
# Windows:
.venv\Scripts\activate
# macOS/Linux:
source .venv/bin/activate
pip install -r requirements.txtThen, in order:
# Phase 1: train both 1-layer and 2-layer models (~20 minutes total on CPU).
python -m src.train
python scripts/plot_emergence.py
# Phase 2: identify the circuit (seconds).
python scripts/plot_attention_patterns.py
python scripts/run_ablation_study.py
# Phase 3: training-time circuit formation (re-runs Phase 1 with extra logging).
python scripts/plot_circuit_formation.pyAll hyperparameters live in configs/default.yaml. Runs are seeded; the default seed is 0. Total CPU runtime is about 30 minutes including the Phase 3 retraining.
induction-heads/
├── src/ # Library code
│ ├── config.py # Dataclass + YAML loader
│ ├── data.py # Random-boundary repeat sequences
│ ├── model.py # Attention-only transformer with ablation hooks
│ ├── train.py # Training with per-position loss + attention logging
│ ├── analysis.py # Ablation and attention extraction
│ └── utils.py # Seeding
├── scripts/ # One-purpose entry points
│ ├── plot_emergence.py
│ ├── plot_attention_patterns.py
│ ├── run_ablation_study.py
│ └── plot_circuit_formation.py
├── configs/
│ └── default.yaml # Hyperparameters
├── results/ # Figures committed; pkls and checkpoints excluded
├── notes/ # Per-phase written discussion
├── requirements.txt
├── LICENSE
└── README.md
Several details worth flagging for anyone reproducing or extending this work:
-
Random boundary positions are essential. An initial fixed-boundary version of the task admitted a 1-layer solution via positional shortcutting ("attend to position
$t - T$ for fixed$T$ "). Randomizing the boundary position per sequence —$T$ uniformly sampled from$[T_{\max}/2, T_{\max}]$ — eliminates this shortcut and forces any successful model to use content-based copying, which is what the induction circuit provides. - Attention-only architecture. No MLP blocks, no LayerNorm. This minimal architecture is the cleanest setting for isolating the induction circuit; adding MLPs and norms does not change the underlying circuit but introduces additional computational pathways that obscure it.
- Head ablation is post-attention, pre-output-projection. We zero the head's contribution to the residual stream after the attention computation but before the W_O projection. This matches the convention in the mechanistic-interpretability literature (e.g., Anthropic's Transformer Circuits work).
- Distributed induction layer. The canonical "two heads do the job" version of induction is an idealization. In this toy model the previous-token role is sharply localized (one head does it almost alone), but the layer-2 induction-reading role is distributed across three of the four available heads. This is consistent with what Olsson et al. observe in larger models.
This is a portfolio practice project. The induction-heads circuit is from Olsson et al. (2022); the toy-model setup is from the same paper. The goal is a clean, reproducible identification of the circuit at the smallest tractable scale, demonstrating that the mechanistic-circuits framework operates on a transformer small enough to be fully analyzable.
Three honest framings: (i) the canonical "two heads do the job" story is an idealization — the layer-1 previous-token role is sharply localized in one head, but the layer-2 induction-reading role is distributed across multiple heads, both here and (per Olsson et al.) in larger models; (ii) the Phase 3 simultaneous-emergence finding is a single-run observation, and multi-seed studies would be needed to test whether the simultaneity holds universally; (iii) an initial experimental-design issue (fixed-boundary positional shortcut) was discovered mid-Phase 1 and corrected by randomizing the boundary — this kind of issue-and-fix loop is included in the project record rather than hidden.
- Olsson, C., et al. (2022). In-Context Learning and Induction Heads. Transformer Circuits Thread.
- Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread.
- Power, A., et al. (2022). Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. arXiv:2201.02177.
MIT. See LICENSE.