Why Attention Is O(n²)

Why Is Attention Quadratic $O(N^2)$?

In a Transformer, self attention computes pairwise similarity scores between all tokens in a sequence of length $N$.

Given Query matrix $Q \in \mathbb{R}^{N \times d}$ and Key matrix $K \in \mathbb{R}^{N \times d}$:

$$\text{Attention Score Matrix } S = Q K^T \in \mathbb{R}^{N \times N}$$

Notice the shape of matrix $S$: $N \times N$.

Every token compares itself to every other token in the sequence.

  Sequence of N = 4 Tokens:                         Pairwise Comparisons (N x N = 16):
  Token 1: ["The"]                                  (1,1) (1,2) (1,3) (1,4)
  Token 2: ["cat"]                                  (2,1) (2,2) (2,3) (2,4)
  Token 3: ["sat"]                                  (3,1) (3,2) (3,3) (3,4)
  Token 4: ["down"]                                 (4,1) (4,2) (4,3) (4,4)

For sequence length $N = 4$, we compute $4^2 = 16$ pairwise dot products.

For sequence length $N = 100,000$, we compute $100,000^2 = 10,000,000,000$ (10 Billion) pairwise dot products!

The Memory Scaling Table

A single attention matrix for a 128k context window requires 65 Gigabytes of VRAM per attention head if stored explicitly!

Solutions to the Quadratic Bottleneck

1. FlashAttention (Tiled Online Softmax): Keeps $O(N^2)$ math theoretically, but avoids storing the $N \times N$ matrix in GPU HBM memory by computing attention in small GPU SRAM tiles.
2. Windowed / Local Attention (Swin / Longformer): Limits attention to local $K$-word neighborhoods ($O(N \cdot K)$ complexity).
3. State Space Models (Mamba / Linear Attention): Replaces quadratic attention with linear $O(N)$ recurrent state dynamics.

Say this out loud

Standard Self Attention has quadratic time and memory complexity O(N^2) because every token Query calculates pairwise dot products with every Key token, producing an N x N matrix. quadrupling sequence length increases memory by 16 times. Systems use FlashAttention GPU tiling or Linear State Space Models like Mamba to handle long context sequences efficiently.

Followups to expect

1. What is the difference between IO bound memory vs FLOP compute bound in attention? Memory bandwidth is the primary bottleneck. Standard PyTorch attention spends most of its time reading and writing the massive $N \times N$ matrix to GPU HBM memory rather than performing floating point math.
2. Why can't we easily replace self attention with linear attention everywhere? Pure linear attention approximations often struggle to recall precise facts buried deep in middle context positions compared to exact $O(N^2)$ dot product attention.