Attention as Content-Based Routing

Five stages. Each one has a distinct failure mode, and each of those failures is documented somewhere in the literature this lesson quotes.

First, project queries, keys and values: map representations into separate spaces for matching and for content transfer. Second, compute compatibility scores: compare each query with the keys it is permitted to see. Third, apply scale and mask: control numeric magnitude and disallow illegal connections. The Transformer paper argues for the scale factor. The mask is where PyTorch and Hugging Face disagree about what True means. Fourth, normalize: convert scores into a distribution-like weighting, commonly with softmax. This is the step that can saturate onto one position — over half of BERT's attention in layers 6–10 lands on a single delimiter. Fifth, aggregate values: produce a content-dependent mixture for each query. This is where weight and effect come apart. A position can take most of the weight and still change almost nothing downstream.

The source of queries and keys determines the architectural role.

Self-attention draws queries, keys and values from one representation set. It mixes information within a sequence or set, can be causal or bidirectional, needs position or structure information, and produces contextual token encodings.

Cross-attention draws queries from one stream and keys or values from another. It is not an abstraction. It is the conditioning interface in a deployed text-to-image system. Latent diffusion was built that way in 2022, and the paper says so plainly: “We turn DMs into more flexible conditional image generators by augmenting their underlying UNet backbone with the cross-attention mechanism”. Q comes from the image latent, K and V from the prompt encoder. The queries are pixels asking what they should become. The keys and values are the words.

That interface was then probed causally rather than admired. Six authors at Google, in the same August, wrote that “we analyze a text-conditioned model in depth and observe that the cross-attention layers are the key to controlling the relation between the spatial layout of the image to each word in the prompt”. Their method, prompt-to-prompt, edits an image by swapping the attention maps and changing nothing else. An intervention, not a visualization.

Neighborhood attention lets queries attend only to selected local or graph-connected elements. It reduces cost and encodes locality or topology, but may miss distant evidence. Shifted windows and graph neighbors are the usual instances.

As key dimension grows, unscaled dot products can have larger variance. Softmax may then become extremely peaked, which leaves small gradients for most alternatives.

The divisor is not a convention someone tuned. The Transformer paper gave the reasoning in 2017: “we suspect that for large values of dk, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by 1/√dk”. A footnote supplies the whole argument in two sentences: “assume that the components of q and k are independent random variables with mean 0 and variance 1. Then their dot product … has mean 0 and variance dk”.

So the divisor tracks the width by construction. Under those assumptions the standard deviation of an unscaled logit is the square root of the key dimension: 8 at the paper's own dk of 64, and 22.63 if a single head ran at the full model width of 512. Dividing by the square root of the key dimension flattens that curve to 1 at every width. It is a design rationale resting on stated assumptions. It is not a promise that every head stays well behaved.

Attention statistics are useful when they are tied to task behavior. The best-documented failure is not hypothetical, and it is a whole model measured head by head. Clark and three co-authors went through all 144 attention heads of BERT-base, 12 layers by 12 heads, and reported this: “For example, over half of BERT’s attention in layers 6-10 focuses on [SEP].” Chance would be about 1/64 for a token occurring twice in a 128-token segment. The second half of the finding is what makes it sharp. The gradient of the loss with respect to that attention goes small from layer 5 onward. The routing that dominates the picture does almost nothing to the output.

An independent group reached the same place. They found the “vertical” attention pattern “is associated predominantly, if not exclusively, with attention to [CLS] and [SEP] tokens”. Disabling single heads raised accuracy by up to 1.2% absolute on MRPC. Disabling a whole layer raised it by 3.2% on RTE. Switching the collapsed routing off made the model better.

Mask errors can produce plausible losses while violating causality or padding rules. They survive review because the conventions genuinely disagree with each other. PyTorch says so in its own reference note: “In scaled_dot_product_attention(), True indicates values to participate in attention. In MultiheadAttention, True indicates values to be masked out (padding).” Same library, opposite polarity. The note then instructs “If migrating from MHA, ensure you invert your boolean mask”. A third convention sits on top of those two. The Hugging Face Transformers glossary documents the keep-polarity, where (for the BertTokenizer) “1 indicates a value that should be attended to, while 0 indicates a padded value” — the same orientation as scaled_dot_product_attention, and therefore inverted relative to MultiheadAttention.

The test that catches this costs a minute. Use a four-element sequence, so that every permitted and forbidden connection can be enumerated by hand. Set distinctive values, chosen so the output reveals exactly which positions were aggregated. Verify each query row separately for causal direction, padding, and batch broadcasting. This is the step where an inverted boolean shows up as a row that saw the future. Test mixed lengths, to confirm that one example's padding cannot affect another. Then repeat in reduced precision, so that masked scores stay numerically safe at deployment precision.

Each head uses its own projections, allowing different compatibility patterns and value channels. Heads can specialize by position, relation, scale, or feature type, but specialization is not guaranteed. More heads also shrink per-head dimension when total width is fixed.

The Transformer's own configuration states the trade: “in this work we employ h = 8 parallel attention layers, or heads. For each of these we use dk = dv = dmodel/h = 64. Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.”

Whether all eight earn their keep was then measured twice, by two groups, with the same answer. Voita and four co-authors report: “For example, on the English-Russian WMT dataset, pruning 38 out of 48 encoder heads results in a drop of only 0.15 BLEU.” A second group observes that “a large proportion of attention heads can be removed at test time without significantly impacting performance, and that some layers can even be reduced to a single head”. In their WMT Transformer, “only 8 (out of 96) heads cause a statistically significant change in performance when they are removed from the model, half of which actually result in a higher BLEU score”.

The asymmetry is the part worth carrying away. That same work finds that reducing the last encoder–decoder attention layer to a single head degrades performance “by at least 13.5 BLEU points”. Self-attention heads are largely fungible after training. The heads that carry conditioning from another stream are not. So ablate head count alongside width, latency, and observed diversity. Run the ablation separately for each attention mode.