Level 4 — ApplicationTransformer Architecture

Transformer Architecture

60 marksprintable — key stays hidden on paper

Level 4 (Application: novel/unseen problems) Time: 60 minutes Total Marks: 60

Instructions: Show all reasoning. Use $...$ for math. Calculators permitted; no external references.

Question 1 — Scaled Dot-Product Attention on Concrete Vectors (12 marks)

A single attention head processes 3 tokens with $d_k = d_v = 2$ . The projected matrices are:

$Q = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix},\quad K = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix},\quad V = \begin{bmatrix} 2 & 0 \\ 0 & 2 \\ 4 & 4 \end{bmatrix}$

(a) Compute the raw score matrix $QK^\top$ and then divide by $\sqrt{d_k}$ . (4) (b) Explain why we divide by $\sqrt{d_k}$ and what would degrade numerically if $d_k$ were 512 and we omitted it. (3) (c) For token 1 (first row), compute the softmax attention weights over the 3 keys (using the scaled scores) and the resulting output vector. Keep 3 decimals. (5)

Question 2 — Designing a Causal Language Model Head (14 marks)

You are asked to adapt an encoder-style block into a decoder-only autoregressive model over a length- $n$ sequence.

(a) Write down the additive mask matrix $M$ (entries $0$ or $-\infty$ ) that must be added to the pre-softmax scores for $n=4$ , and state the rule for entry $M_{ij}$ . (4) (b) A colleague suggests applying the mask after the softmax by zeroing forbidden weights. Prove this is incorrect by showing the row would no longer be a valid probability distribution, and give the correct fix if one insisted on post-softmax masking. (4) (c) Explain precisely why, at inference time, a KV-cache makes causal decoding cheaper, and give the per-new-token attention cost with and without the cache in Big-O terms (sequence so far length $t$ ). (6)

Question 3 — Positional Encoding Under Sequence Extrapolation (12 marks)

(a) For sinusoidal encodings with model dimension $d=4$ , base $10000$ , write the encoding vector $PE(pos)$ as a function of $pos$ (give the 4 component formulas). (4) (b) A model trained on length 512 is asked to process length 4096. Contrast how sinusoidal absolute encodings vs RoPE behave with respect to relative position and extrapolation, and justify why RoPE's rotation formulation preserves relative offsets. (5) (c) Show that in RoPE, the dot product between a query at position $m$ and a key at position $n$ depends only on $(m-n)$ for a single 2D rotary sub-block. (3)

Question 4 — Complexity Trade-off Analysis (12 marks)

A transformer layer has sequence length $n$ , model dim $d$ , FFN hidden dim $4d$ , and $h$ heads.

(a) Derive the asymptotic cost (FLOPs, Big-O) of: (i) the self-attention score+weighting, (ii) the QKV+output projections, (iii) the FFN sublayer. (6) (b) For $n = 8192$ , $d = 1024$ , identify which term dominates and by roughly what ratio (order-of-magnitude, ignoring constants). (3) (c) State the key idea by which FlashAttention reduces memory (not asymptotic FLOP) cost, and explain why it never materialises the full $n \times n$ matrix. (3)

Question 5 — Multi-Head vs Single-Head Reasoning (10 marks)

A model has $d_{model}=512$ . Compare configuration A: 1 head with $d_k=512$ , vs configuration B: 8 heads with $d_k=64$ .

(a) Show both configurations have (approximately) the same number of parameters in the QKV projections. (4) (b) Give a concrete functional advantage of B over A that A structurally cannot replicate, and justify it. (3) (c) After concatenating 8 head outputs of dim 64, what is the role and shape of $W^O$ , and why is it needed even though concatenation already gives dimension 512? (3)

Answer keyMark scheme & solutions

Question 1 (12)

(a) $QK^\top$ : compute row·row dot products. Rows of $Q$ : $q_1=(1,0), q_2=(0,1), q_3=(1,1)$ ; same as keys.

$QK^\top = \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}$ Divide by $\sqrt{d_k}=\sqrt2\approx1.414$ : $\frac{QK^\top}{\sqrt2} = \begin{bmatrix} 0.707 & 0 & 0.707 \\ 0 & 0.707 & 0.707 \\ 0.707 & 0.707 & 1.414 \end{bmatrix}$ Marks: correct matrix (2), correct scaling (2).

(b) Dot products of $d_k$ -dim vectors with unit-variance components have variance $\propto d_k$ ; scaling by $\sqrt{d_k}$ normalises variance to $O(1)$ . Without it at $d_k=512$ , scores grow ~ $\sqrt{512}\approx23\times$ , pushing softmax into saturated regions → near one-hot outputs and vanishing gradients (softmax derivative $\to0$ ), harming training. (variance reasoning 2, gradient consequence 1).

(c) Token 1 scaled scores: $(0.707, 0, 0.707)$ . Exponentials: $e^{0.707}=2.028,\ e^{0}=1,\ e^{0.707}=2.028$ ; sum $=5.056$ . Weights: $(0.401, 0.198, 0.401)$ . Output $= 0.401\,(2,0) + 0.198\,(0,2) + 0.401\,(4,4)$ $= (0.802+1.604,\ 0.396+1.604) = (2.406, 2.000)$ . Marks: exponentials/softmax (3), output vector (2).

Question 2 (14)

(a) Rule: $M_{ij}=0$ if $j\le i$ (key at or before query), $-\infty$ if $j>i$ . $M=\begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0\end{bmatrix}$ (rule 2, matrix 2).

(b) Softmax normalises over all $n$ positions first: $a_j = e^{s_j}/\sum_{k} e^{s_k}$ . Zeroing entries $j>i$ afterward leaves $\sum_{j\le i} a_j < 1$ (mass leaked to now-zeroed future terms), so the row is not a distribution and outputs are mis-scaled. Correct post-softmax fix: zero forbidden weights then renormalise by dividing by the surviving sum $\sum_{j\le i} a_j$ — equivalent to additive $-\infty$ masking. (prove leakage 2, renormalise fix 2).

(c) Without cache: at step $t$ you recompute $K,V$ for all $t$ tokens → attention costs $O(t^2 d)$ per generated token (or $O(t\,d)$ for the single new query against $t$ keys but recomputing all keys is $O(t d)$ projections; recompute of full attention is $O(t^2 d)$ ). With KV-cache: keys/values of past tokens are stored, so only the new token's $q,k,v$ are computed and it attends to $t$ cached keys → $O(t\,d)$ per new token. Over a full sequence: $O(n^2 d)$ cached vs $O(n^3 d)$ recomputing. (cache concept 2, without-cache cost 2, with-cache cost 2).

Question 3 (12)

(a) $d=4$ , indices $i=0,1$ ; frequencies $\omega_i = 10000^{-2i/d}=10000^{-i/2}$ .

$PE(pos)_0 = \sin(pos)$
$PE(pos)_1 = \cos(pos)$
$PE(pos)_2 = \sin(pos/100)$ (since $10000^{-2/4}=1/100$ )
$PE(pos)_3 = \cos(pos/100)$ (each pair/frequency 2 each).

(b) Sinusoidal absolute encodings are added to embeddings; at length 4096 the model sees $pos$ values (and low-frequency phases) never encountered during training → poor extrapolation, though sinusoids are defined for all $pos$ . RoPE instead rotates Q/K by an angle proportional to position; attention scores depend on relative angle $(m-n)$ , so a shift in absolute position leaves relative structure invariant, giving smoother length extrapolation. RoPE encodes relative position implicitly in the dot product rather than as an additive absolute signal. (sinusoidal limitation 2, RoPE relative property 2, justification 1).

(c) For a 2D sub-block, RoPE applies rotation $R(\theta m)$ to $q$ and $R(\theta n)$ to $k$ : $(R(\theta m)q)^\top (R(\theta n)k) = q^\top R(\theta m)^\top R(\theta n) k = q^\top R(\theta(n-m)) k.$ Since $R(\alpha)^\top R(\beta)=R(\beta-\alpha)$ , the result depends only on $(n-m)$ . (rotation composition 2, conclusion 1).

Question 4 (12)

(a) (i) Scores $QK^\top$ : $n\times n$ from $n\times d$ → $O(n^2 d)$ ; weighting $AV$ also $O(n^2 d)$ . Total attention core $O(n^2 d)$ . (ii) QKV projections: 3 matmuls $n\times d$ by $d\times d$ → $O(n d^2)$ ; output projection $O(n d^2)$ . (iii) FFN: $n\times d$ by $d\times4d$ , then back → $O(n d \cdot 4d)=O(n d^2)$ . (each part 2).

(b) Attention term $n^2 d$ vs projection/FFN term $\sim n d^2$ . Ratio attention/FFN $\approx n/d = 8192/1024 = 8$ . So attention (the $n^2d$ term) dominates by roughly $8\times$ at this configuration. (identify crossover 1, ratio 2).

(c) FlashAttention tiles Q,K,V into blocks kept in fast SRAM and computes softmax online (running max + running sum, streaming), so the $n\times n$ score matrix is never written to HBM. This reduces memory from $O(n^2)$ to $O(n)$ (plus recomputation in backward), while FLOPs stay $O(n^2 d)$ . (online softmax/tiling 2, no full matrix 1).

Question 5 (10)

(a) QKV projection params: for config A, three $512\times512$ matrices $=3\cdot512^2$ . For config B, each head has $512\times64$ for Q,K,V; 8 heads → $8\cdot3\cdot(512\cdot64)=3\cdot512\cdot(8\cdot64)=3\cdot512\cdot512=3\cdot512^2$ . Equal. (A count 2, B count and equality 2).

(b) B can attend to different representation subspaces simultaneously — e.g. one head tracks syntactic dependency while another tracks positional/coreference — producing distinct attention distributions per head. A single head yields exactly one attention distribution per query, so it cannot mix multiple relational patterns in one layer. (advantage 2, justification 1).

(c) $W^O$ has shape $512\times512$ ( $(h\cdot d_v)\times d_{model}$ ). Concatenation just stacks head outputs in fixed slots; $W^O$ mixes/recombines information across heads and projects back into the model space, letting the layer learn how head outputs interact rather than treating them as independent fixed partitions. (shape 1, role/mixing 2).

[
  {"claim":"Q1a scaled score matrix correct", "code":"Q=Matrix([[1,0],[0,1],[1,1]]); K=Q; S=Q*K.T/sqrt(2); result = S.equals(Matrix([[1,0,1],[0,1,1],[1,1,2]])/sqrt(2))"},
  {"claim":"Q1c token1 output vector approx (2.406,2.000)", "code":"import math; s=[0.707,0,0.707]; e=[math.exp(x) for x in s]; Z=sum(e); w=[x/Z for x in e]; V=[[2,0],[0,2],[4,4]]; out=[sum(w[i]*V[i][j] for i in range(3)) for j in range(2)]; result = abs(out[0]-2.406)<0.01 and abs(out[1]-2.000)<0.01"},
  {"claim":"Q4b attention/FFN ratio n/d equals 8", "code":"n=8192; d=1024; result = (n*n*d)/(n*d*d)==8"},
  {"claim":"Q5a QKV param counts equal for A and B", "code":"A=3*512**2; B=8*3*(512*64); result = A==B"}
]