Attention Head Specialisation and Probing

A transformer with 12 layers and 12 heads per layer has 144 separate attention mechanisms running at once, each free to look at the sequence however it wants. Nobody assigns any head a job — the training objective only ever sees the final output. Yet when researchers open these models up, they consistently find heads with clean, recognisable specialities: a head that attends from a pronoun back to its noun, a head that tracks the previous token, a head that finds a matching bracket. That division of labour is not designed in — it emerges from gradient descent alone — and understanding why it happens, and how to detect it, is the subject of head specialisation and probing.

The analogy: a newsroom with no assigned beats

Imagine hiring 144 junior reporters and telling them only "help write the best possible final story," with no beat assignments. Left alone, some naturally gravitate toward checking facts, some toward following the money, some toward quoting sources — not because an editor assigned those roles, but because specialising made the group's combined output better than 144 generalists duplicating the same work. Attention heads specialise for the same reason: redundant heads waste capacity, so training pressure nudges different heads toward different, complementary jobs.

What a head actually computes, and how to read it

Each head produces an $n \times n$ attention matrix, where entry $(i,j)$ is how much position $i$ attends to position $j$ (see Scaled Dot-Product Attention and the √d Factor). A head is "specialised" if that matrix follows a consistent, describable pattern across many inputs — always the previous token, or always the syntactic head of the current word — rather than looking diffuse.

Probing checks this systematically. A probe is a small, separately-trained classifier (often just logistic regression) that takes a head's attention weights or output vectors and tries to predict some external label — part of speech, dependency relation, coreference. High probe accuracy from one head's outputs means that head carries the information in an extractable, roughly linear form.

$$\hat{y}_i = \sigma\!\left(w^\top a_i + b\right)$$

In words: take the head's attention weight vector $a_i$ for token $i$, apply a learned weighting $w$ and bias $b$, squash through a sigmoid $\sigma$, and check whether the result predicts label $\hat{y}_i$. High accuracy means the head's output is a strong, decodable signal — not proof the model "uses" it downstream, only that it is present and linearly readable.

Worked example 1: finding a positional head by hand

Take a 6-token sentence and a head's attention matrix (rows = query position, columns = key position, each row summing to 1). Row 4 is $[0.02, 0.03, 0.05, 0.85, 0.03, 0.02]$ — 85% of position 4's mass lands on position 3, one step back. Rows 2, 3, 5, 6 show the same pattern: the largest weight, around 0.8–0.9, sits one column left of the diagonal. That consistent off-by-one pattern, across every row and sentence checked, is the signature of a previous-token head — no probe needed, it's visible directly in the matrix.

Worked example 2: a probe result table

A researcher trains a linear probe on each of 12 heads in one layer to predict "is this token the subject of the sentence's main verb," using 500 labelled sentences and reporting test accuracy against a 50% random-guess baseline:

Head	Probe accuracy
Head 3	91%
Head 7	88%
all other 10 heads	52–58%

Two heads stand far above baseline — strong evidence those two specialise in subject-tracking, while the other ten are doing something unrelated to this label.

Drag the slope and midpoint above and watch the S-curve steepen or shift — that is exactly the sigmoid a probe applies to a head's raw signal; a steep, well-separated curve means the head's output cleanly separates the label, a flat one means it barely helps.

What this means in practice

Probing is how interpretability researchers locate the specific heads responsible for a capability — useful for pruning redundant heads to speed up inference, diagnosing failures, and auditing whether a model learned a shortcut instead of the intended reasoning. In quant contexts, the same technique checks whether a sequence model is attending to genuine price-relevant signal or a spurious positional artefact.

Practice

1. You probe a head and get 95% accuracy predicting sentence length. Is this head "specialised for length," or could this be a confound? What would you check?
2. Sketch what a previous-token head's attention matrix looks like for an 8-token sequence, and contrast it with a head that always attends to position 1 (a "beginning-of-sequence" head).