Why is self-attention O(n^2) and what are the implications?

WHAT THIS TESTS: Whether you understand the self-attention computational graph at a systems level. Interviewers care if you can trace why sequence length creates a quadratic bottleneck, distinguish activation memory from parameter memory, and reason about hardware limits for long documents or high resolution images.

A GOOD ANSWER COVERS: First, the math behind the bottleneck. For a sequence of length n, the query and key matrices each have shape n by d. Multiplying Q by K transpose produces an n by n attention score matrix. Computing softmax over rows and multiplying by the n by d value matrix keeps the operation count proportional to n squared. Second, the memory story. Standard implementations materialize that full n by n matrix in high bandwidth memory to apply softmax, dropout, and masking. For n of 100000 tokens, a float32 attention matrix alone consumes roughly 40 GB, exceeding most single GPU capacities. Third, the practical implication. Very long documents, high resolution images flattened into long patch sequences, or long audio waveforms quickly hit the memory wall. Training becomes impossible without model parallelism or gradient checkpointing, and inference latency grows quadratically. Fourth, the nuance that memory is not fundamentally quadratic. Techniques like memory efficient attention or the approach in Self-attention Does Not Need O(n^2) Memory show that by recomputing attention rather than storing the full matrix, activation memory can be reduced to O(n), though the number of FLOPs remains O(n^2).

COMMON WRONG ANSWERS: Confusing the O(n^2) activation cost with the model parameter count, which is independent of sequence length. Claiming feed forward layers dominate the cost, they scale linearly with n and are not the bottleneck. Asserting that standard attention can be linear time without approximation, which mixes up the exact softmax attention with linear approximations like Performer or Linformer. Ignoring memory bandwidth and only discussing FLOPs, which misses why GPUs actually run out of resources.

LIKELY FOLLOW UPS: How does FlashAttention exploit tiling to reduce high bandwidth memory usage without changing asymptotic compute? What is the complexity of sparse patterns like local sliding window attention? How would you partition a 100k token sequence across multiple GPUs? Can you calculate the exact activation memory for a given batch size, heads, and sequence length?

ONE CONCRETE EXAMPLE: Take a vision transformer on a 1024 by 1024 image with 16 by 16 patches. That yields n of 4096 patches. The attention matrix has 4096 squared or about 16.8 million entries. At float16 that is roughly 32 MB per head. With 16 heads and batch size 32, a single layer needs over 16 GB of activation memory just for the attention scores. A 24 layer model would require hundreds of gigabytes without checkpointing or memory efficient kernels, which is why high resolution ViTs quickly become impractical under standard attention.

Source: huggingface.co