Computational complexity of attention

The fundamental cost breakdown

Memory and operations in standard self-attention

The self-attention mechanism computes three matrices for each token: queries (Q), keys (K), and values (V). The complexity comes from computing the attention scores between all pairs of tokens.

Let's derive the complexity from scratch.

Deriving the time complexity

Step 1: Linear projections

For input $X \in \mathbb{R}^{n \times d}$ and weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$:

$Q = XW_Q$

This is a matrix multiplication: $(n \times d) \times (d \times d_k) = n \times d_k$ output.

Cost: Each output element requires $d$ multiplications, total: $n \cdot d_k \cdot d$ operations.

For all three projections (Q, K, V): $3nd \cdot d_k$ operations.

Why this step? We're transforming each token's embedding into query/key/value space. This is linear in sequence length—not the problem.

Step 2: Computing attention scores $S = QK^T$

$S = QK^T \in \mathbb{R}^{n \times n}$

Matrix dimensions: $(n \times d_k) \times (d_k \times n) = n \times n$

Cost: Each of the $n^2$ output elements requires $d_k$ multiplications.

Total: $n^2 \cdot d_k$ operations.

Why this step? This is where every token compares itself to every other token. The $n \times n$ output is inherently quadratic—this is the bottleneck.

Step 3: Softmax

For each row of $S$ (one query), compute:

$A_{ij} = \frac{\exp(S_{ij}/\sqrt{d_k})}{\sum_{j=1}^n \exp(S_{ij}/\sqrt{d_k})}$

Cost: $O(n)$ per row, $n$ rows total: $O(n^2)$ operations.

Why this step? Normalizing attention weights. Still quadratic in the number of scores.

Step 4: Weighted sum $\text{Output} = AV$

$(n \times n) \times (n \times d_k) = n \times d_k$

Cost: $n^2 \cdot d_k$ operations.

Why this step? Each output token is a weighted combination of all value vectors.

Step 5: Output projection $\text{Out} = ZW_O$

After concatenating heads back into a $d$-dimensional vector, we apply the output projection $W_O \in \mathbb{R}^{d \times d}$:

$(n \times d) \times (d \times d) = n \times d$

Cost: $n \cdot d \cdot d = nd^2$ operations.

Why this step? The output projection mixes information across heads and maps back to model dimension. It contributes another $\mathcal{O}(nd^2)$ term—don't forget it when counting linear-projection cost.

Memory complexity

Why each component matters

Linear projections: $\mathcal{O}(nd^2)$

This scales linearly with sequence length. Doubling $n$ doubles the work. Not a bottleneck for long sequences but dominates for very short sequences where $n \ll d$. This term covers the Q, K, V projections and the output projection $W_O$—together roughly $4nd^2$ operations, but asymptotically $\mathcal{O}(nd^2)$.

Attention matrix: $\mathcal{O}(n^2 d_k)$

This is the core problem. The $n^2$ term means:

• 2× sequence length →4× compute
• 10× sequence length → 100× compute

Why it's unavoidable in standard attention: Every token needs a score with every other token to decide "who to pay attention to."

Multi-head attention scaling

With $h$ attention heads, each of dimension $d_k = d/h$:

• Attention cost: $h$ heads each cost $n^2 d_k$, so total $h \cdot n^2 d_k = n^2 d$ (since $h d_k = d$).
• Projection cost: The Q, K, V projections and output projection each operate on the full $d$-dimensional space once, not once per head. So the projection cost stays $\mathcal{O}(nd^2)$ — it is not multiplied by $h$.

$\boxed{\text{Total} = \mathcal{O}(h\,n^2 d_k + nd^2) = \mathcal{O}(n^2 d + nd^2)}$

Key insight: Using more heads doesn't change the asymptotic complexity (still $\mathcal{O}(n^2 d + nd^2)$), and does not add a factor of $h$ to the projection term. Splitting into heads just reshapes the same $d$-dimensional projections.

Comparison with other architectures

| Architecture | Time complexity | Memory | Receptive field | |------------|-----------------|-----------------| | Self-attention | $\mathcal{O}(n^2 d)$ | $\mathcal{O}(n^2)$ | Global (all tokens) | | RNN (LSTM/GRU) | $\mathcal{O}(nd^2)$ | $\mathcal{O}(d)$ | Sequential (limited) | | CNN (kernel $k$) | $\mathcal{O}(nkd^2)$ | $\mathcal{O}(nd)$ | Local ($k$ tokens) | | Sparse attention | $\mathcal{O}(n \sqrt{n} d)$ | $\mathcal{O}(n \sqrt{n})$ | Structured patterns | | Linear attention | $\mathcal{O}(nd^2)$ | $\mathcal{O}(nd)$ | Global (approximation) |

Trade-off: Self-attention gets global context instantly (depth-1 path between any tokens) but pays quadratic cost. RNNs are linear but require $n$ sequential steps to connect distant tokens.

Efficient attention variants

The $\mathcal{O}(n^2)$ bottleneck has sparked a research area:

1. Sparse attention (Sparse Transformer, Longformer): Only compute attention for subset of token pairs (e.g., local windows + global tokens). Reduces to $\mathcal{O}(n \log n)$ or $\mathcal{O}(n \sqrt{n})$.

2. Linear attention (Performers, Linear Transformers): Approximate attention with kernel methods, achieving $\mathcal{O}(nd^2)$ by never materializing the $n \times n$ matrix.

3. Flash Attention: Keeps $\mathcal{O}(n^2 d)$ complexity but optimizes memory access patterns using tiling, reducing memory from $\mathcal{O}(n^2)$ to $\mathcal{O}(n)$ by not storing the full attention matrix.

Practical implications

Training vs. inference

Training: Requires storing $S$ for backpropagation. Memory is often the limiting factor, not compute.

Inference: Can attention on-the-fly without storing full matrices. KV-caching helps: store past keys/values to avoid recomputing them for each new token.

Why context length is expensive

Quadratic scaling means:

• GPT-3 (2K context): manageable
• GPT-4 (32K context): 16× memory and compute vs. 8K in the attention term
• 100K context: 625× cost vs. 4K

This is why long-context models either use sparse attention or are extremely expensive to run.

Connections

• 4.1.1-Self-attention-mechanism: The mechanism that causes this complexity
• 4.1.7-Multi-head-attention: How splitting heads affects computational cost
• 4.3.2-Sparse-attention-patterns: Solutions to reduce quadratic complexity
• 4.5.1-Memory-optimization-techniques: Flash Attention and gradient checkpointing
• 5.2.3-Positional-encoding-scaling: Why long sequences need efficient attention

#flashcards/ai-ml

What is the time complexity of standard self-attention? :: $\mathcal{O}(n^2 d_k + nd^2)$, dominated by $\mathcal{O}(n^2 d_k)$ for long sequences where the $QK^T$ computation requires $n^2 d_k$ operations.

Compare time complexity of self-attention vs. RNN :: Self-attention is $\mathcal{O}(n^2 d)$ but parallel; RNN is $\mathcal{O}(nd^2)$ but sequential. For long sequences ($n > d$), attention is slower but has global receptive field in one step.

What is the computational bottleneck in self-attention? :: The $QK^T$ matrix multiplication, which is $\mathcal{O}(n^2 d_k)$. This all-to-all comparison between tokens is inherently quadratic and dominates for long sequences.

Concept Map

Hinglish (regional understanding)