Attention Mechanisms: The Foundation of Transformers

AI4Bharat's IndicTrans2 (Gala et al., 2023) translates between English and all 22 scheduled languages of the Indian Constitution using a single encoder-decoder transformer. Give it the English sentence "the UPI payment failed" and ask for Hindi. The encoder reads the whole English sentence once and turns it into a set of vectors, one per source token. The decoder then writes the Hindi output one token at a time, and for every token it writes, it has to decide which of those English vectors matter right now: the token that renders "payment" needs different source information than the token that renders "failed". The decoder cannot pull that information from its own half-written Hindi output, because the relevant English words are not in the Hindi output at all. It has to look outward, into a second sequence that was encoded separately. That outward look, a decoder attending into an encoder's output instead of into itself, is called cross-attention, and it is what makes translation, summarization, and image captioning work in a transformer.

This chapter assumes you already know how a single query, key and value vector combine into a softmax-weighted sum, and how splitting that computation across multiple heads lets a model track several relationships in parallel. It goes past that foundation in two directions that decide whether a transformer works at all outside a lecture slide. The first is architectural: what changes when the query comes from one sequence and the key and value come from another. The second is computational: how that same softmax-weighted sum, run over sequences tens of thousands of tokens long, actually gets computed on a GPU without the hardware running out of memory bandwidth before it finishes.

Cross-attention: a query from one sequence, keys and values from another

Self-attention computes Q, K and V from the same sequence: every token's query is compared against every other token's key, all within one sentence. An encoder-decoder transformer needs a second kind of attention layer, one where that assumption breaks. Inside each decoder block, after the decoder's own masked self-attention has run over the target sequence generated so far, a cross-attention layer takes over. Its query comes from the decoder's own hidden state at the current position. Its key and value come from the encoder's final output: the same fixed set of vectors, reused at every position the decoder will ever generate.

Write it the way the original encoder-decoder transformer (Vaswani et al., 2017) does. Let H_enc be the encoder's output after its own stack of self-attention layers has finished running, one row per source token. Let h_dec(t) be the decoder's hidden state at generation step t, taken after its masked self-attention sublayer. The cross-attention sublayer computes:

Q(t) = h_dec(t) · W_Q_cross
K    = H_enc    · W_K_cross
V    = H_enc    · W_V_cross

CrossAttention(t) = softmax( Q(t)·K^T / sqrt(d_k) ) · V

Three details separate this from self-attention. First, W_Q_cross, W_K_cross and W_V_cross are their own learned weights: not the ones the encoder used for its self-attention, not the ones the decoder used for its masked self-attention. Between the encoder's self-attention, the decoder's masked self-attention and the decoder's cross-attention, the full model carries three independent sets of Q/K/V projections, one per attention sublayer, and cross-attention's set is shared with neither of the other two. Second, K and V do not depend on t. H_enc is computed once, when the encoder finishes its forward pass, and that same K and V feed the cross-attention computation at every decoding step from the first output token to the last. Third, Q depends only on the decoder's own state, which keeps changing, so the attention distribution over the source sentence is recomputed at every output position. Self-attention lets Q, K and V all ride on the same evolving sequence. Cross-attention holds K and V fixed at the encoder's output and lets only Q move underneath them.

Worked example: two decoding steps, one fixed encoder output

Take the toy translation task from the opening: encode "UPI payment failed" and generate a Hindi translation. Suppose that once the encoder finishes, its final hidden states for the three source tokens, already projected through W_K_cross and W_V_cross, are these two-dimensional vectors (small integers, chosen so every step can be checked by hand, not real learned weights):

token       K            V
UPI         [1, 0]       [2, 0]
payment     [0, 1]       [0, 3]
failed      [1, 1]       [1, 1]

These six numbers are now fixed for the rest of decoding. Every future step reuses exactly them. Here d_k = 2, so the scaling factor is sqrt(2) ≈ 1.4142.

Decoder step 1. The decoder has produced no output yet, so its hidden state, projected through W_Q_cross, is some vector Q1 = [1, -1]. Score it against each source token's key:

Q1·K(UPI)     = [1,-1]·[1,0] = 1
Q1·K(payment) = [1,-1]·[0,1] = -1
Q1·K(failed)  = [1,-1]·[1,1] = 1 + (-1) = 0

Scale by 1/sqrt(2): [0.7071, -0.7071, 0.0000]. Exponentiate and normalize:

e^0.7071  = 2.0281
e^-0.7071 = 0.4931
e^0       = 1.0000
sum       = 3.5212

weight(UPI)     = 2.0281 / 3.5212 = 0.5760
weight(payment) = 0.4931 / 3.5212 = 0.1400
weight(failed)  = 1.0000 / 3.5212 = 0.2840

Weighted sum of values: 0.5760·[2,0] + 0.1400·[0,3] + 0.2840·[1,1] = [1.4360, 0.7040]. Step 1 looks mostly at "UPI", the word a translator would need first.

Decoder step 2. The decoder has now emitted one output token, and its hidden state has moved on, giving a different query, Q2 = [-1, 2]. Same encoder K and V as before, nothing recomputed there:

Q2·K(UPI)     = [-1,2]·[1,0] = -1
Q2·K(payment) = [-1,2]·[0,1] = 2
Q2·K(failed)  = [-1,2]·[1,1] = -1 + 2 = 1

scaled: [-0.7071, 1.4142, 0.7071]

e^-0.7071 = 0.4931
e^1.4142  = 4.1133
e^0.7071  = 2.0281
sum       = 6.6345

weight(UPI)     = 0.4931 / 6.6345 = 0.0743
weight(payment) = 4.1133 / 6.6345 = 0.6200
weight(failed)  = 2.0281 / 6.6345 = 0.3057

Context vector: 0.0743·[2,0] + 0.6200·[0,3] + 0.3057·[1,1] = [0.4543, 2.1657]. The weight on "payment" jumped from 14% to 62% between the two steps, and the weight on "UPI" collapsed from 58% to 7%, even though neither K nor V changed at all. Only the query moved, because only the query comes from the decoder's own advancing state. Here is the full computation, reproducing both steps exactly:

import numpy as np

K = np.array([[1,0],[0,1],[1,1]], dtype=float)   # UPI, payment, failed
V = np.array([[2,0],[0,3],[1,1]], dtype=float)
d_k = K.shape[-1]

def cross_attend(Q):
    scores = (K @ Q) / np.sqrt(d_k)
    weights = np.exp(scores) / np.exp(scores).sum()
    return weights, weights @ V

w1, ctx1 = cross_attend(np.array([1.0, -1.0]))
w2, ctx2 = cross_attend(np.array([-1.0, 2.0]))
# w1   ≈ [0.5760, 0.1400, 0.2840]   ctx1 ≈ [1.4359, 0.7041]
# w2   ≈ [0.0743, 0.6200, 0.3057]   ctx2 ≈ [0.4543, 2.1657]

Notice what never appears in this code: nothing about the decoder's earlier output feeds back into K or V. That is the structural signature of cross-attention. A self-attention layer with a new token in the sequence has to recompute keys and values for the whole sequence at that layer. A cross-attention layer just runs the same fixed K and V against a new Q.

FlashAttention: the memory wall standard attention hits

Before FlashAttention (Dao, Fu, Ermon, Rudra and Ré, "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," NeurIPS 2022), training or running a transformer on long sequences hit a wall that had nothing to do with parameter count. The Long Range Arena benchmark's Path-X task requires reasoning over a 16,384-token sequence; standard attention implementations of the time could not fit that computation's memory footprint on a single GPU, so the task was effectively unsolved. The obstacle was not compute. An A100 GPU can do tens of trillions of floating-point operations per second. The obstacle was moving data.

A GPU has two kinds of memory. High-bandwidth memory (HBM) is large: 40 to 80 GB on an A100, enough to hold an entire model's weights and activations. It is also, by on-chip standards, slow: 1.5 to 2.0 TB/s. On-chip SRAM is the opposite. Across an A100's 108 streaming multiprocessors it totals only about 20 MB (192 KB per multiprocessor), but it runs at roughly 19 TB/s, on the order of 10 to 13 times the bandwidth of HBM. Every number a GPU computes has to travel through SRAM at some point; the question is how much it also has to travel through the much slower HBM in between.

A standard implementation of attention moves through HBM three times for no good reason. It computes S = Q · K^T, an N×N matrix, and writes the whole thing to HBM. It reads S back from HBM, applies softmax, and writes the N×N result P back to HBM again. It reads P back from HBM a third time, multiplies by V, and produces the output. Softmax and the scaling step are cheap arithmetically: only a handful of floating-point operations per element. Reading and writing an N×N matrix three times over is not cheap at all, and for a sequence of even a few thousand tokens, that matrix does not fit in SRAM, so every one of those reads and writes has to go through the slow 1.5 to 2.0 TB/s path. The GPU spends most of its time waiting on memory traffic it did not need to generate, not on arithmetic. This is what "memory-bandwidth-bound" means: wall-clock time tracks bytes moved, not floating-point operations performed.

Tiling and the online softmax: computing the same answer without the N×N matrix

FlashAttention never writes S or P to HBM. It loads a block of Q and a block of K and V, small enough that both blocks and all the intermediate numbers for that block fit inside SRAM, computes that block's contribution to the output entirely on-chip, and only then moves to the next block of K and V. The entire chain, score, mask, softmax, weighted sum, runs as one fused GPU kernel instead of three separate ones, so the N×N matrix is never assembled anywhere HBM can see it.

The obstacle to doing this is softmax itself. Ordinary softmax needs the full row of scores before it can normalize: every weight is divided by the sum over every key in the row, and the numerically stable version subtracts the row's maximum before exponentiating. If keys arrive one block at a time, neither the true maximum nor the true sum is known until the last block has been seen. FlashAttention gets around this with the online softmax, a running-statistics update rule (building on Milakov and Gimelshein's online normalizer calculation) that keeps a running maximum m, a running sum of exponentials l, and a running unnormalized output O, and corrects all three every time a new block reveals a larger score than anything seen so far:

m_old, l_old, O_old = -infinity, 0, 0   # before the first block

for each new block of keys/values (scores s_block, values V_block):
    m_new = max(m_old, max(s_block))
    alpha = exp(m_old - m_new)            # rescales everything accumulated so far
    p     = exp(s_block - m_new)
    l_new = alpha * l_old + sum(p)
    O_new = alpha * O_old + p · V_block
    m_old, l_old, O_old = m_new, l_new, O_new

# after the last block:
output = O_old / l_old

The alpha term is the whole trick. Every time a bigger score shows up, the running max moves, and every partial result accumulated under the old, smaller max is multiplied by exp(m_old - m_new), the exact correction needed to rebase it to the new max. Nothing is thrown away and nothing is approximated; the running total is retroactively rescaled to stay consistent with whatever the true maximum turns out to be, one block at a time.

Trace it on one query against four keys, scaled scores s = [1, 2, 4, 0], with (scalar, for clarity) values V = [10, 20, 30, 40], split into two tiles of two keys each.

Tile 1, keys with scores [1, 2] and values [10, 20]: local max is 2, and since nothing has been accumulated yet, m_old is treated as -infinity, so m_new = 2 and the correction alpha = 0 (there is nothing to correct). p = exp([1,2] - 2) = [0.367879, 1.000000]. l = 0·0 + (0.367879+1.000000) = 1.367879. O = 0·0 + (0.367879·10 + 1.000000·20) = 23.678790.

Tile 2, keys with scores [4, 0] and values [30, 40]: local max is 4, larger than the running max of 2, so m_new = 4 and this time the correction is real: alpha = exp(2-4) = 0.135335. Both l and O from tile 1 get multiplied by that 0.135335 before tile 2's contribution is added in. p = exp([4,0] - 4) = [1.000000, 0.018316]. l = 0.135335·1.367879 + (1.000000+0.018316) = 1.203438. O = 0.135335·23.678790 + (1.000000·30 + 0.018316·40) = 33.937209.

Final output: O / l = 33.937209 / 1.203438 = 28.2002. Compare that against ordinary batch softmax over all four scores at once: max is 4, exp(s-4) = [0.049787, 0.135335, 1.000000, 0.018316], sum = 1.203438, weights = [0.041371, 0.112457, 0.830953, 0.015220], output = 0.041371·10 + 0.112457·20 + 0.830953·30 + 0.015220·40 = 28.2002. The two numbers match, because they are the same computation carried out in a different order, with a correction factor that makes the reordering exact.

import numpy as np

def batch_softmax_attention(s, V):
    w = np.exp(s - s.max())
    w = w / w.sum()
    return (w * V).sum()

def tiled_online_attention(s, V, block=2):
    m, l, O = -np.inf, 0.0, 0.0
    for i in range(0, len(s), block):
        sb, vb = s[i:i+block], V[i:i+block]
        m_new = max(m, sb.max())
        alpha = np.exp(m - m_new) if m != -np.inf else 0.0
        p = np.exp(sb - m_new)
        l = alpha * l + p.sum()
        O = alpha * O + (p * vb).sum()
        m = m_new
    return O / l

s = np.array([1.0, 2.0, 4.0, 0.0])
V = np.array([10.0, 20.0, 30.0, 40.0])
# batch_softmax_attention(s, V)  ≈ 28.2002
# tiled_online_attention(s, V)   ≈ 28.2002  (identical to 1e-14)

Real tiles hold far more than two keys, typically 64 to 128, sized to whatever fits alongside Q in one streaming multiprocessor's SRAM, but the recurrence is exactly this one, run more times over bigger blocks. Because the matrix is never materialized, peak extra memory drops from O(N^2), the full attention matrix, to O(N), just the running m and l for each query row. FlashAttention's own IO-complexity result states this precisely: standard attention requires Ω(Nd + N^2) HBM accesses, while FlashAttention requires O(N^2 d^2 / M), where M is the SRAM size. Plugging in real numbers shows why that matters. An A100's 192 KB of SRAM per streaming multiprocessor holds 192 × 1024 / 4 = 49,152 float32 elements. For a typical per-head dimension d = 64, d^2 = 4,096, and 49,152 / 4,096 = 12 exactly, so d^2/M = 1/12. Since standard attention's HBM traffic is dominated by its N^2 term whenever the sequence is much longer than the head dimension, which is the normal case, FlashAttention's HBM accesses come out to roughly a twelfth of standard attention's for these realistic numbers, independent of how long the sequence is. That is a derived illustration of the paper's formula, not a number the paper itself reports; what the paper measured directly, on real hardware, was a 7.6× speedup on the attention computation itself, a 15% end-to-end wall-clock speedup training BERT-large, a 3× speedup training GPT-2, and a 2.4× speedup on the Long Range Arena, the same benchmark whose Path-X task had been out of reach before.

Common misconception: FlashAttention is not a faster approximation

Students who have met approximate or sparse attention variants often assume FlashAttention belongs in the same category: trade some accuracy for speed, the way lowering an image's resolution trades detail for file size. It does not. Nothing in the tiling-and-rescaling recurrence drops a key, rounds a probability early, or skips a term. Every one of the four scores in the worked example above still contributed to the final output; the algorithm just visited them two at a time instead of four at once, and used the alpha correction to keep the running total mathematically equal to what a single batch softmax would have produced. The FlashAttention paper's title says this directly: "Fast and Memory-Efficient Exact Attention." Given identical Q, K and V, FlashAttention and standard attention return the same output up to ordinary floating-point rounding, and the worked example's two numbers, 28.2002 computed both ways, are the proof: that agreement is the rescaling factor's algebra working exactly as derived, to the last digit shown.

Where cross-attention and FlashAttention meet

Tiling and online softmax do not care where Q, K and V came from. The recurrence only ever looks at three matrices and a scale factor; it has no idea whether K and V belong to the same sequence as Q or to an entirely different one. A translation or summarization model whose encoder has processed a long source document, a scanned legal filing running to tens of thousands of tokens, faces exactly the same N×N wall in its cross-attention layers during training as a decoder-only model faces in self-attention: teacher-forced training computes cross-attention for every target position against every source position at once, a matrix sized by target length times source length rather than source length squared, but no smaller in kind. The fix is the same fix. Tile the source sequence into blocks that fit in SRAM, tile the target positions being trained on, and run the identical online-softmax recurrence, with K and V drawn from the encoder's output instead of from the decoder's own sequence. The two ideas in this chapter are not two unrelated topics bundled under one title. One is about which sequence attention reaches into; the other is about how the reaching gets computed without stalling on memory traffic, whichever sequence it reaches into.

Active recall

Work out each answer before reading it.

1. In the cross-attention formula, which of W_Q_cross, W_K_cross and W_V_cross are shared with the decoder's own masked self-attention sublayer?
2. Using this chapter's worked example, a third decoding step arrives with query Q3 = [1, 1]. Without fully normalizing, which source token receives the highest cross-attention weight, and why is computing the full softmax unnecessary to answer that?
3. True or false: during cross-attention, K and V are recomputed at every decoding step, the same way Q is. Justify your answer.
4. Continue the chapter's FlashAttention trace with one more tile: scores [3, -1], values [50, 60], appended after the two tiles already processed (running stats m=4, l=1.203438, O=33.937209). Compute the new running m, l, O, and the final output after this third tile.
5. A classmate says FlashAttention "compresses" attention the way JPEG compresses an image, trading some accuracy for a smaller memory footprint. Using this chapter's evidence, explain what is wrong with that claim.
6. A summarization model's encoder processes a 50,000-token legal document. Explain why cross-attention from the decoder into that encoder output can hit the same memory wall standard self-attention hits, and what fixes it.

Answers

1. None of them. Counting the encoder's self-attention, the decoder's masked self-attention and the decoder's cross-attention, the full model holds three independent sets of Q/K/V projections, one per attention sublayer. Cross-attention's set is shared with neither of the other two. Sharing weights across sublayers is not part of the architecture.
2. Q3·K(UPI) = [1,1]·[1,0] = 1, Q3·K(payment) = [1,1]·[0,1] = 1, Q3·K(failed) = [1,1]·[1,1] = 2. "failed" has the highest raw score. Softmax is built from exp(), a strictly increasing function, and dividing every score by the same positive constant sqrt(d_k) before exponentiating preserves order too, so whichever token has the highest raw score keeps the highest weight after full normalization. Ranking the raw scores answers "which is highest" without computing a single exponential.
3. False. K and V come from H_enc, the encoder's output, which is computed once and does not change during decoding. Only Q, which comes from the decoder's own hidden state, is recomputed at every step. That asymmetry, one side fixed, one side moving, is what makes cross-attention structurally different from self-attention rather than just a renamed copy of it.
4. Local max of the new tile is 3, which is less than the running max of 4, so m_new = max(4,3) = 4, unchanged, and the correction alpha = exp(4-4) = 1: nothing needs rescaling this time. p = exp([3,-1] - 4) = exp([-1,-5]) = [0.367879, 0.006738]. l_new = 1·1.203438 + (0.367879+0.006738) = 1.578055. O_new = 1·33.937209 + (0.367879·50 + 0.006738·60) = 33.937209 + 18.798230 = 52.735439. Final output = 52.735439 / 1.578055 = 33.4180, matching a direct six-key batch softmax over [1,2,4,0,3,-1] and [10,20,30,40,50,60] exactly.
5. JPEG compression discards information: at high compression it loses detail permanently and cannot reconstruct the original pixels. FlashAttention's online-softmax recurrence discards nothing; every key and value in the sequence still contributes its exact term to the final weighted sum, only accumulated in blocks with a correction factor instead of all at once. The worked example computed the same query both ways, batch and tiled, and got 28.2002 either way. A method that returns the identical number is not a compressed or lossy version of the original computation; it is the same computation, reordered.
6. Teacher-forced training computes cross-attention scores between every target position being trained on and every one of the 50,000 source positions, so the score matrix has as many entries as target length times source length, the same quadratic-in-sequence-length shape as self-attention's N×N matrix, just rectangular instead of square. A standard implementation would materialize that whole matrix in HBM the same way standard self-attention does, hitting the same bandwidth wall. The fix is identical to FlashAttention's: tile the source dimension and the target dimension into SRAM-sized blocks and run the online-softmax recurrence over them, with K and V drawn from the encoder's 50,000 source positions instead of from the decoder's own sequence. The algorithm does not need to know which sequence K and V came from to do this.

Think About It

Think about this: How would you explain attention mechanisms: the foundation of transformers to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where attention mechanisms: the foundation of transformers is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting attention mechanisms: the foundation of transformers to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind attention mechanisms: the foundation of transformers, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.