## Self-Attention: Implementation, Complexity, and Efficient Variants This round probes how deeply you understand the attention mechanism — not just the formula, but the implementation, its cost, and the ideas behind the efficient variants used in modern large models. Work through the parts in order; later parts assume the complexity analysis from the first. ### Constraints & Assumptions - Sequence length $n$, model dimension $d$ (per-head dimension $d_k$), batch size $B$, $h$ heads. - Use scaled dot-product attention as the base mechanism; multi-head attention is the standard extension. - "Complexity" means both **time** (FLOPs) and **memory** (peak activation/HBM traffic), since the efficient variants trade these differently. - You may write code in Python with NumPy or PyTorch; correctness and clarity matter more than micro-optimizations. ### Clarifying Questions to Ask - Should the implementation support causal (autoregressive) masking and padding masks, or just full bidirectional attention? - Do you want single-head scaled dot-product attention first, then multi-head, or multi-head directly? - For the complexity discussion, are we interested in the asymptotic order, or the practical bottleneck on a GPU (compute-bound vs. memory-bandwidth-bound)? - For the efficient variants, is the goal exact attention computed more cheaply (e.g., FlashAttention) or an approximation that changes the math (e.g., linear attention)? ### Part 1 — Implement attention and analyze its complexity Write scaled dot-product attention (and sketch multi-head attention), supporting an optional causal mask. Then derive its time and memory complexity in $n$ and $d$, and state the practical consequence for long sequences. ```hint The formula $\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$ — the $\sqrt{d_k}$ scaling keeps the dot products from saturating the softmax. ``` ```hint Where the cost is The two matmuls are $QK^\top$ ($n\times d$ by $d\times n$) and $(\text{scores})V$ ($n\times n$ by $n\times d$); the $n\times n$ score matrix is the source of the quadratic cost. ``` #### What This Part Should Cover ```premium-lock What This Part Should Cover ``` ### Part 2 — FlashAttention: what makes it "flash" Explain how FlashAttention computes **exact** attention faster and with less memory than the naive implementation. Specifically, why is the naive version memory-bandwidth-bound on a GPU, and what does FlashAttention change to fix that? ```hint Name the real bottleneck On a GPU the matmuls are fast; the killer is reading/writing the giant $n\times n$ scores and the softmax intermediates to/from HBM. FlashAttention is an IO-aware algorithm. ``` ```hint The two key tricks Tiling (process Q/K/V in blocks that fit in fast on-chip SRAM) plus an online/streaming softmax (running max and running normalizer) so the full $n\times n$ matrix is never materialized in HBM. ``` #### What This Part Should Cover ```premium-lock What This Part Should Cover ``` ### Part 3 — Linear attention: changing the math to beat quadratic Explain linear (kernelized) attention: how replacing the softmax with a kernel feature map turns the $O(n^2)$ cost into $O(n)$ (or $O(n d^2)$), why associativity is the key, and what you trade away versus exact softmax attention. ```hint The reassociation Softmax couples all pairs, but if $\mathrm{softmax}(QK^\top)$ is replaced by $\phi(Q)\phi(K)^\top$, you can compute $\phi(K)^\top V$ first ($d\times d$) and then multiply by $\phi(Q)$ — the matrix-multiplication order changes the cost from $n^2$ to linear in $n$. ``` #### What This Part Should Cover ```premium-lock What This Part Should Cover ``` ### What a Strong Answer Covers ```premium-lock What a Strong Answer Covers ``` ### Follow-up Questions - FlashAttention gives the *same* result as naive attention, so where exactly does the speedup come from if the FLOP count is essentially unchanged? - During autoregressive decoding, why is the KV cache the memory bottleneck, and how do grouped-query / multi-query attention and linear attention each address it? - Linear attention is $O(n)$ but often underperforms softmax attention on long-context recall. Why, intuitively, and what hybrids try to recover the gap? - How does the causal mask change the cost and the implementation of FlashAttention and of linear attention?