Transformer anatomy

Attention from first principles

Scaled dot-product attention works because learned query-key geometry selects addresses, values carry the retrieved content, square-root scaling cancels dimension-driven logit growth, and temperature sets the sharpness of the route.

The Transformer paper compressed the mechanism into one line: compare queries with keys, normalise the scores, and use the resulting weights to sum values.[5] The line is simple enough to memorise and easy enough to misunderstand. “Attention” sounds like a unit deciding what matters. The equation actually describes a content-addressed router. It does not, by itself, supply meaning, awareness, intention or even a guarantee that the selected information causes the model’s final answer.

This paper rebuilds the router from the puzzle above. The aim is not to survey every attention variant. It is to make one mechanism visible: the route is produced by query-key compatibility, while the payload travels through values. Once that split is clear, scaling and temperature stop looking like arbitrary tricks. They become controls over the statistical regime in which routing occurs.

The address is not the message

Early neural sequence models often compressed an entire input into one fixed-length vector. Bahdanau, Cho and Bengio identified this bottleneck and let a decoder softly search over a sequence of source annotations instead.[1] Luong, Pham and Manning later explored global and local forms of the same basic idea.[2] Related memory systems used content similarity to address stored locations.[3][4] The common move was decisive: do not force every future computation to use one undifferentiated summary. Let the current computation ask for a task-dependent mixture.

The Transformer made this move parallel. For each position, learned linear projections produce three vectors. A query represents what that position is looking for. A key represents how each candidate position should be matched. A value represents what each candidate position can contribute if selected. In self-attention, all three are generated from the same incoming token states, but they are not the same representation and need not encode the same distinctions.

Why three projections are necessary

Suppose one vector had to serve as query, key and value. Any feature that made a token easy to find would also be forced into the content it supplied, and any feature useful in the supplied content would alter how the token was found. The model could still learn, but routing and payload would be entangled. Separate learned projections let one incoming state expose different interfaces to the same head.

Consider the word bank in “the bank approved the loan”. A key projection may emphasise that this token is a candidate organisation or grammatical subject. A value projection may carry number, entity and event features needed by another position. A query at it may emphasise the relation “find an earlier singular organisation”. Nothing requires these vectors to resemble a dictionary definition of bank. Training only requires their interaction to reduce the task loss.

This separation also explains why attention can copy, combine and transform rather than merely point. The weighted sum is taken over value vectors, not over token identities. A value can be a learned transform of the token state, and several values can contribute simultaneously. The head therefore returns a new representation in the value space. Calling the operation “selection” is helpful when one weight dominates, but the general operation is conditional mixture.

There is a hidden boundary here. The projections do not create information absent from the incoming states. If two candidate states are identical in every feature visible to the key projection, the router cannot distinguish them. If the value projection discards a distinction, correct routing cannot recover it. Attention reallocates available information; it does not manufacture missing state.

The distinction is easiest to see by asking which intervention changes which quantity. Change a value while holding its key fixed. The attention weights do not move, but the output changes. Change a key while holding its value fixed. The output changes because the route changes. Change the query. Every candidate is rescored. The three vectors may originate from one token representation, yet they occupy different causal roles.

flowchart TD
  A["The address is not the message"] --> B["How geometry becomes a route"]
  B --> C["The two controls that keep<br/>routing honest"]
  C --> D["Build it, break it, and know<br/>what it proves"]
  D -. "re-test" .-> A
  classDef first fill:#0a1424,stroke:#B96D00,stroke-width:2px,color:#e6f3ff
  classDef step fill:#0a1424,stroke:#3f7da0,stroke-width:1px,color:#e6f3ff
  class A first
  class B,C,D step

How geometry becomes a route

Take one query vector q and one key vector k. Their dot product is the sum of coordinate-wise products. Geometrically, it can also be written as the product of their lengths and the cosine of the angle between them. A large positive score can therefore arise because the vectors point in similar directions, because one or both have large norms, or both.

This matters because the model learns the query and key projections. A key can become attractive by aligning with a query direction, increasing its norm, or changing the upstream representation that produces both. The dot product is not a hand-authored semantic metric. It is a trainable interface through which the loss shapes routing.

With several keys, the model produces one score per candidate. These scores are often called logits because they are unnormalised. They can be negative, do not sum to one, and have no standalone probabilistic meaning. Softmax converts their relative differences into positive weights that sum to one. Adding the same constant to every logit changes nothing, because only differences matter.

Softmax makes a relative claim

Take logits [2, 1, 0]. At temperature 1, softmax produces approximately [0.665, 0.245, 0.090]. The first key receives most of the route, but the other two still contribute. Add 100 to every logit and the weights remain identical. Multiply every logit by ten and the distribution becomes almost one-hot. The first transformation changes no relative evidence; the second changes the expressed margin.

This gives a useful mental model. The logits form a local contest among the candidates available to one query. Softmax does not compare a score with a universal threshold. A score of 2 can be dominant against [1, 0] and weak against [8, 7]. An attention weight therefore cannot be interpreted without the competing keys, the mask and the temperature. Routing confidence is contextual to the candidate set.

The weighted sum adds another dependency. A route weight of 0.6 on a value with norm 0.1 contributes less vector magnitude than a weight of 0.2 on a value with norm 10. Direction matters too: two large contributions can cancel. For this reason, a more complete local diagnostic inspects both the scalar weight and the vector contribution wᵢvᵢ. Even that remains local to the layer because later computation can amplify, rotate or erase the result.

The derivative reveals the saturation problem

For softmax weight wᵢ and logit zⱼ, the derivative is (1/τ)wᵢ(δᵢⱼ − wⱼ). The symbols say something concrete. If a route is already near zero or one, the product contains a near-zero term, so small logit changes produce little weight change. Lower temperature adds a 1/τ factor but also makes extreme probabilities more likely. The net training behaviour depends on the whole distribution.

The original square-root argument targets this regime. If dimension makes raw logits large before the model has learned useful compatibility, softmax can become prematurely sharp and gradients through the routing decision can shrink. Scaling keeps initial comparisons in a range where several candidates can still compete. Once trained, a head may deliberately become sharp, but it should do so through learned geometry rather than an accidental change in dimension.

Softmax has three properties that make it useful here. It produces non-negative weights. It makes the weights sum to one. It is differentiable, so training can adjust queries, keys and upstream representations. Yet none of these properties says that a high weight identifies the unique reason for a later prediction. The value may have a small norm, contain information copied from elsewhere, or be cancelled by subsequent computation.

A realistic configured-system example

Consider a code model processing: approved_limit = 5000; temporary_limit = 7000; charge < limit. At the token limit, the query may be shaped to retrieve the currently relevant definition. Keys at approved_limit and temporary_limit advertise different relations. Values carry features that can update the representation at limit.

If both keys collapse to nearly the same vector, the router cannot distinguish the definitions even when their values are perfectly different. If the keys are distinct but the values omit the amount or scope, the route can be correct while the retrieved content is useless. If a causal mask forbids looking ahead, a future redefinition is not a candidate at all. The configured behaviour therefore depends on address quality, payload quality and route entitlement, not on attention weights alone.

The two controls that keep routing honest

Why divide by the square root of key dimension? The answer is not that larger vectors are intrinsically dangerous. The issue is what happens to a sum of many roughly independent products.

Assume each component of q and k has mean zero and variance one, and components are independent. Each product qᵢkᵢ then has variance one. A dot product adds d_k such products, so its variance is d_k and its standard deviation grows as √d_k. Larger key dimension therefore produces larger-magnitude logits even when no stronger relation has been learned. Softmax sees larger gaps and becomes sharper. The Transformer paper introduced division by √d_k to counteract this effect.[5]

The derivation has a boundary. Trained queries and keys are neither perfectly independent nor guaranteed to have unit variance. Their norms can drift, dimensions can correlate, and normalisation placement changes the distribution. Square-root scaling is therefore a principled baseline, not a theorem that every trained head will remain well-conditioned. Measuring logit variance, attention entropy, query and key norms, and gradient behaviour is more informative than assuming the scale has done its job.

Temperature is a separate control

Temperature divides logits before softmax. At high temperature, score differences shrink and the distribution becomes broader. At low temperature, differences expand and the largest score dominates. Hinton, Vinyals and Dean used this same softmax property to soften class distributions during knowledge distillation.[6] In attention, temperature controls how much the router interpolates among candidates.

Temperature does not normally change the ranking of finite logits. If key A scores above key B, dividing both by the same positive temperature preserves that order. What changes is the margin expressed as probability mass. This makes temperature a competition control rather than a new similarity function.

Scaling and temperature can be written in the same denominator, but they answer different questions. Square-root scaling compensates for a predictable change in logit magnitude caused by representation dimension under stated assumptions. Temperature is an explicit choice or learned parameter that changes how score gaps become mixture weights. Conflating them makes debugging difficult: a head can be sharp because its vectors have large norms, because dimension is uncorrected, because temperature is low, or because the learned relation is genuinely decisive.

Four causes of a sharp route

Suppose one attention row has maximum weight 0.999. That number alone does not tell you why. The winning key may be better aligned with the query. Its norm may be much larger. The logits may have inherited excessive variance from the projections. Or an explicit temperature may be low. These causes predict different responses to intervention.

Normalise query and key vectors to unit length. If the route becomes broad, norm was carrying part of the decision. Hold vectors fixed and raise temperature. If the route broadens while ranking remains stable, competition was the control. Compare otherwise matched heads with different d_k. If unscaled sharpness grows with dimension, statistical scale is implicated. Rotate only the query. If one key remains dominant across unrelated directions, inspect key norms, biases or mask structure.

This produces a practical decomposition:

The decomposition is not a universal explanation of trained models. Correlated dimensions, biases, positional terms and upstream normalisation can produce mixed signatures. Its value is procedural: each hypothesis is tied to a perturbation. The result either narrows the cause or reveals that the simple account is insufficient.

stateDiagram-v2
  [*] --> Claim
  state "Why three projections are<br/>necessary" as Claim
  state "The swapped envelopes" as Test
  state "A key is not a database key" as Revision
  Claim --> Test: expose mechanism
  Test --> Decision: evidence survives
  Test --> Revision: rival survives
  Revision --> Claim: narrow or rebuild
  Decision --> [*]

Build it, break it, and know what it proves

The shortest path to intuition is to manipulate the mechanism directly. The laboratory below contains one query, three key-value pairs, a synthetic dimension model, temperature and an optional mask. Rotate the query and the route moves. Swap values and the route stays while the output changes. Increase dimension without scaling and the route sharpens for statistical rather than semantic reasons. Lower temperature and the same ranking becomes more decisive.

What each intervention establishes

A positive audit does not show that a trained model has learned the intended relation. It only shows that the implementation obeys the mechanism and its invariants. To test learned behaviour, use task-relevant perturbations: replace a pronoun’s antecedent, swap two definitions, move a clause outside a window, or create a distractor with a larger key norm. Observe both the attention distribution and the downstream output.

From equation to production kernel

The direct matrix implementation materialises an n by n score matrix for a sequence of length n. Its arithmetic and memory demands grow quadratically with sequence length. That cost belongs to the standard global attention pattern, not to the conceptual distinction between address and payload. Efficient kernels can preserve the exact mathematical result while changing the order in which blocks are loaded, normalised and multiplied.

FlashAttention is an important example. It tiles the computation to reduce movement between levels of GPU memory while retaining exact attention rather than replacing it with an approximation.[11] The engineering lesson is subtle: an implementation can change its execution schedule without changing its routing contract. Equivalence should be checked numerically within the precision expected from different accumulation orders.

That contract includes more than the formula printed in a paper. Shapes must align. Every query must retain at least one legal key. Boolean masks and additive masks must have declared semantics. The stable softmax must avoid overflow. Mixed-precision execution must not turn a large negative mask into a finite route. Dropout, when used on attention weights during training, changes the realised mixture and should be disabled for deterministic evaluation. Batch and head dimensions must broadcast intentionally rather than accidentally.

A small test harness should therefore include invariants and metamorphic tests. Weights over legal keys sum to one. Masked keys receive zero. Permuting keys and values together permutes weights but preserves the output. Swapping values without keys preserves weights but changes output. Adding a constant to every legal logit preserves weights. Replicating the computation through an optimised kernel matches a simple reference within tolerance. These tests expose errors that an end-task accuracy number can hide.

This failure also clarifies why larger context is not automatically better. Adding candidates changes the relative contest. A plausible distractor can take probability mass from a useful key even when the useful key’s raw score is unchanged. Long-context evaluation should therefore measure route stability under inserted distractors, duplicated facts, superseded facts and irrelevant high-norm tokens, not only whether the answer remains present somewhere in the window.

import numpy as np


def softmax(x: np.ndarray) -> np.ndarray:
    """Numerically stable softmax over the last axis."""
    shifted = x - np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(shifted)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)


def attention(
    q: np.ndarray,
    k: np.ndarray,
    v: np.ndarray,
    *,
    mask: np.ndarray | None = None,
    temperature: float = 1.0,
    scale: bool = True,
) -> tuple[np.ndarray, np.ndarray]:
    """Return (output, routing_weights) for one or more queries."""
    if temperature <= 0:
        raise ValueError("temperature must be positive")
    if q.shape[-1] != k.shape[-1]:
        raise ValueError("query and key dimensions must match")
    if k.shape[-2] != v.shape[-2]:
        raise ValueError("every key must have one value")

    logits = q @ np.swapaxes(k, -1, -2)
    if scale:
        logits = logits / np.sqrt(q.shape[-1])
    logits = logits / temperature

    if mask is not None:
        if mask.shape != logits.shape:
            raise ValueError("mask must match attention logits")
        logits = np.where(mask, logits, -np.inf)
        if np.any(np.all(~mask, axis=-1)):
            raise ValueError("each query needs at least one legal key")

    weights = softmax(logits)
    output = weights @ v
    return output, weights


# Minimal invariants
q = np.array([[1.0, 0.0]])
k = np.array([[1.0, 0.0], [0.8, 0.6], [0.0, 1.0]])
v = np.array([[10.0, 0.0], [0.0, 10.0], [-5.0, 5.0]])

out, w = attention(q, k, v, scale=False)
assert np.allclose(w.sum(axis=-1), 1.0)

swapped_out, swapped_w = attention(q, k, v[[1, 0, 2]], scale=False)
assert np.allclose(w, swapped_w)       # route unchanged
assert not np.allclose(out, swapped_out)  # payload changed

Run with Python 3.10+ and NumPy. The code follows the standard scaled dot-product structure and makes masking, temperature and invariants explicit. PyTorch’s official operator applies the scale before softmax and may select optimised kernels under the same high-level contract.[10]

The strongest boundary: routing weight is not causal responsibility

Attention weights are tempting explanations because they are explicit and easy to plot. Yet empirical studies have found cases where attention distributions correlate poorly with other importance measures, where removing highly weighted items has limited effect, and where very different distributions yield similar predictions.[8][9] Other work argues that the verdict depends on the definition of explanation and on the diagnostic protocol.[7]

The dispute becomes clearer through the address-payload split. A weight measures how strongly one value vector is mixed at one layer and head. Causal responsibility for a final prediction also depends on what the value contains, its norm, interactions with other positions, residual pathways, later layers and the output readout. A routing record can be useful evidence without being a sufficient explanation.

A disciplined audit sequence

Begin with a single query and a tiny candidate set. Verify the reference calculation by hand. Then perform the payload swap, key replacement, query rotation, temperature sweep and mask test. Only after those local mechanisms behave as expected should you inspect a full trained layer. In the trained model, freeze all but one causal object at a time and record both local output change and final task change.

For interpretation, declare the target before running the test. “Which value changed this head’s output?” is narrower than “Which token caused the next word?” The first can often be answered with value replacement and contribution analysis. The second crosses residual connections, other heads, feed-forward transformations and the final readout. A valid local explanation should not be promoted silently into a model-wide one.

For design, decide whether the task requires soft mixture, near-hard selection or an explicit null route. A relation-extraction head may benefit from decisive routing. Evidence synthesis may need several contributors. Safety-sensitive retrieval may require a gate that refuses all candidates when none meets a sufficiency threshold. Temperature and softmax can shape competition, but they cannot define adequacy by themselves.

sequenceDiagram
  participant Q as Question
  participant M as Mechanism
  participant E as Evidence
  participant D as Decision
  Q->>M: The address is not the message
  M->>E: Why three projections are necessary
  E-->>M: compare rival explanations
  E->>D: Separate address, payload and regime before reading the map
  D-->>Q: revise scope or act