Mechanistic Interpretability › Code Implementation
Language AI HandbookAll chapters Mechanistic Interpretability › Code ImplementationOne section · Sections in this chapter8 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement a demonstration of induction head behavior. We will build a minimal two-layer transformer, train it on a task that requires in-context pattern completion, and then verify that induction-head-like behavior emerges in the attention patterns. This implementation mirrors the analytical setup used by Olsson et al. (2022) in their original induction head study, scaled down for clarity.
Setup and Data Generation
import numpy as np
import torch
torch.manual_seed(42)
np.random.seed(42)
# Vocabulary: tokens 0 to VOCAB_SIZE-1
VOCAB_SIZE = 20
SEQ_LEN = 32
D_MODEL = 64
N_HEADS = 2
N_LAYERS = 2
BATCH_SIZE = 128Creating Induction Training Data
The key to surfacing induction head behavior is training on data that rewards pattern completion. We create sequences where a bigram [A, B] appears near the start, then [A] appears again later. The correct next token after the second [A] is [B]. By making the sequences otherwise random, we ensure the model cannot rely on statistical priors; the only reliable signal is the in-context bigram pattern.
from torch.utils.data import DataLoader, TensorDataset
def make_induction_data(n_samples, vocab_size, seq_len, repeat_prob=0.5):
"""
Create sequences with random tokens, then with probability repeat_prob,
insert a repeated subsequence [A, B] where the second A should predict B.
"""
seqs = np.random.randint(1, vocab_size, size=(n_samples, seq_len))
# For half the sequences, insert a repeated bigram
for i in range(n_samples):
if np.random.rand() < repeat_prob:
# Pick a random bigram
a = np.random.randint(1, vocab_size - 1)
b = np.random.randint(1, vocab_size - 1)
# Insert [A, B] at positions 2-3 and [A] at position 14
seqs[i, 2] = a
seqs[i, 3] = b
seqs[i, 14] = (
a # Second occurrence; model should predict b at position 15
)
return seqs
train_seqs = make_induction_data(10000, VOCAB_SIZE, SEQ_LEN)
# Targets are next-token predictions
train_X = torch.tensor(train_seqs[:, :-1], dtype=torch.long)
train_y = torch.tensor(train_seqs[:, 1:], dtype=torch.long)
train_loader = DataLoader(
TensorDataset(train_X, train_y), batch_size=BATCH_SIZE, shuffle=True
)Minimal Transformer Architecture
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.W_Q = nn.Linear(d_model, d_model, bias=False)
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)
self.W_O = nn.Linear(d_model, d_model, bias=False)
def forward(self, x, return_attn=False):
B, T, C = x.shape
Q = self.W_Q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
K = self.W_K(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
V = self.W_V(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
scores = Q @ K.transpose(-2, -1) / (self.d_k**0.5)
# Causal mask
mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
scores = scores.masked_fill(mask, float("-inf"))
attn = F.softmax(scores, dim=-1)
out = (attn @ V).transpose(1, 2).contiguous().view(B, T, C)
out = self.W_O(out)
if return_attn:
return out, attn
return out
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.attn = MultiHeadAttention(d_model, n_heads)
self.mlp = nn.Sequential(
nn.Linear(d_model, 4 * d_model),
nn.GELU(),
nn.Linear(4 * d_model, d_model),
)
self.ln1 = nn.LayerNorm(d_model)
self.ln2 = nn.LayerNorm(d_model)
def forward(self, x, return_attn=False):
if return_attn:
attn_out, attn_weights = self.attn(self.ln1(x), return_attn=True)
x = x + attn_out
x = x + self.mlp(self.ln2(x))
return x, attn_weights
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
class MiniTransformer(nn.Module):
def __init__(self, vocab_size, d_model, n_heads, n_layers, seq_len):
super().__init__()
self.embed = nn.Embedding(vocab_size, d_model)
self.pos_embed = nn.Embedding(seq_len, d_model)
self.blocks = nn.ModuleList(
[TransformerBlock(d_model, n_heads) for _ in range(n_layers)]
)
self.ln_f = nn.LayerNorm(d_model)
self.unembed = nn.Linear(d_model, vocab_size, bias=False)
def forward(self, x, return_attn_layer=None):
B, T = x.shape
positions = torch.arange(T, device=x.device).unsqueeze(0)
h = self.embed(x) + self.pos_embed(positions)
attn_all = []
for i, block in enumerate(self.blocks):
if return_attn_layer is not None:
h, attn_weights = block(h, return_attn=True)
attn_all.append(attn_weights)
else:
h = block(h)
h = self.ln_f(h)
logits = self.unembed(h)
if return_attn_layer is not None:
return logits, attn_all
return logitsTraining
We train the model with a next-token prediction objective. Because the training data contains the induction pattern, the model has an incentive to develop induction heads to solve the task efficiently. The training is deliberately simple: no learning rate schedule, no dropout, just vanilla AdamW. In practice, the induction circuit forms reliably even under these minimalist conditions, which shows how strongly gradient descent favors this particular solution.
model = MiniTransformer(VOCAB_SIZE, D_MODEL, N_HEADS, N_LAYERS, SEQ_LEN - 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)
loss_fn = nn.CrossEntropyLoss()
losses = []
for epoch in range(50):
epoch_loss = 0.0
for xb, yb in train_loader:
logits = model(xb)
loss = loss_fn(logits.view(-1, VOCAB_SIZE), yb.view(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
losses.append(epoch_loss / len(train_loader))Final training loss: 2.7654 Loss reduction: 2.9565 -> 2.7654
The training loss decreases consistently, indicating the model is learning the pattern-completion structure. We can now inspect the attention patterns to check whether induction-like behavior emerges. The test we construct is strict: we embed the induction pattern in a random-token sequence, so the model's only path to the correct prediction is through the induction mechanism.
Inspecting Attention Patterns
model.eval()
# Create a test sequence with a clear induction pattern: [A, B, ...random..., A, ?]
test_token_a = 5
test_token_b = 12
test_seq = np.random.randint(1, VOCAB_SIZE, size=(1, SEQ_LEN - 1))
test_seq[0, 2] = test_token_a
test_seq[0, 3] = test_token_b
test_seq[0, 14] = test_token_a # Trigger for induction
test_tensor = torch.tensor(test_seq, dtype=torch.long)
with torch.no_grad():
logits, attn_all = model(test_tensor, return_attn_layer=True)
# Get predictions at position 14 (after the repeated A)
probs = F.softmax(logits[0, 14], dim=-1).numpy()
predicted_token = int(np.argmax(probs))
predicted_prob = float(probs[predicted_token])
b_prob = float(probs[test_token_b])
# Attention weights: shape [batch, n_heads, seq_len, seq_len]
layer1_attn = attn_all[1][0].numpy() # [n_heads, T, T]
layer0_attn = attn_all[0][0].numpy() # [n_heads, T, T]Token A = 5, Token B = 12 Predicted next token: 7 (prob: 0.106) Probability of correct token B=12: 0.096
The model's prediction at position 14 (after the second occurrence of token A) assigns higher probability for token B, showing induction-head behavior. The fact that the model assigns meaningful probability to the correct token despite the random background tokens in the sequence confirms that it has learned the in-context pattern structure, not just statistical regularities in the training distribution.
Visualizing Attention Patterns
The attention heatmaps show the characteristic induction pattern: head 0 in layer 1 shows a strong bright cell at row 14, column 3, exactly where we would expect an induction head to attend. Row 14 is the current position (the second A), and column 3 is the position of B in the early part of the sequence.
Measuring the Previous Token Head
One head in layer 0 should specialize as the previous token head: it should consistently attend from position to position . This is the sub-diagonal pattern in the attention matrix. A head that implements this pattern reliably across all positions is providing the "ingredient enrichment" that the induction head in layer 1 depends on.
def mean_subdiag_weight(attn_matrix, T=20):
"""Mean attention weight from position i to position i-1 (subdiagonal)."""
weights = []
for i in range(1, T):
weights.append(attn_matrix[i, i - 1])
return float(np.mean(weights))
head0_prev = mean_subdiag_weight(layer0_attn[0])
head1_prev = mean_subdiag_weight(layer0_attn[1])
prev_token_head = 0 if head0_prev > head1_prev else 1Layer 0, Head 0: mean previous-token attention = 0.124 Layer 0, Head 1: mean previous-token attention = 0.019 Previous-token head: Head 0
One of the two layer-0 heads shows substantially higher previous-token attention than the other, consistent with the induction head circuit's requirement for a previous token head. This functional asymmetry between the two heads in the same layer is a direct consequence of the division of labor described earlier: one head specializes in predecessor-enrichment, while the other handles a different computation.
OV Circuit Analysis
Finally, we probe the OV circuit of the putative induction head to verify it acts as a copy circuit. A copy circuit writes a vector to the residual stream that promotes the same token it attended to, effectively forwarding the attended token's identity to the output. The full OV circuit matrix should look approximately diagonal if the head is truly copying.
with torch.no_grad():
block = model.blocks[1]
d_k = D_MODEL // N_HEADS
# Reshape weight matrices to [n_heads, d_k, d_model] and [d_model, n_heads, d_k]
W_V = block.attn.W_V.weight.T.view(N_HEADS, d_k, D_MODEL)
W_O = block.attn.W_O.weight.view(D_MODEL, N_HEADS, d_k).permute(1, 0, 2)
# OV circuit for head 0: maps d_model -> d_model
# W_V: [d_k, d_model], W_O: [d_model, d_k]
W_V_h0 = W_V[0] # [d_k, d_model]
W_O_h0 = W_O[0] # [d_model, d_k]
OV = W_O_h0 @ W_V_h0 # [d_model, d_model]
W_E = model.embed.weight # [vocab_size, d_model]
W_U = model.unembed.weight # [vocab_size, d_model]
# OV effect matrix: ov_effect[i, j] = how much attending to token j promotes logit i
ov_effect = (W_U @ OV @ W_E.T).detach().numpy() # [vocab_size, vocab_size]
self_promotion_strength = float(np.mean(np.diag(ov_effect)))
n_off = VOCAB_SIZE * (VOCAB_SIZE - 1)
off_diag_mean = float(
(np.sum(ov_effect) - np.sum(np.diag(ov_effect))) / n_off
)
copy_ratio = self_promotion_strength / (abs(off_diag_mean) + 1e-8)OV circuit self-promotion (diagonal mean): -0.1304 OV circuit off-diagonal mean: 0.0321 Copy ratio (diagonal / off-diagonal magnitude): -4.06
If the OV circuit is acting as a copy circuit, the diagonal of the OV effect matrix should be much larger than the off-diagonal entries: each attended token promotes its own logit. The copy ratio quantifies this. A ratio well above 1.0 indicates that the head implements a copy operation, which is the functional signature of an induction head. This is a weight-level verification of what the attention pattern analysis suggested: the induction head does not just happen to attend to the right position; it is structurally wired to promote the token it attends to.
OV Circuit Heatmap
We visualize the full OV effect matrix directly. A strong copy circuit should look like a near-diagonal matrix: each token (column) should most strongly promote itself (the matching row). Off-diagonal patterns, where one token promotes a different token, would suggest a more complex transformation rather than pure copying.
Training Loss Curve
Previous Token Head Visualization
The bar chart reveals that one head in layer 0 consistently attends to the immediately preceding position across all sequence positions, while the other head distributes attention more broadly. This is the functional signature of the previous token head, and it provides the compositional substrate that the layer-1 induction head requires. The consistency of the pattern across positions, rather than just at a few positions, is important: it demonstrates that this is a learned strategy rather than a coincidence at a specific sequence location.