COMP5318 — Week 9 Supplement
Mathematical Foundations
Scaled dot-product attention, the \(\mathrm{softmax}(QK^\top/\sqrt{d_k})V\) formula, multi-head concatenation, and sinusoidal positional encoding.
Scaled dot-product self-attention
Let the input be a matrix \(X \in \mathbb{R}^{n \times d_{\text{model}}}\) — one row per token. Self-attention learns three projections, scores all query-key pairs, normalises with softmax, and returns a weighted sum of value vectors.
Q, K, V projections
\[ Q = X W^Q, \qquad K = X W^K, \qquad V = X W^V \]
where \(W^Q, W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}\) and \(W^V \in \mathbb{R}^{d_{\text{model}} \times d_v}\). \(Q\) and \(K\) lie in the same space so dot products are meaningful.
Scaled dot-product attention
\[ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left( \frac{Q K^\top}{\sqrt{d_k}} \right) V \]
The matrix \(QK^\top \in \mathbb{R}^{n \times n}\) holds all pairwise scores. Row \(i\) is softmaxed into a distribution over source positions, and the output for token \(i\) is the corresponding convex combination of value rows of \(V\).
If \(q\) and \(k\) have independent entries of variance 1, then \(q \cdot k\) has variance \(d_k\). For large \(d_k\), un-scaled scores reach extreme magnitudes, the softmax saturates and gradients vanish. Dividing by \(\sqrt{d_k}\) keeps the score variance \(O(1)\).
Multi-head attention
Run \(h\) parallel attention heads. Each head has its own projection matrices \(W_i^Q, W_i^K, W_i^V\) of width \(d_k = d_v = d_{\text{model}}/h\).
Per-head attention
\[ \text{head}_i = \mathrm{Attention}(X W_i^Q,\; X W_i^K,\; X W_i^V) \]
Concatenate and project
\[ \mathrm{MultiHead}(X) = \mathrm{Concat}(\text{head}_1, \text{head}_2, \ldots, \text{head}_h)\, W^O \]
with \(W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}\). The result has the same shape as the input, so encoder layers can be stacked.
Positional encoding
Self-attention is permutation-equivariant, so the model needs an explicit signal of token position. The original Transformer uses fixed sinusoidal encodings added to the input embeddings.
Sinusoidal positional encoding
For position \(p \in \{0, 1, \ldots, n-1\}\) and embedding dimension index \(2i\) or \(2i+1\):
\[ PE_{(p, 2i)} = \sin\!\left( \frac{p}{10000^{2i/d_{\text{model}}}} \right), \qquad PE_{(p, 2i+1)} = \cos\!\left( \frac{p}{10000^{2i/d_{\text{model}}}} \right) \]
Each dimension is a sinusoid of a different wavelength, ranging from \(2\pi\) up to \(10000 \cdot 2\pi\). The encoding is then added to the input embedding: \(\tilde x_p = x_p + PE_p\).
Because \(\sin(a+b)\) and \(\cos(a+b)\) are linear combinations of \(\sin a, \cos a, \sin b, \cos b\), the encoding for position \(p+k\) is a fixed linear function of the encoding for position \(p\). The model can therefore learn to attend by relative offset.
Putting it together: one encoder/decoder layer
Each encoder layer is two sub-layers with residual + layer norm wrappers:
Encoder layer
\[ z = \mathrm{LayerNorm}\!\bigl(x + \mathrm{MultiHead}(x)\bigr) \]
\[ \mathrm{out} = \mathrm{LayerNorm}\!\bigl(z + \mathrm{FFN}(z)\bigr) \]
where \(\mathrm{FFN}(z) = \max(0, zW_1 + b_1)W_2 + b_2\) is a position-wise two-layer MLP with ReLU.
Three sub-layers: masked self-attention (causal mask so position \(t\) attends only to positions \(\le t\)), then encoder-decoder attention with \(Q\) from the decoder and \(K, V\) from the top encoder output, then the same FFN. Each is wrapped by residual + layer norm.