Attention Mechanism

Dot-product attention, multi-head attention, Q/K/V, complexity

The RNN/LSTM module ended on a ceiling: even with the forget gate holding a gradient highway open, reaching back reliably still degrades past a couple hundred steps, and even a *ten*-step reach has already threaded through several forget gates along the way. Is there a mechanism where reaching an earlier word costs exactly the same — one step — no matter how far back it is?

Start with a smaller, more concrete version of the same question. Take "The bank by the river was steep." To resolve *which* "bank" this is — the sloped edge of a river, not a financial institution — a model has to notice that "river" sits nearby, and lean on it heavily. Now take an unrelated sentence: "The bank raised its interest rate." Here the correct reading needs "interest rate" instead, and nothing about geography. Different sentences, different neighbor — but the same underlying move both times: look at *some* other words, decide how much each one matters for the word you're resolving right now, and lean on the ones that matter.

Picture each word not as a fixed dot but as a color. Understanding "bank" in context means mixing in a lot of "river"'s color and almost none of "interest"'s — a blend, weighted by relevance, not a copy of any single neighbor. That blend is what attention computes, formally, and it is the single idea behind every Transformer.

The old way (encoder-decoder RNNs, circa 2014) could not do this at all — it crammed the whole sentence into one fixed-size summary vector before translation even began. By the time the model reached "bank," that summary had been overwritten word by word and the "river" clue was diluted or gone. Attention replaces the single squeezed summary with a direct, weighted look back at every earlier word, no matter how many sit in between.

Q, K, V — the mechanism

Put a name on the shape of that leaning: "bank"’s version of the question would be answered by "river"’s version of the tag — a comparison the model has to make, not a copy it can shortcut. Think of it like a library search: you show up with a question (a query), every book carries a catalog tag describing what it's about (a key), and the book itself is what you walk away with (its value). You compare your question against each tag and lean on the books whose tags match best. In self-attention, every word does all three jobs at once — question-asker, catalog-tag, and payload — for every other word, simultaneously.

None of this is hand-designed. Each token's embedding is multiplied by three learned matrices — W_Q, W_K, W_V — trained by the exact same backpropagation that trains every other weight in the network; the network decides for itself, from data, what makes a good "question" and a good "tag" for the task at hand. To decide how much word *i* should attend to word *j*, compare *i*'s query with *j*'s key (a dot product) — a big match means "this one is relevant." Run all those scores through a softmax [reminder: softmax turns any list of numbers into positive weights that sum to 1] so they become attention weights, then take the weighted average of the *values*. That is the whole operation.

Watch it work on the sentence above, with tiny stand-in numbers for what a trained model would actually produce (real weights run to hundreds of numbers; two numbers per word is enough to walk the arithmetic by hand). Give "The," "bank," and "river" 2-number embeddings: x_The = [1, 0], x_bank = [0, 1], x_river = [2, 1]. Say the network has already learned W_Q = [[2, 0], [1, 3]], W_K = [[1, 1], [2, 0]], W_V = [[1, 0], [0, 2]].

"Bank"'s query is x_bank · W_Q = [1, 3] — call this q_bank, the question "bank" is asking of every other word. Every word's key comes from the same W_K: k_The = [1, 1], k_bank = [2, 0], k_river = [4, 2] — the catalog tags q_bank will be compared against. The raw relevance score for each word is the dot product q_bank · k: against "The," 1×1 + 3×1 = 4; against itself, 1×2 + 3×0 = 2; against "river," 1×4 + 3×2 = 10 — already the largest of the three.

Divide each raw score by √d_k (here d_k = 2, so √2 ≈ 1.41) to get the scaled relevance score: 2.83, 1.41, 7.07. Run those three through softmax and you get attention weights of about 1.4% on "The," 0.3% on itself, and 98.2% on "river" — "bank" has decided, almost entirely, to lean on "river." Compute the values: v_The = [1, 0], v_bank = [0, 2], v_river = [2, 2]. Blend them by the attention weights: 0.014×[1,0] + 0.003×[0,2] + 0.982×[2,2] ≈ [1.98, 1.97] — call this "bank"'s attention-updated representation. It lands almost exactly on top of v_river = [2, 2]: after one attention step, "bank"'s vector has absorbed nearly all of "river" and almost none of "The" or itself. The disambiguation the opening paragraph asked for has happened, numerically, inside this arithmetic.

Why divide by √d_k at all? In this 2-dimensional toy it barely matters — leaving the scores unscaled (4, 2, 10) pushes "river"'s weight from 98.2% up to 99.7%, since a gap of 10 versus 2 already dominates the softmax either way. But real attention runs with d_k up in the tens or low hundreds, not 2 — and the *variance* of a dot product between random-looking query and key components grows with d_k (each of the many summed terms adds its own variance). At d_k = 64, typical raw scores can land dozens of points apart instead of single digits like this toy's 10 vs. 2. Feed a spread that wide into softmax and it collapses to a near one-hot pick — one word gets essentially all the weight, every other path's gradient goes to zero, and the model stops being able to learn from them. Dividing by √d_k rescales the spread back down to roughly what this toy example shows, keeping the softmax soft enough that gradients keep flowing through more than one path.

Many heads, many kinds of relationship

A single attention pass, with one W_Q/W_K/W_V, can only capture one *kind* of relationship at a time. So Transformers run several in parallel — multi-head attention — each head with its own learned W_Q, W_K, W_V, and therefore its own queries, keys, and values. One head might specialize in grammatical links (verb ↔ subject), another in word meaning, another in position. Their outputs are combined, giving the model several different lenses on the same sentence instead of forcing everything through one.

One more choice: how far is each word allowed to look? Everything worked out above let "bank" look at every other word in the sentence, before and after it — that’s *bidirectional* attention, right for tasks like classification where the whole sequence is sitting there to read at once. Generation is different: predicting the next word one token at a time, the model must not be allowed to peek at words that come after the one it’s producing. The fix is a causal mask — before the softmax, every score for a "future" position is forced to −∞, so softmax turns it into exactly 0 weight. That single restriction is the entire difference between encoder-style (bidirectional) attention and decoder-style (causal, or masked) attention.

The catch is cost: comparing every word with every other word is O(n²) — double the sequence length and you quadruple the work and memory. That quadratic cost is the single biggest constraint on long-context Transformers, and a whole family of tricks (FlashAttention, sparse attention, and others) exists to tame it.

Key points

• Use self-attention when the task requires modeling relationships between any two positions in a sequence, especially when those positions are far apart. Encoder-only (bidirectional) attention for classification and understanding tasks; causal (masked) attention for generation. For sequences longer than ~8k tokens, O(n²) memory becomes the bottleneck — use FlashAttention (exact, 2–4× faster, O(n) memory via tiling) as the first-line solution before considering approximate methods. Cross-attention applies when you need one sequence to query another — translation, image captioning, conditioning in diffusion models.
• The production trap: missing causal mask in autoregressive models. Without masking future positions to −∞ before softmax, the model during training can attend directly to the token it is predicting — achieving near-zero training loss while learning nothing about language. The symptom is excellent training loss and near-random test generation. Always verify the causal mask is applied before the softmax, not after. Multi-head attention in decoder-only models must use upper-triangular masking for every head.
• The diagnostic: visualise attention weights on a known example before trusting any trained model. For an encoder model, check whether the attention distribution for a given word is concentrated on related words (e.g., the subject attends to its verb) or diffuse noise. Uniform attention weights across all positions indicate the model has not learned meaningful relationships — either the query-key projections are not trained or the softmax temperature is too high. Log the entropy of attention distributions per head per layer: low entropy = head is attending specifically; high entropy = head is attending uniformly (potentially wasted capacity).

Recap

• Ceiling this module removes: LSTM's forget-gate highway (previous module) still degrades past ~200 steps; attention gives an O(1) reach to *any* earlier word regardless of distance.
• Core need, demonstrated twice: "bank" needs "river" nearby to read as riverbank, not the financial sense; "bank raised its interest rate" needs "interest rate" instead — same move both times: weigh other words by relevance, lean on the ones that matter.
• Attention = a weighted blend, not a copy of one neighbor — mix in a lot of "river"'s representation, almost none of "interest"'s. Old encoder-decoder RNNs couldn't do this: one fixed-size summary vector, overwritten word by word, diluting exactly the clue that mattered.
• Q, K, V are computed, not hand-designed: x·W_Q, x·W_K, x·W_V — three matrices learned by the same backprop as every other weight — give each token a query ("what am I looking for"), key ("what do I offer"), value ("what I hand over if picked").
• Mechanism: score = Q·K (dot product, big match = relevant) → scale by √d_k → softmax [outputs positive weights summing to 1] → weighted average of the *values*.
• Worked toy (d=2): q_bank·k_river=10 vs. 4 ("The") and 2 (itself) → scaled ≈7.07/2.83/1.41 → softmax ≈98.2%/1.4%/0.3% → blended output ≈[1.98,1.97], almost exactly v_river=[2,2]. Disambiguation happens inside the arithmetic.
• Why √d_k: dot-product variance grows with d_k; at real d_k (tens–hundreds) unscaled scores spread far enough apart to push softmax to near one-hot, zeroing gradients on every other path. Scaling keeps it soft enough to keep learning (barely visible at this toy's d_k=2, decisive at d_k=64+).
• Multi-head attention runs several in parallel, each with its own W_Q/W_K/W_V — one head can track grammar, another meaning, another position — several lenses instead of forcing everything through one.
• The catch is O(n²) cost: doubling sequence length quadruples work and memory — the single biggest constraint on long-context Transformers; FlashAttention, sparse attention, and similar tricks exist to tame it.

Check your understanding