Checklist USA-NA-AIO 2025 Round 2 · Task 2
Multi-Head Attention and Its Variants (GQA and MLA)
Analyse and implement multi-head attention, grouped-query attention and multi-head latent attention, and relate GQA to MLA.
The task
Problem 2 (100 points, 14 parts) sets up notation for attending and attended sequences and asks for the shapes of the projection matrices, the attention scores and outputs, and a from-scratch multi-head attention module.
It then introduces grouped-query attention (a module MyGQA without copying key/value projections and without loops), shows that MHA is a special case, proves that GQA can be represented as multi-head latent attention (MLA) and implements the conversion GQA_2_MLA, gives a counterexample showing MLA is not always representable as GQA, derives reduced projection matrices for MLA, and discusses key/value caching in autoregressive generation.
Abridged by SOTA from the official materials. The official statement has the exact rules, and it wins wherever this summary differs.
At a glance
- You submit
- Written answers and code in a notebook.
- Rules
- Only torch, torch.nn and numpy as imported in the starter cell.
- Format
- 2025 USA-NA-AIO Round 2, 27 April 2025 at MIT (IOAI news article). Statements and official solutions were posted on the USAAIO forum on 14 May 2025.