Identifying circuits and components

import torch, numpy as np

@torch.no_grad()
def attention_patterns(model, tok, text):
    """Attention weights for every layer and head: [layers][B, heads, T, T]."""
    ids = tok(text, return_tensors="pt")
    out = model(**ids, output_attentions=True)
    pieces = tok.convert_ids_to_tokens(ids["input_ids"][0])
    return out.attentions, pieces

def classify_heads(attentions):
    """Cheap structural summaries that identify well-known head types."""
    rows = []
    for L, A in enumerate(attentions):
        a = A[0]                                  # [heads, T, T]
        T = a.size(-1)
        for H in range(a.size(0)):
            m = a[H]
            rows.append({
                "layer": L, "head": H,
                # Attends to itself: often a no-op / null head.
                "self": float(m.diagonal().mean()),
                # Attends one position back: 'previous token' head.
                "prev": float(m.diagonal(-1).mean()) if T > 1 else 0.0,
                # Attends to position 0: attention sink, extremely common.
                "first": float(m[:, 0].mean()),
                # Entropy: low = focused on few positions, high = diffuse.
                "entropy": float(-(m * m.clamp_min(1e-12).log()).sum(-1).mean()),
            })
    return rows

# Also worth profiling: the norm of each head's OUTPUT contribution, which is
# what actually reaches the residual stream. A head with a sharp pattern but a
# tiny output norm is not doing much, and the pattern alone would mislead you.
def head_output_norms(model, tok, text, layer):
    captured = {}
    h = model.transformer.h[layer].attn.register_forward_hook(
        lambda m, i, o: captured.__setitem__("o", (o[0] if isinstance(o, tuple) else o).detach()))
    model(**tok(text, return_tensors="pt"))
    h.remove()
    dim = model.config.n_embd
    n_heads = model.config.n_head
    per_head = captured["o"][0].view(-1, n_heads, dim // n_heads)
    return per_head.norm(dim=-1).mean(0)          # one number per head

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score

def sparse_probe(acts, labels, C=0.01):
    """L1 penalty drives most coefficients to exactly zero."""
    clf = LogisticRegression(penalty="l1", C=C, solver="liblinear", max_iter=5000)
    clf.fit(acts, labels)
    nz = np.flatnonzero(clf.coef_[0])
    return nz, clf.coef_[0][nz], cross_val_score(clf, acts, labels, cv=5).mean()

def stability_selection(acts, labels, C=0.01, n_boot=100, frac=0.7):
    """Which neurons are selected consistently, not just once?"""
    rng = np.random.default_rng(0)
    counts = np.zeros(acts.shape[1])
    n = int(frac * len(labels))
    for _ in range(n_boot):
        idx = rng.choice(len(labels), n, replace=False)
        nz, _, _ = sparse_probe(acts[idx], labels[idx], C)
        counts[nz] += 1
    return counts / n_boot

freq = stability_selection(activations, labels)
stable = np.flatnonzero(freq > 0.8)
print(f"stable neurons: {stable.tolist()}")

# The accuracy-vs-sparsity curve is the informative output: sweep C and plot
# cross-validated accuracy against the number of non-zero weights. A sharp
# elbow means the information really is concentrated; a gradual slope means it
# is distributed and no small set is privileged.
for C in [0.001, 0.01, 0.1, 1.0]:
    nz, _, acc = sparse_probe(activations, labels, C)
    print(f"C={C:<6} neurons={len(nz):4}  acc={acc:.3f}")

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

class SparseAutoencoder(nn.Module):
    def __init__(self, d_in, expansion=16, l1=1e-3):
        super().__init__()
        d_hidden = d_in * expansion            # overcomplete: the whole point
        self.enc = nn.Linear(d_in, d_hidden)
        self.dec = nn.Linear(d_hidden, d_in, bias=False)
        self.pre_bias = nn.Parameter(torch.zeros(d_in))
        self.l1 = l1
        # Unit-norm decoder columns, or the model shrinks features and inflates
        # decoder weights to cheat the sparsity penalty.
        with torch.no_grad():
            self.dec.weight /= self.dec.weight.norm(dim=0, keepdim=True)

    def forward(self, x):
        centred = x - self.pre_bias
        f = F.relu(self.enc(centred))          # sparse feature activations
        recon = self.dec(f) + self.pre_bias
        return recon, f

    def loss(self, x):
        recon, f = self(x)
        mse = (recon - x).pow(2).sum(-1).mean()
        sparsity = f.abs().sum(-1).mean()
        return mse + self.l1 * sparsity, mse, (f > 0).float().sum(-1).mean()

sae = SparseAutoencoder(d_in=768, expansion=16)
opt = torch.optim.Adam(sae.parameters(), lr=1e-3)

for batch in activation_loader:                # batch: [B, 768]
    total, mse, l0 = sae.loss(batch)
    opt.zero_grad(); total.backward()
    # Re-normalise decoder columns after every step.
    with torch.no_grad():
        sae.dec.weight /= sae.dec.weight.norm(dim=0, keepdim=True)
    opt.step()

print(f"mse {mse.item():.4f}  active features per input {l0.item():.1f}")

# Interpreting a feature: find its top-activating corpus examples, exactly as
# with a neuron — but now they usually share one clear property. Resample any
# feature that never fires, or that dictionary slot is wasted.

import numpy as np
from scipy.linalg import eigh

def ged(A, B, reg=0.05):
    """Directions maximising var(A)/var(B). A, B: [n_samples, n_dims]."""
    Ca = np.cov(A - A.mean(0), rowvar=False)
    Cb = np.cov(B - B.mean(0), rowvar=False)

    # Shrinkage regularisation — required when n_samples < n_dims, and it also
    # stops near-singular Cb producing enormous meaningless eigenvalues.
    Cb = (1 - reg) * Cb + reg * np.trace(Cb) / Cb.shape[0] * np.eye(Cb.shape[0])

    vals, vecs = eigh(Ca, Cb)              # generalised, not standard
    order = np.argsort(vals)[::-1]
    return vals[order], vecs[:, order]     # eigenvalue = variance ratio

vals, vecs = ged(acts_category_a, acts_category_b)
print("top variance ratios:", vals[:5].round(2))

top = vecs[:, 0]
proj_a, proj_b = acts_category_a @ top, acts_category_b @ top
print(f"separation: {abs(proj_a.mean() - proj_b.mean()) / (proj_a.std() + proj_b.std()):.2f}")

# The control that decides whether it is real: fit the direction on a training
# split and evaluate the separation OUT OF SAMPLE. GED will always find a
# separating direction in-sample when dimensions exceed samples.
def ged_out_of_sample(A, B, folds=5):
    scores = []
    for tr_a, te_a, tr_b, te_b in split_folds(A, B, folds):
        _, vecs = ged(tr_a, tr_b)
        w = vecs[:, 0]
        scores.append(abs((te_a @ w).mean() - (te_b @ w).mean()) /
                      ((te_a @ w).std() + (te_b @ w).std() + 1e-9))
    return float(np.mean(scores))

# Run it per layer for a laminar profile of where the category is separable.
for L in range(n_layers):
    print(L, round(ged_out_of_sample(a_by_layer[L], b_by_layer[L]), 3))