Core day 39 of 60 · Week 8 · Part II · application depth

Attention and Transformers

A ticket-triage model reads "bag delayed refund" and a colleague shows an attention heat map in which "refund" gets 40% of one token's weight as "the reason" for the prediction. Can you compute such weights for three tokens by hand, give the shapes through the layer, and say what the map does and does not show?

By the end: You compute query, key and value projections, scaled scores, softmax weights and outputs for three tokens with d_k = 2, apply a causal mask, state the shape at every step, and explain why attention weights are not an explanation.

At a glance

Each token mixes all value vectors with softmax weights from query-key matches; shapes are fixed by n and d_k, cost grows with n squared, and weights are not reasons.

Query key value
Three learned projections of the same token vectors, each (n, d_k).
Softmax over keys
Scores QK^T / sqrt(d_k) are (n, n); each row of weights sums to 1.
Mask before softmax
Future positions get minus infinity, so decoders cannot peek.
n squared cost
Doubling the context quadruples the score matrix per head and layer.
Not an explanation
A heat map shows how values were mixed; test it by masking tokens.

Recall · 5 min

No notes: (1) the shape of $A B$ when $A$ is $(3, 2)$ and $B$ is $(2, 3)$ ; (2) what softmax does to a vector of scores and what it is applied over; (3) one reason a recurrent encoder struggles with a sentence of 200 tokens (Day 37). Today's layer answers the third.

Open due reviews

Concept · 10 min

A recurrent network pushes everything through one small state. In attention, every token looks directly at every other token. Think of a soft lookup table: each token asks with a query, offers a key to be matched on and a value to hand over. The picture breaks in one place: the roles are learned linear projections of the same vector, and mean only what training made them mean.

Attention is a matrix product with a softmax over keys - change a score, a scale or the mask and watch the weights move

Predict first

With the causal mask on, what does token 0 attend to in a three-token sequence?

Text description of the figure

A three-token example (flight, was, late) showing query, key and value vectors, the scaled score matrix, the row-wise softmax weights and the output vectors. Controls edit the query, key and value entries, the scaling and a causal mask, with presets for large scores and identical keys.

Four parts of one attention head Overview

Query
What this token is looking for: q = x W_Q.
Key
What this token can be matched on: k = x W_K.
Value
What this token hands over if matched: v = x W_V.
Softmax over keys
Turns one query's scores into weights that sum to 1.

Shapes through one head (n = 3 tokens, d_model = 3, d_k = 2) Diagram

X(3, 3): one row per token
- to K = X W_K
- to V = X W_V
Q = X W_Q(3, 2)
- to Q K^T / sqrt(d_k)
K = X W_K(3, 2)
- to Q K^T / sqrt(d_k)
V = X W_V(3, 2)
- to Output = A V
Q K^T / sqrt(d_k)(3, 3) scores
Softmax per row(3, 3): rows sum to 1
Output = A V(3, 2): mixed values

The only (n, n) objects are the scores and the weights; everything else has one row per token.

One head on six tokens - no mask, causal mask, and no scaling Illustration

Encoder: every token reads all six. Each row sums to 1 (checked); columns do not. Average spread 4.8 of 6 keys.

Decoder: future keys are set to minus infinity before the softmax, so the upper triangle is exactly 0 and the first token reads only itself.

Same Q and K without dividing by sqrt(16) = 4: scores 4x larger, rows sharper (average spread 2.3). Nothing was learned; only the scale changed.

Seeded random token vectors (d_model = 16) and random projections (d_k = 16); no training, so which token attends to which is arbitrary. Spread = 2 to the power of the row's entropy in bits: 1 = one key, 6 = uniform.

Rows are queries, columns keys. The mask zeroes the upper triangle; dropping the scale sharpens every row without anything being learned.

One transformer block, bottom to top Stack

Token vectors plus positionsEmbeddings (Day 38) plus sinusoidal or learned position vectors.
Multi-head self-attentionh heads in parallel, d_k = d_model / h, concatenated and mixed by W_O.
Residual plus LayerNormAdd the input back, then normalize each token's vector.
Feed-forward networkThe same two-layer MLP (often GELU, Day 32) applied to each position.
Residual plus LayerNormAdd and normalize again; the output has the input's shape.

Stacks of these form encoders and decoders (Vaswani et al., 2017). Many current models normalize before each sublayer instead (pre-LN).

When does the n-squared part dominate a layer's compute? Illustration

d_model = 1024: the per-token parts grow like n, the score parts like n squared, and they meet at n = 6 d = 6,144 tokens. Below that, longer context costs roughly linearly.

d_model = 768 (47% at 4k tokens)
d_model = 1,024 (40% at 4k tokens)
d_model = 4,096 (14% at 4k tokens)

Share of a layer's multiply-adds spent on the n x n part. Wider models push the crossover out: at 4,096 tokens a d_model = 4,096 layer spends only a seventh of its compute there.

Exact multiply-add counts for the forward pass of one layer, batch 1, ignoring softmax, normalisation and memory traffic. Memory is a separate story: a naive implementation also stores the n x n weights per head.

The per-token projections and feed-forward network grow linearly; scores and the weighted sum grow with n squared and take over at about six times d_model.

Recurrence versus self-attention Comparison

Criterion	Recurrent network (Day 37)	Self-attention
Steps between two distant positions	Up to n steps, a product of factors	One step: every pair is scored directly
Parallel over positions	No: step t waits for step t - 1	Yes, within a layer
Memory of the past	One fixed-size state	Every position, in the context window
Cost growth with length n	Linear in n	Quadratic in n per head and layer

Project the token matrix $X$ $(n, d_{model})$ into $Q$ , $K$ , $V$ , each $(n, d_k)$ ; scores $QK^\top / \sqrt{d_k}$ are $(n, n)$ ; softmax over keys makes each row sum to 1; the output is weights times $V$ .
The $\sqrt{d_k}$ keeps scores in a range where the softmax is not saturated.
Multi-head attention runs $h$ maps in parallel and mixes them with $W_O$ . Self-attention ignores order, so positions are added; blocks add residuals, Layer normalization and a feed-forward network.
A Causal mask sets future scores to $-\infty$ before the softmax.
Cost: $n^2$ score entries per head and layer.
Weights are not explanations by default (Jain and Wallace, 2019; Wiegreffe and Pinter, 2019, show it depends on the definition): treat a heat map as a hypothesis and test it by perturbation.

Assumptions

Pairwise interactions between all positions are worth their $n^2$ cost; cheap for short sequences, not for very long ones.
Order information is supplied (positions) and, for generation, future tokens are masked.
The training objective rewards the roles you hope heads play; interpretation of a head is a hypothesis.

Application and workplace decision · 10 min

Three tokens "bag", "delayed", "refund" are one-hot rows of $X$ ( $3 \times 3$ ), so each token's query is the matching row of $W_Q$ (hand-set, not learned). With $d_k = 2$ : $W_Q$ rows $(1,0), (0,1), (1,1)$ ; $W_K$ rows $(1,0), (1,1), (0,1)$ ; $W_V$ rows $(2,0), (0,2), (1,1)$ .

By hand, for token 0:

(a) the raw scores $q_0 \cdot k_j$ and (b) the scaled scores;
(c) the softmax weights and (d) the output $\sum_j a_{0j} v_j$ .

Then run the code, compare, and read the causal weights.

Your answer to the task (decision, method, baseline). Write it before reading the interpretation.

Python

import numpy as np
np.set_printoptions(precision=3, suppress=True)
X = np.eye(3)                      # 3 tokens, d_model = 3: bag, delayed, refund (one-hot)
Wq = np.array([[1, 0], [0, 1], [1, 1]], dtype=float)   # (d_model, d_k) = (3, 2), hand-set
Wk = np.array([[1, 0], [1, 1], [0, 1]], dtype=float)
Wv = np.array([[2, 0], [0, 2], [1, 1]], dtype=float)
Q, K, V = X @ Wq, X @ Wk, X @ Wv   # each (3, 2)
d_k = Q.shape[1]

def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

scores = Q @ K.T / np.sqrt(d_k)    # (3, 3): query i against key j
A = softmax(scores)                # softmax over keys: each row sums to 1
out = A @ V                        # (3, 2)
print("scores / sqrt(d_k):\n", scores)
print("weights (rows sum to 1):\n", A)
print("output:\n", out)

mask = np.triu(np.ones((3, 3), dtype=bool), k=1)   # True above the diagonal = future positions
A_causal = softmax(np.where(mask, -np.inf, scores))
print("causal weights:\n", A_causal)
print("token 0 attends only to itself:", bool(A_causal[0, 0] == 1.0))
print("shapes: Q", Q.shape, "scores", scores.shape, "out", out.shape)

This code is executed in the repository's checks (pnpm lesson:py / lesson:sql). It is not run on this site; copy it into your notebook.

Technical depth · optional · 32 min in total

Open only what you need. Nothing here is required to complete the day; the assessment covers the core above.

The layer, defined3 minScaled dot-product attention as a matrix formula and row by row, what is learned, and the assumptions written as properties of the equation.

Scaled dot-product attention Formula

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

$QK^\top$: (n, n): dot product of query i with key j
$\sqrt{d_k}$: rescales scores so their spread does not grow with d_k
$\operatorname{softmax}$: applied to each row, that is over keys
$V$: (n, d_k): the vectors being mixed

Attention for one query, written row by row Formula

o_i = \sum_{j=1}^{n} a_{ij}\, v_j, \qquad a_{ij} = \frac{\exp(s_{ij})}{\sum_{l} \exp(s_{il})}, \qquad s_{ij} = \frac{q_i \cdot k_j}{\sqrt{d_k}}, \qquad q_i = x_i W_Q

$x_i$: the d_model-vector of token i (embedding plus position)
$q_i, k_j, v_j$: x_i W_Q, x_j W_K, x_j W_V: three learned linear views of the same tokens
$s_{ij}$: how well query i matches key j, scaled so its spread does not grow with d_k
$a_{ij}$: softmax over j: non-negative, sums to 1 across keys for each query i
$o_i$: a weighted average of the values: it always lies inside their convex hull

The matrix form in the formula figure is these n rows stacked. What is learned: W_Q, W_K, W_V (and W_O with several heads); the weights a_ij are recomputed for every input and carry no parameters.

One head as an algorithm, with the cost of each step Sequence

1. Project
Q = X W_Q, K = X W_K, V = X W_V: one row per token, about 3 n d_model d_k multiply-adds
2. Score
S = Q K^T / sqrt(d_k), an (n, n) matrix: about n^2 d_k multiply-adds
3. Mask
Set forbidden entries of S (future keys, padding) to minus infinity
4. Softmax per row
A = softmax(S) over keys: every row is non-negative and sums to 1
5. Mix
O = A V: each output row is a weighted average of value rows, about n^2 d_k multiply-adds

Steps 2 and 5 are the n-squared part; steps 1 and the output projection are per token.

Three bets the formula makes Overview

Pairwise relevance
Every pair (i, j) is scored: any token may matter to any other. The price is n^2.
Order from outside
The formula never reads i or j as numbers, so positions must enter through x_i.
A learned similarity
q_i . k_j is a dot product in a learned space: 'relevant' means only what training rewarded.

Softmax weighting by hand on three tokens4 minOne query pushed through scores, scale, softmax and value mixing on a three-event bank session, then the application's own weights to check your hand work.

Softmax weighting by hand: a bank session of three events (d_k = 4) Sequence

Tokens
login, new-device, transfer. Keys: k0 = (1, 0, 0, 0), k1 = (1, 1, 0, 1), k2 = (0, 1, 1, 0)
Query of transfer
q2 = (1, 1, 0, 1). Values (2-d risk features): v0 = (1, 0), v1 = (0, 1), v2 = (1, 1)
Raw scores
q2 . k0 = 1, q2 . k1 = 1 + 1 + 0 + 1 = 3, q2 . k2 = 0 + 1 + 0 + 0 = 1
Scale by 1 / sqrt(4)
(0.5, 1.5, 0.5)
Exponentiate
e^0.5 = 1.649, e^1.5 = 4.482, e^0.5 = 1.649; sum 7.779
Normalise
Weights (0.212, 0.576, 0.212): transfer reads mostly from new-device, and the row sums to 1
Mix the values
o2 = 0.212 (1, 0) + 0.576 (0, 1) + 0.212 (1, 1) = (0.424, 0.788)
Without the scale
softmax(1, 3, 1) = (0.107, 0.787, 0.107): sharper from the same vectors

Hand-set vectors (not a trained fraud model), verified in numpy; on paper only the exponentials need a calculator. Read o2 as an average: mostly v1, a little of v0 and v2. Nothing in the weights says which event made the session risky. The application starts one step earlier, from one-hot tokens and hand-set projections.

Check your hand result: today's weights without a mask (rows sum to 1) Grid

key →
query ↓	bag	delayed	refund
bag	0.401	0.401	0.198
delayed	0.198	0.401	0.401
refund	0.248	0.503	0.248

Printed by the application code. Outputs: bag (1.000, 1.000), delayed (0.797, 1.203), refund (0.745, 1.255). 'refund' attends most to 'delayed', not to itself.

Why divide by the square root of d_k4 minThe variance of a dot product derived term by term, what it does to the softmax and its gradient, then measured on random vectors.

The square-root argument, derived Sequence

Assume
Components q_c and k_c independent, mean 0, variance 1 (roughly true at initialisation)
One term
E[q_c k_c] = E[q_c] E[k_c] = 0 and Var(q_c k_c) = E[q_c^2] E[k_c^2] = 1
Sum of d_k terms
q . k has mean 0 and variance d_k (independent terms add their variances)
Spread
SD(q . k) = sqrt(d_k): 2 at d_k = 4, 4 at d_k = 16 (20,000 random pairs in numpy: 3.99)
Rescale
Divide by sqrt(d_k): the scaled score has variance 1 at every width
Why it matters
Softmax slopes are a_i (1 - a_i) and -a_i a_j: max 0.22 for a uniform 3-way row, 0.02 for (0.98, 0.01, 0.01)

Saturated rows pass almost no gradient to the scores, so the projections barely learn. Vaswani et al. (2017) give this variance argument for the scale.

Spread (standard deviation) of q . k for unit-variance components Chart

d_k = 42
d_k = 164
d_k = 648
d_k = 25616

A sum of d_k unit-variance products has variance d_k, so its spread is sqrt(d_k). Dividing by sqrt(d_k) brings it back to 1.

Without the scale, rows collapse toward one key as d_k grows Illustration

A dot product of d_k unit-variance products has standard deviation sqrt(d_k): measured 1.3, 2.1, 2.9, 4.0, 5.7, 8.4, 11.3, 16.0.

Scaled rows stay spread over 5.5 to 6.2 of 8 keys at every d_k. Unscaled rows collapse to 1.2 keys at d_k = 256: near one-hot, so almost no gradient reaches the other keys.

Seeded random queries and keys with independent unit-variance components (the situation at initialisation), 40 draws of 8 tokens per d_k. Effective keys = 2 to the power of the row entropy in bits.

Measured on random queries and keys, the situation at initialisation. The scale keeps the softmax in its responsive range at every width.

What a larger score spread does to the softmax Sequence

Scaled scores (0.707, 0.707, 0)
Weights (0.401, 0.401, 0.198): today's token 0
Unscaled (1, 1, 0)
Weights (0.422, 0.422, 0.155): slightly sharper
Eight times larger (8, 8, 0)
Weights (0.500, 0.500, 0.0002): the third key is ignored
Consequence
Near-one-hot rows pass almost no gradient to the other keys

The causal mask by hand, and the library defaultsadvanced4 minMasking worked on the same three events, causal and padding masks, the two opposite boolean conventions in PyTorch, and the head-averaged weight map.

The causal mask by hand, on the same session Sequence

Row of new-device
Query q1 = (0, 1, 1, 0); scaled scores against keys 0, 1, 2: (0, 0.5, 1.0)
Unmasked softmax
(0.186, 0.307, 0.506): the largest weight goes to transfer, an event that has not happened yet
Mask the future
Set s_12 to minus infinity; exp(minus infinity) = 0, so the key drops out of the sum
Masked softmax
Over keys 0 and 1 only: e^0 = 1, e^0.5 = 1.649, so (0.378, 0.622, 0)
Row of login
Only key 0 is allowed, so its weight is exactly 1 whatever the scores
Row of transfer
The last row sees every earlier key and is unchanged by the mask

Verified in numpy. The mask is applied to the scores before the softmax; zeroing weights after it would leave rows that no longer sum to 1.

The causal mask on today's three tokens: weights after masking Grid

key →
query ↓	bag	delayed	refund
bag (token 0)	1.000	0.000	0.000
delayed (token 1)	0.330	0.670	0.000
refund (token 2)	0.248	0.503	0.248

Masked scores become minus infinity BEFORE the softmax, so the upper triangle is exactly 0 and every row still sums to 1. Token 0 gives itself weight 1; the last row sees every key and is unchanged. Computed by the application code.

Two masks, one mechanism: minus infinity before the softmax Comparison

Criterion	Causal mask	Padding mask
Removes	Keys after the query position	Padded keys
Protects against	A decoder reading the token it must predict	Outputs that depend on what a sentence was batched with
Unit test	Change a later token; earlier outputs stay fixed	Add padding; real tokens' outputs stay fixed

PyTorch attention defaults that change results Table

Setting	Default	Why it matters
SDPA scale	None = 1 / sqrt(E)	Pass your own only on purpose; the default is the paper's scale
SDPA is_causal	False	Decoders must ask for the mask; passing it with attn_mask is an error
SDPA boolean attn_mask	True = may attend	The opposite of MultiheadAttention's key_padding_mask
MHA key_padding_mask	True = padding, ignore	Invert with ~mask when moving between the two APIs
MHA batch_first	False: (seq, batch, feature)	A silent transpose if your tensors are batch first
MHA need_weights, average_attn_weights	True, True	The returned map is averaged over heads, not one head

Read from the PyTorch 2.14 pages for scaled_dot_product_attention and MultiheadAttention. A row whose keys are all masked gives a softmax of nothing, which comes out as NaN. Because returned maps are head averages, a 'head that attends to refund' may not exist.

Properties and the tests that check them3 minPermutation equivariance, convex mixing, input-dependent weights, parameter count and quadratic cost, each written in symbols and tied to a cheap test.

Properties of self-attention, each with the test that checks it Table

Property	In symbols	Consequence	Test
Permutation equivariant	Attn(PX) = P Attn(X) without positions	Word order is invisible until positions are added	Permute tokens; outputs permute (the lab)
Convex mixing	o_i in the convex hull of the values	One head cannot output beyond its values; W_O and the MLP must	Check o_i against min and max of V
Input-dependent weights	a_ij is a function of X	The map changes per example; parameters do not	Recompute the map on two inputs
Parameters independent of n	W_Q, W_K, W_V, W_O sized by d_model only	Longer context costs activations, not weights	Count parameters at two context lengths
Quadratic scores	n x n per head and layer	Time and naive memory grow with n squared	Profile at n and 2n
Masked rows	a_ij = 0 for masked j, rows still sum to 1	No leak from the future or from padding	Change a later token; earlier outputs fixed

Run these tests before reading any heat map. Equivariance follows from the row formula: permuting X's rows permutes Q, K and V alike, so S and A permute on both axes and O on its rows. Convex mixing is why blocks need the residual path and the MLP: attention alone only averages what is there.

Multiple heads, and sharing keys and valuesadvanced3 minHow heads split the width, what concatenation and W_O do, and why decoders share key-value heads.

Multi-head attention: shapes (d_model = 8, two heads) Diagram

X(n, 8)
- to Head 2
Head 1own W_Q, W_K, W_V; output (n, 4)
- to Concatenate
Head 2own W_Q, W_K, W_V; output (n, 4)
Concatenate(n, 8)
Mix with W_O(8, 8); output (n, 8)

With d_k = d_model / h, all heads together cost about as much as one full-width head. Heads can specialise; nothing forces them to.

How heads share the width Comparison

Criterion	One full-width head	Multi-head, h heads	Grouped-query
Query heads	1, width d_model	h, width d_model / h each	h
Key-value heads	1, width d_model	h, one per query head	g < h, each shared by h / g query heads
Attention weights (no biases)	4 d_model^2	4 d_model^2: the width is split, not multiplied	W_K and W_V shrink to g / h of their width
Decoder cache	Keys and values at full width	Same as one head	Shrinks by h / g, at a small quality cost

Heads change how the width is used, not how many weights there are. Grouped-query attention is Ainslie et al. (2023); the decoder cache is Day 42's topic.

Positions and order3 minSelf-attention is order-blind until positions are added; what sinusoidal codes encode and one alternative.

Sinusoidal positions - similarity that depends only on distance Illustration

Strong diagonal band: nearby positions have similar codes, distant ones less so (colour runs from cosine 0.45 to 1). The pattern repeats along every diagonal.

start at p = 0
start at p = 25
start at p = 200

Similarity depends only on the distance k, not on where you start: the three curves differ by at most 8e-16. The original paper chose sinusoids so that relative offsets would be easy to attend to.

Exact sinusoidal encodings (Vaswani et al., 2017) with d = 64; no data. Other schemes differ: learned position vectors have no built-in structure, and ALiBi (Press et al., 2022) adds a distance penalty to the scores instead.

Without positions, attention is order-blind. Sinusoidal codes make the similarity of two positions a function of their offset alone.

Sinusoidal positions (Vaswani et al., 2017) Formula

PE_{(pos, 2i)} = \sin\!\left(\frac{pos}{10000^{2i/d_{model}}}\right), \quad PE_{(pos, 2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d_{model}}}\right)

$pos$: position of the token in the sequence
$i$: index of the dimension pair: low i varies fast, high i slowly

Without positions, permuting the input tokens only permutes the outputs (the lab checks this), so word order would be invisible. ALiBi (Press et al., 2022) adds no position vectors: it subtracts a distance-growing penalty from each score, and extrapolated to longer inputs in their experiments.

Cost, memory and the n-squared part3 minScore-matrix memory per head, and why exact attention can still be fast.

Score-matrix entries per head and layer, stored at 2 bytes each Chart

1,024 tokens1,048,576 entries2.1 MB
2,048 tokens4,194,304 entries8.39 MB
4,096 tokens16,777,216 entries33.55 MB

Each doubling of the length quadruples the entries. FlashAttention (Dao et al., 2022) computes exact attention in tiles without storing the matrix; compute still grows with n squared, but dominates a layer only past about 6 d_model tokens. Generation adds a key-value cache that grows with context and layers (Day 42).

Attention weights and explanations3 minWhat a heat map shows, and the perturbation test that checks it.

Testing a heat map instead of believing it Diagram

Highlighted tokenfor example refund
Mask or replace itkeep everything else
Change in probabilityover many tickets
Large and consistent?
- yes: to Evidence of reliance
no
Decorationdo not present it as why
Evidence of reliancenot causation

Compare with a simple attribution method as well. Jain and Wallace (2019): weights often disagree with gradient-based importance, and different weights can give the same prediction. Wiegreffe and Pinter (2019): whether weights explain depends on the definition of explanation; test, do not assume.

What a staff-level answer addsadvanced2 minNuances and misreadings that separate an experienced answer from a textbook one.

The textbook answer and what a staff-level answer adds Comparison

Criterion	Common reading	Staff-level reading
What attention is	A learned map of what matters	A data-dependent mixing matrix; projections are the parameters
Cost	One n-squared number	Compute is n squared always; memory only without a fused kernel
Heads	'Multi-head means more parameters'	The width is split, not multiplied
The heat map	'The model attends to X, so X is the reason'	One head, one layer, averaged maps and residuals sit in between
Raising the context	'More context cannot hurt'	First measure where the label evidence sits in your inputs

Memory anchor

Softmax over keys, scaled by sqrt(d_k), masked before the softmax, n squared in cost - and a weight is not a reason.

Useful for

Reading a transformer diagram or a paper and knowing what each matrix, shape and mask is doing.
Reasoning about context-length cost before changing a limit (long contracts, claim files, call transcripts), and about why decoder models need a causal mask.
Setting honest expectations for attention heat maps shown to users (agents triaging tickets, clinicians reading risk flags, analysts reviewing flagged transactions).

Not suitable for

Explaining individual predictions - attention weights alone do not do this.
Training or tuning a production transformer; this day is the mechanism and its shapes at three tokens.
Judging model quality from the architecture name; a transformer must still beat strong baselines on your task (Days 27 and 35).

Traps

Presenting attention heat maps as the reason

A highlighted word looks like an explanation, and users, auditors and sometimes the team itself treat it as one, although the same prediction can come from different weights.

How to avoid it

Label it as a diagnostic, run a perturbation test (mask or remove the highlighted tokens), and compare with a second attribution method before showing it as a reason.

A missing or mismatched mask

Without a causal mask a decoder leaks the target and reports excellent loss; with padding unmasked, padded positions take attention weight and results change with batch composition.

How to avoid it

Unit-test that changing a later token does not change earlier outputs, and that adding padding does not change real tokens' outputs.

Senior reasoning

A senior reads attention as a data-dependent mixing matrix with a known cost and known failure modes: missing masks, order blindness, quadratic memory, and interpretation traps. They reduce a claim about a model to a shape check, a unit test (mask, padding, permutation) or a perturbation experiment, and they ask where the label evidence sits in the input before paying for context.

Interview question

"Walk me through scaled dot-product attention for three tokens. Why scale, why mask, and why not trust the heat map?"

Outline your answer first.

Recall cards scheduled after completion

Give the shapes in single-head attention for n tokens, and say what the softmax is taken over.
What does quadratic complexity mean for context length, and what is the explanation caveat?

Reviewed on a 1 · 3 · 7 · 14 · 30-day schedule; a miss returns the card to 1 day.

Prerequisites and related concepts

References

Attention Is All You Need — Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin, 2017core
Defines scaled dot-product and multi-head attention, positional encodings and the transformer block used in this lesson.
Attention is not Explanation — Sarthak Jain, Byron C. Wallace, 2019core
Shows that attention weights often do not track feature importance, the basis for treating heat maps as a hypothesis rather than an explanation.
Attention is not not Explanation — Sarah Wiegreffe, Yuval Pinter, 2019optional
Argues that whether attention explains depends on the definition of explanation, and proposes tests such as a uniform-weights baseline.
Neural Machine Translation by Jointly Learning to Align and Translate — Dzmitry Bahdanau, Kyunghyun Cho, Yoshua Bengio, 2015optional
Attention over encoder states as the answer to the fixed-length bottleneck of encoder-decoder models.
Layer Normalization — Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton, 2016optional
Per-example layer normalisation independent of batch size, used in every transformer block.
Deep Residual Learning for Image Recognition — Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun, 2016optional
Residual connections as the source of the skip connections used around attention and feed-forward sublayers.
On Layer Normalization in the Transformer Architecture — Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, Tie-Yan Liu, 2020optional
Normalisation placement (pre-LN versus post-LN) and warmup in transformer training.
torch.nn.functional.scaled_dot_product_attention (PyTorch documentation) — PyTorchoptional
Default scale, causal flag and the boolean mask convention (True takes part), the inverse of MultiheadAttention's key_padding_mask.
MultiheadAttention (PyTorch documentation) — PyTorchoptional
key_padding_mask marks keys to ignore; returned weights are averaged over heads by default.
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai, 2023optional
Grouped-query attention shares key-value heads across query heads to speed decoder inference.
Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation — Ofir Press, Noah A. Smith, Mike Lewis, 2022optional
Distance-based biases on attention scores instead of position embeddings, with length extrapolation.
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Re, 2022optional
States that attention time and memory are quadratic in sequence length and describes exact attention that reduces memory traffic by tiling.

Optional sections

Deep dive · 10 min

Why $\sqrt{d_k}$ . If the components of $q$ and $k$ are independent with zero mean and unit variance, $q \cdot k$ has mean 0 and variance $d_k$ . For $d_k = 64$ the standard deviation is 8, so raw scores differ by many units, the softmax approaches a one-hot vector, and its gradients vanish. Dividing by $\sqrt{d_k}$ restores unit variance.

Parameter counts. For $d_{model}$ with $h$ heads of size $d_{model}/h$ , the four attention matrices $W_Q, W_K, W_V, W_O$ hold $4 d_{model}^2$ weights (without biases) whatever $h$ is; a feed-forward block with a 4x hidden expansion holds $8 d_{model}^2$ . At $d_{model} = 512$ that is about 1.0 million and 2.1 million per block.

Masks. A padding mask removes padded keys so that batch composition does not change results; a causal mask removes future keys. Both are added to the scores before the softmax. Cross-attention takes queries from the decoder and keys and values from the encoder.

Where this goes. Decoder-only models stacked from these blocks and trained on next-token prediction are the LLMs of Days 42-45; the cost of attention is why context length, caching and retrieval are engineering decisions (Days 42, 48 and 52).

Hands-on lab: One transformer block in numpy - heads, positions, residual + LayerNorm, mask tests, scaling · 20 min

Build a two-head block for six tokens and test its properties: output shape, normalised rows, permutation behaviour without positions, loss of that property with positions, causal-mask isolation of earlier outputs, the effect of the $\sqrt{d_k}$ scale, and the size of the score matrix at three context lengths (2 bytes per entry is an assumption for half precision). Then remove the residual connection and describe what changes in the output statistics.

import numpy as np

rng = np.random.default_rng(39)
n, d_model, heads = 6, 8, 2
d_k = d_model // heads

def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

def layer_norm(x, eps=1e-5):
    return (x - x.mean(-1, keepdims=True)) / np.sqrt(x.var(-1, keepdims=True) + eps)

Wq, Wk, Wv = (rng.normal(0, 0.5, (heads, d_model, d_k)) for _ in range(3))
Wo = rng.normal(0, 0.5, (heads * d_k, d_model))

def block(X, causal=False):
    outs = []
    for h in range(heads):
        Q, K, V = X @ Wq[h], X @ Wk[h], X @ Wv[h]            # each (n, d_k)
        s = Q @ K.T / np.sqrt(d_k)                           # (n, n)
        if causal:
            s = np.where(np.triu(np.ones((len(X), len(X)), bool), 1), -np.inf, s)
        outs.append(softmax(s) @ V)                          # (n, d_k)
    mixed = np.concatenate(outs, axis=-1) @ Wo               # (n, d_model)
    return layer_norm(X + mixed)                             # residual, then LayerNorm

def sinusoid(n, d):
    pos, i = np.arange(n)[:, None], np.arange(d // 2)[None]
    ang = pos / 10000 ** (2 * i / d)
    return np.concatenate([np.sin(ang), np.cos(ang)], axis=-1)

X = rng.normal(size=(n, d_model))
Y = block(X)
print("block output shape:", Y.shape, "| row means ~0:", np.allclose(Y.mean(-1), 0, atol=1e-8))
perm = rng.permutation(n)
print("no positions - permuting tokens permutes outputs:", np.allclose(block(X[perm]), Y[perm]))
P = sinusoid(n, d_model)
print("with positions - still only permutes:", np.allclose(block(X[perm] + P), block(X + P)[perm]))
X2 = X.copy()
X2[-1] += 5.0                                                # change only the last token
print("causal: earlier outputs unchanged:", np.allclose(block(X2, True)[:-1], block(X, True)[:-1]))
print("non-causal: earlier outputs change:", not np.allclose(block(X2)[:-1], Y[:-1]))
q, k = rng.normal(size=(2000, 64)), rng.normal(size=(2000, 64))
dots = (q * k).sum(1)
print("std of q.k for d_k=64: raw %.1f, scaled %.1f" % (dots.std(), (dots / 8).std()))
for n_tok in (1_000, 8_000, 32_000):
    print(f"n={n_tok:6d}: {n_tok**2:.1e} score entries per head = {n_tok**2 * 2 / 1e6:7.0f} MB at 2 bytes")

Interview practice · 10 min

Your team's insurance claims-document classifier gets a 4x longer context. Latency and memory jump by more than 4x. Explain why and what you would try before buying bigger GPUs.

Answer as you would in the interview.

Content version 6 · last verified 2026-10-18 · depends on numpy