AI-ML interleaved practice
Instructions. Work each problem on its own. These are deliberately mixed across subtopics — before computing, decide which concept and method applies. Show all work. Marks in brackets. Use for any math.
1. [4 marks] A self-attention layer processes a sequence of length with model dimension . State the time and memory complexity of standard scaled dot-product attention in terms of and , and explain why this scaling is exactly the limitation that Flash Attention targets (and what Flash Attention does NOT change asymptotically).
2. [3 marks] Given query vector and two key vectors , with head dimension , compute the two scaled dot-product scores (before softmax). Give exact values.
3. [4 marks] A BPE tokenizer starts from a character vocabulary. Given the training corpus word-with-counts: low:5, lower:2, newest:6, widest:3. Perform the first merge step: identify the most frequent adjacent symbol pair and state the merged token. Show the pair counts you used.
4. [3 marks] For the sinusoidal positional encoding , explain why this specific functional form was chosen over learned absolute positions in one respect that matters for generalization, and contrast briefly with what RoPE encodes instead.
5. [4 marks] A language model assigns the following probabilities to a 4-token test sequence: . Compute the perplexity. Show the formula you selected and why perplexity (not raw cross-entropy) is being reported.
6. [3 marks] You must choose an architecture for: (a) machine translation, (b) autoregressive text generation, (c) sentence-embedding for classification. Assign each to encoder-only, decoder-only, or encoder-decoder and justify each choice in one sentence.
7. [4 marks] Multi-head attention uses heads with model dimension . What is the dimension per head under the standard convention? If instead you used a single head of full dimension , state one representational advantage multi-head gives that the single head loses.
8. [3 marks] Write the causal (masked) attention mask logic: for a length-4 sequence, give the mask matrix (1 = allowed, 0 = blocked) used in decoder self-attention, and explain in one sentence why this is required for the causal language modeling objective but NOT for BERT's masked language modeling.
9. [4 marks] A model has vocabulary and embedding dimension . (a) How many parameters in the input embedding matrix? (b) If input and output (unembedding) weights are tied, how many parameters are saved compared to untied? (c) Name one tradeoff of increasing .
10. [4 marks] Consider "Pre-LN" vs "Post-LN" transformer blocks. Given a residual sublayer computing (Pre-LN) versus (Post-LN), state which placement gives more stable gradients in very deep networks and explain the why in terms of the residual gradient path.
Answer keyMark scheme & solutions
1. Tests 4.1.13 (complexity) + 4.1.14 (Flash attention). Time: ; Memory: for the attention matrix (plus activations). The term dominates because every query attends to every key. Flash Attention reduces memory to (no materialized matrix; tiling + online softmax, recomputation) and improves wall-clock via reduced HBM I/O, but the compute remains asymptotically — it is IO-aware, not sub-quadratic. Why method: recognize this is an asymptotic-analysis question, not a numeric one; the trap is claiming Flash changes the FLOP order.
2. Tests 4.1.4 (scaled dot-product). . Scaled: . . Scaled: . Both scaled scores . Why method: apply scaling by ; the dot products happen to be equal here.
3. Tests 4.2.2 (BPE). Split into characters (with counts): pairs and their frequencies:
l o: 5+2 = 7 (low, lower)o w: 5+2 = 7w e: 2 (lower) +6 (newest) +3 (widest)... let's count carefully across words:newest→ n e w e s t: pairsn e:6,e w:6,w e:6,e s:6,s t:6widest→ w i d e s t:w i:3,i d:3,d e:3,e s:3,s t:3lower→ l o w e r:l o:2,o w:2,w e:2,e r:2low→ l o w:l o:5,o w:5
Aggregate top pairs: e s = 6+3 = 9, s t = 6+3 = 9, l o = 7, o w = 7.
Most frequent = e s or s t (tie at 9). Standard tie-break picks one (e.g., first encountered) — first merge: es (or st; either accepted with the count 9 shown).
Why method: count adjacent symbol pairs across the whole weighted corpus; the trap is forgetting to weight by word count.
4. Tests 4.1.6 (sinusoidal PE) + 4.1.7 (RoPE contrast). Sinusoidal encodings let the model represent relative positions via linear combinations: is a linear function of (rotation), and they extrapolate to sequence lengths unseen in training (fixed, not learned). RoPE instead injects position by rotating the query/key vectors so that the attention dot-product depends only on the relative offset , embedding relative position directly into rather than adding to inputs. Why method: conceptual contrast question; identify "additive absolute-but-relative-friendly" vs "multiplicative rotation in attention."
5. Tests 4.2.9 (perplexity). , . probs: , , , . Sum . Mean . . (Equivalently geometric mean of inverse probs: .) Why method: perplexity = exp of average negative log-likelihood; reported because it's an interpretable "effective branching factor."
6. Tests 4.1.10 (architecture choice). (a) MT → encoder-decoder: encoder builds full bidirectional source representation, decoder generates target conditioned on it. (b) Autoregressive generation → decoder-only: causal masking, next-token prediction. (c) Sentence embedding for classification → encoder-only (BERT-style): bidirectional context, no generation needed. Why method: map task structure (input→output type, need for generation, need for bidirectionality) to the correct block layout.
7. Tests 4.1.5 (multi-head). . Multi-head lets different heads attend to different subspaces/relations simultaneously (e.g., syntactic vs positional patterns) and averages/concatenates them; a single 512-d head must collapse all relations into one attention distribution, losing this representational diversity. Why method: standard convention divides dimension across heads; advantage is subspace specialization.
8. Tests 4.1.11 (masked/causal attention) + contrast 4.2.7. Lower-triangular mask: Position attends only to positions . Required for causal LM so predicting token cannot "cheat" by seeing future tokens; BERT's MLM is bidirectional (predicts masked tokens using full context) so no causal mask. Why method: recognize causal-vs-bidirectional distinction drives the mask.
9. Tests 4.2.4 (vocab tradeoffs) + 4.2.5 (embeddings/tied weights). (a) parameters. (b) Tying reuses the same matrix for input and output, saving one full copy: 38,400,000 parameters saved. (c) Larger : fewer/shorter token sequences (shorter effective context in tokens) but a much larger embedding + softmax layer (more params, rarer tokens undertrained) — a memory/coverage vs efficiency tradeoff. Why method: param count = ; tying saves exactly one such matrix.
10. Tests 4.1.9 (residual + LN placement). Pre-LN is more stable in very deep nets. In Pre-LN, the residual path has an identity term whose gradient passes through unnormalized, giving a clean skip connection so gradients don't vanish/explode with depth. In Post-LN, the LayerNorm sits on the residual sum, scaling gradients along the skip path and requiring careful warmup/learning-rate tuning to remain stable. Why method: trace gradient through the identity branch to argue stability.
[
{"claim":"Scaled dot-product scores in Q2 both equal 1/sqrt(3)",
"code":"import sympy as sp\nq=sp.Matrix([1,0,1]); k1=sp.Matrix([1,0,0]); k2=sp.Matrix([0,1,1])\ndk=sp.sqrt(3)\ns1=(q.dot(k1))/dk; s2=(q.dot(k2))/dk\nresult = sp.simplify(s1-1/sp.sqrt(3))==0 and sp.simplify(s2-1/sp.sqrt(3))==0"},
{"claim":"Perplexity of probs [0.5,0.25,0.5,0.1] is 160**0.25 ~ 3.557",
"code":"import sympy as sp\nps=[sp.Rational(1,2),sp.Rational(1,4),sp.Rational(1,2),sp.Rational(1,10)]\nN=len(ps)\nPP=(1/(ps[0]*ps[1]*ps[2]*ps[3]))**sp.Rational(1,N)\nresult = abs(float(PP)-3.5566)<1e-3"},
{"claim":"Embedding matrix params V*d = 38.4M and tied saving equals same amount",
"code":"V=50000; d=768\nemb=V*d\nsaved=V*d\nresult = (emb==38400000) and (saved==38400000)"}
]