Watch an induction circuit form
A two-layer transformer trains in your browser on letter sequences that repeat themselves. Copying is the only thing it can learn.
Every head's attention on one sequence
Rows are the letter doing the looking; columns are the letters it looks at. Darker means more attention. Select a head to enlarge it. Switch heads off to test what depends on them.
All heads are on.
Layer 2, head 1: induction score n/a
A previous-letter head lights up the line just left of the diagonal. An induction head lights up a line further left: each letter in the second copy looks at the letter after its first appearance.
Top row: your letters. Bottom row: the model's guess for the letter that comes next. Repeat part of your sequence to test copying.
Challenges
0 of 4 complete
- Train until the phase change: the loss on copied letters suddenly drops below half of the loss on random ones. (not yet)
- After the phase change, switch off just one layer-2 head. Does copying stop? (not yet)
- Without touching layer 2, find the one layer-1 head whose removal more than doubles the loss on copied letters. (not yet)
- Type a sequence that repeats itself and get the model to predict at least 5 of the repeated letters. (not yet)
Try this
- Train at full speed and wait for the phase-change marker. Then compare the two charts.
- Pause, then switch off one layer-2 head at a time. Does copying stop?
- Keep training to about step 1,200, then look for the single layer-1 head whose removal breaks copying.
- Try another seed. The phase change lands at a similar step, but different heads end up doing the work.
Measure a head's induction scorePython with NumPy
import numpy as np
# A start marker, the letters Q J F K, then the same four again.
tokens = ["start", "Q", "J", "F", "K", "Q", "J", "F", "K"]
T = len(tokens)
first, second, length = 1, 5, 4 # the copy begins at position 5
gap = second - first
def causal_softmax(scores):
"""Turn scores into attention: each position sees itself and earlier ones."""
masked = np.where(np.tril(np.ones((T, T))) == 1, scores, -np.inf)
w = np.exp(masked - masked.max(axis=1, keepdims=True))
return w / w.sum(axis=1, keepdims=True)
# A layer-1 head that strongly prefers the position just before it.
prev_scores = np.zeros((T, T))
for i in range(1, T):
prev_scores[i, i - 1] = 4.0
# A layer-2 head that, inside the copy, prefers the letter after the
# earlier occurrence of the current letter.
ind_scores = np.zeros((T, T))
for i in range(second, T - 1):
ind_scores[i, i - gap + 1] = 2.0
prev_attn = causal_softmax(prev_scores)
ind_attn = causal_softmax(ind_scores)
# The demo's scores: average over queries inside the copy (but its last letter).
queries = range(second, second + length - 1)
prev_score = np.mean([prev_attn[i, i - 1] for i in queries])
ind_score = np.mean([ind_attn[i, i - gap + 1] for i in queries])
print(f"previous-letter score: {prev_score:.2f}")
print(f"induction score: {ind_score:.2f}")
print("letter F at position 7 looks most at:", tokens[int(ind_attn[7].argmax())], "at position", int(ind_attn[7].argmax()))previous-letter score: 0.90
induction score: 0.55
letter F at position 7 looks most at: K at position 4