Self-Attention
Every token asks a question (query), every token advertises what it holds (key), and each token's new representation is a softmax-weighted mix of what the relevant tokens carry (value). The formula fits on one line; the weights it produces are a computation, not an explanation.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
How does a token decide which other tokens matter, what is the formula, and what can and cannot be read off the attention weights?
A compliance team is using a transformer classifier to flag messages that may contain a regulated claim. Auditors have asked "why did the model flag this one?", and an engineer has proposed showing the attention weights as the explanation. Someone needs to say what those weights are and whether they answer the auditor's question.
Attention weights say how much each token attended to each other token. Show the auditor the tokens with the highest weight for the flagged message and call it the explanation.
The weights are per head per layer, and different heads attend to different things — position, syntax, punctuation, the previous token. Averaging them produces a blur; picking one produces whichever story you chose.
- The weights are per head per layer, and different heads attend to different things — position, syntax, punctuation, the previous token. Averaging them produces a blur; picking one produces whichever story you chose.
- A token with a high attention weight contributes its *value* vector, and that value has already been mixed with other tokens by earlier layers. The weight says where information was pulled from at that step, not which input word caused the decision.
- Interventions expose the gap. Zero out the highest-weight attention and the prediction often does not change; the same output can be produced with quite different attention patterns. Presented to an auditor as "why", this is a story that happens to be available, and the compliance risk is now the explanation rather than the model (Explainability).
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- Predict whether a message contains a claim of a regulated kind; the label is a compliance officer's decision on a sample of past messages, and the action is routing to review.
- The auditor's question is a different target — a faithful account of why this output was produced — and this lesson is about whether attention supplies it.
- One example is a message of a few dozen to a few hundred tokens with a binary label. The claim, when present, is often a short phrase whose meaning depends on words elsewhere — "guaranteed" is a claim near "returns" and not near "delivery".
- The model is a pretrained encoder with a classification head, fine-tuned on a few thousand labelled messages (Fine-Tuning).
- Attention weights are available for every head in every layer; there are hundreds of them per token, and they disagree with each other.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- Each position's vector x is projected three ways: a query q = xW_Q (what am I looking for), a key k = xW_K (what do I hold), and a value v = xW_V (what I contribute if selected). Relevance between positions i and j is the dot product q_i · k_j — large when the query and key point the same way in the projected space.
- The scores for position i over all j are divided by √d_k so their variance does not grow with the projection width, then passed through a softmax so they sum to one. The output for position i is Σ_j softmax(q_i·k_j / √d_k) v_j — a convex combination of the value vectors, weighted by relevance. In matrix form: Attention(Q, K, V) = softmax(QKᵀ / √d_k) V.
- A single head computes one notion of relevance. Multi-head attention runs h independent projections with d_k = d/h each, concatenates the h outputs and projects once more, so one layer can attend by syntax in one head and by lexical similarity in another. In a decoder a causal mask sets scores for future positions to −∞ before the softmax, so a token cannot attend to what comes after it.
Query, key, value
Three linear maps turn each token's vector into a query, a key and a value. The names are the mechanism: a query is what this position is looking for, a key is what each position advertises, and a value is what a position hands over if it is selected. The relevance of j to i is the dot product of i's query with j's key.
The scores are scaled by √d_k — without it a wide projection produces dot products with large variance, the softmax saturates to a near-one-hot, and the gradient through it vanishes. After the softmax, position i's output is a weighted average of the values, and that output goes on to the feed-forward layer as the new representation of position i.
1import numpy as np2 3def attention(X, W_q, W_k, W_v, mask=None):4 Q, K, V = X @ W_q, X @ W_k, X @ W_v # n x d_k each5 scores = Q @ K.T / np.sqrt(K.shape[1]) # n x n — the quadratic term6 if mask is not None:7 scores = np.where(mask, scores, -np.inf) # causal: hide the future8 w = np.exp(scores - scores.max(axis=1, keepdims=True))9 w = w / w.sum(axis=1, keepdims=True) # softmax, rows sum to 110 return w @ V, w # n x d_k, and the weights11 12# Attention(Q, K, V) = softmax(Q Kᵀ / sqrt(d_k)) V13# multi-head: run h of these with d_k = d / h, concatenate, project onceThe returned w is what gets drawn as a heat-map. Notice that the output is w @ V, and V at layer twelve is a function of every token via layers one to eleven. The weight tells you which *row of V* was mixed in, not which input word.
A four-token worked example
Take four tokens and a two-dimensional projection so the numbers can be read. Queries and keys are chosen so that "guaranteed" asks a question that "returns" answers loudly, "delivery" answers weakly, and the others barely at all. The softmax turns the scores into weights that sum to one across the row.
The point of the example is the shape: a row of the weight matrix is a distribution over positions, computed from dot products, and the token's new representation is a mix of values in those proportions. The same message with "delivery" instead of "returns" produces a different row, which is how the same word can be a claim in one context and not in another.
The query/key geometry learned during pretraining and fine-tuning still places the tokens that make a phrase a claim near each other in production messages.
holds when Claims are phrased the way the fine-tuning set phrased them; message lengths are similar; the vocabulary of regulated terms has not moved.
breaks when Senders adopt new euphemisms, split a claim across sentences the heads do not bridge, or pad messages so the softmax mass over the claim is diluted.
respond Collect the new phrasings as labelled examples and fine-tune; do not adjust the threshold to recover recall, which trades a phrasing problem for a precision one.
tokens: [we] [guaranteed] [returns] [delivery] d_k = 2
queries q_i keys k_j
we ( 0.2, 0.1) we ( 0.1, 0.2)
guaranteed ( 1.0, 1.0) guaranteed ( 0.2, 0.2)
returns ( 0.3, 0.9) returns ( 1.0, 1.2)
delivery ( 0.4, 0.2) delivery ( 0.6, 0.1)
row for "guaranteed": q = (1.0, 1.0), scaled by sqrt(2)
score(we) = (0.1 + 0.2)/1.41 = 0.21
score(guaranteed) = (0.2 + 0.2)/1.41 = 0.28
score(returns) = (1.0 + 1.2)/1.41 = 1.56
score(delivery) = (0.6 + 0.1)/1.41 = 0.50
softmax -> we 0.14 guaranteed 0.15 returns 0.54 delivery 0.17
new("guaranteed") = 0.14 v_we + 0.15 v_guaranteed + 0.54 v_returns + 0.17 v_delivery
full weight matrix (rows = query token, columns = key token):
we guar ret deliv
we 0.24 0.25 0.29 0.22
guaranteed 0.14 0.15 0.54 0.17
returns 0.15 0.17 0.51 0.17
delivery 0.21 0.22 0.33 0.24Weights are not explanations
The heat-map is seductive because it is a real quantity from inside the model, drawn over the input words. But the auditor asked a counterfactual question — what made the model flag this — and the weights answer a mechanistic one: where, at this layer, information was pulled from. Those coincide only when the values are still close to the input words, which after a few layers they are not.
The test is intervention. If removing the tokens an explanation names does not move the prediction, the explanation is not describing the decision. Attention heat-maps fail that test often enough that they should be presented as a diagnostic and not as an account; attribution methods that pass it are the honest answer, and a human-in-the-loop decision is the most honest one (Explainability).
| Trigger | Symptom | Cause | Response |
|---|---|---|---|
| Auditor asks why a message was flagged | A heat-map highlighting "guaranteed" is filed as the explanation | Weights show where information flowed at one layer, not which input caused the output; ablating the word often does not change the decision | File an ablation- or gradient-based attribution with its faithfulness check, and the reviewer's decision as the explainable one |
| Model update changes attention patterns | All historical "explanations" now look different for the same decisions | Attention is one of many internal routes to the same output; it is not stable under retraining | Re-run the faithfulness check; treat explanation drift as a monitoring signal, not as a change in model reasoning |
| Sender pads a message with neutral text | Score falls below threshold with the claim intact | Softmax rows sum to one; more positions dilute the mass on the claim | Chunk long messages and score chunks; add padded examples to the robustness test set (Robustness Testing) |
How to build it
Most important first.
- Use attention weights for what they are: a diagnostic of where information flowed inside one layer. They are useful for debugging a model that ignores its input, and not as a faithful account of the decision.
- For the auditor, give an explanation method whose faithfulness can be tested — an input-ablation or gradient-based attribution checked by removing the attributed tokens and confirming the prediction moves (Explainability, Attribution Is Not Causality).
- Better still, design the decision so the explanation is the process: route the flagged message with the model's score and the matched phrase from a rule set to a human reviewer, and let the human's decision be the explainable one (Human Oversight).
- Keep the attention math in mind when sizing: scores are n×n per head per layer, which is the quadratic cost and the memory ceiling (Transformer Fundamentals).
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- For the model: precision and recall at the routing threshold on held-out messages, sliced by claim type — the numbers that map to the decision.
- For any explanation offered: a faithfulness check — the fraction of predictions that change when the tokens the explanation names are removed. An explanation that survives its own ablation is not explaining.
- Do not measure "explanation quality" by whether a human finds it plausible. Plausible and faithful are different properties, and attention weights are reliably plausible.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- The relevance the model learned between tokens — encoded in W_Q and W_K — still matches how claims are phrased in production; a new phrasing convention is a distribution shift attention cannot signal on its own.
- The production message length stays within the range the attention was trained over; dilution across an unusually long message shifts the softmax mass whether or not the claim is present.
- Whatever is shown to the auditor as an explanation has a faithfulness check that is re-run when the model changes, so the explanation and the model do not drift apart.
- Offline: for a sample of flagged messages, ablate the top-attended tokens and record how often the decision flips; compare against an attribution method on the same sample. The one whose ablation flips decisions is the one that is explaining.
- Online: log the model's score and the reviewer's decision together, and monitor the disagreement rate per claim type; that is the number that says whether the model still understands the phrasing (Prediction Logging).
- Over time: keep a fixed set of messages with known claims and re-run the faithfulness check after each model update; an explanation method that used to pass and now fails is a finding.
What can go wrong
- The auditors accept attention heat-maps; a later model update changes the attention pattern without changing the decisions; the "explanations" all change and nobody can say what changed in the model.
- A message contains a claim spread across two sentences that a single head does not bridge; the classifier misses it and the heat-map shows attention firmly on the wrong sentence — the diagnostic worked, the decision did not.
- Adversarial phrasing: a sender learns that padding a message with neutral text spreads the attention out and lowers the score. The softmax normalises to one, so dilution is a real attack surface (Adversarial Inputs).
- Faithful attribution methods cost several forward passes per explanation and produce fuzzier pictures than a heat-map; the heat-map is cheap and pretty and answers a different question.
- Multi-head attention buys multiple simultaneous notions of relevance at the price of h separate score matrices per layer, all of them memory at long context.
- Routing every flag to a human reviewer gives an explainable decision and a review queue whose cost scales with the false-positive rate (Threshold Selection).
- "The model attended to the word 'guaranteed', so that is why it flagged the message." The weight says information from that position was mixed into this one at that layer. Remove the word and the model may flag it anyway — the value vectors of neighbouring positions already carried it.
- "Attention is interpretable by design." Attention is inspectable by design. Interpretable would mean the inspection predicts the model's behaviour under intervention, and it often does not.
- "Averaging attention across heads gives the overall attention." Heads compute different relevance functions; their average is a quantity no part of the model computes.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- SIMPLIFIEDThe worked 4-token example uses a tiny projection width and hand-picked numbers so the softmax is readable; real heads have d_k of 64 or 128, and the scores are for the shape of the computation, not measurements from any model.
- GENERALThe formula and the argument that weights are not explanations apply to every attention-based model regardless of task; what differs across tasks is whether a faithful attribution method exists that an auditor will accept.
- CONTESTEDA serious position holds that attention weights, read carefully — per head, with the value norms taken into account, and aggregated through the layers by attention rollout or flow — are a genuinely informative account of the model's computation, and that dismissing them wholesale discards a useful diagnostic. That is fair as a diagnostic for engineers; the claim contested here is narrower — that a raw heat-map is a faithful explanation of a single decision to an auditor, which the ablation evidence does not support.
Where the depth lives
This domain teaches the model and hands the rest off by name.