Investigating neurons and dimensions

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("gpt2-medium")
model = AutoModelForCausalLM.from_pretrained("gpt2-medium").eval()

class Recorder:
    """Context manager so hooks cannot outlive their scope."""
    def __init__(self, modules):
        self.modules, self.acts, self.handles = modules, {}, []

    def __enter__(self):
        for name, module in self.modules.items():
            self.handles.append(module.register_forward_hook(self._make(name)))
        return self

    def _make(self, name):
        def hook(module, inputs, output):
            t = output[0] if isinstance(output, tuple) else output
            self.acts[name] = t.detach().cpu()      # detach or you leak
        return hook

    def __exit__(self, *exc):
        for h in self.handles:
            h.remove()

targets = {f"mlp{i}": model.transformer.h[i].mlp for i in range(6)}
with Recorder(targets) as rec:
    out = model(**tok("The cat sat on the mat", return_tensors="pt"),
                output_hidden_states=True)

# hidden_states is the RESIDUAL STREAM after each block, plus the embeddings
# at index 0. So len(hidden_states) == n_layers + 1.
print(len(out.hidden_states), out.hidden_states[0].shape)

# A hook on .mlp captures only the MLP's CONTRIBUTION, before it is added to
# the residual stream. These are different tensors and confusing them is the
# most common misreading in this area:
#   hidden_states[i+1] == hidden_states[i] + attn_out_i + mlp_out_i

# And hidden_states[-1] for a causal LM has already had the final layernorm
# applied — so it is not simply the last block's output.

import torch

def maximise_neuron(model, layer, neuron, seq_len=8, steps=300, lr=0.1):
    """Optimise a continuous embedding to drive one unit high."""
    emb_layer = model.transformer.wte
    dim = emb_layer.weight.size(1)
    x = torch.randn(1, seq_len, dim, requires_grad=True) * 0.02
    opt = torch.optim.Adam([x], lr=lr)

    captured = {}
    h = model.transformer.h[layer].mlp.register_forward_hook(
        lambda m, i, o: captured.__setitem__("a", o))

    for _ in range(steps):
        model(inputs_embeds=x)
        loss = -captured["a"][0, :, neuron].mean()      # ascent = descend -a
        opt.zero_grad(); loss.backward(); opt.step()
    h.remove()

    # Which real tokens is the optimised vector nearest to? Usually the
    # similarities are all low, which is the point: the optimum is not a
    # plausible input, so the result is hard to interpret linguistically.
    W = emb_layer.weight.detach()
    Wn = W / W.norm(dim=1, keepdim=True)
    xn = x.detach()[0] / x.detach()[0].norm(dim=1, keepdim=True)
    sims = xn @ Wn.T
    return sims.max(dim=1)          # values and token ids

import torch, numpy as np
from collections import defaultdict

@torch.no_grad()
def scan_corpus(model, tok, texts, layer, neuron):
    """Activation of one unit at every token across a corpus."""
    records = []
    captured = {}
    h = model.transformer.h[layer].mlp.register_forward_hook(
        lambda m, i, o: captured.__setitem__("a", o.detach()))

    for text in texts:
        ids = tok(text, return_tensors="pt", truncation=True, max_length=256)
        model(**ids)
        acts = captured["a"][0, :, neuron]
        pieces = tok.convert_ids_to_tokens(ids["input_ids"][0])
        for i, (piece, a) in enumerate(zip(pieces, acts.tolist())):
            context = "".join(pieces[max(0, i - 6): i + 1])
            records.append((a, piece, context))
    h.remove()
    return records

recs = scan_corpus(model, tok, corpus, layer=8, neuron=1423)
vals = np.array([r[0] for r in recs])

print(f"mean {vals.mean():+.3f}  std {vals.std():.3f}  "
      f"p99 {np.percentile(vals, 99):+.3f}  max {vals.max():+.3f}")
print(f"fraction above zero: {(vals > 0).mean():.2%}")

recs.sort(reverse=True)
for a, piece, ctx in recs[:15]:
    print(f"{a:+7.3f}  {piece!r:15} ...{ctx}")

# Also print a slice from the middle of the distribution. If those look like
# the top examples, the unit is not selective — it is just generally active.
mid = len(recs) // 2
for a, piece, ctx in recs[mid: mid + 5]:
    print(f"{a:+7.3f}  {piece!r:15} ...{ctx}")

import torch, numpy as np

@torch.no_grad()
def word_activation(model, tok, sentence, word, layer, neuron, how="last"):
    ids = tok(sentence, return_tensors="pt")
    captured = {}
    h = model.transformer.h[layer].mlp.register_forward_hook(
        lambda m, i, o: captured.__setitem__("a", o.detach()))
    model(**ids)
    h.remove()

    pieces = tok.convert_ids_to_tokens(ids["input_ids"][0])
    # Find the contiguous span whose concatenation matches the word.
    span = find_span(pieces, word, tok)      # returns (start, end)
    a = captured["a"][0, span[0]:span[1], neuron]

    return {"first": a[0], "last": a[-1], "mean": a.mean(), "max": a.max()}[how]

# The control that decides whether a result is real:
def by_token_length(model, tok, items, layer, neuron):
    groups = {}
    for sentence, word, label in items:
        n = len(tok.tokenize(" " + word))
        a = word_activation(model, tok, sentence, word, layer, neuron)
        groups.setdefault(n, []).append((label, float(a)))

    for n, rows in sorted(groups.items()):
        pos = [a for lab, a in rows if lab == 1]
        neg = [a for lab, a in rows if lab == 0]
        if pos and neg:
            print(f"{n} tokens: n={len(rows):3}  "
                  f"diff {np.mean(pos) - np.mean(neg):+.3f}")
# A consistent difference WITHIN each length group is evidence.
# An effect that only appears when lengths are pooled is a confound.

import numpy as np
from scipy import stats
from statsmodels.stats.multitest import multipletests

def find_tuned_neurons(acts, labels, alpha=0.05):
    """acts: [n_stimuli, n_neurons]; labels: 0/1 per stimulus."""
    a, b = acts[labels == 1], acts[labels == 0]

    t, p = stats.ttest_ind(a, b, axis=0, equal_var=False)   # Welch

    # Effect size matters more than p when n is large.
    pooled = np.sqrt((a.var(0, ddof=1) + b.var(0, ddof=1)) / 2)
    d = (a.mean(0) - b.mean(0)) / np.maximum(pooled, 1e-9)   # Cohen's d

    reject, p_adj, *_ = multipletests(p, alpha=alpha, method="fdr_bh")

    hits = np.where(reject & (np.abs(d) > 0.8))[0]           # large effect only
    order = hits[np.argsort(-np.abs(d[hits]))]
    return [(int(i), float(t[i]), float(p_adj[i]), float(d[i])) for i in order]

for idx, t_, p_, d_ in find_tuned_neurons(activations, labels)[:10]:
    print(f"neuron {idx:5}  t={t_:+6.2f}  p_adj={p_:.2e}  d={d_:+.2f}")

# Then the step that actually matters: re-test the survivors on stimuli held
# out from the search. Most will not replicate, and that is the useful result.

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score, StratifiedKFold
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

def probe(acts, labels, C=1.0, folds=5):
    pipe = make_pipeline(
        StandardScaler(),
        LogisticRegression(C=C, max_iter=2000, penalty="l2"),
    )
    cv = StratifiedKFold(folds, shuffle=True, random_state=0)
    return cross_val_score(pipe, acts, labels, cv=cv).mean()

def permutation_baseline(acts, labels, n=50):
    rng = np.random.default_rng(0)
    return [probe(acts, rng.permutation(labels)) for _ in range(n)]

real = probe(activations, labels)
null = permutation_baseline(activations, labels)
print(f"probe {real:.3f}   permuted {np.mean(null):.3f} "
      f"(95th pct {np.percentile(null, 95):.3f})")
# Only the margin over that 95th percentile is evidence.

# Laminar profile: where does the information first become linearly available?
for L in range(n_layers):
    print(L, round(probe(acts_by_layer[L], labels), 3))
# A sharp rise at some depth is the interesting result — it localises where
# the feature is computed, which is what makes this worth doing per layer.