Attention Mechanisms

Every recurrent model from RNNs, LSTMs, and GRUs answers the same way: compress the history into one hidden vector and hope what mattered survived the squeeze. Attention refuses the compression. At each output position it asks the full input set "how relevant are you to me right now?", scores every position, and mixes them by those scores. The three-projection formulation (Bahdanau et al., 2014, formalized by Sutskever-era seq2seq work) is the vocabulary of that question:

• Query (Q) — what this output position is looking for. Built from the current position's representation.
• Key (K) — what each input position advertises. Used only for matching.
• Value (V) — what each input position actually delivers if it is matched.

The library metaphor holds precisely: the query is your search term, keys are the index cards, values are the documents. You match against cards, but you retrieve documents — which is why K and V are separate projections instead of one.

import torch, torch.nn as nn

d_model, d_k = 64, 16
proj_q = nn.Linear(d_model, d_k, bias=False)
proj_k = nn.Linear(d_model, d_k, bias=False)
proj_v = nn.Linear(d_model, d_model, bias=False)

x = torch.randn(2, 6, d_model)          # (batch, timesteps, d_model)
Q, K, V = proj_q(x), proj_k(x), proj_v(x)     # each (2, 6, 16 or 64)

scores = Q @ K.transpose(-2, -1)        # (2, 6, 6): every position vs every position
print(scores.shape)                     # torch.Size([2, 6, 6])

In self-attention each position supplies its own query from its own representation. For the sentence "the cat sat on the mat", the position holding "sat" produces a query whose best key match is the noun it acts on — say the resulting weights come out [0.08, 0.61, 0.10, 0.07, 0.09, 0.05], so its output is 61% the value of "cat" plus small shares of the rest. The "mat" position builds a different query and gets a different weight vector; the mixture is rebuilt per output position, every timestep, from scratch.

The score q_t · k_j is a dot product, and it behaves like a similarity: aligned vectors score large, orthogonal score near zero, opposing score negative. Each row of scores is then turned into a distribution over input positions, and the final output is the value-vector mixture those weights describe — z_t = Σ_j α_{t,j} v_j, with α ≥ 0 and Σα = 1, so the output is a convex combination of inputs, an interpolator you can read: 0.7 of one token, 0.3 of another.

Why learn W_Q and W_K rather than scoring raw x_t · x_j? Because raw vectors encode topic and frequency — common tokens have typical magnitudes and the similarity you need ("holds the subject of my sentence", "is the noun my verb acts on") is a task-specific relation, not surface similarity. Learned projections let different positions be findable along different axes: one projection space answers syntactic questions, another semantic ones. Why a separate W_V? Matching relevance and carrying content are different roles; keys must be mutually comparable across all positions, while values only need to be useful once selected.

Two variants show up immediately in reading lists. Self-attention uses Q, K, V from the same sequence — every position querying its own sentence. Cross-attention takes Q from one sequence (the decoder's current position) and K, V from another (the encoder's output) — the exact mechanism by which a translation decoder looks back at the source sentence. The formula that consumes these three tensors, the scaling, the masking, and the multi-head split are the next chapter's worked example, and the whole structure becomes an architecture in Transformer Architecture Overview.

The complete mechanism fits on one line, and you should be able to reproduce it numerically without a framework:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V

Take d_k = 4, one query q = [1, 0, 1, 0], three keys and values:

k1 = [1, 0, 1, 0]   ->  q·k1 = 1 + 0 + 1 + 0 = 2      v1 = [2, 0]
k2 = [0, 1, 0, 1]   ->  q·k2 = 0                         v2 = [0, 2]
k3 = [1, 1, 0, 0]   ->  q·k3 = 1 + 1 + 0 + 0 = 1        v3 = [1, 1]

scale:  √d_k = 2   ->  logits [1.0, 0.0, 0.5]
softmax ->  e: [2.718, 1.000, 1.649], sum 5.367
           α = [0.506, 0.186, 0.307]

z = 0.506*[2, 0] + 0.186*[0, 2] + 0.307*[1, 1] = [1.32, 0.68]

The query aligns with k1, so v1 gets just over half the weight, v3 (half-aligned) gets 31%, and the orthogonal v2 still keeps 19% — attention is soft, never a hard pick. Now the part candidates stumble on: why divide by √d_k? If query and key components are roughly independent with variance 1, each product in the dot has variance 1 and the sum of d_k products has variance d_k — so the logit's spread grows as √d_k. At d_k = 64, raw logits of magnitude ~8 are unremarkable, and [8, 4, 0] softmaxes to [0.982, 0.018, 0.000]: an effectively one-hot weight whose gradient (α(1-α) ≈ 0.018) barely moves the projections. Divide those by √64 = 8 and [1, 0.5, 0] softmaxes to [0.506, 0.307, 0.186] — a soft distribution sitting where softmax still learns. Scaling keeps logits in the trainable region of the softmax, and that is the entire answer to the most-asked attention question.

The batched implementation, complete with the two masking conventions:

import torch

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / d_k ** 0.5      # (B, T, T)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = torch.softmax(scores, dim=-1)            # rows sum to 1
    return weights @ V, weights                        # (B, T, d_v), (B, T, T)

A causal mask zeroes everything above the diagonal so position t may only attend to ≤ t — required for autoregressive generation, where reading the future is cheating. A padding mask blocks −inf onto pad positions so they collect zero weight. And multi-head attention splits d_model into h heads of d_k = d_model/h, runs the code above h times in parallel (each head with its own W_Q/W_K/W_V), concatenates, and projects back — different heads specialize (one tracking adjacent-word relations, another coreference), at the cost of h separate T×T score matrices.

B, T, d_model, h = 2, 6, 64, 8
d_k = d_model // h                                   # 8 per head

Q = q_proj(x).view(B, T, h, d_k).transpose(1, 2)     # (B, h, T, 8)
K = k_proj(x).view(B, T, h, d_k).transpose(1, 2)
V = v_proj(x).view(B, T, h, d_k).transpose(1, 2)

scores = Q @ K.transpose(-2, -1) / d_k ** 0.5        # scale by PER-HEAD d_k: 8, not 64
w = scores.softmax(-1)                               # (B, h, T, T)
out = (w @ V).transpose(1, 2).reshape(B, T, d_model) # concat the heads back

Two details in that snippet earn interview credit. The scaling uses the per-head dimension because each head's logits are d_k-term sums — scaling by full d_model would over-cool them. And the memory bill multiplies by heads: 8 heads x 512² scores is 2.1M entries per layer per batch element before anything else happens, which is why sequence length, not parameter count, is usually the binding constraint. Cost, and what the single-vector bottleneck it replaces could never do, is where this topic closes.

Suppose a 200-token sentence whose verb at position 190 depends on a noun at position 3. Compare how the two mechanisms handle it, feature by feature — this comparison is the standard "why did attention beat RNNs?" interview question, and each row comes with a number.

Two rows deserve elaboration. Capacity: an LSTM improves what it stores, not where — everything still routes through h_t ∈ R^256, a few kilobytes describing 200 tokens. Attention keeps the input itself as memory: token 3 is still sitting there, retrievable with weight α_{190,3} whenever a query matches its key. And gradient: backprop to token 3's score crosses one softmax edge instead of 187 multiplicative steps, so learning a 180-token-away dependency does not require surviving a long product. The empirical signature was visible the year attention arrived: Bahdanau's translation model's BLEU score degraded gracefully as sentences lengthened while no-attention baselines fell off a cliff past ~30-40 words, and the learned alignments printed as plausible diagonal-and-near-diagonal soft traces.

The costs are real, so state them before an interviewer does. Scores grow as T²: 512 tokens is 262,144 entries per head per layer, and at 8,192 tokens a single head holds 67M floats — which is why production systems use FlashAttention-style kernels, KV caches, and sliding windows to attack memory before accuracy ever becomes the limit. And attention is permutation-invariant: shuffle the keys and the mixture for a given query is identically shuffled output — nothing knows that "cat" preceded "sat". Recurrence supplied order for free; attention needs it injected explicitly.

Two honest caveats round out the comparison. The parallelism row is really the teacher-forcing row from Sequence Modeling Basics: because every position conditions on true previous targets, the whole score matrix computes in one shot — attention did not remove that assumption, it scales it. And attention weights look like explanations without being one: a head can route decisive information along paths whose visualized weight is small, and studies find weak correlation between attention patterns and gradient-based importance. Read them as a diagnostic of where the model looked, not as proof of how it reasoned — an interview that distinguishes the two reads as someone who has actually interpreted the plots.

Streaming deserves one more honest note: a recurrent model updates one fixed-size state in constant work per token, while each new token under attention must score every key in the history — O(T) work per token. The KV cache keeps those keys and values so generation does not re-run the entire prefix at every step, but the cache still grows linearly with context, which is the engineering wall every long-context system runs into.

The last row is the handoff. With scoring, masking, multi-heading, and the recurrence comparison in place, the only remaining questions are structural — how many layers, where the positional signal goes, what surrounds the attention block — and those are answered by the final topic of this sequence, Transformer Architecture Overview.