Self-Attention: Implementation, Complexity, and Efficient Variants

Scenario

You are optimizing the training and inference pipeline for a large language model with context window length N and hidden dimension D across H attention heads (per-head dimension d = D / H). The standard scaled dot-product self-attention mechanism becomes a critical bottleneck as N scales to 32k tokens and beyond.

Tasks

1. Mathematical Formulation and Complexity Derivation

Given input representation X of shape (N, D) and projection weight matrices W_Q, W_K, W_V of shape (D, D):

• Write down the mathematical definition of standard multi-head self-attention for a single batch.
• Derive the exact FLOP count and peak memory consumption (in terms of N, D, and d) for the forward pass, explicitly accounting for the attention weight matrix S = (Q K^T) / sqrt(d) and output projection.
• Identify the memory footprint bottleneck during backpropagation when storing activations for gradient computation.

2. Algorithmic Optimization: FlashAttention

Standard attention is memory-bound rather than compute-bound on modern accelerators (e.g., NVIDIA H100/A100) due to repeated reads/writes of the (N, N) attention matrix to high-bandwidth memory (HBM).

• Explain how FlashAttention addresses this bottleneck using GPU SRAM tiling and the online softmax trick.
• Derive the online softmax formula used to compute running normalization factors across tiled blocks of keys and values without materializing the full N x N matrix.
• Explain how FlashAttention reduces backward pass memory from O(N^2) to O(N) without storing the full intermediate attention matrix.

3. Linear Attention Variants

An alternative line of research replaces standard softmax attention with kernel-based linear attention: Attention(Q, K, V) = phi(Q) (phi(K)^T V).

• Explain how changing the matrix multiplication associativity reduces time complexity from O(N^2 d) to O(N d^2).
• Discuss the trade-offs of linear attention regarding causal autoregressive decoding, expressive power, and why standard softmax attention is still preferred for frontier models.