Level 2 — RecallTransformer Architecture

Transformer Architecture

⏱ 30 minutes40 marksprintable — key stays hidden on paper

Subject: AI-ML Chapter: Transformer Architecture Time limit: 30 minutes Total marks: 40 Level 2 — Recall: definitions, standard textbook problems, short derivations

Answer all questions. Show working where required.

Q1. (4 marks) State two limitations of RNNs that motivated the transformer architecture. Briefly explain how self-attention addresses each.

Q2. (4 marks) Write the formula for scaled dot-product attention given query matrix $Q$ , key matrix $K$ , and value matrix $V$ . Explain the role of the scaling factor $\frac{1}{\sqrt{d_k}}$ .

Q3. (5 marks) In a self-attention layer the model dimension is $d_{model}=512$ and the number of heads is $h=8$ . Assuming $d_k=d_v=d_{model}/h$ , compute: (a) the per-head key/value dimension $d_k$ ; (2) (b) the total number of parameters in the three projection matrices $W^Q, W^K, W^V$ (ignore biases) across all heads. (3)

Q4. (4 marks) Define multi-head attention. Why is using multiple heads beneficial compared to a single attention head?

Q5. (5 marks) The sinusoidal positional encoding is defined as $PE_{(pos,2i)} = \sin\!\left(\frac{pos}{10000^{2i/d}}\right),\quad PE_{(pos,2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d}}\right).$ For $d=4$ and $pos=1$ , compute the 4-dimensional positional encoding vector $[PE_{(1,0)}, PE_{(1,1)}, PE_{(1,2)}, PE_{(1,3)}]$ . Give values to 4 decimal places.

Q6. (4 marks) Explain the purpose of masked (causal) attention in the decoder. What values are placed in the masked positions before the softmax, and why?

Q7. (4 marks) Compare the computational and memory complexity of self-attention with respect to sequence length $n$ . State the complexity and explain briefly why.

Q8. (3 marks) Distinguish between encoder-only, decoder-only, and encoder–decoder transformer architectures, giving one example task each.

Q9. (4 marks) Describe the two sublayers of a standard transformer encoder block. State where residual connections and layer normalization are applied in the original "Attention is All You Need" design.

Q10. (3 marks) What is the key idea behind FlashAttention? State one advantage it provides over the naive attention implementation.

End of paper

Answer keyMark scheme & solutions

Q1. (4 marks)

Sequential computation prevents parallelization over time steps (slow training) — self-attention processes all positions in parallel. (2)
Difficulty modelling long-range dependencies / vanishing gradients over long sequences — self-attention gives constant path length $O(1)$ between any two tokens. (2) Why: attention connects every pair directly, removing recurrence.

Q2. (4 marks) $\text{Attention}(Q,K,V)=\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V \quad (2)$ Scaling by $1/\sqrt{d_k}$ counteracts the growth of dot-product magnitude with dimension; without it large values push softmax into regions of tiny gradients (saturation). (2)

Q3. (5 marks) (a) $d_k = 512/8 = 64$ . (2) (b) Each of $W^Q, W^K, W^V$ maps $d_{model}\to d_{model}$ (all heads combined) $=512\times512$ params. Total $=3\times512\times512 = 3\times262144 = 786432$ . (3) Why: full projection matrix is $d_{model}\times(h\cdot d_k)=512\times512$ .

Q4. (4 marks) Multi-head attention runs $h$ attention functions in parallel on different learned linear projections of Q,K,V, concatenates outputs, and projects with $W^O$ : $\text{MHA}=\text{Concat}(head_1,...,head_h)W^O.$ (2) Benefit: allows the model to jointly attend to information from different representation subspaces / positions, capturing multiple relationship types a single head would average out. (2)

Q5. (5 marks) For $d=4$ , $pos=1$ :

$PE_{(1,0)}$ : $i=0$ , $\sin(1/10000^0)=\sin(1)=0.8415$ . (1.25)
$PE_{(1,1)}$ : $\cos(1)=0.5403$ . (1.25)
$PE_{(1,2)}$ : $i=1$ , $\sin(1/10000^{2/4})=\sin(1/100)=\sin(0.01)=0.0100$ . (1.25)
$PE_{(1,3)}$ : $\cos(0.01)=0.99995$ . (1.25) Vector $\approx[0.8415, 0.5403, 0.0100, 0.99995]$ .

Q6. (4 marks) Masked attention prevents a position from attending to future positions, preserving the autoregressive property so predictions depend only on known (past) tokens. (2) Masked positions are set to $-\infty$ (large negative) before softmax, so their softmax weights become 0. (2)

Q7. (4 marks) Self-attention is $O(n^2 d)$ time and $O(n^2)$ memory in sequence length $n$ . (2) Because computing $QK^\top$ produces an $n\times n$ attention matrix — every token attends to every other token. (2)

Q8. (3 marks)

Encoder-only (e.g., BERT) — classification/understanding. (1)
Decoder-only (e.g., GPT) — autoregressive text generation. (1)
Encoder–decoder (e.g., original Transformer/T5) — machine translation / seq2seq. (1)

Q9. (4 marks) Two sublayers: (1) multi-head self-attention, (2) position-wise feed-forward network. (2) In the original design each sublayer uses a residual connection followed by layer norm: $\text{LayerNorm}(x+\text{Sublayer}(x))$ (post-norm). (2)

Q10. (3 marks) FlashAttention computes exact attention in a tiled/block-wise manner using online softmax, keeping intermediate blocks in fast SRAM and never materializing the full $n\times n$ matrix in HBM. (2) Advantage: reduced memory (linear in $n$ ) and faster wall-clock time via fewer memory reads/writes (IO-awareness). (1)

[
  {"claim":"d_k = 512/8 = 64","code":"result = (512//8 == 64)"},
  {"claim":"Total QKV projection params = 786432","code":"result = (3*512*512 == 786432)"},
  {"claim":"PE(1,0)=sin(1)~0.8415, PE(1,1)=cos(1)~0.5403","code":"import sympy as sp; s=sp.N(sp.sin(1)); c=sp.N(sp.cos(1)); result = (abs(s-0.8415)<1e-3 and abs(c-0.5403)<1e-3)"},
  {"claim":"PE(1,2)=sin(0.01)~0.0100, PE(1,3)=cos(0.01)~0.99995","code":"import sympy as sp; s=sp.N(sp.sin(sp.Rational(1,100))); c=sp.N(sp.cos(sp.Rational(1,100))); result = (abs(s-0.0100)<1e-3 and abs(c-0.99995)<1e-4)"}
]