Attention is a weighted sum with a precise mask
Calculate attention numerically, derive its softmax backward pass, and test causal and permutation properties.
Original publication date unrecordedUpdated
Lesson preparation & details
Level: intermediate
By the end, you should be able to
- Calculate scaled dot-product attention by hand
- Implement causal masking before softmax
- Distinguish permutation equivariance from invariance
Bring with you
Sources checked: . A source check does not change the publication or update date.
Editorial review: · What review means
On the workbench / Attention
Change the score. Watch the mixture.
Two fixed values: 2 and 10. The second score is zero; adjust the first score before softmax. The output is a weighted average, not a vote for just one value.
Output = 0.500 × 2 + 0.500 × 10 = 6.000
The mask excludes the second score before normalization, so its weight is exactly zero. This fixed two-value instrument does not learn projections or demonstrate model understanding. Compare it with the full matrix calculation below.
In this article · 24 sections
Attention turns a hard lookup into a learned mixture. Start with a dictionary, build the weighted sum, then test the shapes, gradients and causal boundary with small numerical examples.
Prerequisites
You should be able to multiply (n, d) by (d, m) and explain which axis a sum removes. Review matrix products and broadcasting if these shapes are unfamiliar. Prior additive-attention intuition is useful but not required: the soft-lookup construction below starts from a dictionary.
1 · Start with something you already understand: a dictionary lookup
Forget attention for a moment. Let's talk about Python dictionaries.
d = {"cat": [1.0, 0.0, 0.0],
"dog": [0.0, 1.0, 0.0],
"fish": [0.0, 0.0, 1.0]}
d["cat"] # → [1.0, 0.0, 0.0]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:
- Replace exact-string match with a similarity score (dot product).
- Turn the scores into a probability distribution (soft-max).
- 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.
Attention as a room full of raised hands
Imagine you walk into a crowded room with a question in mind. You look around, and every person in the room has a label pinned to their chest (that's their key) and a thing they'd tell you in their pocket (that's their value). You compare your question against each label and mentally score how relevant that person seems — similarity scoring is the dot product. Then instead of picking just one person, you turn those scores into a distribution of "how much attention I'll pay to each" (that's the softmax), and take a weighted average of everyone's pocket-thing. That's one attention pass. Self-attention is the same room, except every person is simultaneously asking their own question and answering everyone else's — the whole room exchanges weighted value information through score computation, normalization and value aggregation.
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:
The dot product is the product of the vector norms and the cosine of their angle. Aligned nonzero vectors have a positive score, perpendicular vectors have zero, and opposite vectors have a negative score. Alignment alone does not determine magnitude: scaling a vector also scales its score.
Dot products let us calculate every query–key score with one matrix multiplication, Q @ K.T, and differentiate the result through the projections. Other scoring functions are possible. This is not a commutativity claim: matrix products generally do not commute.
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→ largerw_j. Order is preserved. - Controllable concentration. For fixed scores, multiplying by an increasing positive scale concentrates mass on their maximum set as the scale tends to infinity; equal maxima retain equal weights. Shrinking the scale toward zero approaches a uniform distribution. A factor of ten is not a universal guarantee of either extreme.
def softmax(x):
e = np.exp(x - x.max()) # subtract max for numerical stability
return e / e.sum()
w = softmax(scores)
print(w) # [0.4355, 0.3960, 0.1684] approximately
y = w @ V # weighted average of value vectors
print(y) # [7.9196, 2.0804] approximatelyThat output y 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.
Numerical stability trick: always subtract the max before exponentiating. exp(1000) overflows to infinity; exp(0) = 1 doesn't. Soft-max is invariant to adding a constant to all inputs, so subtracting the max is free.
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 all q_i, k_i are mutually independent, unit-variance and zero-mean, the dot product has variance d. So typical dot products have magnitude ≈ √d.
Large differences between scores can make softmax nearly one-hot, producing small derivatives for many entries. Absolute magnitude alone is not decisive: equal large scores still produce uniform weights after a stable softmax.
Fix: divide the scores by √d_k before soft-max. This is the "scaled" in scaled dot-product attention.
Section 7 derives this scaling under an explicit independence model and explains its limits.
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 ∈ ℝ^{d_{model} \times d_k} and W_V ∈ ℝ^{d_{model} \times d_v} 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.
Read the equation from left to right: compare queries and keys, normalize each query's scores, then combine the values.
What each piece is doing
- Q — 'what am I looking for?' — one query vector per position, derived from that position's input.
- K — 'what do I offer?' — one key vector per position, derived from that position's input.
- V — 'here's the content I provide if you attend to me' — one value vector per position.
- Q Kᵀ — all-pairs similarity matrix, shape (n, n). Entry (i, j) is how much position i wants information from position j.
- / √d_k — normalize score variance under the independence model; saturation can still occur after learning.
- softmax along rows — turn similarities into a probability distribution per query.
- · V — weighted average of value vectors, using those probabilities as weights.
The equation has two paths: Q and K determine the weights, while V supplies the content those weights combine. Follow those paths from the three projections below. The diagram includes a batch axis B; within each batch item, T is the sequence length we called n above.
Figure text equivalent — labels and relationships
- Node “X: (B, T, d_model)” (square).
- Node “× W_Q” (square).
- Node “× W_K” (square).
- Node “× W_V” (square).
- Node “Q: (B, T, d_k)” (square).
- Node “K: (B, T, d_k)” (square).
- Node “V: (B, T, d_v)” (square).
- Node “scores = Q Kᵀ / √d_k (B, T, T)” (square).
- Node “weights = softmax(scores, dim=-1) (B, T, T)” (square).
- Node “Y = weights · V (B, T, d_v)” (square).
- “X: (B, T, d_model)” leads to “× W_Q”. Line: normal; ends: arrow point.
- “X: (B, T, d_model)” leads to “× W_K”. Line: normal; ends: arrow point.
- “X: (B, T, d_model)” leads to “× W_V”. Line: normal; ends: arrow point.
- “× W_Q” leads to “Q: (B, T, d_k)”. Line: normal; ends: arrow point.
- “× W_K” leads to “K: (B, T, d_k)”. Line: normal; ends: arrow point.
- “× W_V” leads to “V: (B, T, d_v)”. Line: normal; ends: arrow point.
- “Q: (B, T, d_k)” leads to “scores = Q Kᵀ / √d_k (B, T, T)”. Line: normal; ends: arrow point.
- “K: (B, T, d_k)” leads to “scores = Q Kᵀ / √d_k (B, T, T)”. Line: normal; ends: arrow point.
- “scores = Q Kᵀ / √d_k (B, T, T)” leads to “weights = softmax(scores, dim=-1) (B, T, T)”. Line: normal; ends: arrow point.
- “weights = softmax(scores, dim=-1) (B, T, T)” leads to “Y = weights · V (B, T, d_v)”. Line: normal; ends: arrow point.
- “V: (B, T, d_v)” leads to “Y = weights · V (B, T, d_v)”. Line: normal; ends: arrow point.
Q and K meet at the scores node; V bypasses that node and joins after softmax. The position-by-position weights therefore change which values contribute to each output, rather than adding another feature projection. In the worked example next, compute these same intermediates for three tokens.
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 = XSo 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.248, 0.248, 0.503]]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.4011·[1,0] + 0.1978·[0,1] + 0.4011·[1,1] = [0.802, 0.599].
Row 1: 0.1978·[1,0] + 0.4011·[0,1] + 0.4011·[1,1] = [0.599, 0.802].
Row 2: 0.2483·[1,0] + 0.2483·[0,1] + 0.5035·[1,1] = [0.7517, 0.7517].
Y ≈ [[0.802, 0.599],
[0.599, 0.802],
[0.752, 0.752]]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.
Practice: Do this by hand
Redo the calculation with W_Q = [[1, 0], [0, 0]] (queries look only at dim 0) and W_K = W_V = I. Predict what happens to the attention weights, then verify. For the query from [0,1], the new query is zero and all valid keys receive equal weights. The other rows prefer large dim-0 keys but need not collapse to one-hot probabilities.
5.1 · A two-key scalar check
For , keys and , scores are approximately . The weights are approximately . With scalar values 2 and 10, the result is . The result can mix both values; it need not retrieve exactly one of them.
import numpy as np
q = np.array([[1., 0.]])
k = np.eye(2)
v = np.array([[2.], [10.]])
scores = q @ k.T / np.sqrt(2)
weights = np.exp(scores-scores.max(axis=-1, keepdims=True))
weights /= weights.sum(axis=-1, keepdims=True)
out = weights @ v
np.testing.assert_allclose(weights.sum(axis=-1), 1.)
assert abs(out.item()-4.641907) < 1e-6
print('weights', weights.tolist(), 'output', out.item())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 the numbers differ, inspect the score matrix and normalization axis before the final multiply. The helper above defaults to axis=-1; explicitly passing axis=None normalizes across the entire array instead of separately for each query.
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: (d_model, d_k); W_V: (d_model, d_v) — weights shared across positions
Q, K: (B, T, d_k); V: (B, T, d_v)
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 @ VThis framework reference illustrates the batch axes; it has not been executed here. The NumPy fixtures provide the bounded local checks.
7 · Why the √d_k? A deeper derivation
Assume all query and key components are mutually independent with mean 0 and variance 1. This initialization-style model is not guaranteed for learned projections, especially a query and key derived from the same token. Then:
So the standard deviation of a raw dot product is √d_k, about 22.6 for d_k = 512 under this model. This raises the risk of large score differences and softmax saturation; it does not imply that every row is concentrated or that all training necessarily stalls.
Under those assumptions, dividing by √d_k restores unit score variance. It reduces a dimension-driven source of saturation; it does not guarantee that trained logits never saturate.
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. Without a restricting mask, each query can use every key in that 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 ∈ (B, T_enc, d_k)andV ∈ (B, T_enc, d_v)- 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.
A conventional text-only decoder stack uses causal self-attention rather than an encoder–decoder cross-attention sublayer. Multimodal and retrieval-augmented designs can add cross-attention, so inspect the architecture rather than inferring its sublayers from a product label.
9 · Shape debugging cheatsheet
Use the following shape checks to localize a mismatch.
If self-attention throws a shape error, check in this order
-
- Print X.shape. Should be (B, T, d_model). If it's (B, d_model, T) you fed it channels-first, transpose.
-
- Print Q.shape, K.shape, V.shape. Q and K should be (B,T,d_k); V may be (B,T,d_v). Key and value sequence lengths must agree.
-
- Print scores.shape after Q @ K.transpose(-2, -1). Must be (B, T, T). If it's (B, d_k, d_k), you transposed the wrong axes.
-
- Print weights.sum(dim=-1). Every element must equal 1.0 within floating-point error. If not, you may have soft-maxed on the wrong axis, or encountered a non-finite or entirely masked row.
-
- Print output.shape. Should match (B, T, d_v). If it lost the batch dim, you called .squeeze() somewhere.
10 · The three big pitfalls
Wrong normalization axis. Every query needs a distribution over keys. A softmax across queries produces a different operator, even if the resulting array shape looks correct. Check row sums, then test a non-square cross-attention example so a coincidental square shape cannot conceal transposition.
Masking after softmax. Zeroing prohibited probabilities after normalization leaves rows summing below one unless renormalized. Prefer setting prohibited logits to negative infinity before softmax, with at least one allowed key per query. A fully masked row makes the usual subtract-maximum implementation undefined; reject it or explicitly define a zero output. A large finite sentinel is not a universal equivalent for every dtype.
Unjustified explanation. Attention weights show one intermediate routing distribution. Values, output projections, residual paths and later layers affect predictions. A bright attention cell does not prove causal importance. Perturbation tests can help investigate a particular model but require their own assumptions.
10.1 · Causality is an invariant you can test
When position predicts token , it may use positions up to and including , not future positions. Set forbidden scores to negative infinity before softmax. Multiplying probabilities by a mask afterward without renormalizing changes the row sum and is not the same operation.
Every query needs at least one allowed key. An all-masked row produces undefined normalization; a generic padding implementation must handle that case explicitly. Our causal example always permits the diagonal.
import numpy as np
rng = np.random.default_rng(2)
x = rng.normal(size=(4, 3))
Wq, Wk, Wv = [rng.normal(size=(3, 2)) for _ in range(3)]
def attention(x, causal=True):
q, k, v = x @ Wq, x @ Wk, x @ Wv
s = q @ k.T / np.sqrt(q.shape[-1])
if causal:
s = np.where(np.tril(np.ones(s.shape, dtype=bool)), s, -np.inf)
a = np.exp(s-s.max(axis=-1, keepdims=True))
a /= a.sum(axis=-1, keepdims=True)
return a @ v, a
out, a = attention(x)
changed = x.copy(); changed[-1] += 100
np.testing.assert_allclose(attention(changed)[0][:-1], out[:-1])
assert np.all(a[np.triu_indices(4, 1)] == 0)
np.testing.assert_allclose(a.sum(axis=-1), 1)
perm = np.array([2, 0, 3, 1])
np.testing.assert_allclose(attention(x[perm], False)[0], attention(x, False)[0][perm])
print('causal isolation and unmasked permutation equivariance verified')Without positional information and without the causal mask, permuting all input positions permutes outputs: this is equivariance, not invariance. Adding a position-dependent mask or position embeddings changes the argument. “Attention ignores order” is too imprecise unless those assumptions are stated.
10.2 · Projection roles, symmetry and biases
Q, K and V are learned projections of the same input in self-attention, not guaranteed semantic labels such as “intent,” “identity” and “meaning.” Their roles come from the computation. Tying query and key projections makes the raw score matrix symmetric, but does not generally make row-normalized attention symmetric: different rows can have different normalization sums. Separate projections permit a nonsymmetric bilinear score without tying the two feature maps.
import numpy as np
x = np.array([[1., 0.], [1., 2.]])
scores = x @ x.T / np.sqrt(2)
weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights /= weights.sum(axis=-1, keepdims=True)
assert np.allclose(scores, scores.T)
assert not np.allclose(weights, weights.T)
# Adding the same key bias adds a constant to each score row.
bias = np.array([.3, -.7])
biased = x @ (x + bias).T / np.sqrt(2)
bp = np.exp(biased - biased.max(axis=-1, keepdims=True))
bp /= bp.sum(axis=-1, keepdims=True)
np.testing.assert_allclose(bp, weights)
print('symmetric scores need not give symmetric weights; shared key-bias cancellation checked')For a fixed query, a shared additive key bias contributes the same scalar to every allowed score, so it cancels under softmax. A query bias contributes a term that varies across keys and does not generally cancel. A value bias is added to the output when the allowed weights sum to one. These facts do not imply that every Transformer omits Q/K/V biases or that omitting them always improves stability. The reference PyTorch class makes a deliberate bias-free choice; it is not a universal architecture requirement.
10.3 · Backward pass and a shape-level derivation
Let P be row-wise softmax of S, S=QK-transpose/sqrt(d_k), and Y=PV. For upstream G=dL/dY, dV=P-transpose G and dP=G V-transpose. For each row, dS=P times (dP-sum(P*dP)), where the sum is over keys. Therefore dQ=dS K/sqrt(d_k) and dK=dS-transpose Q/sqrt(d_k). Propagate these through Q=XW_Q, K=XW_K and V=XW_V, adding all three contributions to dX. A masked score has zero derivative when its mask is fixed. This is a vector–Jacobian product, not a need to allocate a full softmax Jacobian.
For a cross-attention example with 2 queries, 5 keys, key width 3 and value width 4, S and P have shape (2,5) and Y has shape (2,4). The formulas produce dQ (2,3), dK (5,3) and dV (5,4). Write those shapes before implementing the backward pass.
11 · Modern-2025 twist — exact attention, cache compression and recurrence
The equation is an operator; its implementation and architecture determine cost. FlashAttention tiles exact softmax attention and carries a running row maximum, normalization sum and weighted-value accumulator. If a new tile raises the maximum from m to m', rescale the old denominator and numerator by exp(m-m') before incorporating the tile. This avoids materializing all scores in high-bandwidth memory. Pairwise arithmetic remains quadratic; wall time depends on device, dtype, head dimensions and workload.
Multi-query and grouped-query attention share keys/values across query heads. They reduce cache size without reducing the number of query outputs; see the multi-head and KV-cache lessons for shape accounting. MLA compresses the stored representation with learned low-rank structure and may retain additional positional components. A memory ratio requires the actual model dimensions, not only a head count.
A recurrent state such as KDA has a different capacity/cost tradeoff: bounded state supports constant per-token storage for that recurrent component, but does not preserve a separate exact key/value vector for every prior token. Hybrid models retain global-attention layers for selective access. Do not assert that every recurrent architecture uses softmax attention in half its layers; that is architecture-specific. The Kimi architecture study derives one pinned hybrid rather than generalizing a headline to all models.
Practice: calculate the bytes of a materialized attention matrix for batch 2, 8 heads, 4096 tokens and two-byte elements. The result is 536,870,912 bytes (512 MiB), excluding projections, output and gradients. Reducing KV heads does not by itself shrink this query-head score matrix. FlashAttention avoids storing the whole matrix, but that byte estimate is not its actual peak memory.
Different implementation improvements solve different bottlenecks
FlashAttention-2, v1 reduces non-matrix-multiply work and changes how thread blocks and warps share work to improve occupancy and reduce shared-memory traffic. FlashAttention-3, v2 targets Hopper hardware: it overlaps movement and compute using asynchrony and warp specialization, interleaves matmul with softmax, and studies FP8 processing. Their reported speedups belong to specified hardware, shapes, precision and baselines—not every attention call.
Ring Attention, v4 instead distributes sequence blocks across devices and circulates key/value blocks while computing blockwise attention. This addresses per-device memory placement and communication overlap; it does not make full pairwise attention linear in total sequence length. Good overlap requires adequate computation relative to communication and a suitable interconnect. A paper's million-token experiment does not establish that a named proprietary model uses that implementation.
Choosing a position-mixing operator
| Operator | What one output can use | Cost and state, under fixed widths | Important boundary |
|---|---|---|---|
| Full self-attention | Any permitted key in one layer, with input-dependent weights | O(n²) pairwise arithmetic; a naive score matrix uses O(n²) space, while tiled implementations avoid storing it | Short graph distance does not guarantee successful long-range learning or interpretability |
| Local convolution of width k | A local neighborhood per layer, with learned shared kernel weights | O(nk) work; receptive field grows with stacked layers or dilation | Locality and translation structure can be useful inductive biases; nonlinearity still makes the network's overall behavior input-dependent |
| Fixed-state recurrence | Earlier inputs through the carried state | O(n) sequential state updates and fixed-size inference state for the stated architecture | Standard backpropagation retains/recomputes states; training memory is not automatically constant, and some structured recurrences admit parallel scans |
A hybrid can spend cheap local or recurrent computation on most positions and use global attention selectively. This is a design option, not proof that any one mixture dominates every task. Quantify storage and information access separately: a fixed-size recurrent state cannot preserve an arbitrary collection of independent past values exactly.
Worked comparison: for n=4096, full attention has 16,777,216 query–key pairs per head; a width-31 local stencil has at most 126,976 position–offset pairs before boundary handling. The latter cannot directly choose an arbitrary far-away token in one local layer. Fewer pairs alone therefore does not imply equivalent behavior.
12 · Where this fits in the roadmap
The next concepts build on this operator. With model width d, total query/key/value head width d, and four dense projection matrices, multi-head attention has 4d² parameters excluding biases. Splitting that fixed total width into more heads does not multiply this count by the number of heads; architectures with shared KV heads or different widths need their own accounting.
- S037 — Multi-head attention. Run several attention heads in parallel with different projections, concatenate. This provides multiple learned subspaces; interpretable roles such as syntax or coreference are not guaranteed.
- S038 — Positional encodings. Without position information or an order-dependent mask, the derived operator is permutation-equivariant: permuting inputs permutes outputs. Language cares about order; position features and masks change that symmetry.
- 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.
Exercises and answer sketches
- If every allowed score is zero, what are the weights? Answer: uniform over allowed positions, not over all positions including masked ones.
- Why can values have width different from keys? Answer: keys determine matching dimensions with queries; values are combined only after weights are computed.
- How do you detect a transposed causal mask? Answer: perturb a future input and assert that earlier outputs do not change; also inspect zero probability above the diagonal.
- Does an attention heatmap establish why a model answered incorrectly? Answer: it visualizes one intermediate routing pattern. Output projections, residual paths and later layers also matter; causal explanations require interventions and a specified question.
🧠 Retention scaffold
Q1 (recall) — Write Attention(Q, K, V) = softmax(Q Kᵀ / √d_k) V from memory and label every symbol including tensor shapes.
Q2 (recall) — Reproduce the length-3, d=2 worked example from §5 in your head: what were the final attention weights for query row 2?
Q3 (recall) — Why √d_k and not d_k or √d_model? Give the variance argument in one sentence.
Q4 (recall) — What tests distinguish score saturation from normalization along the wrong axis?
Q5 (recall) — Shape check: for input X ∈ (B, T, d_model), give the shape of Q, K, V, scores, weights, and output.
Q6 (stretch) — Read the DeepSeek-V3 MLA section (§2.1) and explain in your own words how compressing K/V into a shared latent breaks the naive KV-cache memory scaling law while preserving the softmax(QKᵀ) structure. Which of the three steps we derived does MLA modify?
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 §11 (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. Different heads can learn different relations, without a guarantee that each has a distinct interpretable role. 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.
Source and execution boundary
The Transformer paper, revision 7, sections 3.2 and 4 supports scaled dot-product attention, head concatenation, causal masking and the comparison with recurrence/convolution. Projection-symmetry and bias conclusions above follow directly from the displayed algebra and finite NumPy checks. The cited FlashAttention and Ring Attention abstracts support their named implementation mechanisms, not measured performance on this machine. The worked arithmetic is checked in the corrective CPU artifact. PyTorch snippets and GPU kernels were not run; the NumPy tests are not a benchmark or model-quality claim.
Pause / Recall / Apply
Can you explain it without the page?
Close the example. Reconstruct the core idea, then change one assumption. Mark complete when you’re ready; you can always undo it.
Stored in this browser only. No account, no sync. Clearing browser data removes your record.