ary/index_
ary/ projects/ induction-probe

induction-probe

the notebook from my induction-heads essay, packaged as reproducible code. one-command setup.

// about

This is the companion code for my essay on the emergence of in-context learning. If you want to reproduce the experiment — a 4-layer transformer that develops induction heads around step 4,100 — this repo gives you a one-command path to running it.

what’s here

  • train.py — the training loop, logging both loss and induction score
  • probe.py — the per-head induction-score probe, essentially the one from Olsson et al.
  • notebook.ipynb — the same experiment in a more exploratory form, with the plots from the essay
  • configs/ — the three seeds I used, so the results in the essay are reproducible

install (from source)

Requires Python 3.11+, PyTorch 2.2+, and a CUDA-capable GPU. I ran it on a single 3090 in about 40 minutes per seed.

git archive --remote=ssh://git@ary.example.com:induction-probe.git HEAD | tar -x
# or, if you got a tarball:
tar -xzf induction-probe-main.tar.gz && cd induction-probe-main

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

# reproduce the three seeds from the essay
./run-all-seeds.sh

# the plots go here:
open figures/

things I want to try next

  • Scan dropout from 0 to 0.5 in 0.05 steps — does it always kill induction-head formation, or just slow it?
  • Do the same experiment on a 1-layer model. Should be impossible, but I want to see the failure mode.
  • Check whether the n-gram matcher is actually a phase that precedes induction, or whether they emerge together.

If you run this and find something interesting, email me.