Writing · · 5 min read
Deviating Attention
Attention heads specialize, and they do not all specialize on the same thing. What the division of labor looks like, and what breaks when it is ignored.
#attention#transformers#interpretability#machine-learning
A transformer attention is not monolithic. It is divided into heads, and the heads divide the work. Some attend to the previous token. Some attend to punctuation. Some attend to syntax. Some appear to do nothing, and then turn out to matter.
The division is emergent. The model trains on language, and the gradient creates specialists because specialization is what the loss rewards. The structure below starts from a single monolithic attention operation, derives the multi-head split, builds head specializations step by step, and ends with the measurements that reveal what pruning, distillation, and averaging destroy.
The naive attention: one operation for everything
The simplest attention is a single head that computes one set of weights per position. Every query attends to the same keys with one learned projection. The implementation is the standard scaled dot-product.
def attention_single(q, k, v):
# q, k, v: sequence of vectors, shape (n, d).
# Output: weighted sum of values per query.
n = len(q)
d = len(q[0])
scale = d ** -0.5
out = []
for i in range(n):
scores = [sum(q[i][t] * k[j][t] for t in range(d)) * scale
for j in range(n)]
weights = softmax(scores)
out_vec = [0.0] * d
for j in range(n):
for t in range(d):
out_vec[t] += weights[j] * v[j][t]
out.append(out_vec)
return out
def softmax(xs):
m = max(xs)
exps = [2.718281828 ** (x - m) for x in xs]
total = sum(exps)
return [e / total for e in exps]The naive operation has one bottleneck: one projection must serve every kind of relationship. Positional adjacency, syntactic dependency, and lexical matching all compete for the same query and key spaces. The model needs separate subspaces so each relationship can be learned without interference.
The multi-head split: many small attentions
The fix splits the model dimension into heads. Each head receives its own query, key, and value projections and computes a separate attention. The outputs concatenate back into the model dimension.
def multi_head(q, k, v, heads):
# Split each vector into `heads` chunks.
# Each head computes attention on its chunk.
d = len(q[0])
chunk = d // heads
head_outs = []
for h in range(heads):
qh = [row[h * chunk:(h + 1) * chunk] for row in q]
kh = [row[h * chunk:(h + 1) * chunk] for row in k]
vh = [row[h * chunk:(h + 1) * chunk] for row in v]
head_outs.append(attention_single(qh, kh, vh))
# Concatenate head outputs per position.
out = []
for i in range(len(q)):
row = []
for h in range(heads):
row.extend(head_outs[h][i])
out.append(row)
return out, head_outsThe split multiplies the number of attention patterns without changing total compute much. Each head still costs where is the head dimension. Total cost stays near because . The gain is representational: each head can learn a different relation in its own subspace.
The model dimension divides cleanly in most architectures. A model with 24 layers and 32 heads per layer carries 768 heads. A model with 16 heads per layer still carries hundreds of specialists. The population size is the point: the model is a society of operations, not one operation repeated.
What heads do: the specializations
The empirical literature finds a recurring set of behaviors. Each specialization solves a different bottleneck of single-head attention.
Positional heads attend to a fixed offset: the previous token, the next token, the token two positions back. The implementation is a sharp pattern on the diagonal or one step off it.
def positional_pattern(n, offset):
pattern = [[0.0] * n for _ in range(n)]
for i in range(n):
j = i + offset
if 0 <= j < n:
pattern[i][j] = 1.0
return patternPositional heads are the simplest specialists. They implement local structure that the single operation would have to share with everything else.
Syntactic heads attend to grammatically related tokens: the subject of a verb, the antecedent of a pronoun. The pattern is content-based rather than position-based. The same offset gives different targets in different sentences.
def syntactic_example():
# Tokens: "The cat that chased the mouse slept."
# A subject head maps "slept" to "cat" (distance 5)
# and "chased" to "cat" (distance 3).
n = 8
targets = {"slept": "cat", "chased": "cat"}
return targetsSyntactic heads are harder to characterize and more important. Their function is conditional on the sentence. The pattern cannot be read from a single weight matrix in isolation.
Induction heads attend to previous occurrences of the current token and copy what followed. The two-layer circuit is the mechanism behind in-context learning: a prior-token head marks where the token last appeared, an induction head copies the following content.
def induction_lookup(context, current):
# Return what followed the previous occurrence of `current`.
last_seen = -1
for i, tok in enumerate(context):
if tok == current:
last_seen = i
if last_seen + 1 < len(context):
return context[last_seen + 1]
return None
ctx = ["the", "cat", "sat", "the", "cat", "slept"]
print(induction_lookup(ctx, "cat")) # "slept"Induction heads run in pairs across layers. Layer provides the previous-token signal. Layer uses the signal to copy forward. The circuit requires at least two layers and cannot be assigned to a single head.
Punctuation heads attend to commas, periods, and structural markers. The pattern looks useless and is not. Punctuation provides a reference frame for the sentence, which other heads reuse.
def punctuation_anchors(tokens):
return [i for i, t in enumerate(tokens) if t in {",", ".", ";", ":"}]Orphan heads lack a clear function. Some are near-uniform, some are noise. Ablating them sometimes does nothing. Sometimes it destroys performance. The distinction is not visible from the attention pattern alone. The output must be inspected: a head can attend diffusely yet contribute a precise value transformation.
Why the distribution of work is uneven
The total compute divides among heads. The division is not equal. A few heads carry much of the load. Many carry little. The loss function does not require equal contributions, and training does not produce them.
def head_importance(head_outputs, baseline_loss, layer):
# Ablate one head at a time and measure the loss increase.
scores = {}
for h, out in enumerate(head_outputs):
zeroed = zero_head(layer, h)
scores[h] = loss(layer, zeroed) - baseline_loss
return scoresThe importance distribution is long-tailed. A small fraction of heads account for a large fraction of total effect. The tail matters because importance is input-conditional. A head that averages unimportant can be critical on rare inputs. Average-case ranking misses the tail.
def tail_importance(head, inputs):
# Measure worst-case rather than average effect.
worst = 0.0
for x in inputs:
effect = ablation_delta(head, x)
worst = max(worst, effect)
return worstWhat breaks when specialization is ignored
Pruning. Removing heads by importance ranking compresses the model up to a point. Past that point, the removed heads are the ones whose contribution is conditional on rare inputs. The pruned model passes average benchmarks and fails on the tail.
Interpretability. A head that attends to a syntactic pattern in one example attends to something else in another. The interpretation is conditional. The statement "head 14 in layer 7 is the subject head" is a summary over a distribution, not a fact about the head.
Distillation and averaging. Treating attention as one mechanism — distilling the mean pattern, averaging head outputs, assuming interchangeability — loses the division of labor. The information lives in the division. Averaging smooths the specialists into a generalist that performs worse than the population.
Robustness. Adversarial inputs hijack attention. A spurious pattern that triggers a strong head dominates the representation. Other heads cannot always override it. Defenses that inspect only the mean attention miss the trigger.
What to measure
Head importance needs three measurements. Ablation delta measures causal effect: loss change when the head output is zeroed. Distribution across inputs measures conditionality: importance on average versus worst case. Circuit membership measures coordination: whether the head acts alone or as part of a multi-layer circuit such as induction.
def circuit_check(layer, head, test_inputs):
base = run_with_head(layer, head, test_inputs)
ablated = run_without_head(layer, head, test_inputs)
return [b - a for b, a in zip(base, ablated)]Redundancy is the fourth measurement. Pairs and groups of heads may back each other up. Single-head ablation underestimates the effect of removing a pair. Group ablations over random subsets expose the redundancy structure.
A checklist
- Report ablation deltas per head, not attention entropy. Patterns alone mislead.
- Measure worst-case importance across inputs alongside average importance.
- Identify circuits before pruning. Induction heads are pairs, not singles.
- Test group ablations to expose redundancy before removing heads.
- Inspect value and output projections, not only attention weights. Orphan heads transform values.
- Distill per-head patterns, not the mean pattern. The division is the mechanism.
The head is not the neuron. The model is not one attention repeated. It is a population of operations that learned together to do the work, with uneven load, conditional roles, and partial redundancy. The population matters, and any compression or interpretation that ignores the population pays for the omission.