Modifying the MLP

import torch, numpy as np, torch.nn.functional as F

@torch.no_grad()
def neuron_medians(model, tok, corpus, layer):
    """Per-neuron median activation over a corpus — the neutral value."""
    collected = []
    box = {}
    h = model.transformer.h[layer].mlp.register_forward_hook(
        lambda m, i, o: box.__setitem__("a", o.detach()))
    for text in corpus:
        model(**tok(text, return_tensors="pt", truncation=True, max_length=128))
        collected.append(box["a"][0])
    h.remove()
    return torch.cat(collected, dim=0).median(dim=0).values     # [d_mlp]

def replace_neurons(indices, medians):
    idx = torch.as_tensor(indices, dtype=torch.long)
    def fn(x):
        x = x.clone()
        x[..., idx] = medians[idx]
        return x
    return fn

@torch.no_grad()
def degradation_curve(model, prompt, target_id, layer, order, medians, steps=None):
    steps = steps or [0, 1, 2, 5, 10, 20, 50, 100, 200, 500]
    rows = []
    for n in steps:
        if n == 0:
            loss = -F.log_softmax(model(**prompt).logits[0, -1], -1)[target_id].item()
        else:
            with intervene(model.transformer.h[layer].mlp,
                           replace_neurons(order[:n], medians)):
                loss = -F.log_softmax(model(**prompt).logits[0, -1], -1)[target_id].item()
        rows.append((n, loss))
    return rows

med = neuron_medians(model, tok, corpus, layer=8)

# The comparison that means something: selective ordering vs. random ordering.
selective = degradation_curve(model, prompt, target_id, 8, tuned_order, med)
rng = np.random.default_rng(0)
random_order = rng.permutation(med.numel())
baseline = degradation_curve(model, prompt, target_id, 8, random_order, med)

for (n, a), (_, b) in zip(selective, baseline):
    print(f"n={n:4}  selective {a:6.3f}   random {b:6.3f}   gap {a - b:+.3f}")
# A gap means selectivity predicts causal importance. No gap means the tuned
# neurons were not privileged — a genuinely useful negative result.

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

def select_by_ttest(acts_a, acts_b, alpha=0.05, min_d=0.5):
    t, p = stats.ttest_ind(acts_a, acts_b, axis=0, equal_var=False)
    pooled = np.sqrt((acts_a.var(0, ddof=1) + acts_b.var(0, ddof=1)) / 2)
    d = (acts_a.mean(0) - acts_b.mean(0)) / np.maximum(pooled, 1e-9)
    reject, _, *_ = multipletests(p, alpha=alpha, method="fdr_bh")
    hits = np.flatnonzero(reject & (np.abs(d) > min_d))
    return hits[np.argsort(-np.abs(d[hits]))]          # strongest first

def laminar_lesion_profile(model, tok, layers, stimuli, behaviour_metric):
    rows = []
    for L in layers:
        a = collect_activations(model, tok, stimuli.condition_a, L)
        b = collect_activations(model, tok, stimuli.condition_b, L)
        selected = select_by_ttest(a, b)

        med = neuron_medians(model, tok, corpus, L)
        before = behaviour_metric(model)
        with intervene(model.transformer.h[L].mlp, replace_neurons(selected, med)):
            after = behaviour_metric(model)

        # Matched control: the same NUMBER of random neurons. Without this,
        # any lesion of 200 units looks impressive.
        rng = np.random.default_rng(L)
        control_idx = rng.choice(med.numel(), len(selected), replace=False)
        with intervene(model.transformer.h[L].mlp, replace_neurons(control_idx, med)):
            control = behaviour_metric(model)

        rows.append({"layer": L, "n_selected": len(selected),
                     "effect": before - after,
                     "control_effect": before - control})
    return rows

for r in laminar_lesion_profile(model, tok, range(12), stimuli, metric):
    print(f"layer {r['layer']:2}  n={r['n_selected']:4}  "
          f"effect {r['effect']:+.3f}  control {r['control_effect']:+.3f}")

import torch, numpy as np

def remove_directions(directions):
    """directions: [k, dim]. Projects the activation onto the orthogonal
    complement of their span, removing a distributed feature that no single
    neuron ablation could touch."""
    D = torch.as_tensor(np.atleast_2d(directions), dtype=torch.float32)
    # Orthonormalise: the directions need not be orthogonal to begin with.
    Q, _ = torch.linalg.qr(D.T)                        # [dim, k]
    P = Q @ Q.T                                        # projection onto the span

    def fn(x):
        return x - x @ P.T                             # keep the complement
    return fn

# Test a feature found observationally, with the specificity control.
for name, direction in candidate_features.items():
    with intervene(model.transformer.h[8].mlp, remove_directions(direction)):
        target = target_behaviour(model)
        control = control_behaviour(model)
    print(f"{name:24} target {target:.3f}  control {control:.3f}")
# A real feature: target drops, control barely moves. If both drop, you removed
# something the model needs generally.

# Sweep the dimensionality: how many directions must go before the behaviour
# breaks? A sharp threshold means the feature is genuinely low-dimensional.
for k in [1, 2, 4, 8, 16, 32]:
    with intervene(model.transformer.h[8].mlp, remove_directions(top_directions[:k])):
        print(k, round(target_behaviour(model), 3), round(control_behaviour(model), 3))