Mechanistic Interpretability: Looking Inside the Black Box

import datetime, numpy as np, pandas as pd, torch
import matplotlib.pyplot as plt
from importlib.metadata import version as pkg_version
from transformer_lens import HookedTransformer, utils

torch.manual_seed(0); np.random.seed(0)
torch.set_grad_enabled(False)                       # inference, except the two toy trainings

device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
model = HookedTransformer.from_pretrained("gpt2", device=device)
NL, NH, DM, DH = model.cfg.n_layers, model.cfg.n_heads, model.cfg.d_model, model.cfg.d_head

print("executed:", datetime.date.today().isoformat())
print(f"torch {torch.__version__} · transformer_lens {pkg_version('transformer_lens')} · device {device}")
print(f"GPT-2 small: {NL} layers · {NH} heads · d_model {DM} · d_head {DH}")

prompt = "Michael Jordan plays the sport of"
logits, cache = model.run_with_cache(prompt)
print("the model's next word:", repr(model.to_string(logits[0, -1].argmax())))

final = cache["resid_post", NL - 1][0, -1]                  # last layer, last token
hand = model.ln_final(final) @ model.W_U + model.b_U        # unembed it ourselves
print("hand-built logits == the model's own:", torch.allclose(hand, logits[0, -1], atol=1e-3))

ans_dir = model.tokens_to_residual_directions(model.to_single_token(" basketball"))
parts, labels = cache.decompose_resid(layer=-1, pos_slice=-1, apply_ln=True, return_labels=True)
dla = (parts[:, 0, :] @ ans_dir)                            # each component's push on the answer

top = torch.argsort(dla, descending=True)[:6]
print("who writes ' basketball'? top contributors:")
for i in top:
    print(f"  {labels[i]:>12}  {dla[i]:+.2f}")

plt.figure(figsize=(7, 3))
vals = dla.cpu().numpy()
plt.bar(range(len(vals)), vals, color=["#2c5142" if v >= 0 else "#c4521e" for v in vals])
plt.axhline(0, color="#5b6a62", lw=.8)
plt.xticks(range(len(labels)), labels, rotation=90, fontsize=6)
plt.ylabel("push on ' basketball'"); plt.tight_layout(); plt.show()

head_results = cache.stack_head_results(layer=-1, apply_ln=True)   # [n_layers·n_heads, batch, pos, d]
head_dla = head_results[:, 0, -1, :] @ ans_dir
print("split the attention outputs into heads — top writers of ' basketball':")
for i in torch.argsort(head_dla, descending=True)[:6]:
    print(f"  L{int(i // NH)}H{int(i % NH)} : {head_dla[i]:+.2f}")

from datasets import load_dataset
from sklearn.linear_model import LogisticRegression

tr = load_dataset("nyu-mll/glue", "sst2", split="train[:800]")     # train the probe + direction here
te = load_dataset("nyu-mll/glue", "sst2", split="train[800:1200]") # …score it on 400 held-out sentences

def reps_of(rows):
    R = []
    for s in rows["sentence"]:
        _, c = model.run_with_cache(s, names_filter=lambda n: n.endswith("resid_post"))
        R.append(torch.stack([c["resid_post", l][0, -1] for l in range(NL)]).cpu().numpy())
    return np.stack(R), np.array(rows["label"])                    # [n, NL, d], [n]  (1 = positive)

reps, y = reps_of(tr)
reps_te, y_te = reps_of(te)
base = max(y_te.mean(), 1 - y_te.mean())

test_acc = []
for l in range(NL):
    clf = LogisticRegression(max_iter=2000).fit(reps[:, l], y)
    test_acc.append(clf.score(reps_te[:, l], y_te))
print(f"linear-probe HELD-OUT accuracy by layer on SST-2 (majority baseline = {base:.2f}):")
print("  " + "  ".join(f"L{l}:{a:.2f}" for l, a in enumerate(test_acc)))
LAYER = int(np.argmax(test_acc))
print(f"\nsentiment generalizes out of sample — best at layer {LAYER} (test acc {max(test_acc):.2f}).")

def mean_diff(labels):                                        # unit direction = positive mean − negative mean
    d = torch.tensor(reps[labels == 1, LAYER].mean(0) - reps[labels == 0, LAYER].mean(0), dtype=torch.float32)
    return (d / d.norm()).to(device)
direction = mean_diff(y)                                      # the real sentiment direction (train labels)
control = mean_diff(np.random.default_rng(0).permutation(y))  # same recipe, labels shuffled → a control

def hooks(alpha, d):
    def add(resid, hook): return resid + alpha * d.to(resid.device, resid.dtype)
    return [] if alpha == 0 else [(utils.get_act_name("resid_post", LAYER), add)]

pairs = [(" great", " terrible"), (" good", " bad"), (" wonderful", " awful"),
         (" excellent", " horrible"), (" pleasant", " unpleasant")]
pt = torch.tensor([model.to_single_token(p) for p, _ in pairs])
nt = torch.tensor([model.to_single_token(n) for _, n in pairs])
prompts = ["The restaurant downtown was", "The movie was", "The book felt"]   # unseen by the probe

def gap(alpha, d):                                            # mean (pos − neg) gap over 5 pairs × 3 prompts
    tot = 0.0
    for p in prompts:
        with model.hooks(fwd_hooks=hooks(alpha, d)):
            f = model(p)[0, -1].cpu()
        tot += (f[pt] - f[nt]).mean().item()
    return tot / len(prompts)

print("mean positive − negative logit gap (5 pairs × 3 unseen prompts)")
print("    α  |  sentiment dir  |  shuffled control")
for a in [-15, -10, -5, 0, 5, 10, 15]:
    print(f"  {a:+3d}  |     {gap(a, direction):>+6.2f}      |      {gap(a, control):>+6.2f}")

def train_toy(density, m=5, d=2, steps=2500, seeds=6):
    best, imp = None, torch.tensor([0.9 ** i for i in range(m)])
    for seed in range(seeds):
        torch.manual_seed(seed)
        with torch.enable_grad():
            W = torch.nn.Parameter(torch.randn(d, m) * 0.1); b = torch.nn.Parameter(torch.zeros(m))
            opt = torch.optim.Adam([W, b], lr=1e-2)
            for _ in range(steps):
                x = (torch.rand(1024, m) < density).float() * torch.rand(1024, m)
                loss = (imp * (x - torch.relu(x @ W.t() @ W + b)) ** 2).sum(1).mean()
                opt.zero_grad(); loss.backward(); opt.step()
        if best is None or float(loss) < best[1]:
            best = (W.detach(), float(loss))
    return best[0]                                              # W : [d, m]

sparse_W = None
for dens, label in [(1.0, "dense  (every feature always on)"), (0.05, "sparse (features rarely on)")]:
    W = train_toy(dens); norms = W.norm(dim=0)
    print(f"{label}: {int((norms ** 2 > 0.05).sum())}/5 features survive   norms = {np.round(norms.numpy(), 2)}")
    if dens == 0.05:
        sparse_W = W

G = sparse_W.t() @ sparse_W                                     # Gram matrix: diagonal = feature energy, off-diagonal = overlap
print("\nGram matrix WᵀW (sparse) — 5 columns can't be orthogonal in 2-D, so the")
print("off-diagonals are nonzero: that's the interference a pentagon packing accepts.")
print(np.round(G.numpy(), 2))

plt.figure(figsize=(3.2, 3.2))
U = sparse_W / sparse_W.norm(dim=0, keepdim=True)
for i in range(U.shape[1]):
    plt.arrow(0, 0, float(U[0, i]), float(U[1, i]), head_width=0.05, length_includes_head=True, color="#2c5142")
    plt.text(float(U[0, i]) * 1.18, float(U[1, i]) * 1.18, f"f{i}", ha="center", va="center", fontsize=9)
plt.axhline(0, color="#c9c9c9", lw=.6); plt.axvline(0, color="#c9c9c9", lw=.6)
plt.xlim(-1.35, 1.35); plt.ylim(-1.35, 1.35); plt.gca().set_aspect("equal")
plt.title("5 sparse features packed in 2-D", fontsize=9); plt.tight_layout(); plt.show()

from datasets import load_dataset

ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train")
ids = []
for line in ds["text"]:
    if len(line) > 200:
        ids += model.to_tokens(line)[0, 1:].tolist()           # drop the BOS token
    if len(ids) >= 300 * 128:
        break
toks = torch.tensor(ids[:300 * 128], device=device).view(-1, 128)   # reused by the SAE in §4
flat = toks.reshape(-1).cpu()

mlp_name = utils.get_act_name("post", 6)
mlp_chunks = []
for i in range(0, toks.shape[0], 16):
    _, c = model.run_with_cache(toks[i:i + 16], names_filter=mlp_name)
    mlp_chunks.append(c[mlp_name].reshape(-1, model.cfg.d_mlp).cpu())
mlp_acts = torch.cat(mlp_chunks)

# reproducible pick: among the highest-variance neurons, take the one whose top
# firings are the LEAST concentrated on a single token — i.e. genuinely diverse
cand = mlp_acts.var(0).topk(20).indices.tolist()
def token_spread(n):
    tp = mlp_acts[:, n].topk(20).indices
    return torch.bincount(flat[tp]).max().item() / len(tp)      # lower = fires on more distinct tokens
neuron = min(cand, key=token_spread)
print(f"MLP neuron #{neuron} — the contexts it fires hardest on, across {toks.shape[0]} lines of WikiText:")
for pos in mlp_acts[:, neuron].topk(8).indices.tolist():
    start = pos - pos % 128
    print(f"  …{model.to_string(flat[max(pos - 6, start):pos + 1])!r}")

LSAE = 6
name = utils.get_act_name("resid_post", LSAE)                  # same `toks` from §3, a different site
chunks = []
for i in range(0, toks.shape[0], 16):
    _, c = model.run_with_cache(toks[i:i + 16], names_filter=name)
    chunks.append(c[name].reshape(-1, DM).cpu())               # keep the pile on the CPU
acts = torch.cat(chunks)
print(f"harvested {acts.shape[0]:,} activation vectors (width {DM}) at resid_post {LSAE}")

class TopKSAE(torch.nn.Module):
    def __init__(self, m, k):
        super().__init__(); self.k = k
        self.b_pre = torch.nn.Parameter(torch.zeros(DM))
        self.enc = torch.nn.Linear(DM, m); self.dec = torch.nn.Linear(m, DM, bias=False)
    def encode(self, x):
        z = torch.relu(self.enc(x - self.b_pre))
        val, idx = z.topk(self.k, -1)
        return torch.zeros_like(z).scatter_(-1, idx, val)      # keep K, zero the rest
    def forward(self, x):
        a = self.encode(x); return self.dec(a) + self.b_pre, a

NTR = 252                                                      # train on these; hold the rest back
mean = acts[:NTR * 128].mean(0)                                # mean from TRAIN only — no eval leakage
Xc = acts - mean                                               # centre once, on the CPU
train_c = Xc[:NTR * 128]                                       # …so the eval below is genuinely out-of-sample
def train_sae(m, k, steps=1500):
    torch.manual_seed(0)
    sae = TopKSAE(m, k).to(device); opt = torch.optim.Adam(sae.parameters(), 1e-3)
    with torch.enable_grad():
        for _ in range(steps):
            x = train_c[torch.randint(0, train_c.shape[0], (2048,))].to(device)   # stream a batch to the GPU
            recon, _ = sae(x); loss = ((recon - x) ** 2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
            with torch.no_grad():
                sae.dec.weight /= sae.dec.weight.norm(dim=0, keepdim=True) + 1e-6
    return sae

mean_d = mean.to(device); eval_toks = toks[NTR:]               # the held-out sequences

def ce(hook=None):
    fwd = [] if hook is None else [(name, hook)]
    with model.hooks(fwd_hooks=fwd):
        return model(eval_toks, return_type="loss").item()

def recon_hook(sae):
    def h(resid, hook):
        r, _ = sae(resid.reshape(-1, DM) - mean_d); return (r + mean_d).reshape(resid.shape)
    return h
def ablate_hook(resid, hook):
    return mean_d.expand(resid.reshape(-1, DM).shape).reshape(resid.shape)

eval_c = Xc[NTR * 128:]                                        # held-out activations
ce_clean, ce_ablate = ce(), ce(ablate_hook)
print(f"loss: clean {ce_clean:.2f}   mean-ablated {ce_ablate:.2f}  (the gap the SAE must refill)\n")

rows, sae32 = [], None                                         # a tiny frontier: 3 widths × 3 sparsities
for mult in [2, 4, 8]:
    for k in [16, 32, 64]:
        sae = train_sae(mult * DM, k)
        lr = (ce_ablate - ce(recon_hook(sae))) / (ce_ablate - ce_clean)
        with torch.no_grad():
            l0 = (sae.encode(eval_c.to(device)) > 0).float().sum(-1).mean().item()
        rows.append({"width": f"{mult}×", "K": k, "L0": round(l0, 1), "loss_recovered": round(lr, 3)})
        if (mult, k) == (4, 32):
            sae32 = sae
frontier = pd.DataFrame(rows)
print(frontier.to_string(index=False))

plt.figure(figsize=(4.4, 3))
for w, g in frontier.groupby("width", sort=False):
    plt.plot(g["L0"], g["loss_recovered"], marker="o", label=w)
plt.xlabel("average L0 (features active per token)"); plt.ylabel("loss recovered")
plt.legend(title="dict width"); plt.tight_layout(); plt.show()

with torch.no_grad():
    codes = torch.cat([sae32.encode(Xc[i:i + 4096].to(device)).cpu()
                       for i in range(0, Xc.shape[0], 4096)])
freq = (codes > 0).float().mean(0)

def contexts(scores, k=5):
    out = []
    for pos in scores.topk(k).indices.tolist():
        start = pos - pos % 128
        out.append(model.to_string(flat[max(pos - 6, start):pos + 1]))
    return out

fid = 0                                                        # chosen by browsing the band: fires on coordinated adjectives ("… and …")

print(f"RAW MLP neuron #{neuron} (from §3) — a grab-bag:")
for cx in contexts(mlp_acts[:, neuron]): print(f"    …{cx!r}")
print(f"\nSAE feature #{fid} — fires on {freq[fid] * 100:.1f}% of tokens, and it coheres:")
for cx in contexts(codes[:, fid]): print(f"    …{cx!r}")

effect = (sae32.dec.weight[:, fid] @ model.W_U).cpu()
print("\nfeature promotes (decoder direction through the unembedding):",
      [model.to_string(t).strip() for t in effect.topk(6).indices])

fvec = sae32.dec.weight[:, fid].to(device)
tt = effect.topk(4).indices                                   # the feature's own top-promoted token ids
labels = [model.to_string(int(t)).strip() for t in tt]
prompt = "The impact was social and"
print(f"steering feature #{fid} on {prompt!r} — output logits of the tokens it promotes:")
print("    α  | " + " | ".join(f"{lb:>13}" for lb in labels))
for a in [-30, 0, 30]:
    def add(resid, hook): return resid + a * fvec.to(resid.device, resid.dtype)
    with model.hooks(fwd_hooks=[] if a == 0 else [(name, add)]):
        f = model(prompt)[0, -1].cpu()
    print(f"  {a:+3d}  | " + " | ".join(f"{f[t]:>+13.1f}" for t in tt))

l, h = 5, 5                                               # L5H5 — a known induction head (caught in §6)
W_QK = model.W_Q[l, h] @ model.W_K[l, h].T                 # [d_model, d_model]
W_OV = model.W_V[l, h] @ model.W_O[l, h]                   # [d_model, d_model]
print(f"head L{l}H{h}")
print(f"  W_QK : {tuple(W_QK.shape)}  rank {torch.linalg.matrix_rank(W_QK.cpu()).item():>2}   (≤ d_head = {DH})")
print(f"  W_OV : {tuple(W_OV.shape)}  rank {torch.linalg.matrix_rank(W_OV.cpu()).item():>2}   (≤ d_head = {DH})")
print("The head's learnable linear maps factor through two rank-64 products — far more")
print("inspectable than a full 768×768 map (softmax, LayerNorm and the inputs still matter).")

torch.manual_seed(0); S = 40
rep = torch.cat([torch.randint(0, model.cfg.d_vocab, (1, S))] * 2, dim=1).to(device)   # [A…][A…]
logits, cache = model.run_with_cache(rep)
nll = -torch.log_softmax(logits[0, :-1], -1)[torch.arange(2 * S - 1, device=device), rep[0, 1:]]
print(f"loss on the first copy  : {nll[:S - 1].mean():.2f}  (it's guessing noise)")
print(f"loss on the second copy : {nll[S:].mean():.2f}  (it's copying — off a cliff)")

dst = torch.arange(S, 2 * S - 1, device=device); src = dst - S + 1
score = torch.zeros(NL, NH)
for l in range(NL):
    score[l] = cache["pattern", l][0][:, dst, src].mean(-1).cpu()

top_idx = torch.topk(score.flatten(), 4).indices
print("induction heads, found by their attention stripe:")
for i in top_idx:
    print(f"  L{int(i // NH)}H{int(i % NH)} : {score.flatten()[i]:.2f}")

li, hi = int(top_idx[0] // NH), int(top_idx[0] % NH)
fig, ax = plt.subplots(1, 2, figsize=(7, 3))
ax[0].imshow(score.numpy(), cmap="Greens", origin="lower", aspect="auto")
ax[0].set_xlabel("head"); ax[0].set_ylabel("layer"); ax[0].set_title("where induction lives")
ax[1].imshow(cache["pattern", li][0, hi].cpu().numpy(), cmap="magma")
ax[1].set_title(f"L{li}H{hi} attention"); ax[1].set_xlabel("attends to"); ax[1].set_ylabel("from")
plt.tight_layout(); plt.show()

copied = rep[0, S + 1:2 * S]
dirs = model.tokens_to_residual_directions(copied)         # [S-1, d_model]
ph = cache.stack_head_results(layer=-1, apply_ln=True)[:, 0, S:2 * S - 1, :]   # [head, pos, d]
dla_heads = (ph * dirs[None]).sum(-1).mean(-1)
ph_top = torch.topk(dla_heads, 4).indices
print("heads that actually write the copied token's logit:")
for i in ph_top:
    print(f"  L{int(i // NH)}H{int(i % NH)} : {dla_heads[i]:+.2f}")
print("consistent with composition: stripe heads locate the match; slightly later heads appear to cash it in.")

def make_ablation(heads):                                      # hooks that zero these heads' output
    def ablate(z, hook):
        z = z.clone()
        for lyr, hd in heads:
            if lyr == hook.layer():
                z[:, :, hd, :] = 0
        return z
    return [(utils.get_act_name("z", lyr), ablate) for lyr in sorted({l for l, _ in heads})]

def second_copy_loss(fwd=None):
    with model.hooks(fwd_hooks=fwd or []):
        lg = model(rep)
    nll = -torch.log_softmax(lg[0, :-1], -1)[torch.arange(2 * S - 1, device=device), rep[0, 1:]]
    return nll[S:].mean().item()

top_heads = [(int(i // NH), int(i % NH)) for i in top_idx[:4]]                          # highest stripe score
low_heads = [(int(i // NH), int(i % NH)) for i in torch.argsort(score.flatten())[:4]]   # lowest stripe score
print(f"second-copy loss — intact               : {second_copy_loss():.2f}")
print(f"                   4 LOW-score heads off : {second_copy_loss(make_ablation(low_heads)):.2f}  (control)")
print(f"                   4 stripe heads off    : {second_copy_loss(make_ablation(top_heads)):.2f}")

import transformer_lens.patching as patching
clean = "When John and Mary went to the store, John gave a drink to"
corrupt = "When John and Mary went to the store, Mary gave a drink to"
ct, rt = model.to_tokens(clean), model.to_tokens(corrupt)
MARY, JOHN = model.to_single_token(" Mary"), model.to_single_token(" John")

assert ct.shape == rt.shape, "clean and corrupt must tokenize to the same length for position-wise patching"
def probs(prompt):
    p = torch.softmax(model(prompt)[0, -1], -1)
    return round(p[MARY].item(), 3), round(p[JOHN].item(), 3)
print("sanity — clean   p(' Mary'), p(' John'):", probs(clean))
print("sanity — corrupt p(' Mary'), p(' John'):", probs(corrupt))

def logit_diff(logits): return logits[0, -1, MARY] - logits[0, -1, JOHN]
_, clean_cache = model.run_with_cache(ct)
clean_ld, corr_ld = logit_diff(model(ct)), logit_diff(model(rt))
print(f"logit diff (Mary − John): clean {clean_ld:+.2f}  →  names swapped {corr_ld:+.2f}")

def metric(logits): return (logit_diff(logits) - corr_ld) / (clean_ld - corr_ld)   # 0 = corrupt, 1 = clean
heat = patching.get_act_patch_resid_pre(model, rt, clean_cache, metric)            # [layer, pos]
best = int(torch.argmax(heat)); bl, bp = divmod(best, heat.shape[1])   # a single argmax, so layer & token match
print(f"best single-spot recovery: {heat[bl, bp]:.2f}  at layer {bl}, token {model.to_str_tokens(corrupt)[bp]!r}")

plt.figure(figsize=(7, 3.2))
plt.imshow(heat.cpu().numpy(), cmap="RdBu_r", origin="lower", aspect="auto", vmin=-1, vmax=1)
plt.colorbar(label="recovery")
xt = [t.replace("Ġ", " ").strip() or "·" for t in model.to_str_tokens(clean)]
plt.xticks(range(len(xt)), xt, rotation=55, ha="right", fontsize=7)
plt.ylabel("layer"); plt.title("where the indirect-object computation lives"); plt.tight_layout(); plt.show()

logits, cache = model.run_with_cache("Michael Jordan plays the sport of")
print("layer → the word the model would commit to if it stopped there")
for l in range(0, NL, 2):
    v = model.ln_final(cache["resid_post", l][0, -1]) @ model.W_U + model.b_U
    print(f"  L{l:>2}  {model.to_string(v.argmax()).strip():<12} (p={torch.softmax(v, -1).max():.2f})")
print(f"  final {model.to_string(logits[0, -1].argmax()).strip()}")

basket = model.to_single_token(" basketball")
p_basket = [torch.softmax(model.ln_final(cache["resid_post", l][0, -1]) @ model.W_U + model.b_U, -1)[basket].item()
            for l in range(NL)]
plt.figure(figsize=(4.2, 3))
plt.plot(range(NL), p_basket, marker="o", color="#2c5142")
plt.xlabel("layer"); plt.ylabel("naive logit-lens p(' basketball')"); plt.tight_layout(); plt.show()