Attention From Scratch

We left off with token embeddings: a sequence of vectors, one per token. But each vector so far knows only about its own token. Attention is the mechanism that lets tokens look at each other — and it’s the single component whose cost grows quadratically with sequence length, which is why long context is expensive.

The intuition: each token attends to the others

To predict the next token after “The cat sat on the ___”, the model needs “cat” and “sat” to inform the blank. Attention lets every position gather information from every other position, weighted by relevance. “Relevance” is learned and computed on the fly.

The mechanism gives each token three roles, each a learned linear projection of its embedding:

• Query (Q) — what I am looking for.
• Key (K) — what I offer, as an index others match against.
• Value (V) — the actual content I pass along if matched.

A token’s new representation is a weighted sum of everyone’s Values, where the weights come from how well its Query matches each Key.

The formula, piece by piece

$$\text{Attention}(Q,K,V) = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d}}\right)V$$

Read it inside-out:

1. QKᵀ — every Query dotted with every Key. For a sequence of n tokens, this produces an n×n matrix of raw scores: entry (i, j) is how much token i attends to token j.
2. ÷√d — scale by the square root of the head dimension so the dot products don’t blow up and saturate the softmax.
3. softmax — turn each row into a probability distribution: token i’s attention weights over all tokens sum to 1.
4. ·V — multiply those weights by the Value vectors to produce each token’s output: a relevance-weighted blend of the sequence.

        Keys →   k1   k2   k3   k4
      ┌─────────────────────────────┐
 q1   │  s11  s12  s13  s14         │   each row = one query's
 q2   │  s21  s22  s23  s24         │   scores over all keys,
 q3   │  s31  s32  s33  s34         │   softmax'd to weights
 q4   │  s41  s42  s43  s44         │   that sum to 1
      └─────────────────────────────┘
            the n × n attention matrix

In decoder LLMs this matrix is causally masked: token i may only attend to tokens ≤ i (no peeking at the future), so the upper triangle is zeroed. It’s still an n×n object.

Multi-head attention

One attention computation captures one kind of relationship. Multi-head attention runs h of them in parallel, each with its own smaller Q/K/V projections of dimension d_head = d_model / h. One head might track syntax, another long-range references. Their outputs are concatenated and projected back to d_model. Same quadratic shape, just h times in parallel.

The cost: O(n²) in sequence length

Here is the throughline. That n×n matrix is the problem — in both compute and memory:

• Compute: forming QKᵀ is ~n² × d multiply-adds; applying the weights to V is another ~n² × d. Doubling the sequence length roughly quadruples attention’s FLOPs.
• Memory: the scores matrix has n² entries per head, per layer. Materializing it naively is what blows up GPU memory on long sequences.

Why long context is expensive

Everything else in a transformer (embeddings, the FFN) scales linearly with n — twice the tokens, twice the work. Attention scales quadratically. So as context grows, attention’s share of the cost grows with it. A 1M-token context isn’t “125× a 8k context” — in attention terms it’s closer to 125² ≈ 15,000×. That gap is the entire reason FlashAttention and other long-context tricks exist.

Worked example: the attention matrix at n = 8,192

Take a single sequence of n = 8,192 tokens. The raw scores matrix for one head has:

8,192 × 8,192 = 67,108,864 entries ≈ 67M numbers.

In FP16 (2 bytes each) that’s 67M × 2 = ~134 MB — for one head, one layer, one request.

Now scale to the real model. Llama-3-8B has 32 layers and 32 query heads:

134 MB × 32 heads × 32 layers ≈ ~137 GB if you naively materialized every attention matrix at once.

That’s more than an 80 GB H100 holds — for a single request’s intermediate scores. The model weights are only ~16 GB; the transient attention matrices dwarf them at long context.

The lever this motivates

You never actually need the full n×n matrix in memory at once — you only need its product with V. FlashAttention (Part 3) computes that product in tiles, never materializing the matrix, cutting attention’s memory from O(n²) to O(n) and avoiding slow round-trips to GPU memory. And the KV cache avoids recomputing K and V every step. Both are direct responses to the quadratic cost you just sized.

Under the hood — half the scores are masked, and that’s a free 2×

The causal mask isn’t only a correctness detail; it’s a performance lever. In a decoder, entry (i, j) is discarded whenever j > i — the entire upper triangle above the diagonal. A naive dense kernel therefore computes roughly twice the QKᵀ scores it actually keeps. A mask-aware kernel can skip any tile of the matrix that lies wholly in the future region and never compute it at all.

        keys →   k1   k2   k3   k4
 q1   │  ✓    ·    ·    ·   │   · = masked (future): never needed
 q2   │  ✓    ✓    ·    ·   │
 q3   │  ✓    ✓    ✓    ·   │   only the lower triangle is real work
 q4   │  ✓    ✓    ✓    ✓   │

This is exactly the optimization FlashAttention-2 (2023) added over the original: for causal attention it stops once a block is fully masked and never touches it — one of several changes (alongside better GPU work-partitioning) that together roughly doubled attention throughput over the first version. The lesson threads forward: the cheapest score is the one a mask lets you skip computing.

The architect’s lens

Attention is the transformer’s one irreducible component — and its dominant scaling cost. The five questions:

• Why does it exist? Because after embedding, each token vector knows only about its own token; attention lets every position gather information from every other, weighted by a learned Query–Key match and blended over Values.
• What problem does it solve? Contextualization — turning isolated token vectors into representations that depend on the whole sequence, so “the cat sat on the ___” can route “cat” and “sat” into the blank.
• What are the trade-offs? Cost scales as O(n²) in sequence length: the n×n score matrix is n² entries per head, per layer, so doubling context quadruples attention’s FLOPs, and at n=8,192 the transient matrices (~137 GB across 32 heads × 32 layers) can dwarf the ~16 GB of weights.
• When should I avoid it? You never drop it, but you avoid its naive form — materializing the full n×n matrix (use FlashAttention) and recomputing K/V each step (use the KV cache); the quadratic term is what makes very long context the expensive regime to design around.
• What breaks if I remove it? Tokens can no longer see one another — the FFN would process each position in isolation and all cross-token reasoning disappears; this is exactly the recurrence-free mixing that Attention Is All You Need introduced.

Check your understanding

1. In one sentence each, what are the Query, Key, and Value roles for a token?
2. Walk through softmax(QKᵀ/√d)·V: what does each of the four operations produce?
3. Why does multi-head attention exist, and how does it change the cost shape (if at all)?
4. Attention scales as O(n²) while the FFN scales as O(n). Why does that make long context disproportionately expensive?
5. From the worked example, why are the transient attention matrices potentially larger than the model weights themselves at long context — and what later technique avoids materializing them?

Show answers

1. Query = what this token is looking for; Key = what this token offers as a match target; Value = the content this token contributes if it’s matched.
2. QKᵀ produces the n×n raw score matrix (every query dotted with every key); ÷√d scales the scores to keep softmax well-behaved; softmax turns each row into attention weights summing to 1; ·V produces each token’s output as a weighted blend of all Value vectors.
3. It runs h attention computations in parallel, each capturing a different relationship, then concatenates and projects them. The cost shape stays O(n²) — it’s just done h times in parallel with smaller per-head dimensions.
4. Doubling n doubles the FFN work but quadruples attention’s, so attention’s share grows with context; very long contexts are dominated by the quadratic attention term.
5. The scores matrix is n² entries per head per layer, which at n=8,192 across 32 heads × 32 layers (~137 GB) far exceeds the ~16 GB of weights. FlashAttention avoids materializing it by computing the QKᵀ·V product in tiles.