Transformer Architecture

Question 1

(a) [6 marks] Each term $q_{im}k_{jm}$: since $q,k$ independent, mean $= E[q]E[k] = 0$. (1) Variance of one term: $\text{Var}(q_{im}k_{jm}) = E[q^2]E[k^2] - 0 = 1\cdot 1 = 1$. (2) Sum of $d_k$ independent terms: $E[s_{ij}] = 0$, $\text{Var}(s_{ij}) = d_k$. (2) To restore unit variance divide by standard deviation $\sqrt{d_k}$, not $d_k$ (which would over-shrink to variance $1/d_k$). Hence scaling $1/\sqrt{d_k}$ keeps logits $O(1)$. (1)

(b) [8 marks] Softmax: $a_i = e^{z_i}/\sum_l e^{z_l}$. For $i=j$: $\frac{\partial a_i}{\partial z_i} = a_i(1-a_i)$. (2) For $i\ne j$: $\frac{\partial a_i}{\partial z_j} = -a_i a_j$. (2) Compactly: $\frac{\partial a_i}{\partial z_j} = a_i(\delta_{ij} - a_j)$. (1) Failure mode: without scaling, $\text{Var}(s)=d_k$ grows, so softmax inputs become large in magnitude → softmax saturates → one $a_i \to 1$, others $\to 0$. Then $a_i(1-a_i)\to 0$ and $a_ia_j\to 0$, so the whole Jacobian $\to 0$: vanishing gradients, no learning signal through attention weights. (2) Scaling by $1/\sqrt{d_k}$ keeps logits $O(1)$, keeping the softmax in its non-saturated regime and the Jacobian well away from zero. (1)

(c) [4 marks]

• $QK^\top$: $n\times d_k$ times $d_k\times n$ → $O(n^2 d_k)$ time, $O(n^2)$ memory. (1.5)
• Softmax over $n\times n$ matrix: $O(n^2)$ time and memory. (1)
• $\times V$: $n\times n$ times $n\times d_v$ → $O(n^2 d_v)$ time, output $O(n d_v)$. (1.5) Overall $O(n^2 d)$ time, $O(n^2)$ memory (dominant).

(d) [4 marks] Idea: tiling / blockwise computation — never materialise the full $n\times n$ score matrix; process $K,V$ in blocks and accumulate the output using an online (streaming) softmax so only $O(n)$ statistics (running max and running sum/output) are kept in fast SRAM. (2) Numerical technique: the online softmax with running max subtraction (log-sum-exp / max-rescaling), rescaling accumulated partial outputs as new blocks arrive to avoid overflow. (2)

Question 2

(a) [6 marks] Each of $W^Q,W^K,W^V$ maps $d_{\text{model}} \to h\cdot d_k = d_{\text{model}}$, so each is $d_{\text{model}}\times d_{\text{model}}$ = $d_{\text{model}}^2$ params. (3) $W^O$ maps concatenated $h\cdot d_v = d_{\text{model}}$ back to $d_{\text{model}}$: $d_{\text{model}}^2$. (1) Total $= 4 d_{\text{model}}^2$, independent of $h$ because $h\,d_k = d_{\text{model}}$ cancels the $h$ dependence. (2)

(b) [8 marks] Let $\omega_i = 10000^{-2i/d}$. The pair is $\begin{pmatrix}\sin(\omega_i pos)\\ \cos(\omega_i pos)\end{pmatrix}$. (2) Using angle addition: $\sin(\omega_i(pos+k)) = \sin(\omega_i pos)\cos(\omega_i k) + \cos(\omega_i pos)\sin(\omega_i k)$ $\cos(\omega_i(pos+k)) = \cos(\omega_i pos)\cos(\omega_i k) - \sin(\omega_i pos)\sin(\omega_i k)$. (3) Hence $M_k^{(i)} = \begin{pmatrix} \cos(\omega_i k) & \sin(\omega_i k)\\ -\sin(\omega_i k) & \cos(\omega_i k)\end{pmatrix},$ independent of $pos$ — a rotation by angle $\omega_i k$. (3)

(c) [6 marks] Rotate query at position $m$: $\tilde q_m = R(m\theta) q$; key at $n$: $\tilde k_n = R(n\theta) k$. (2) Score $= \tilde q_m^\top \tilde k_n = q^\top R(m\theta)^\top R(n\theta) k$. (2) Using $R(a)^\top R(b) = R(b-a)$: $= q^\top R((n-m)\theta) k$, depending only on $n-m$. (2)

Question 3

(a) [8 marks]

def causal_self_attention(X, Wq, Wk, Wv, dk):
    Q = X @ Wq                       # (n, dk)
    K = X @ Wk
    V = X @ Wv
    scores = (Q @ K.T) / np.sqrt(dk) # (n, n)
    n = X.shape[0]
    mask = np.triu(np.ones((n, n)), k=1).astype(bool)  # True above diagonal
    scores[mask] = -np.inf           # future positions masked
    A = softmax(scores, axis=-1)      # row-wise
    return A @ V

Marks: Q/K/V projection (2), scaled scores (1), correct upper-triangular mask (2), $-\infty$ before softmax (1). (6) Why before softmax: softmax normalises over all positions; inserting $-\infty$ makes $e^{-\infty}=0$ so masked positions get exactly zero weight and the remaining weights renormalise to sum to 1. Masking after softmax (zeroing) would leave the denominator including future tokens, so probabilities would not sum to 1 and would leak information about future magnitudes. (2)

(b) [6 marks] Post-LN: $x_{l+1} = \text{LN}(x_l + F(x_l))$ — the residual is inside the normalisation, so gradients must pass through every LN, and the identity path is not clean; at large depth this causes gradient magnitudes to explode/vanish, requiring warmup. (3) Pre-LN: $x_{l+1} = x_l + F(\text{LN}(x_l))$ — the residual stream is an unnormalised identity highway; gradients flow additively/unimpeded from top to bottom (each layer contributes an $I + \partial F$ term), keeping gradient norms bounded and stable, so training proceeds without warmup. (3)

(c) [4 marks]

1. Encoder self-attention: Q, K, V all from the encoder input sequence. (1)
2. Decoder masked self-attention: Q, K, V all from the (partial) decoder sequence, causally masked. (1.5)
3. Encoder-decoder cross-attention: Q from the decoder, K and V from the encoder output. (1.5)

[
  {"claim":"Raw dot-product of two iid unit-variance d_k-dim vectors has variance d_k",
   "code":"dk=symbols('d_k',positive=True); var=dk*1; result = (var==dk)"},
  {"claim":"Softmax Jacobian diagonal is a_i(1-a_i)",
   "code":"z1,z2=symbols('z1 z2'); a1=exp(z1)/(exp(z1)+exp(z2)); J=diff(a1,z1); expr=a1*(1-a1); result = simplify(J-expr)==0"},
  {"claim":"MHA total projection params = 4*d_model^2 independent of h",
   "code":"dm,h=symbols('d_model h',positive=True); dk=dm/h; per=dm*(h*dk); total=3*per+ (h*dk)*dm; result = simplify(total-4*dm**2)==0"},
  {"claim":"Rotation composition R(a)^T R(b) = R(b-a) gives relative dependence",
   "code":"a,b=symbols('a b'); R=lambda t: Matrix([[cos(t),-sin(t)],[sin(t),cos(t)]]); lhs=simplify(R(a).T*R(b)); rhs=R(b-a); result = simplify(lhs-rhs)==zeros(2,2)"}
]