Worked examples — The original - Attention is All You Need - architecture

Every symbol below is built from the ground up. If you have never seen a matrix, read it as a grid of numbers: a row is a horizontal strip, a column a vertical strip, and the entry at "row $i$, column $j$" sits where strip $i$ crosses strip $j$.

---

The scenario matrix

Here $d_k$ = the size of each key/query vector (how many numbers describe one token in an attention head). $n$ = number of tokens. "Score" = the raw dot product $Q{K}^{\top }$ before softmax.

#	Case class	What makes it special	Covered by
A	Standard forward pass	full attention, all positions	Ex 1
B	Scaling sign & size	why divide by $\sqrt{d_k}$; large vs small $d_k$	Ex 2
C	Softmax saturation (limiting)	scores $\to \infty$, gradient $\to 0$	Ex 3
D	Degenerate input: single token ($n=1$)	attention becomes identity	Ex 4
E	Causal mask, every row	future = $-\infty$, all triangle rows	Ex 5
F	Cross-attention (asymmetric $Q$ vs $K$)	$Q$ from decoder, $K,V$ from encoder	Ex 6
G	Positional encoding: all four sign quadrants	$\sin /\cos$ positive & negative	Ex 7
H	Residual + LayerNorm edge (zero variance)	$\sigma =0$ input, why $ϵ$ exists	Ex 8
I	Real-world word problem	translation, count the FLOPs-ish shape	Ex 9
J	Exam twist: multi-head splitting	$d_{\text{model}}/h$, concat back	Ex 10

We now hit every cell.

---

Example 1 — Cell A: the standard forward pass

Forecast: Guess — will token 1's output lean toward itself or toward token 2? Write your guess.

1. Scores $Q{K}^{\top }$. Row $i$, col $j$ = dot product of row $i$ of $Q$ with row $j$ of $K$. $Q{K}^{\top }=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right].$ Why this step? The dot product measures overlap: identical directions give $1$, perpendicular give $0$. Each token here is perpendicular to the other, so off-diagonal is $0$.

2. Scale by $\sqrt{d_k}=\sqrt{2}\approx 1.414$. $\frac{Q{K}^{\top }}{\sqrt{2}}=\left[\begin{array}{cc}0.7071& 0\\ 0& 0.7071\end{array}\right].$ Why this step? Keeps the numbers small so softmax stays in its sensitive region (see Ex 3).

3. Softmax row 1 over $[0.7071,0]$. Exponentiate: $e^{0.7071}=2.0281$, $e^0=1$, sum $=3.0281$. $\text{row 1}=[0.6698,0.3302].$ Why this step? Softmax turns scores into a probability row that sums to $1$ — "how much attention token 1 pays to each token."

4. Weighted sum with $V=X$: output for token 1 $=0.6698V_1+0.3302V_2=[0.6698,0.3302]$. Why this step? The output is a blend of value vectors weighted by attention.

Verify: Row weights sum to $0.6698+0.3302=1.0000$ ✓. Token 1 leans toward itself ($0.67>0.33$) because it matched itself with score $1$ — matching the geometry in the figure (red arrow longest onto itself).

---

Example 2 — Cell B: why the $\sqrt{d_k}$ sign and size matter

Forecast: Does scaling make the two dimensions give the same number? Guess yes/no.

1. Raw score = $q\cdot k=\sum 1\cdot 1=d_k$. So $d_k=4$ gives $4$; $d_k=64$ gives $64$. Why this step? A dot product of two all-ones vectors is just the count of dimensions — it grows linearly with $d_k$.

2. Scaled score $=d_k/\sqrt{d_k}=\sqrt{d_k}$. For $d_k=4$: $\sqrt{4}=2$. For $d_k=64$: $\sqrt{64}=8$. Why this step? Scaling by $\sqrt{d_k}$ turns linear growth into square-root growth — much gentler.

Verify: Unscaled ratio $64/4=16$; scaled ratio $8/2=4=\sqrt{16}$ ✓. Scaling shrinks the blow-up from $16\times$ down to $4\times$. See Scaled Dot-Product Attention — this is exactly the "variance control" argument.

---

Example 3 — Cell C: the saturation limit (why big scores kill gradients)

Forecast: As $s$ grows huge, does attention become $[1,0]$, $[0.5,0.5]$, or $[0,1]$?

1. Weight on token 1 $=\frac{e^s}{e^s+e^0}=\frac{1}{1+e^{-s}}$ (the sigmoid). Why this step? Divide top and bottom by $e^s$ to get the standard logistic form.

2. Limit $s\to +\infty$: $e^{-s}\to 0$, weight $\to 1$. Limit $s\to -\infty$: $e^{-s}\to \infty$, weight $\to 0$. Why this step? Covers both signs of $s$ — the two flat tails of the curve.

3. Slope of the curve is $\text{weight}\cdot (1-\text{weight})$. At weight $=1$ or $0$ the slope $=0$. Why this step? Zero slope = zero gradient = the network stops learning. That is saturation, and it's exactly why we scale (Ex 2) to keep $s$ moderate.

Verify (numeric): at $s=10$, weight $=1/(1+e^{-10})=0.99995$ and slope $=0.99995\cdot 0.00005=4.5\times 10^{-5}\approx 0$ ✓. Nearly dead gradient — the accent-red flat region in the figure.

---

Example 4 — Cell D: degenerate single token ($n=1$)

Forecast: Does attention change the token at all?

1. Score matrix is $1\times 1$: just $q\cdot k$, one number. Why this step? With $n=1$ there is nothing else to compare to.

2. Softmax of a single number is always $1$ (since $e^s/e^s=1$). Why this step? A probability row of length one must sum to $1$, so it is $1$.

3. Output $=1\cdot V_1=V_1$. Attention is the identity. Why this step? The weighted sum has one term with weight $1$.

Verify: For any score $s$, $e^s/e^s=1$ exactly ✓. This is the degenerate edge: attention cannot mix information when there is only one token — it just passes it through. (In a real model the FFN and residual still transform it; see Residual Connections.)

---

Example 5 — Cell E: causal mask, every triangle row

Forecast: Row 1 attends to how many tokens? Row 3?

1. Mask rule: entry $(i,j)$ becomes $-\infty$ if $j>i$ (future), else stays. Why this step? Prevents a token from peeking at tokens that come after it — the core of Attention Masking and autoregressive generation.

2. Row 1 = $[0,-\infty ,-\infty ]$. Softmax: $e^0=1$, $e^{-\infty }=0$. Weights $=[1,0,0]$. Why this step? Token 1 has no past, so it attends only to itself.

3. Row 2 = $[0,0,-\infty ]$. Weights $=[0.5,0.5,0]$. Why this step? Token 2 splits equally across positions 1–2.

4. Row 3 = $[0,0,0]$. Weights $=[\frac{1}{3},\frac{1}{3},\frac{1}{3}]$. Why this step? Token 3 sees all three past-or-present positions equally.

Verify: each row sums to $1$: $1$, $0.5+0.5=1$, $3\times \frac{1}{3}=1$ ✓. The lower-triangular red region in the figure is exactly the set of allowed attentions.

---

Example 6 — Cell F: cross-attention (asymmetric shapes)

Forecast: Is the score matrix $2\times 2$, $3\times 3$, or $2\times 3$?

1. $Q$ comes from the decoder ($2$ rows), $K,V$ from the encoder ($3$ rows). See Encoder-Decoder Architecture. Why this step? Cross-attention lets each output token query the whole source sentence.

2. Score $=Q{K}^{\top }$ has shape $(2\times d_k)(d_k\times 3)=2\times 3$. Why this step? Matrix-product shapes: inner dims $d_k$ cancel, outer dims survive.

3. Softmax runs across the $3$ encoder columns (one distribution per decoder token). Why this step? Each decoder token must decide how to split its attention over the $3$ source tokens; the row must sum to $1$.

Verify: shape check $(2\times 2)\cdot (2\times 3)=2\times 3$ ✓; each of the $2$ rows sums to $1$ over $3$ entries. Output shape $=(2\times 3)\cdot (3\times d_k)=2\times d_k$, back to decoder length ✓.

---

Example 7 — Cell G: positional encoding across all sign quadrants

Forecast: At which position does $\sin$ first go negative?

1. $\sin (1)=0.8415$ — positive, rising (quadrant I).
2. $\sin (2)=0.9093$ — positive, past the peak, falling (quadrant II).
3. $\sin (4)=-0.7568$ — negative, falling (quadrant III).
4. $\sin (5)=-0.9589$ — negative, near trough (quadrant IV).

Why this step? Positional Encoding uses $\sin /\cos$ precisely because they take both signs and repeat smoothly, so the model can read a position's phase. Covering all four quadrants shows the encoding is never "stuck" positive.

Verify: signs are $+,+,-,-$; first negative at $\text{pos}=4$ (since $\pi \approx 3.14$, $\sin$ crosses zero just after pos $3$) ✓.

---

Example 8 — Cell H: LayerNorm at the zero-variance edge

Forecast: What is $x-\mu$ here? What breaks if $ϵ=0$?

1. Mean $\mu =(5+5+5+5)/4=5$. So $x-\mu =[0,0,0,0]$. Why this step? A constant vector has all entries equal to its mean.

2. Variance ${\sigma }^2=0$. Without $ϵ$ we'd divide by $\sqrt{0}=0$ — undefined. Why this step? This is the degenerate corner; $ϵ$ guarantees the denominator is never zero.

3. With $ϵ$: output $=0/\sqrt{0+10^{-5}}=0$ for every entry. Why this step? The numerator is already $0$, so the output is a clean zero vector, not a NaN.

Verify: $\sqrt{10^{-5}}=0.0031623$; $0/0.0031623=0$ ✓. The $ϵ$ turned a "divide-by-zero crash" into a safe $0$.

---

Example 9 — Cell I: real-world translation word problem

Forecast: Guess the per-head dimension before computing.

1. (a) Per-head dim $=d_{\text{model}}/h=512/8=64$. Why this step? Multi-Head Self-Attention splits the $512$ features evenly across $8$ heads.

2. (b) Score matrices per layer = one per head = $8$. Across $6$ encoder layers: $8\times 6=48$. Why this step? Each layer runs $h=8$ independent attention operations, each producing its own $n\times n$ score matrix ($n=3$ here).

3. Sanity on shape: each score matrix is $3\times 3$ (source attends to source), consistent with Teacher Forcing feeding the full target only on the decoder side, not here. Why this step? Encoder self-attention is symmetric over the $3$ source tokens.

Verify: $512/8=64$ ✓; $8\times 6=48$ ✓.

---

Example 10 — Cell J: the exam twist (split, attend, concat)

Forecast: Does concat give $3\times 8$ or $6\times 4$?

1. Concat along the feature axis: stack the two $3\times 4$ head outputs side by side → $3\times 8$. Why this step? Concatenation glues columns so each token regains all $d_{\text{model}}=8$ features. Row count $n=3$ is unchanged.

2. $W^O\in {\mathbb{R}}^{8\times 8}$ mixes the concatenated heads: $(3\times 8)(8\times 8)=3\times 8$. Why this step? Without $W^O$ the heads never talk; $W^O$ blends their specialised views into one representation.

3. Check total head width $=h\cdot d_k=2\cdot 4=8=d_{\text{model}}$. Why this step? Confirms the split-and-rejoin conserves dimension — a favourite exam gotcha.

Verify: concat shape $3\times (2\cdot 4)=3\times 8$ ✓; after $W^O$: $3\times 8$ ✓; $h\cdot d_k=8=d_{\text{model}}$ ✓.

---

Recall Self-test (reveal after answering)

Single-token self-attention output ::: the token's own value vector $V_1$ (attention weight $=1$). Decoder row 2 causal weights on 3 tokens with equal scores ::: $[0.5,0.5,0]$. Cross-attention score shape for 2 decoder tokens and 3 encoder tokens ::: $2\times 3$. Per-head dimension for $d_{\text{model}}=512,h=8$ ::: $64$. Why does LayerNorm need $ϵ$? ::: to avoid dividing by zero when the input has zero variance.