DL S036 · Self-Attention Derived from Scratch

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 36 / 80 · Module M07 · ~2 hours

The story — a hallway at Google Brain, spring 2017

Picture a whiteboard in Building 40 at Google's Mountain View campus, sometime in early 2017. Ashish Vaswani, Noam Shazeer, and Niki Parmar are standing in front of it. They've been trying to make sequence-to-sequence translation faster. The dominant architecture — stacked LSTMs with additive attention (Bahdanau, Luong) — was state of the art for machine translation, but it was slow: every timestep waited on the previous one, so you couldn't parallelise across the sequence. Training a big EN→DE model took weeks on eight P100 GPUs.

Shazeer, who had been quietly rebellious about recurrence for a while, made a claim on that whiteboard that sounded almost silly at the time: "We don't need the RNN. Attention alone is enough." The paper that came out of that hallway six months later (Vaswani et al., "Attention Is All You Need," NeurIPS 2017) had exactly that title. It trained in 12 hours on 8 P100s and beat every LSTM baseline on WMT14 EN-DE (28.4 BLEU vs the previous 26.3). Not "a bit better." Better AND an order of magnitude cheaper. That combination is what makes a paper detonate.

Eight years later, in 2026, essentially every large model you use every day — GPT-4o, Claude 3.7 Sonnet, Gemini 2.5, Llama 4, DeepSeek-V3, Qwen 2.5, Mistral Large, Grok 3 — is a direct descendant of the equations on that whiteboard. Even AlphaFold's Evoformer, Whisper, Stable Diffusion's text encoder, ESM protein models, and RT-2 for robotics all use the same primitive. The primitive is self-attention. And self-attention is really just three things stacked in a specific way: a dot product, a soft-max, and a weighted average.

For thirty-five sessions we've been walking. Now we start running.

Everything we've built — MLPs, convnets, RNNs, LSTMs, seq2seq with additive attention — was a warm-up. This session is where the modern era of deep learning actually begins. Every model you've heard of in the last five years — GPT-4, Claude, Llama, Gemini, Whisper, Stable Diffusion's text encoder, AlphaFold's Evoformer, ESM, DALL·E, SAM — is built on top of one primitive. That primitive is self-attention. And self-attention is really just three things stacked in a specific way: a dot product, a soft-max, and a weighted average.

If you internalise the derivation in this session, the rest of M07 is downhill. Multi-head, positional encodings, the full block, nanoGPT — all of it is engineering polish on top of the equation you're about to derive. So take your time. Do the numeric worked example on paper. When you can write Attention(Q, K, V) = softmax(Q Kᵀ / √d) V without hesitating and explain what each piece is doing in plain English, you are ready for the next four sessions.

Here's the puzzle we're going to solve. Given a sequence of vectors x₁, x₂, …, x_n (say, word embeddings), we want each position i to produce a new vector y_i that mixes in information from other positions. RNNs solved this by walking left-to-right and squashing everything into a hidden state. That works, but it's sequential (can't parallelise), it forgets long-range detail, and gradients decay through time. Convnets solved it by using local windows — great for images, bad for language where the relevant word might be 40 tokens back. Attention solves it by letting each position reach out to every other position and pull in a weighted mix of their content, in a single parallel matrix multiply.

The trick is: how do we decide the weights? That's the whole game. And the answer — the answer that ate the world — is let the model learn the weights from a similarity score between learned query and key projections.

Prerequisites

• Session 002 — dot products as similarity, matrix multiplication shapes.
• Session 003 — you'll be reading (B, T, T) tensors constantly this module.
• Session 035 — Bahdanau additive attention. This session generalises it.

1 · Start with something you already understand: a dictionary lookup

Forget attention for a moment. Let's talk about Python dictionaries.

This is a hard lookup. You give it a query key, and it returns exactly one value — the one whose key matched exactly. No match → error.

Now imagine your query is fuzzy. You want to ask "give me the value for something roughly like a cat" and get back a mixture — mostly the cat vector, a little bit of dog (also mammal), and almost none of the fish. That's a soft lookup.

We're going to build attention as three generalisations of this soft lookup:

1. Replace exact-string match with a similarity score (dot product).
2. Turn the scores into a probability distribution (soft-max).
3. Learn the keys, values, and queries as projections of the input rather than hardcoding them.

Do those three, and out falls scaled dot-product attention. Let's go.

2 · Step 1 — from string keys to vector similarity

Instead of keys being strings like "cat", let each key be a vector k_j ∈ ℝ^d. Let each value v_j also be a vector. And let the query q be a vector too.

The natural "how much does q match k_j?" score for two vectors is the dot product:

Two vectors pointing the same direction → large positive score. Perpendicular → zero. Opposite → negative. This is exactly the geometric intuition from S002.

So the soft lookup becomes: compute the similarity of the query to every key, then use those similarities as weights on the values.

import numpy as np
 
# 3 items in our "dictionary", each key/value is 2-D
K = np.array([[1.0, 0.0],   # key for "cat"
              [0.9, 0.1],   # key for "dog"
              [0.0, 1.0]])  # key for "fish"
 
V = np.array([[10.0,  0.0],
              [ 9.0,  1.0],
              [ 0.0, 10.0]])
 
q = np.array([1.0, 0.05])   # a fuzzy "cat-ish" query
 
scores = K @ q              # shape (3,)  — one score per key
print(scores)               # [1.00, 0.905, 0.05]

The score vector [1.00, 0.905, 0.05] says: query is very similar to key 0 (cat), quite similar to key 1 (dog), barely similar to key 2 (fish). Good — that matches our intuition. Now we need to turn those raw scores into weights that sum to 1.

3 · Step 2 — from scores to a probability distribution

We want weights w_j ≥ 0 with Σ_j w_j = 1 so we can take a convex combination of the values. The canonical way to turn arbitrary real numbers into a probability distribution is soft-max:

Two important properties:

• Monotone. Larger s_j → larger w_j. Order is preserved.
• Peaky, controllable by scale. Multiply all scores by 10 and soft-max concentrates almost all mass on the max. Divide by 10 and it becomes nearly uniform. Temperature.

def softmax(x):
    e = np.exp(x - x.max())   # subtract max for numerical stability
    return e / e.sum()
 
w = softmax(scores)
print(w)                      # [0.46, 0.42, 0.12]  approximately
 
y = w @ V                     # weighted average of value vectors
print(y)                      # [8.62, 1.62]  approximately

That output y ≈ [8.62, 1.62] is the attention output for query q: mostly the "cat" value [10, 0], with a chunky contribution from "dog" [9, 1], and a tiny bit of "fish" [0, 10]. Exactly what a soft lookup should do.

3.1 · A subtle problem — variance grows with dimension

Here's a quirk that will bite us in step 3. When keys and queries are high-dimensional (say d = 512, typical for real transformers), the dot product q · k = Σ q_i k_i is a sum of d random terms. If q_i, k_i are unit-variance and zero-mean, the dot product has variance d. So typical dot products have magnitude ≈ √d.

Plug big numbers into soft-max and it becomes a one-hot: all mass on the single largest score, zero everywhere else. The gradient through soft-max in that regime is essentially zero — the model can't learn.

Fix: divide the scores by √d_k before soft-max. This is the "scaled" in scaled dot-product attention.

We'll re-derive why the √d_k factor is exactly right in §7. For now, believe it and move on.

4 · Step 3 — learn the keys, values, and queries

In our toy example, we hand-wrote K, V, and q. In a real model, we don't have hand-labelled keys and values sitting around — we have a sequence of input embeddings x_1, …, x_n, and we want the model to learn how to turn each x_i into a query, a key, and a value.

The trick: three learned linear projections.

where X ∈ ℝ^{n \times d_{model}} is the input (one row per token), and W_Q, W_K, W_V ∈ ℝ^{d_{model} \times d_k} are learnable weight matrices.

Now the beauty. For each position i, we take q_i (the i-th row of Q) and dot it with every row of K — but we can do this for all queries simultaneously by computing Q Kᵀ, which is an (n × n) matrix of all pairwise scores.

Divide by √d_k, soft-max along each row, multiply by V. Done.

That's the equation. Six symbols. Read it out loud until it sounds like English.

5 · A numeric worked example — do this on paper

Let's compute self-attention by hand for a length-3 sequence with d_model = d_k = 2. This is small enough to fit on one page and big enough to be honest.

Input X (three tokens, 2-dim embeddings):

X = [[1, 0],
     [0, 1],
     [1, 1]]

Weights (chosen to be simple, not learned):

W_Q = [[1, 0], [0, 1]]     # identity — Q = X
W_K = [[1, 0], [0, 1]]     # identity — K = X
W_V = [[1, 0], [0, 1]]     # identity — V = X

So Q = K = V = X in this toy case.

Step A · scores.

Check: entry (i, j) = x_i · x_j. E.g. (0, 2) = 1·1 + 0·1 = 1. ✓

Step B · scale. d_k = 2, so √d_k = 1.414. Divide every entry:

scaled = [[0.707, 0.000, 0.707],
          [0.000, 0.707, 0.707],
          [0.707, 0.707, 1.414]]

Step C · row-wise soft-max. For row 0: exp([0.707, 0, 0.707]) = [2.028, 1.000, 2.028], sum = 5.056, so weights ≈ [0.401, 0.198, 0.401].

Doing all three rows:

W ≈ [[0.401, 0.198, 0.401],
     [0.198, 0.401, 0.401],
     [0.288, 0.288, 0.424]]

Sanity checks: every row sums to 1 (✓), row 0 attends most to positions 0 and 2 (the ones with a 1 in dim 0 that matches x_0), row 2 attends fairly evenly with a slight bias to itself. Reasonable.

Step D · output. Y = W · V = W · X.

Row 0: 0.401·[1,0] + 0.198·[0,1] + 0.401·[1,1] = [0.802, 0.599].

Row 1: 0.198·[1,0] + 0.401·[0,1] + 0.401·[1,1] = [0.599, 0.802].

Row 2: 0.288·[1,0] + 0.288·[0,1] + 0.424·[1,1] = [0.712, 0.712].

Y ≈ [[0.802, 0.599],
     [0.599, 0.802],
     [0.712, 0.712]]

Compare to the input X = [[1,0], [0,1], [1,1]]. The output at each position has been "smoothed" by mixing in the other positions, weighted by similarity. Position 0 used to be pure "dim-0"; now it has a bit of "dim-1" mixed in (from position 1 and 2). That's attention doing its job — moving information between positions.

6 · Code it up in ten lines of NumPy

import numpy as np
 
def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)
 
def self_attention(X, W_Q, W_K, W_V):
    Q = X @ W_Q            # (n, d_k)
    K = X @ W_K            # (n, d_k)
    V = X @ W_V            # (n, d_v)
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)   # (n, n)
    weights = softmax(scores, axis=-1)
    return weights @ V, weights       # (n, d_v), (n, n)

Run it on our toy input:

X = np.array([[1., 0.], [0., 1.], [1., 1.]])
I = np.eye(2)
Y, W = self_attention(X, I, I, I)
print(Y.round(3))
print(W.round(3))

You should get the numbers from §5. If you don't, the bug is almost always: forgot axis=-1 on soft-max (default axis=None computes soft-max over the whole flattened tensor — a classic gotcha), or forgot to transpose K.

6.1 · Batched version — the shape you'll actually use

Real code processes a batch of B sequences at once, each of length T. Shapes:

X:  (B, T, d_model)
W_Q, W_K, W_V:  (d_model, d_k)  — same weights applied to every batch element and every position
Q, K, V:  (B, T, d_k)
scores = Q @ K.transpose(-1, -2):  (B, T, T)
weights = softmax(scores, axis=-1):  (B, T, T)
output = weights @ V:  (B, T, d_v)

In PyTorch this becomes:

import torch, torch.nn as nn, torch.nn.functional as F, math
 
class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
 
    def forward(self, x):
        Q, K, V = self.W_Q(x), self.W_K(x), self.W_V(x)
        scores = Q @ K.transpose(-2, -1) / math.sqrt(K.size(-1))
        weights = F.softmax(scores, dim=-1)
        return weights @ V

Fifteen lines. That is the primitive that ate deep learning.

7 · Why the √d_k? A deeper derivation

Assume the components of q and k are i.i.d. with mean 0 and variance 1. Then:

So the standard deviation of a raw dot product is √d_k. If we don't scale, then for d_k = 512 the typical score is ≈ ±23. Feeding scores of that magnitude into soft-max concentrates virtually all mass on the top one or two positions, and the Jacobian of soft-max becomes vanishingly small everywhere else. Gradients die. Training stalls.

Dividing by √d_k restores unit variance on the scores, which keeps soft-max in the "reasonable slope" regime.

8 · Self-attention vs cross-attention

Everything above was self-attention: Q, K, V all come from the same input X. The sequence attends to itself. Every position can look at every other position of the same sequence.

Cross-attention is the same math, but Q comes from one sequence and K, V come from another:

Shapes:

• X_dec ∈ (B, T_dec, d) → Q ∈ (B, T_dec, d_k)
• X_enc ∈ (B, T_enc, d) → K, V ∈ (B, T_enc, d_k)
• scores (B, T_dec, T_enc), output (B, T_dec, d_v)

This is how the original transformer decoder pulls information from the encoder in machine translation: the decoder position asks "what's relevant in the source sentence for the word I'm about to generate?" and the encoder positions answer.

Decoder-only transformers (GPT, Llama, most modern LMs) don't have cross-attention — they only use self-attention with a causal mask. We'll build that in S040.

9 · Shape debugging cheatsheet

You will get shape errors. Everyone does. Here's how to diagnose them in 30 seconds.

10 · The three big pitfalls

11 · Mermaid diagram — the attention data flow

Read left to right. Three projections fan out from X, meet in a scores matrix, get soft-maxed into weights, and pull V through to produce the output. Nothing more, nothing less.

12 · Modern-2025 twist — attention is not "solved"

The equation softmax(QKᵀ/√d) V is timeless, but the implementation has been rewritten three times since 2022, each time saving billions of dollars of compute:

• FlashAttention (Dao et al., 2022, arXiv:2205.14135) — tiled the attention computation to fit in SRAM. Same math, ~2–4× wall-clock speedup on A100/H100.
• FlashAttention-2 (Dao, 2023) — reordered loops for better GPU occupancy. ~2× on top of FA-1.
• FlashAttention-3 (Shah, Bikshandi, Zhang et al., 2024) — exploits Hopper H100's asynchronous WGMMA and TMA units + FP8 e4m3 support. Reaches 75% of H100 theoretical peak (740 TFLOPs FP16, 1.2 PFLOPs FP8). This is what makes long-context Llama 3 405B training economically possible.
• Multi-head Latent Attention (MLA) in DeepSeek-V3 tech report (2024) compresses K and V into a low-rank latent, cutting KV-cache memory by 93.3% vs vanilla MHA. That's why DeepSeek-V3 can serve 128k-token contexts on far fewer H800s than its peers. We revisit MLA in S042/S043.
• Ring Attention (Liu, Zaharia, Abbeel, 2023) and Blockwise Parallel Transformer shard the sequence across devices, enabling Gemini 1.5's 1M-token context and Llama 4 Scout's 10M-token context (April 2025 release).

The research frontier as of 2026 is linear-attention hybrids (Mamba-2, Jamba 1.5, Griffin, RWKV-7) that mix attention with structured state-space models to get sub-quadratic scaling without giving up in-context recall. Attention isn't dying — it's being composed with recurrence. But every one of those hybrids still uses scaled dot-product attention in at least half its layers. The primitive isn't going anywhere.

Further reading:

• Tri Dao's FlashAttention-3 blog (Together AI, Jul 2024) — the tiling story with pictures.
• DeepSeek-V3 technical report — sections 2.1 (MLA) and 3.1 (auxiliary-loss-free MoE).
• Llama 4 Scout & Maverick announcement (Meta AI, Apr 2025) — 10M-context via interleaved chunked attention.

13 · Where this fits in the roadmap

We just built the atom. The rest of M07 builds molecules:

• S037 — Multi-head attention. Run several attention heads in parallel with different projections, concatenate. Lets one head learn syntax, another semantics, another position-relative patterns.
• S038 — Positional encodings. Attention as-derived is permutation-invariant (shuffle the input tokens and you get shuffled output tokens). Language cares about order. We add position info.
• S039 — The full transformer block. Attention + FFN + LayerNorm + residual = one transformer block. Stack N of them = a transformer.
• S040 — nanoGPT. Rebuild karpathy/nanoGPT line by line using everything above.

🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________________________

Spaced review: re-read §5 (numeric example) and §7 (√d derivation) in 24 hours. Revisit the full session on day 7. In week 2 revisit §12 (modern twist) as the bridge into S042/S043.

Next session (S037): why one attention head when you can have eight? We split d_model across h heads, run them in parallel, concat and project. Each head learns different relations — syntax, coreference, position-relative offsets. In 2025 the story continues with Grouped-Query Attention (Llama 3), Multi-Query Attention (PaLM), and Multi-head Latent Attention (DeepSeek-V3).

Sticky note (keep on your desk): Attention = weighted sum where weights come from learned similarity. Q asks, K matches, V is pooled. Divide by √d_k to keep softmax honest.

Previous: ← DL S035 · Next: DL S037 →