Module 3 · Reading 3.2 · read before S14: Attention

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

In the previous section you measured the core weakness of recurrence: information from token 3 reaches token 300 only by surviving a 297-step relay, and the gradients that would teach the model to use that information decay along the same path. Attention is the mechanism that deletes the relay. Instead of compressing history into one crawling state vector, the model keeps every position’s representation around and, at each step, computes a weighted average over all of them — with weights the model calculates on the fly, per query, based on relevance.

The idea entered deep learning as a modest patch to machine translation models in 2014. Within three years it had eaten the whole architecture. This section builds attention from its historical starting point (the seq2seq bottleneck) to the exact form used inside every modern language model: scaled dot-product attention over queries, keys, and values. You will implement it in NumPy, read real attention weights off a worked example, and see with concrete numbers why the innocuous-looking division by \(\sqrt{d_k}\) is load-bearing. Then the ledger: what this mechanism costs. You will measure the quadratic growth that makes context length expensive, and build the KV cache that makes generation affordable — the two facts about attention that shape practical work with large models more than any other.

How to read this

The bottleneck: one vector per sentence

The 2014 sequence-to-sequence recipe (Sutskever, Vinyals & Le; and concurrently Cho et al.) translated a sentence in two phases. An encoder RNN read the source sentence and produced a final hidden state — a single vector, perhaps 1000 numbers. A decoder RNN then generated the target sentence, seeing the source only through that vector. Whatever the decoder needed to know about the source (every word, every nuance, every agreement constraint) had to fit through that fixed-size pipe.

For short sentences this worked strikingly well. For long ones it degraded: Cho et al. reported in 2014 that translation quality fell off markedly as source sentences grew longer, exactly what you would expect if a fixed-size summary were being asked to hold more than it can. The encoder also had to decide, up front and blind, what the decoder would later need — but which source words matter depends on which target word is being produced. Translating the verb needs the source verb; translating a name needs the name. A single static summary serves all of these requests with the same vector.

The fix is to stop summarizing. Keep all \(T\) encoder hidden states, one per source token, and let the decoder look back at them: every step, afresh.

Learned alignment: Bahdanau and Luong

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio proposed exactly this in “Neural Machine Translation by Jointly Learning to Align and Translate” (published at ICLR 2015; the preprint circulated from September 2014). At each decoding step, their model computed a relevance score between the decoder’s current state and every encoder state, turned the scores into weights with a softmax, and fed the weighted average of encoder states, the context vector, into the decoder. The scoring function was a small learned MLP, so this variant is called additive attention. The paper’s alignment heatmaps became famous: without any alignment supervision, the weights lit up on source words that human translators would consider the counterpart of the word being produced.

Historical note. The word “attention” stuck because the mechanism resembles selective visual attention: glance at what is currently relevant. Bahdanau et al. actually motivated it as alignment, the classical machine-translation problem of matching source and target words, which statistical MT systems had handled with separate alignment models since the IBM models of the early 1990s. Attention folded alignment into the network and learned it end-to-end from the translation objective alone — no aligned word pairs, no separate model. That is the part contemporaries found remarkable.

A year later, Luong, Pham, and Manning (EMNLP 2015) systematically compared scoring functions and showed that the simplest one (a plain dot product between decoder state and encoder state) worked as well as the learned MLP while being much cheaper: scoring all positions at once is a single matrix multiply. Dot-product scoring is what survived. This 2014–2015 lineage still ran on RNNs; attention was a bolt-on that let a recurrent decoder peek at a recurrent encoder. The realization that attention could replace the recurrence entirely came in 2017 and is the subject of the next section.

Two scoring functions, side by side

The additive-versus-multiplicative fork is worth seeing in code, both because you will meet the terms in papers and because the comparison clarifies what a scoring function is: any function that maps a (query, key) pair to a single relevance number. Everything downstream (softmax, weighted average) is identical regardless of how the scores were produced. The script below runs both scorers over the same keys and values. Dot-product (multiplicative, Luong-style) scoring is \(q \cdot k\), parameter-free. Additive (Bahdanau-style) scoring is a one-hidden-layer MLP: project query and key into a shared space, add, squash with tanh, and project to a scalar with a learned vector \(v\):

\[\operatorname{score}(q, k) = v^\top \tanh(W_q q + W_k k)\] 
import numpy as np

np.random.seed(0)

def softmax(z, axis=-1):
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

# Two scoring functions over the same keys and values. Dot-product
# (Luong-style) scoring is parameter-free; additive (Bahdanau-style)
# scoring runs query and key through a tiny one-hidden-layer MLP with
# its own weights W_q, W_k, v. Everything after the scores -- softmax,
# weighted average of values -- is identical.

d_k, d_att = 4, 8   # key dim; hidden size of the additive scorer

K = np.array([[1.0, 0.0, 1.0, 0.0],    # key 0: "animal-ish"
              [0.0, 1.0, 0.0, 1.0],    # key 1: "place-ish"
              [1.0, 0.0, 0.9, 0.1]])   # key 2: also "animal-ish"
V = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [0.5, 0.5]])
q = np.array([2.0, 0.0, 2.0, 0.0])     # the "animal" query

def dot_scores(q, K):
    return K @ q / np.sqrt(d_k)

W_q = np.random.randn(d_att, d_k) * 0.5
W_k = np.random.randn(d_att, d_k) * 0.5
v_a = np.random.randn(d_att) * 0.5

def additive_scores(q, K):
    # score(q, k_i) = v^T tanh(W_q @ q + W_k @ k_i)
    return np.tanh(W_q @ q + K @ W_k.T) @ v_a

for name, scores in [("dot-product", dot_scores(q, K)),
                     ("additive   ", additive_scores(q, K))]:
    w = softmax(scores)
    out = w @ V
    print(f"{name}: scores = {np.round(scores, 4)}  weights = {np.round(w, 4)}"
          f"  output = {np.round(out, 4)}")

n_params = W_q.size + W_k.size + v_a.size
print(f"\nadditive scorer parameters: {n_params} (dot-product scorer: 0)")

# Cost of scoring T queries against T keys. Dot-product: one length-d_k
# dot per pair, a single matmul. Additive (with W_q @ q and W_k @ k
# precomputed for all positions): ~2*d_att ops per pair, plus a
# (T, T, d_att) intermediate if vectorized.
for T in [512, 4096]:
    dot = T * T * d_k
    add = T * T * 2 * d_att + 2 * T * d_att * d_k
    print(f"T = {T:>4}: dot-product ~{dot/1e6:>7.1f}M ops, "
          f"additive ~{add/1e6:>7.1f}M ops, "
          f"additive intermediate: {T*T*d_att*4/1e9:.2f} GB float32")

Running this prints:

dot-product: scores = [2.  0.  1.9]  weights = [0.4902 0.0663 0.4435]  output = [0.7119 0.2881]
additive   : scores = [ 0.8151 -0.2195  0.7713]  weights = [0.4324 0.1537 0.4139]  output = [0.6394 0.3606]

additive scorer parameters: 72 (dot-product scorer: 0)
T =  512: dot-product ~    1.0M ops, additive ~    4.2M ops, additive intermediate: 0.01 GB float32
T = 4096: dot-product ~   67.1M ops, additive ~  268.7M ops, additive intermediate: 0.54 GB float32

Read the two weight rows first. Both scorers produce a valid attention distribution, and both happen to favor keys 0 and 2 here — keys 0 and 2 are nearly identical vectors, and any smooth function scores near-identical inputs near-identically. But the resemblance is superficial. The dot-product scores read relevance directly off the geometry of the space (key 0 scores 2.0 because it aligns with the query); the additive scores are whatever a randomly initialized MLP happens to compute, and only training would align them with anything meaningful. That is the real difference in spirit: dot-product attention assumes similarity-in-representation-space is relevance and asks the surrounding network to learn representations that make it so, while additive attention learns the relevance function itself, at the price of 72 extra parameters here (and correspondingly more at real scale).

The cost lines explain why scaled dot-product attention became the mainstream Transformer choice. Scoring all pairs is one matrix multiply, while a naive additive scorer uses more arithmetic and can materialize a (T, T, d_att) tensor. Additive attention has not disappeared (it remains useful in some encoder–decoder and smaller recurrent systems), but the rest of this chapter focuses on scaled dot products.

People. Bahdanau was a visiting student in Yoshua Bengio’s lab in Montreal when he worked on the 2014 paper, with Cho then a postdoc there — attention began as, by most accounts, an internship project on a known engineering problem. Cho has since co-authored much of the foundational seq2seq literature; Bengio shared the 2018 Turing Award with Hinton and LeCun for deep learning broadly. The mechanism their paper introduced is, by citation count, among the most influential ideas in the field’s history.

Queries, keys, and values

Strip away the RNN scaffolding and a clean abstraction remains, crystallized in the 2017 Transformer paper’s terminology. Attention is a soft dictionary lookup involving three roles:

  • a query \(q\): what I am looking for (one vector per position doing the looking);
  • keys \(K\): one vector per stored item, advertising what that item is about;
  • values \(V\): one vector per stored item, the content actually returned.

A hash map matches a query to exactly one key and returns its value. Attention softens this: score the query against every key, softmax the scores into weights that are positive and sum to 1, and return the weighted average of all values:

\[\operatorname{attention}(q, K, V) = \sum_i \operatorname{softmax}_i(q \cdot k_i / \sqrt{d_k})\, v_i\]

or, for a whole batch of queries stacked as rows of \(Q\):

\[\operatorname{attention}(Q, K, V) = \operatorname{softmax}(QK^\top / \sqrt{d_k})\, V\]

In Bahdanau’s setting the query was the decoder state and the keys and values were both the encoder states. Separating keys from values is the extra generality: an item may advertise itself one way (key) and contribute something different (value). In practice all three are produced by learned linear projections of token representations, so the model learns what to ask, what to advertise, and what to hand over.

Three readings of the same equation are worth carrying around, because different corners of the literature favor different ones. As a soft dictionary, attention retrieves by similarity instead of equality, useful when nothing will match exactly. As an expectation, the output is \(E[v]\) under a distribution the model computes over stored items, which is why attention outputs are smooth and differentiable while a hard lookup is not. As content-based addressing, it is memory access by what is stored rather than where, an idea that also drove the Neural Turing Machine line of work (Graves et al., 2014) contemporaneously with translation attention, a sign of how much the field was converging on differentiable memory access from several directions at once. The dictionary reading explains the API; the expectation reading explains the gradients; the addressing reading explains why attention generalizes far beyond translation — the stored items can be source words, image patches, retrieved documents, or, in the S16 reading, every token of the model’s own input.

Scaled dot-product attention in NumPy

The mechanism is a dozen lines. The example below builds a tiny, fully hand-inspectable instance: three stored items whose keys live in a 4-dimensional space where the first/third coordinates mean roughly “animal-ish” and the second/fourth mean “place-ish”. Keys 0 and 2 point in similar directions; key 1 points elsewhere. Each key carries a distinctive 2-dimensional value so you can see exactly what the output mixes.

import numpy as np

np.random.seed(0)

def softmax(z, axis=-1):
    z = z - z.max(axis=axis, keepdims=True)  # numerical stability
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

def scaled_dot_product_attention(Q, K, V):
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)      # (n_q, n_k)
    weights = softmax(scores, axis=-1)   # each row sums to 1
    return weights @ V, weights

# A hand-built example with d_k = 4. Keys 0 and 2 point in similar
# directions; key 1 points somewhere else entirely.
K = np.array([[1.0, 0.0, 1.0, 0.0],    # key 0: "animal-ish"
              [0.0, 1.0, 0.0, 1.0],    # key 1: "place-ish"
              [1.0, 0.0, 0.9, 0.1]])   # key 2: also "animal-ish"
V = np.array([[1.0, 0.0],              # value carried by key 0
              [0.0, 1.0],              # value carried by key 1
              [0.5, 0.5]])             # value carried by key 2

queries = {
    "query ~ animal": np.array([[2.0, 0.0, 2.0, 0.0]]),
    "query ~ place":  np.array([[0.0, 2.0, 0.0, 2.0]]),
    "query ~ mixed":  np.array([[1.0, 1.0, 1.0, 1.0]]),
}

for name, Q in queries.items():
    out, w = scaled_dot_product_attention(Q, K, V)
    print(f"{name}: weights = {np.round(w[0], 4)}  ->  output = {np.round(out[0], 4)}")

Running this prints:

query ~ animal: weights = [0.4902 0.0663 0.4435]  ->  output = [0.7119 0.2881]
query ~ place: weights = [0.1053 0.7783 0.1164]  ->  output = [0.1635 0.8365]
query ~ mixed: weights = [0.3333 0.3333 0.3333]  ->  output = [0.5 0.5]

Walk through the first row. The animal-flavored query dots strongly with keys 0 and 2 (both animal-ish) and weakly with key 1, so after the softmax it spends 49% of its weight on key 0, 44% on key 2, and only 7% on key 1. The output, [0.7119, 0.2881], is correspondingly a blend of values 0 and 2 with a trace of value 1. The place-flavored query flips the pattern: 78% of its weight lands on key 1 and the output is dominated by that key’s value. The mixed query scores every key identically (each dot product is 2.0), so the weights are exactly uniform and the output is the plain average of the values.

Two properties are worth pausing on. First, nothing here depends on the number of stored items or on their order: attention handles any number of keys, and permuting the key/value rows permutes nothing in the output. That order-blindness becomes a real issue in the S16 reading, where positional encodings must reintroduce word order. Second, the entire computation is differentiable, so “where to look” is trained by ordinary backprop, jointly with everything else — this is the property that let Bahdanau’s model learn alignment for free.

Why divide by sqrt(d_k)

Now the scaling factor, which looks like a fussy detail and is not. Real models use key dimensions of 64–128 per attention head (the S16 reading runs several heads in parallel), and dot products grow with dimension: if \(q\) and \(k\) have independent entries with mean 0 and variance 1, each term of \(q \cdot k = \sum_j q_j k_j\) has mean 0 and variance 1, and the \(d_k\) terms are independent, so their variances add: \(q \cdot k\) has mean 0 and variance \(d_k\), hence typical magnitude \(\sqrt{d_k}\). At 64–128 dimensions that is a logit standard deviation of 8–11. Already at \(d_k = 64\) the top key of 16 takes 0.83 of the softmax weight on average; at \(d_k\) of 256–1024 (standard deviation 16–32) it takes 0.92–0.97, nearly all of the probability mass. The experiment below measures this: for each d_k, it draws a random query and 16 random keys 200 times, and reports the spread of raw dot products plus the average winning softmax weight, with and without scaling.

import numpy as np

np.random.seed(0)

def softmax(z):
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()

# Dot products of random unit-variance vectors grow with dimension:
# if q, k have i.i.d. N(0,1) entries, q . k has variance d_k.
# Without scaling, softmax over 16 keys saturates as d_k grows.
n_keys, trials = 16, 200

print(f"{'d_k':>5}  {'std of q.k':>11}  {'max weight (raw)':>17}  {'max weight (scaled)':>20}")
for d_k in [4, 16, 64, 256, 1024]:
    max_raw, max_scaled, dots_all = [], [], []
    for _ in range(trials):
        q = np.random.randn(d_k)
        K = np.random.randn(n_keys, d_k)
        dots = K @ q
        dots_all.append(dots)
        max_raw.append(softmax(dots).max())
        max_scaled.append(softmax(dots / np.sqrt(d_k)).max())
    print(f"{d_k:>5}  {np.concatenate(dots_all).std():>11.2f}  "
          f"{np.mean(max_raw):>17.4f}  {np.mean(max_scaled):>20.4f}")

# Saturated softmax means dead gradients. The Jacobian of softmax has
# entries involving p_i * (1 - p_i); when one p_i ~ 1, all entries ~ 0.
print("\ngradient factor p*(1-p) at the winning key:")
for d_k in [4, 256]:
    np.random.seed(0)
    q = np.random.randn(d_k)
    K = np.random.randn(n_keys, d_k)
    for label, logits in [("raw", K @ q), ("scaled", K @ q / np.sqrt(d_k))]:
        p = softmax(logits)
        i = p.argmax()
        print(f"  d_k = {d_k:>4}, {label:>6}: max p = {p[i]:.6f}, p*(1-p) = {p[i]*(1-p[i]):.6f}")

Running this prints:

  d_k   std of q.k   max weight (raw)   max weight (scaled)
    4         1.93             0.4271                0.2341
   16         4.01             0.7050                0.2501
   64         7.86             0.8271                0.2422
  256        16.03             0.9164                0.2460
 1024        31.95             0.9651                0.2499

gradient factor p*(1-p) at the winning key:
  d_k =    4,    raw: max p = 0.404117, p*(1-p) = 0.240806
  d_k =    4, scaled: max p = 0.234212, p*(1-p) = 0.179357
  d_k =  256,    raw: max p = 0.999238, p*(1-p) = 0.000762
  d_k =  256, scaled: max p = 0.180193, p*(1-p) = 0.147724

The first table confirms the theory: the standard deviation of raw dot products tracks \(\sqrt{d_k}\) (1.93, 4.01, 7.86, 16.03, 31.95; roughly doubling as d_k quadruples). Without scaling, the average winning weight climbs from 0.43 at d_k = 4 to 0.97 at d_k = 1024: at large d_k (256 and above), unscaled attention is nearly a hard argmax at initialization, before the model has learned anything about what deserves attention. With the \(\sqrt{d_k}\) division, the winning weight sits near 0.24 regardless of dimension — softly peaked, dimension-independent.

Why does saturation matter? Every entry of the softmax Jacobian carries probability factors. At d_k = 256 unscaled, the winning key has p = 0.9992 and diagonal factor 0.0008, about 200 times smaller than the scaled case’s 0.148. This makes gradients through the attention weights very small, though not necessarily zero and not the only gradient path in a full block. The Vaswani et al. scaling addresses this saturation tendency.

Pitfall. The single most common attention implementation bug is applying the softmax along the wrong axis. The weights for one query must sum to 1 across keys (axis=-1 when scores have shape (n_q, n_k)). Softmax over the query axis also produces a plausible-looking matrix of positive numbers — shapes check out, the code runs, and the model trains to a mediocre loss. Test your implementation the way this section’s first script does: build a case where you know which key should win, read the weights, and assert that weights.sum(axis=-1) is all ones. With several queries stacked, a softmax over the wrong axis can still let the right keys win, so reading the weights alone does not catch it.

One mechanism, two wirings

Nothing in the QKV formulation says where the queries, keys, and values come from, and the two standard wirings deserve names now, because the rest of this chapter uses them constantly.

Cross-attention is the original 2015 arrangement: queries come from one sequence, keys and values from another. In translation, the decoder’s state asks, the encoder’s states answer — information flows across from source to target. Cross-attention is how any model conditions generation on something else: a translation on its source sentence, a caption on an image’s features, an answer on retrieved documents. The two sequences can have different lengths and even different modalities; all the mechanism needs is queries in the same space as the keys.

Self-attention sets all three projections loose on the same sequence: every position issues a query against every position’s key, including its own. Each token’s new representation becomes a mixture of all tokens’ values, weighted by learned relevance — the word “bank” can absorb “river” from five tokens away in one step. Self-attention was the later, more radical idea: it repurposes a lookup mechanism as a representation-building layer, and it is the move that let the 2017 Transformer delete recurrence entirely. A useful contrast with the previous section: an RNN relates positions 3 and 300 through 297 intermediate states, while self-attention relates them through one dot product (maximum path length 1, at the price the next experiment measures).

When you read architecture diagrams, the wiring is always identifiable from where K and V originate: from the same stream as Q, it is self-attention; from a different stream, cross-attention. The Transformer’s decoder, as you will see next section, uses both in every layer.

The price of looking everywhere

Deleting the relay has a cost, and you should know it precisely, because it shapes nearly every engineering decision around modern models. With \(T\) queries scoring \(T\) keys, the score matrix has \(T^2\) entries; computing it, softmaxing it, and using it to mix values all scale with \(T^2\). Double the sequence, quadruple the attention cost — in both arithmetic and memory. Compare the recurrent bargain: an RNN’s cost is linear in \(T\), but its steps are serial. Attention buys constant path length and full parallelism, and pays quadratically for it.

The script below makes the accounting concrete and then checks it against reality. The table counts exact multiply-accumulate operations (MACs) for the two matrix products and the exact size of the float32 score matrix. The measurement times the full attention computation at doubling lengths; because absolute milliseconds depend on the machine, the script reports only the observed growth rate, classified against the candidates: linear would double per doubling, quadratic would quadruple, cubic would octuple.

import time
import numpy as np

np.random.seed(0)

# Self-attention cost versus sequence length. For T tokens the score
# matrix Q @ K^T has T^2 entries, so both the arithmetic and the memory
# of the attention step grow quadratically -- the table counts them
# exactly. The measurement at the end times the real computation at
# doubling lengths and classifies the observed growth; absolute times
# depend on the machine, so only the growth rate per doubling is
# reported.

def softmax(z, axis=-1):
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V):
    w = softmax(Q @ K.T / np.sqrt(Q.shape[-1]), axis=-1)
    return w @ V

d = 64
lengths = [256, 512, 1024, 2048, 4096]

print(f"{'T':>5}  {'score+mix MACs':>14}  {'vs prev':>7}  {'score matrix':>12}")
prev = None
for T in lengths:
    macs = 2 * T * T * d            # Q@K^T plus weights@V: d MACs per entry each
    mem = T * T * 4 / 1e6           # float32 score matrix, in MB
    ratio = f"{macs / prev:.1f}x" if prev else "-"
    print(f"{T:>5}  {macs / 1e6:>13.0f}M  {ratio:>7}  {mem:>9.1f} MB")
    prev = macs

# Now measure wall-clock time (best of 5 runs at each length).
best = []
for T in lengths:
    Q = np.random.randn(T, d).astype(np.float32)
    K = np.random.randn(T, d).astype(np.float32)
    V = np.random.randn(T, d).astype(np.float32)
    attention(Q, K, V)              # warm-up
    reps = []
    for _ in range(5):
        t0 = time.perf_counter()
        attention(Q, K, V)
        reps.append(time.perf_counter() - t0)
    best.append(min(reps))

# Consecutive timing ratios bounce with machine noise, so take the
# median growth factor per doubling and bucket it against the candidates:
# linear -> ~2x, quadratic -> ~4x, cubic -> ~8x.
ratios = sorted(b / a for a, b in zip(best, best[1:]))
med = ratios[len(ratios) // 2]
if med < 3.0:
    verdict = "linear-like (about 2x per doubling)"
elif med < 6.0:
    verdict = "quadratic-like (about 4x per doubling)"
else:
    verdict = "cubic-or-worse (8x or more per doubling)"
print(f"\nmeasured median time growth per doubling of T: {verdict}")

Running this prints:

    T  score+mix MACs  vs prev  score matrix
  256              8M        -        0.3 MB
  512             34M     4.0x        1.0 MB
 1024            134M     4.0x        4.2 MB
 2048            537M     4.0x       16.8 MB
 4096           2147M     4.0x       67.1 MB

measured median time growth per doubling of T: quadratic-like (about 4x per doubling)

The arithmetic quadruples per doubling by construction, and the wall clock, measured, follows suit — on the machine used to build these readings, the per-doubling factors hovered around 4, drifting a bit above at small \(T\) (fixed overheads) and settling toward 4 as the matrices grow. Two readings of the same table matter. Per token, cost grows linearly with context: each token attends to \(T\) others, so doubling the context doubles the work each token does. Per sequence, it is the quadratic total that hits your budget, and the score matrix’s memory grows the same way: 67 MB at T = 4096 for a single head of this small width, before any batch dimension. This is why “context length” is a headline specification of every model, why long-context support is engineering news rather than a config change, and why an entire subfield works on cheaper attention (a story taken up in Going deeper).

Attention at generation time: the KV cache

The quadratic bill above is for processing a whole sequence at once, the training regime. Generation adds its own wrinkle, and a piece of vocabulary you will meet in every discussion of LLM serving: the KV cache. A causal model produces one token at a time, and each new token’s query must score against the keys and values of all earlier tokens. Recomputed naively from scratch each step, the projections repeat almost all of their work — token 1’s key gets recomputed at step 2, 3, 4, and so on. The fix is to compute each token’s key and value once, the first time that token exists, and append them to a growing cache.

import numpy as np

np.random.seed(0)

# The KV cache. During generation, a causal model produces one token at a
# time, and the new token's attention needs the keys and values of every
# earlier token. Recomputing them from scratch each step repeats work the
# model already did; caching them makes each step's projection cost
# constant. This script verifies that the cached computation produces the
# same numbers, then counts the work saved.

d, T = 32, 16
W_q, W_k, W_v = (np.random.randn(d, d) / np.sqrt(d) for _ in range(3))
X = np.random.randn(T, d)          # the token representations, step by step

def softmax(z):
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()

def attend(q, K, V):
    return softmax(K @ q / np.sqrt(d)) @ V

# Version 1: no cache. At each step t, re-project ALL tokens so far.
outputs_naive, proj_macs_naive = [], 0
for t in range(1, T + 1):
    K = X[:t] @ W_k                # recomputed from scratch every step
    V = X[:t] @ W_v
    q = X[t - 1] @ W_q
    proj_macs_naive += (2 * t + 1) * d * d
    outputs_naive.append(attend(q, K, V))

# Version 2: KV cache. Project only the NEW token; append to the cache.
K_cache = np.zeros((0, d)); V_cache = np.zeros((0, d))
outputs_cached, proj_macs_cached = [], 0
for t in range(1, T + 1):
    k_new = X[t - 1] @ W_k         # one row, not t rows
    v_new = X[t - 1] @ W_v
    q = X[t - 1] @ W_q
    proj_macs_cached += 3 * d * d
    K_cache = np.vstack([K_cache, k_new])
    V_cache = np.vstack([V_cache, v_new])
    outputs_cached.append(attend(q, K_cache, V_cache))

diff = max(np.abs(a - b).max() for a, b in zip(outputs_naive, outputs_cached))
print(f"max |naive - cached| over all {T} steps: {diff:.2e}")
print(f"projection MACs, naive:  {proj_macs_naive:>9,}")
print(f"projection MACs, cached: {proj_macs_cached:>9,}")
print(f"savings factor at T = {T}: {proj_macs_naive / proj_macs_cached:.1f}x "
      f"(grows with T: the naive version repeats old work every step)")
print(f"cache size at T = {T}: 2 * {T} * {d} floats per attention layer")

Running this prints:

max |naive - cached| over all 16 steps: 6.66e-16
projection MACs, naive:    294,912
projection MACs, cached:    49,152
savings factor at T = 16: 6.0x (grows with T: the naive version repeats old work every step)
cache size at T = 16: 2 * 16 * 32 floats per attention layer

The first line is the license for everything else: the cached and uncached computations agree to float64 round-off, because caching changes when numbers are computed, never what is computed. The savings are structural. Naive projection work over \(T\) steps sums \(1 + 2 + \cdots + T\) re-projections (quadratic in \(T\)), while the cached version does exactly \(T\), one per new token; the 6.0x factor at T = 16 scales up with sequence length. What the cache does not remove is the attention itself: step \(t\) still scores against \(t\) keys, so per-step cost grows linearly and total generation cost stays quadratic. And the trade is memory. Two vectors per token, per layer, per head add up: the cache is the reason long-context generation is memory-hungry, the reason serving systems obsess over cache size, and (recall the previous section) the exact mirror image of an RNN’s constant-size state. Recurrence compresses history into a fixed vector and pays in fidelity; attention keeps history verbatim and pays in memory.

What attention still lacks

Attention as built here is a mechanism, not yet a model. Three gaps remain. It has no notion of order: you saw that permuting key and value rows together leaves each query’s output unchanged. One weighted average per query lets each position express only a single “kind” of lookup at a time, when a word might simultaneously need its subject, its modifiers, and its coreferents. And a single attention layer is linear in its value vectors once the softmax weights are fixed, and no nonlinearity is applied to the values themselves; it needs depth and nonlinearity around it to compute anything rich. The Transformer answers all three (positional encodings, multi-head attention, and stacked blocks with feedforward layers) and is the subject of the S16 reading.

Going deeper

The quadratic wall, and three ways past it. The \(T^2\) cost you measured spawned a research program of its own. One family sparsifies: restrict each query to a subset of keys (a sliding local window, a few global tokens, strided patterns) so cost scales with \(T\) times the window size (Longformer and Big Bird, both 2020, are representative). A second family linearizes: replace the softmax with kernel feature maps so the \((QK^\top)V\) product can be reassociated as \(Q(K^\top V)\), collapsing the \(T^2\) term entirely (the Performer, 2020, is the best-known). Both families trade some quality or generality for the asymptotics, and neither displaced exact attention in mainstream LLMs. What actually shipped at scale was a third move: keep the exact computation but reorganize it. FlashAttention (Dao et al., NeurIPS 2022) computes exact attention in tiles that never materialize the full \(T \times T\) score matrix in slow GPU memory, cutting memory from quadratic to linear and speeding up wall-clock time substantially — a reminder that on modern hardware, memory movement, not arithmetic, is often the real cost. Its descendants are now the default kernel in most training and serving stacks.

Shrinking the cache. The KV cache experiment showed memory growing as \(2\,T\,d\) per layer: multiply by 32-plus layers, a batch of concurrent users, and long contexts, and the cache dominates serving memory. Two widely deployed variants attack it by sharing: multi-query attention (Shazeer, 2019) keeps many query heads but a single shared key/value head, shrinking the cache by the head count; grouped-query attention (Ainslie et al., EMNLP 2023) interpolates, giving each group of query heads one KV pair, and recovers most of multi-head quality at a fraction of the cache — the compromise adopted by many current open-weight models. That production architectures are being reshaped around the memory footprint of attention’s bookkeeping is the clearest sign of how central this one mechanism became.

Attention weights are not explanations. The Bahdanau heatmaps made attention look interpretable, and reading weights as “what the model considered important” became a standard analysis move. Treat it with care. A weight is one link in a deep composition: values are mixed, projected, and re-mixed across layers and heads, so a large weight on a token does not establish that the token drove the prediction, and research on the question (notably a 2019 back-and-forth under the titles “Attention is not Explanation” and “Attention is not not Explanation”) found that attention patterns can sometimes be altered substantially without changing outputs. Weights are a legitimate descriptive signal (these readings use them repeatedly), but causal claims need causal tests, like the perturbation experiments you will run in the S16 reading.

Common pitfalls

  • Attention output looks like an uninformative average and training is sluggish. Cause: missing \(\sqrt{d_k}\) scaling can leave logits either saturated (large \(d_k\)) or, with very small init, near-zero and uniform; either regime slows learning of where to look. Fix: scale by \(\sqrt{d_k}\) and check the entropy of attention weights early in training.
  • Model attends confidently to padding tokens. Cause: scores for padded key positions were left finite, so the softmax assigns them mass. Fix: add -inf (or a very large negative number) to masked positions before the softmax, never after.
  • Softmax overflows to NaN for long inputs. Cause: exponentiating large logits directly. Fix: subtract the row max before exponentiating, as in this section’s softmax; every production implementation does this.
  • Weights sum to 1 over the wrong dimension. Cause: softmax axis confusion (see the Pitfall callout). Fix: assert weights.sum(axis=-1) is all ones for score shape (n_q, n_k), and unit-test against a hand-checkable case.

Try it yourself

  1. In the NumPy example, scale the animal query from [2, 0, 2, 0] up to [8, 0, 8, 0] and rerun. What happens to the weight distribution, and what does this tell you about how a trained model can sharpen its own attention?
  2. Add a fourth key/value pair that duplicates key 0 exactly. How does the animal query’s weight redistribute, and why does the output change even though you added no new “information”?
  3. Reproduce the saturation table with 4 keys and with 256 keys instead of 16. Does the number of keys change how badly unscaled attention saturates, or only the baseline 1/n_keys weight?
  4. The additive-scorer example used random weights. Train them: with the animal/place keys and values fixed, define a target output for each query and fit W_q, W_k, v_a by gradient descent (autograd via PyTorch, or finite differences). How closely can the trained additive scorer reproduce the dot-product scorer’s weights?
  5. In the complexity script, hold T = 1024 fixed and double d from 16 to 256. Predict the growth rate of the MAC count before running, then check whether measured time follows — and explain why the score-matrix memory column does not move at all.
  6. Extend the KV-cache script to also count the score and mix MACs per step. Plot (or print) per-step totals with and without the cache: which part of generation does the cache make constant, and which part still grows with t?

Further reading

  • Bahdanau, Cho & Bengio, “Neural Machine Translation by Jointly Learning to Align and Translate,” ICLR 2015. arXiv:1409.0473 — the paper that introduced attention to translation; the alignment heatmaps in its appendix repay a look.
  • Luong, Pham & Manning, “Effective Approaches to Attention-based Neural Machine Translation,” EMNLP 2015. arXiv:1508.04025 — the systematic scoring-function comparison behind this section’s additive-versus-dot-product story.
  • Cho et al., “Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation,” EMNLP 2014. arXiv:1406.1078 — the seq2seq bottleneck documented firsthand, including the degradation on long sentences.
  • Vaswani et al., “Attention Is All You Need,” NeurIPS 2017. arXiv:1706.03762 — source of the QKV terminology and the \(\sqrt{d_k}\) scaling argument you verified above.
  • Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need,” 2019. arXiv:1911.02150 — multi-query attention: the KV cache’s memory bill, and the first mainstream fix for it.
  • Dao et al., “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” NeurIPS 2022. arXiv:2205.14135 — how exact attention escapes the quadratic memory cost by never materializing the score matrix.
  • Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,” EMNLP 2023. arXiv:2305.13245 — grouped-query attention, the multi-head/multi-query compromise in most current open-weight models.
  • d2l.ai, Chapter 11: Attention Mechanisms and Transformers, especially the attention scoring functions and Bahdanau attention sections — code-first walkthroughs of both scorers with trained models attached.
  • PyTorch docs: torch.nn.functional.scaled_dot_product_attention — the production entry point, including which fused kernel (FlashAttention among them) it dispatches to.

Before class you should be able to

  • State what the fixed-size encoder vector bottlenecked, and how attention removes it
  • Name the roles of query, key, and value in one sentence each
  • Compute attention weights by hand for a small Q, K, V example (softmax over the key axis)
  • Explain why the \(\sqrt{d_k}\) scaling exists and what goes wrong without it
  • Tell self-attention from cross-attention by where K and V originate
  • State how attention cost grows with sequence length, per token and per sequence
  • Explain what the KV cache stores, what it saves, and what it costs
  • Spot the wrong-axis softmax bug in a code snippet

Connects to

  • This session is covered by Quiz 3 (Module 3: S13–S22), given at S23.
  • HW3 (Transformer from scratch, out at S18): the scaled dot-product function built here is its innermost component.
  • The KV-cache and quadratic-cost facts return in S26 (efficiency) and matter for any project serving a generative model.

← S13 · RNNs and LSTMs · All readings · S16 · Transformer I →