Problems

Problem II.3.B01

Attention on three tokens, by hand

numeric▲▲△

All values rounded to 4 d.p. Intermediate quantities carry full precision; only what is printed is rounded.

STATEMENT

Three tokens attend to one another under scaled dot-product attention. Compute the whole forward pass by hand and confirm that every attention row is a probability distribution.

GIVEN

\mat{Q} = \mat{K} = \mat{V} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix} \in \R^{3 \times 2}, \qquad \dk = 2 .

Row $i$ of each matrix belongs to token $i$ . Row-major throughout, as fixed in the Apparatus.

FIND

The score matrix $\mat{S}$ , the scaled scores $\mat{Z}$ , the attention weights $\mat{A}$ , and the output $\mat{O}$ — each a $3 \times 3$ matrix except $\mat{O}$ , which is $3 \times 2$ .

STRATEGY

Name each intermediate rather than nesting them, so a wrong digit can be traced to the step that produced it.

SOLUTION

Step 1 — the scores. $\mat{S} = \mat{Q}\mat{K}^{\mathsf T}$ , so $S_{ij}$ is the inner product of query row $i$ with key row $j$ .

\mat{S} = \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}

Token 3’s query is $(1,1)$ , which agrees with both basis directions, so its row is the largest. Nothing about position enters: $S_{ij}$ is a function of content alone.

Step 2 — the scaling. Divide by $\sqrt{\dk} = \sqrt 2 \approx 1.4142$ .

\mat{Z} = \begin{bmatrix} 0.7071 & 0 & 0.7071 \\ 0 & 0.7071 & 0.7071 \\ 0.7071 & 0.7071 & 1.4142 \end{bmatrix}

Step 3 — the softmax, row by row. Take row 1, $\vec{z} = (0.7071, 0, 0.7071)$ . Subtract the row maximum first, which changes nothing and removes the overflow (Log-sum-exp 0.NU.02):

\vec{z} - \max = (0,\; -0.7071,\; 0), \qquad e^{0} = 1, \quad e^{-0.7071} = 0.4931 .

The denominator is $1 + 0.4931 + 1 = 2.4931$ , so

\vec{a}_1 = \left(\tfrac{1}{2.4931},\; \tfrac{0.4931}{2.4931},\; \tfrac{1}{2.4931}\right) = (0.4011,\; 0.1978,\; 0.4011) .

Row 2 is row 1 with the first two entries exchanged, by the symmetry of the inputs. Row 3 has $\vec{z} = (0.7071, 0.7071, 1.4142)$ ; subtracting the maximum gives $(-0.7071, -0.7071, 0)$ , exponentials $0.4931, 0.4931, 1$ , denominator $1.9862$ , and

\vec{a}_3 = (0.2483,\; 0.2483,\; 0.5035) .

Step 4 — the aggregation. $\mat{O} = \mat{A}\mat{V}$ . For row 1, column 1:

O_{11} = 0.4011(1) + 0.1978(0) + 0.4011(1) = 0.8022 ,

and column 2 is $0.4011(0) + 0.1978(1) + 0.4011(1) = 0.5989$ .

Answer

\mat{A} = \begin{bmatrix} 0.4011 & 0.1978 & 0.4011 \\ 0.1978 & 0.4011 & 0.4011 \\ 0.2483 & 0.2483 & 0.5035 \end{bmatrix}, \qquad \mat{O} = \begin{bmatrix} 0.8022 & 0.5989 \\ 0.5989 & 0.8022 \\ 0.7517 & 0.7517 \end{bmatrix}

$\mat{A}$ is $3 \times 3$ and dimensionless; $\mat{O}$ is $3 \times 2$ and carries the units of $\mat{V}$ .

Check — numeric · ii-3-b01-attention-by-hand.py

S = [[sum(q[i] * k[i] for i in range(d_k)) for k in K] for q in Q]
Z = [[s / sqrt(d_k) for s in row] for row in S]
A = [softmax(row) for row in Z]
O = [[sum(a[j] * V[j][c] for j in range(3)) for c in range(2)] for a in A]

The full twenty-line script is scripts/snippets/ii-3-b01-attention-by-hand.py, and its output is diffed against the digits above on every build.

Executed in CI. The digits above are the digits it printed.

Check — sanity

Three independent reasons the answer is right.

Rows sum to one. $0.4011 + 0.1978 + 0.4011 = 1.0000$ , and likewise for the other two. A softmax row that does not sum to one is an arithmetic slip, not a modelling choice.

Every output lies inside the convex hull of $\mat{V}$ . The value rows are $(1,0)$ , $(0,1)$ and $(1,1)$ ; every entry of $\mat{O}$ lies in $[0,1]$ , as a convex combination must. This is the whole content of Proposition 1: attention mixes, and cannot leave the hull.

The symmetry survives. $\mat{Q} = \mat{K}$ here, so $\mat{S}$ is symmetric, and rows 1 and 2 of the answer are exchanged copies of one another. They are.

Where this breaks

The convex-hull argument holds only because the weights are non-negative and sum to one. Remove the softmax — score-weighted sums, or normalisation by anything other than the total — and the output can leave the hull entirely. That is why the feed-forward block after attention is not decoration.

Variation

Recompute with $\mat{V} = \begin{bmatrix} 2 & 0 \\ 0 & 2 \\ 2 & 2 \end{bmatrix}$ and predict $\mat{O}$ before doing any arithmetic. Then set $\dk = 8$ with the same $\mat{Q}$ and $\mat{K}$ and say what happens to the sharpness of $\mat{A}$ , and why.

II.3.X01Softmax at three temperatureslimit▲△△

Compute $\softmax(\vec{s}/\tau)$ for $\vec{s} = (2, 1, 0)$ at $\tau = 1$ , $\tau = 0.5$ and $\tau = 2$ . State what each result says about the temperature’s role, and name the two limits.

Hint

Divide before exponentiating, and subtract the largest entry of the scaled vector first. The three denominators are all you need.

Solution

τ = 1. Scaled scores $(2,1,0)$ . Subtract the maximum: $(0,-1,-2)$ , giving exponentials $1$ , $0.3679$ , $0.1353$ , and a denominator of $1.5032$ .

$\vec{a} = (0.6652,\; 0.2447,\; 0.0900)$

τ = 0.5. Scaled scores $(4,2,0)$ . Shifted: $(0,-2,-4)$ , exponentials $1$ , $0.1353$ , $0.0183$ , denominator $1.1536$ .

$\vec{a} = (0.8668,\; 0.1173,\; 0.0159)$

τ = 2. Scaled scores $(1, 0.5, 0)$ . Shifted: $(0,-0.5,-1)$ , exponentials $1$ , $0.6065$ , $0.3679$ , denominator $1.9744$ .

$\vec{a} = (0.5065,\; 0.3072,\; 0.1863)$

Each row sums to $1.0000$ .

What it says. Lowering $\tau$ multiplies every score gap by $1/\tau$ before exponentiating, so the distribution concentrates on the largest entry; raising $\tau$ shrinks the gaps toward zero and flattens it. The mass on the top token moves from $0.5065$ to $0.8668$ as $\tau$ falls from 2 to 0.5, on unchanged scores.

The two limits. As $\tau \to 0^{+}$ the distribution approaches $\text{one-hot}(\arg\max \vec{s})$ ; as $\tau \to \infty$ it approaches the uniform distribution over the three entries, $(1/3, 1/3, 1/3)$ . Neither limit is reached at any finite $\tau$ .

The contract

Attention on three tokens, by hand

Exercises