Build an LLMMiddleLesson 397 min read

Step 3 · Attention, derived slowly

The one idea the whole field rests on. Three vectors per token, one dot product, one softmax. Click the diagram to walk through it.

🧪 Lab · 25 min 📝 Lesson test

Lesson in motion

In 60 seconds

Step 3 · Attention, derived slowly

The one idea the whole field rests on. Three vectors per token, one dot product, one softmax. Click the diagram to walk through it.

1/4

In simple words

Every word asks a question ("what am I looking for?"), every word wears a label ("here is what I am"), and every word carries a message ("here is what I would tell you"). Each word listens most to the words whose label best answers its question.

Formally, each token produces three vectors by multiplying its embedding by three learned matrices:

Vector	Name	Plain meaning
`q`	Query	What I am looking for
`k`	Key	What I am, advertised to others
`v`	Value	What I will contribute if you attend to me

Tap any box in the diagram

Query

Produced by multiplying this token's current vector by a learned matrix W_q. Think of it as the question this position is asking of the rest of the sequence. In "the animal did not cross the street because it was tired", the query at "it" is roughly "which noun am I referring to?"

Click each stage. Attention is one matrix multiply, one scale, one mask, one softmax, one more matrix multiply. Everything else in a transformer is plumbing around it.

Causal self-attention, completepython

import torch, torch.nn as nn, torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head, dropout=0.0):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head = n_head
        self.d_head = d_model // n_head
        self.qkv  = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # (B, T, C) -> (B, n_head, T, d_head)
        q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)

        # fused, memory-efficient, and applies the causal mask for us
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.drop(self.proj(y))

Why multiple heads

One attention operation can only average one way. Splitting the vector into heads lets the model run several attention patterns at once — one head tracking syntax, another tracking which entity a pronoun refers to, another tracking position. They are computed in parallel and concatenated, so heads are close to free.

Watch out

The cost you cannot escape: attention compares every token to every token, so compute and memory grow with the square of sequence length. Doubling context quadruples the work. This single fact drives FlashAttention, sliding windows, GQA and every long-context trick in Module 48.

The thing worth remembering for Track B

Do this

Attention is why the system has no trust boundary. Every token in the context can influence every other token, weighted only by learned relevance — never by origin. There is no place in this mechanism to write "but that part came from a web page". The security problem is architectural, not a missing feature.

Watch and read more

**Attention in transformers, step by step**3Blue1Brown · video

**Let's build GPT: from scratch**Andrej Karpathy · attention section · video

Attention Is All You Needthe original paper · reading

Lab

Attention implemented twice, verified identical.

~25 min

The problem

Write causal self-attention from scratch with explicit loops. Verify against F.scaled_dot_product_attention to within floating-point tolerance. Then measure memory and time at sequence lengths 128, 512, 2048 and plot.

Starter codepython

def attention_by_hand(q, k, v):
    d = q.size(-1)
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(d)
    mask = torch.triu(torch.ones_like(scores, dtype=torch.bool), diagonal=1)
    scores = scores.masked_fill(mask, float("-inf"))
    return torch.softmax(scores, dim=-1) @ v

assert torch.allclose(attention_by_hand(q,k,v),
                      F.scaled_dot_product_attention(q,k,v,is_causal=True), atol=1e-4)

You are done when

Your implementation matches the reference to 1e-4
You measured quadratic memory growth
Removing the 1/sqrt(d) scale visibly destabilised training

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Remove the 1/sqrt(d) scaling. Predict the failure precisely, then verify.Reveal

Dot products of d-dimensional vectors grow like sqrt(d), so at d=64 the logits are ~8x larger. Softmax saturates: one weight goes to ~1, the rest to ~0, gradients through the softmax vanish, and the layer stops learning — loss plateaus early rather than diverging. Verify by printing softmax entropy per layer with and without the scale; you will see it collapse toward zero.

Please sign in to continue.

Questions people ask

Why divide by the square root of the head dimension?

Dot products of high-dimensional random vectors grow with the square root of the dimension. Without scaling, the softmax saturates — one weight goes to 1, the rest to 0 — and gradients vanish. It is a small detail that makes training possible.

Is attention the same as memory?

It is more like a lookup over the current context. Nothing persists between forward passes. Everything the model "remembers" is either in its weights or in the tokens currently in front of it.

Why is the value matrix separate from the key?

So relevance and content can differ. A token can be very findable for a certain query while contributing something quite different once found. Tying them together measurably reduces capability.

Do I need to implement attention myself?

Write it once by hand to understand it. Then use F.scaled_dot_product_attention in production — it dispatches to FlashAttention kernels and is dramatically faster and more memory-efficient than a naive implementation.

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.