Exercises — Transformers — attention mechanism, self-attention

• $q_i=W_Q{x}_i$ — the Query: what am I looking for?
• $k_i=W_K{x}_i$ — the Key: what do I advertise to others?
• $v_i=W_V{x}_i$ — the Value: what content do I pass on if selected?

"Self"-attention means all three come from the same sequence.

Row-wise, i.e. across keys / positions. Each row of $Q{K}^{\top }$ holds one query's scores against every key. Softmax on that row turns those scores into weights that sum to 1, giving a probability distribution over "which tokens to read." It is not applied over the feature dimension.

• Learned (fixed after training, same for every input): $W_Q$, $W_V$ (also $W_K$, $W_O$).
• Computed on the fly (change with each input): ${\alpha }_{ij}$, $z_i$.

Raw dot product: $q_1\cdot k_2=3\cdot 0+4\cdot 2=8$. Scale: $s_{12}=\frac{8}{\sqrt{2}}=\frac{8}{1.41421}=5.65685$.

Scaled: $s_{11}=1/\sqrt{2}=0.70711$, $s_{12}=0$. ${\alpha }_{11}=\frac{e^{0.70711}}{e^{0.70711}+e^0}=\frac{2.02812}{3.02812}=\mathrm{0.66975.}$ ${\alpha }_{12}=1-0.66975=\mathrm{0.33025.}$ Check: they sum to 1, as attention weights must.

$z_1=0.66975(2,0)+0.33025(0,3)=(1.33950,0.99074).$ Interpretation: token 1 mostly grabs $v_1$ (its query matched key 1) but mixes in a third of $v_2$ — a soft lookup, not a hard pick. See the geometry below.

Raw scores: $s_{21}=q_2\cdot k_1=0$, $s_{22}=q_2\cdot k_2=1$. Scale by $\sqrt{2}$: $s_{21}=0$, $s_{22}=0.70711$. Softmax: ${\alpha }_{22}=0.66975$, ${\alpha }_{21}=0.33025$ (mirror of L2.2). Output: $z_2=0.33025(2,0)+0.66975(0,3)=(0.66049,2.00924).$

$q\cdot k={\sum }_{m=1}^{d_k}{q}_m{k}_m$ is a sum of $d_k$ independent products, each with variance $1$ (since $\mathrm{Var}(q_m{k}_m)=\mathrm{Var}(q_m)\mathrm{Var}(k_m)=1$ for zero-mean unit-variance independent factors). Variances add: $\mathrm{Var}(q\cdot k)=d_k=64\Rightarrow \text{std}=\sqrt{64}=8.$ Dividing by $\sqrt{d_k}=8$ scales the std by $1/8$, giving std $\approx 1$ — exactly the "keep variance near 1" goal from the parent note. Linked idea: Gradient vanishing/saturation.

The outputs swap identically: the model returns $z_B,z_A$ — the same vectors in swapped slots. Because every score $q_i\cdot k_j$ depends only on which vectors meet, not where they sit, self-attention is permutation-equivariant. Consequence: it has no built-in notion of order, which is why Positional Encoding must be added to inject "who came first."

The dot product $q\cdot k=\Vert q\Vert \Vert k\Vert \cos \theta$, where $\theta$ is the angle between them. It is large and positive when the two vectors point the same way ($\cos \theta \to 1$) — precisely "these two are relevant." It is smooth and differentiable (unlike a hard match) so gradients flow, and it's one cheap matrix multiply $Q{K}^{\top }$ for all pairs at once (parallelism). Distance would need extra normalisation and doesn't factor into a single matmul as cleanly. See Dot product & cosine similarity.

$d_k=\frac{d_{\text{model}}}{h}=\frac{512}{8}=64.$ Each head outputs a $d_v=64$-wide vector per token. Concatenating $h=8$ heads gives width $8\times 64=512$. Then $W_O$ maps $512\to 512$, i.e. $W_O\in {\mathbb{R}}^{512\times 512}$, returning to $d_{\text{model}}$. Note the total compute stays comparable to one big head — that's the elegance. See Multi-Head Attention.

(a) One hop. Query $q_{180}$ dots directly with key $k_{120}$; the path length is 1 regardless of the 60-step gap. (b) An RNN must thread the $t=120$ signal through 60 sequential hidden-state updates, each blurring/diluting it (long-range forgetting). Self-attention's single direct connection preserves the signal with no information loss, so the model can put a large ${\alpha }_{180,120}$ and flag the true cause.

• $Q,K\in {\mathbb{R}}^{3\times 4}$ (rows = tokens, cols = $d_k$).
• $V\in {\mathbb{R}}^{3\times 5}$.
• $Q{K}^{\top }\in {\mathbb{R}}^{3\times 3}$ — all pairwise scores; row $i$ = token $i$'s scores.
• softmax (row-wise) keeps shape ${\mathbb{R}}^{3\times 3}$.
• $Z=(\text{weights})V\in {\mathbb{R}}^{3\times 5}$ — same width as the values, one output row per token.

(a) Unscaled: ${\alpha }_1=\frac{e^8}{e^8+e^0}=\frac{2980.958}{2981.958}=0.99966$, so ${\alpha }_2=0.00034$ — nearly one-hot. (b) Scaled scores: $8/8=1$ and $0$. Then ${\alpha }_1=\frac{e^1}{e^1+1}=\frac{2.71828}{3.71828}=0.73106$, ${\alpha }_2=0.26894$ — a healthy blend. (c) In the unscaled case the second weight is $\approx 0$, so $\partial (\text{output})/\partial s_2\approx 0$ — gradients to that token vanish; scaling by $\sqrt{d_k}$ keeps both weights sizeable, so learning signal flows to both. This is the Gradient vanishing/saturation fix in action.

Token 1. Raw scores: $s_{11}=q_1\cdot k_1=1+1=2$, $s_{12}=q_1\cdot k_2=1-1=0$. Scale by $\sqrt{2}$: $s_{11}=1.41421$, $s_{12}=0$. Softmax: ${\alpha }_{11}=\frac{e^{1.41421}}{e^{1.41421}+1}=\frac{4.11325}{5.11325}=0.80443$, ${\alpha }_{12}=0.19557$. Output: $z_1=0.80443(1,0)+0.19557(0,1)=(0.80443,0.19557)$.

Token 2. Raw scores: $s_{21}=q_2\cdot k_1=1-1=0$, $s_{22}=q_2\cdot k_2=1+1=2$. By the same arithmetic (mirror): ${\alpha }_{22}=0.80443$, ${\alpha }_{21}=0.19557$. Output: $z_2=0.19557(1,0)+0.80443(0,1)=(0.19557,0.80443)$.

Each token attends mostly to itself (its query matches its own key), mixing in ~20% of the other — a clean symmetric example.

Set the future scores to $-\infty$ before softmax (add a mask of $-\infty$ to those entries). Since $e^{-\infty }=0$, those weights become exactly 0 and the surviving weights renormalise. Check on token 1 with $j=2$ masked: scaled scores $s_{11}=1.41421$, $s_{12}=-\infty$. Softmax: ${\alpha }_{11}=\frac{e^{1.41421}}{e^{1.41421}+0}=1.00000$, ${\alpha }_{12}=0$. Output: $z_1=1\cdot (1,0)+0\cdot (0,1)=(1.00000,0.00000)$ — token 1 now reads only itself, as required.