AI-ML interleaved practice
Instructions: Solve all problems. Show your reasoning. Each problem draws from a different subtopic — part of the challenge is recognizing which framework or formula applies before you compute. Use notation for math. Total: 100 marks.
1. (12 marks) A team has a fixed compute budget of FLOPs (where = parameters, = training tokens). Under Chinchilla scaling, the compute-optimal allocation follows and . If you increase your compute budget by a factor of , by what factor should parameters and data each grow to stay compute-optimal? A rival trains a model with the Chinchilla-optimal parameter count on the same compute — qualitatively, what problem does their model suffer from?
2. (10 marks) In a Mixture-of-Experts layer with experts and top- routing, each token is dispatched to 2 experts. If the total expert parameters are but only are active per token, state (a) the ratio of total-to-active parameters, and (b) why this decouples model capacity from inference FLOPs.
3. (8 marks) You are designing an agent that must answer "What was the temperature in Paris yesterday, and is that warmer than the historical average?" Sketch the ReAct trace (Thought / Action / Observation cycles) needed to solve this. Identify how many tool calls are required and why interleaving reasoning with acting beats pure chain-of-thought here.
4. (10 marks) A transformer with hidden dim , layers processes a sequence of length in fp16 (2 bytes). Compute the KV-cache memory for a single sequence. Then explain how Grouped-Query Attention with 8 KV-heads instead of 32 query-heads reduces this, giving the reduction factor.
5. (10 marks) A load-balancing auxiliary loss in MoE is defined as , where is the fraction of tokens routed to expert and is the average routing probability for expert . Explain why this loss is minimized under uniform routing, and what failure mode occurs without it.
6. (12 marks) Speculative decoding uses a small draft model to propose tokens, which a large target model verifies in one forward pass. If the draft model's token acceptance rate is , the expected number of tokens accepted per verification step is . Compute this expected value. If a normal decode step costs the same as one verification step, estimate the speedup factor.
7. (8 marks) State-space models like Mamba claim complexity vs. transformer attention's . For , express the ratio of compute. Give one reason Mamba trades away, compared to full attention.
8. (10 marks) In ZeRO Stage 3 (full sharding, i.e. FSDP), across GPUs, the optimizer states, gradients, AND parameters are all partitioned. If a naive data-parallel setup needs bytes per GPU for these three, roughly what does each GPU hold under ZeRO-3? Name the communication cost you pay for this memory saving.
9. (10 marks) An "emergent ability" is often reported as a sharp jump in a metric at some scale. A critic argues these jumps are partly an artifact of the metric choice (e.g., exact-match accuracy). Explain the critic's argument and how switching to a continuous/token-level metric changes the picture.
10. (10 marks) You deploy a code-generation agent that must never output shell commands that delete files. Describe two complementary guardrail mechanisms — one at generation time (constrained decoding) and one post-hoc — and explain why interpretability (6.3) matters for auditing why the model wanted to emit a dangerous command.
Answer keyMark scheme & solutions
1. — Subtopic 6.1.1 (Neural scaling laws) Since and , a compute increase gives factor for both parameters and data.
- Parameters grow , data grows (jointly, since and ✓).
- The rival with optimal params on same compute is under-trained (too few tokens per parameter). Their model is oversized/data-starved — worse loss than the compute-optimal point, and wasteful at inference.
Why this method: The tell is "compute-optimal" and — pure Chinchilla, not a MoE or systems question.
2. — Subtopic 6.1.4 (MoE architecture) (a) Total-to-active ratio . (b) MoE decouples capacity from FLOPs: the model stores parameters (large capacity, more knowledge) but each token only computes through (top-2). Inference FLOPs scale with active params, not total — so you get a bigger model at the cost of a smaller model's compute.
Why this method: "total vs active parameters" + top- signals capacity/compute decoupling, distinct from #5's balancing.
3. — Subtopic 6.2.2 (ReAct) Trace:
- Thought: Need yesterday's Paris temp → use weather tool.
- Action:
get_weather(city="Paris", date="yesterday") - Observation: e.g. 14°C.
- Thought: Need historical average → another lookup.
- Action:
get_climate_avg(city="Paris") - Observation: e.g. 11°C.
- Thought: 14 > 11, so yes, warmer.
- Answer: Warmer than average.
2 tool calls required. ReAct beats pure CoT because the answer depends on external live/historical facts the model can't hallucinate reliably; interleaving lets each Observation ground the next reasoning step.
Why this method: Real-world factual grounding + acting = ReAct, not planning decomposition (6.2.4).
4. — Subtopic 6.1.13 (KV-cache optimization) KV-cache size (2 for K and V). bytes bytes GiB. GQA with 8 KV-heads vs 32 query-heads shares KV across groups: KV memory scales with KV-heads, so reduction factor . New cache GiB.
Why this method: Memory-per-sequence formula and head-sharing = KV-cache, not model parallelism.
5. — Subtopic 6.1.6 (Load balancing) is a dot product of the token-fraction vector and probability vector , scaled by . By Cauchy-Schwarz / convexity, given and , the product is minimized when both are uniform (), and the loss reaches its minimum . Non-uniform routing correlates high with high , raising the loss. Failure mode without it: expert collapse — the router sends most tokens to a few favored experts; others get no gradient and die, wasting capacity.
Why this method: , , auxiliary loss = load balancing specifically (not routing mechanics of 6.1.5).
6. — Subtopic 6.1.12 (Speculative decoding) Expected accepted tokens . So ~3.36 tokens per verification step. If a verification step ≈ one normal decode step in cost, speedup (minus draft-model overhead, so realistically ~2–2.5×).
Why this method: draft/verify + acceptance rate formula = speculative decoding.
7. — Subtopic 6.1.11 (State-space models) Ratio of transformer to Mamba compute . So attention does ~ more work in the sequence-mixing term at this length. Trade-off: Mamba's fixed-size recurrent state compresses history, losing the exact random-access recall that full attention provides — worse at some in-context copy/retrieval tasks.
Why this method: vs + "trade away" = SSM.
8. — Subtopic 6.1.9 (FSDP / ZeRO-3) Under ZeRO-3, params + grads + optimizer states are each split across GPUs, so each GPU holds bytes for these components. Communication cost: you must all-gather parameter shards before each layer's forward/backward and reduce-scatter gradients — extra communication volume replaces the saved memory.
Why this method: "all three sharded" = Stage 3 / FSDP, not Stage 1/2 or plain data parallel.
9. — Subtopic 6.1.3 (Emergent abilities) Critic's argument: metrics like exact-match are discontinuous/nonlinear — a model can steadily improve token-level correctness, but the composite task only "flips" to success once all subtokens are right, producing a sharp apparent jump. Under a continuous metric (per-token log-likelihood or partial credit), the underlying capability improves smoothly with scale, and the "emergence" cliff largely disappears. So some emergence is a measurement artifact, not a phase transition in the model.
Why this method: metric-artifact critique is the emergent-abilities debate.
10. — Subtopics 6.2.10 (Guardrails) + 6.3.1 (Interpretability)
- Generation-time (constrained decoding): grammar/regex-constrained or allow-list decoding that masks tokens forming forbidden commands (e.g., block
rm -rfpatterns via a finite-state constraint), so unsafe sequences are unreachable. - Post-hoc filter: a separate classifier/validator scans the completed output and rejects/rewrites destructive commands before execution. They're complementary: constrained decoding prevents generation, the filter catches paraphrases/obfuscations the constraint missed. Interpretability matters because auditing why the model proposed a dangerous command (via attention/feature attribution or circuit analysis) lets you fix the root cause (e.g., a spurious feature triggered by user phrasing) rather than only patching symptoms.
Why this method: two-layer safety + auditing = guardrails combined with the interpretability motivation.
[
{
"claim": "16x compute gives 4x params and 4x data under Chinchilla, and 4*4=16",
"code": "factor = 16**0.5\nresult = (abs(factor - 4) < 1e-9) and (factor*factor == 16)"
},
{
"claim": "Speculative decoding expected accepted tokens with beta=0.8, gamma=4 is ~3.3616",
"code": "beta=0.8; gamma=4\nexp = (1 - beta**(gamma+1))/(1-beta)\nresult = abs(exp - 3.3616) < 1e-4"
},
{
"claim": "KV cache = 2*L*S*d*2 bytes = 8 GiB for L=32,S=8192,d=4096",
"code": "L=32; S=8192; d=4096; b=2\nkv = 2*L*S*d*b\ngib = kv/(1024**3)\nresult = abs(gib - 8.0) < 1e-6"
}
]