Interfering with attention

import torch

def ablate_head_output(head, n_heads, dim):
    """Zero one head's slice of the concatenated attention output.
    Attach to attn.c_proj as a FORWARD PRE-hook, so you edit the input to the
    output projection — that is where heads are still separable."""
    head_dim = dim // n_heads
    lo, hi = head * head_dim, (head + 1) * head_dim

    def pre_hook(module, inputs):
        x = inputs[0].clone()
        x[..., lo:hi] = 0
        return (x,)
    return pre_hook

handle = model.transformer.h[9].attn.c_proj.register_forward_pre_hook(
    ablate_head_output(6, model.config.n_head, model.config.n_embd))
with torch.no_grad():
    ablated = model(**prompt).logits
handle.remove()

# The other kind: force the pattern uniform (causal-masked), leaving the head
# active but non-selective. Requires patching inside the attention computation,
# which for GPT-2 means overriding the module's forward or using a library
# such as TransformerLens that exposes the pattern directly.

def uniform_causal_pattern(T, device):
    m = torch.tril(torch.ones(T, T, device=device))
    return m / m.sum(-1, keepdim=True)

import torch, numpy as np

@torch.no_grad()
def head_importance(model, prompt, answer_id, wrong_id):
    """Effect of ablating each head individually. The output is the map that
    every head-level finding in the literature comes from."""
    n_layers, n_heads = model.config.n_layer, model.config.n_head
    dim = model.config.n_embd

    def logit_diff():
        lg = model(**prompt).logits[0, -1]
        return (lg[answer_id] - lg[wrong_id]).item()

    base = logit_diff()
    effects = np.zeros((n_layers, n_heads))

    for L in range(n_layers):
        for H in range(n_heads):
            h = model.transformer.h[L].attn.c_proj.register_forward_pre_hook(
                ablate_head_output(H, n_heads, dim))
            effects[L, H] = logit_diff() - base
            h.remove()
    return effects

eff = head_importance(model, ioi_prompt, mary_id, john_id)

flat = np.dstack(np.unravel_index(np.argsort(np.abs(eff).ravel())[::-1], eff.shape))[0]
for L, H in flat[:10]:
    print(f"L{L}H{H}  Δlogit_diff {eff[L, H]:+.3f}")

print(f"\nheads with |effect| > 0.5: {(np.abs(eff) > 0.5).sum()} of {eff.size}")
# Typically a small handful. The overwhelming majority of heads are irrelevant
# to any given behaviour, which is what makes circuit analysis tractable at all.

# Negative values are as interesting as positive: a head whose removal IMPROVES
# the correct answer was suppressing it. IOI has exactly such heads.

import torch, numpy as np

@torch.no_grad()
def representational_impact(model, prompt, layer, head, probe_layer=-1):
    """How much does ablating this head reorganise a later layer's structure?"""
    base = model(**prompt, output_hidden_states=True).hidden_states[probe_layer][0].numpy()

    h = model.transformer.h[layer].attn.c_proj.register_forward_pre_hook(
        ablate_head_output(head, model.config.n_head, model.config.n_embd))
    ablated = model(**prompt, output_hidden_states=True).hidden_states[probe_layer][0].numpy()
    h.remove()

    return {
        # 1 - RSA: how much the similarity STRUCTURE changed.
        "structure_change": 1 - rsa(base, ablated),
        # Mean per-token cosine displacement: how far vectors MOVED.
        "mean_displacement": float(np.mean([
            1 - np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12)
            for a, b in zip(base, ablated)])),
    }

# Compare the two views for every head. Heads that score low on logit effect
# and high on structure change are the ones a logit-only sweep would miss.

import torch, numpy as np

@torch.no_grad()
def capture_head_input(model, prompt, layer):
    """The concatenated per-head output, before the output projection."""
    box = {}
    h = model.transformer.h[layer].attn.c_proj.register_forward_pre_hook(
        lambda m, inputs: box.__setitem__("x", inputs[0].detach().clone()))
    model(**prompt)
    h.remove()
    return box["x"]              # [B, T, dim]

def patch_head_slice(donor, head, n_heads, dim, position=None):
    head_dim = dim // n_heads
    lo, hi = head * head_dim, (head + 1) * head_dim

    def pre_hook(module, inputs):
        x = inputs[0].clone()
        if position is None:
            x[..., lo:hi] = donor[..., lo:hi]
        else:
            x[:, position, lo:hi] = donor[:, position, lo:hi]
        return (x,)
    return pre_hook

@torch.no_grad()
def head_patch_map(model, clean, corrupted, answer_id, wrong_id):
    n_layers, n_heads = model.config.n_layer, model.config.n_head
    dim = model.config.n_embd
    ld = lambda p: (lambda lg: (lg[answer_id] - lg[wrong_id]).item())(
        model(**p).logits[0, -1])

    base_clean, base_corr = ld(clean), ld(corrupted)
    out = np.zeros((n_layers, n_heads))

    for L in range(n_layers):
        donor = capture_head_input(model, clean, L)
        for H in range(n_heads):
            h = model.transformer.h[L].attn.c_proj.register_forward_pre_hook(
                patch_head_slice(donor, H, n_heads, dim))
            out[L, H] = (ld(corrupted) - base_corr) / (base_clean - base_corr + 1e-9)
            h.remove()
    return out

# Add a position loop for the full (layer, head, position) tensor. Slower, and
# it is what turns 'these heads matter' into 'this head reads from position i
# and writes to position j' — an actual mechanism.