attention-analysis

Attention Analysis

Two different objects get called "attention", and you need both:

• The pattern — softmax(QK^T/√d), a [batch, heads, query, key] matrix saying where each head looks. Under eager the attention module hands it back as attn.output[1]; .source reaches the raw softmax inside.
• The per-head output — what each head writes into the residual stream. A slice of the output projection's input.

Patterns tell you about routing; outputs tell you about content.

import torch
import nnsight
from nnsight import TransformersModel

# eager is required — sdpa/flash never materialize the probability matrix
model = TransformersModel("openai-community/gpt2", dispatch=True,
                          attn_implementation="eager")

prompt = "When Mary and John went to the store, John gave a drink to"
n_layers = len(model.transformer.h)
n_heads = model.config.n_head
head_dim = model.transformer.h[0].attn._module.head_dim   # not n_embd // n_heads
tokens = [model.tokenizer.decode([i]) for i in model.tokenizer(prompt).input_ids]

Getting patterns

Every layer in one pass (model.output.attentions also works, but only with output_attentions=True; the loop lets you pick layers):

with torch.no_grad():                     # a trace runs with autograd on
    with model.trace(prompt):
        patterns = nnsight.save([block.attn.output[1] for block in model.transformer.h])

print(len(patterns), patterns[0].shape)   # 12 x [batch, heads, query, key]

If the weights are None, the model was not loaded with attn_implementation="eager" — sdpa and flash_attention_2 never materialize the matrix, and model.output.attentions is the empty tuple rather than an error.

Anchor before you score. Two properties are cheap and both must hold, or you are holding something that is not a probability matrix:

first_layer = patterns[0]
seq = first_layer.shape[-1]
upper = torch.triu(torch.ones(seq, seq, dtype=torch.bool), diagonal=1)

assert torch.allclose(first_layer.sum(-1), torch.ones_like(first_layer.sum(-1)))
assert (first_layer[..., upper] == 0).all()      # causal mask: exactly zero, not small