Attention Complexity › The Quadratic Bottleneck

Self-attention proceeds through three distinct operations, each contributing to the overall computational cost. These three stages correspond exactly to the formula you have seen before in this book, but here we examine them through the lens of cost rather than function. Each stage has a well-defined computational profile, and understanding them separately makes the total complexity easy to build up from first principles.

The three stages are:

1. Computing attention scores: Each token asks "how relevant is every other token to me?" This requires comparing all $n$ tokens against all $n$ tokens, producing an $n \times n$ grid of relevance scores.

2. Applying softmax: The raw scores are converted into proper attention weights that sum to 1, allowing us to interpret them as a probability distribution. This normalizes along each row of the score matrix.

3. Computing outputs: Each token gathers information from all other tokens, weighted by the attention scores, to produce its contextual representation. This is a weighted sum over value vectors.

The first and third steps are where the quadratic cost emerges. Step two is also quadratic, but it is dominated by the other two when the model dimension is large. Let's examine each stage carefully, counting operations precisely.

The heart of self-attention is measuring pairwise relevance between tokens. Given a sequence of $n$ tokens, we need to compute how strongly each token should attend to every other token. This is done by comparing query vectors against key vectors using the dot product. You can think of the query as a question and the key as an index: the dot product measures how well a token's question matches another token's index entry.

The reason we use a matrix multiplication here, rather than computing each score individually in a loop, is efficiency. Modern hardware, especially GPUs and TPUs, is designed to execute large matrix multiplications at peak throughput. Expressing the operation as a single $QK^T$ multiplication allows the hardware to exploit its full parallelism. The result is still the same $n^2$ numbers, but computed much faster than $n^2$ separate dot products would be.

Mathematically, we arrange all queries into a matrix $Q$ and all keys into a matrix $K$, then compute:

$$S = QK^T$$

where:

• $S$: the $n \times n$ attention score matrix, where $S_{ij}$ measures how much token $i$ should attend to token $j$
• $Q$: the query matrix with shape $n \times d_k$, containing one query vector per token
• $K^T$: the transposed key matrix with shape $d_k \times n$, rearranging the key vectors for efficient matrix multiplication
• $n$: the sequence length (number of tokens)
• $d_k$: the dimension of each query/key vector (typically 64 in standard transformers)

Why does this produce $n^2$ values? Each row of $Q$ (one token's query) gets multiplied against every column of $K^T$ (every token's key), producing one score per pair. With $n$ rows and $n$ columns, we get an $n \times n$ grid of scores.

The computational cost of this matrix multiplication is substantial. Each entry in the output matrix requires $d_k$ multiplications and $d_k - 1$ additions. With $n^2$ entries, the total operation count is $n^2 \cdot d_k$.

The critical observation is that this score computation already scales as $O(n^2 d_k)$. The $n^2$ term is unavoidable: we cannot know which token pairs are important without examining all of them. Even in cases where the final attention weights are sparse (most tokens barely attending to most others), we still must compute all the raw scores before we can discover which ones are large. The sparsity only emerges after the computation, not before it.

Raw dot products can be any real number, positive or negative, with no upper bound. To use them as weights for a weighted average, we need to transform them into a valid probability distribution. The softmax function accomplishes this by exponentiating each score and normalizing within each row.

The key insight is that softmax operates independently on each row of the score matrix. Each row corresponds to one query token, and its $n$ scores express how much that token should attend to each of the $n$ positions. Softmax converts these $n$ arbitrary numbers into $n$ non-negative weights that sum to exactly 1. Think of it as converting raw relevance signals into a proper distribution over positions to attend to.

For each token $i$, softmax converts its $n$ raw scores into $n$ attention weights that sum to 1. This involves:

1. Computing $\exp(S_{ij})$ for each of the $n$ scores (making them all positive)
2. Summing these exponentials across the row
3. Dividing each exponential by the sum

Each row requires approximately $3n$ operations (exponentiate, sum, divide), and with $n$ rows, the total is $O(n^2)$. While still quadratic, this is dominated by the $O(n^2 d_k)$ cost of score computation when $d_k$ is large. In practice, the softmax step is rarely the performance bottleneck, but it does introduce a numerical subtlety: computing $\exp$ of large values causes overflow. The standard remedy is to subtract the row maximum before exponentiating, which leaves the softmax output unchanged (the constant cancels in the numerator and denominator) but keeps all values in a numerically safe range.

The final step uses the attention weights to compute a weighted combination of value vectors. Each token's output is a blend of all tokens' values, weighted by attention. You can think of this as a soft lookup: instead of retrieving a single value at a specific position, each token retrieves a smooth mixture of all values in the sequence, weighted by how relevant each position was determined to be in Stage 1.

The operation computes a new representation for each of the $n$ tokens. Each token's new representation aggregates information from every other token in the sequence, scaled by attention weights. Tokens that received high attention scores contribute more to the output; tokens with near-zero attention weight contribute almost nothing.

The formula for this aggregation is:

$$\text{Output} = \text{softmax}(S) \cdot V$$

where:

• $\text{softmax}(S)$: the $n \times n$ attention weight matrix, with each row summing to 1
• $V$: the value matrix with shape $n \times d_v$, containing one value vector per token
• $d_v$: the dimension of each value vector (typically equal to $d_k$)

This is another matrix multiplication: $(n \times n) \cdot (n \times d_v)$ produces an $n \times d_v$ output. Each of the $n \cdot d_v$ output elements requires $n$ multiply-accumulate operations (summing across the $n$ positions), giving $n^2 \cdot d_v$ total operations.

This step is just as expensive as score computation. Notice that both of the dominant operations involve an $n \times n$ intermediate matrix: in Stage 1 we create this matrix, and in Stage 3 we multiply by it. This $n \times n$ structure is the crux of attention's quadratic scaling, and it is also the primary target of memory-efficient implementations like FlashAttention, which avoid materializing the full $n \times n$ matrix in memory at once.

Now that we have analyzed each stage individually, we can combine the costs to get the total complexity of a single self-attention operation. This summary is important because it makes the $n^2 d$ scaling explicit and shows that neither the softmax step nor the $d$ dimension can offset the quadratic term.

Summing the costs from all three stages:

$$\text{Total FLOPs} = \underbrace{O(n^2 d_k)}_{\text{scores}} + \underbrace{O(n^2)}_{\text{softmax}} + \underbrace{O(n^2 d_v)}_{\text{output}} = O(n^2(d_k + d_v)) = O(n^2 d)$$

where:

• $n$: sequence length
• $d_k$: query/key dimension
• $d_v$: value dimension
• $d$: model dimension (typically $d_k = d_v = d/h$ for $h$ attention heads)

The simplification to $O(n^2 d)$ reflects that $d_k$ and $d_v$ are both proportional to the model dimension $d$. In standard transformers, each head operates on $d/h$ dimensions, but since we have $h$ heads processing in parallel and concatenating their outputs, the total complexity across all heads scales with the full model dimension.

It is worth pausing here to appreciate what this formula tells us. Two factors control the cost: the sequence length squared ($n^2$) and the model dimension ($d$). These two factors contribute very differently to the total. Doubling $d$ doubles the cost. Doubling $n$ quadruples it. For any sequence where $n > 1$, the sequence length term dominates. A model with $n = 2048$ and $d = 768$ spends roughly $2048^2 / 768 \approx 5,461$ times more computation on the $n^2$ factor than on the $d$ factor. This asymmetry is exactly what makes sequence length the central bottleneck in transformer scaling.

Let's trace through a concrete small example to make the formula visceral. Consider a single-head attention layer with sequence length $n = 4$ (four tokens), query/key dimension $d_k = 3$, and value dimension $d_v = 3$.

Stage 1: Score computation. The query matrix $Q$ has shape $4 \times 3$ and the key matrix $K$ has shape $4 \times 3$. Computing $S = QK^T$ produces a $4 \times 4$ matrix. Each of the 16 entries requires 3 multiplications and 2 additions (a dot product of length 3), totaling $16 \times 3 = 48$ multiplications.

Stage 2: Softmax. We apply softmax to each of the 4 rows of $S$. Each row has 4 elements, requiring 4 exponentiations, 1 sum (over 4 values, so 3 additions), and 4 divisions. Approximately $4 \times (4 + 3 + 4) = 44$ operations.

Stage 3: Output computation. The softmax weight matrix has shape $4 \times 4$ and the value matrix $V$ has shape $4 \times 3$. The product has shape $4 \times 3$. Each of the 12 output entries requires 4 multiplications and 3 additions, totaling $12 \times 4 = 48$ multiplications.

The total is dominated by Stages 1 and 3: approximately $48 + 48 = 96$ multiplications (plus roughly the same number of additions). Using our formula $2n^2 d_k = 2 \times 16 \times 3 = 96$, we recover exactly this number. The formula works.

Now compare with scaling to $n = 8$ (doubling the sequence length). Stage 1 becomes $64 \times 3 = 192$ multiplications (4x more). Stage 3 becomes $64 \times 3 = 192$ multiplications (4x more). The total scales by exactly 4x, as the quadratic formula predicts. Every doubling of $n$ multiplies the operation count by 4, regardless of what $d$ is. This is the essential arithmetic that makes long sequences expensive.

This analysis reveals the fundamental scaling behavior of self-attention. The two key findings are:

• Quadratic in sequence length: Doubling $n$ quadruples the computation. This is the defining characteristic that limits context length.
• Linear in model dimension: Doubling $d$ only doubles the computation. Making models wider is relatively cheap.

The quadratic term dominates for long sequences. At $n = 1000$ tokens with $d = 768$ (GPT-2 dimensions), the $n^2 = 1,000,000$ factor far exceeds the $d = 768$ factor. Understanding this asymmetry helps when reasoning about transformer scalability.

The key insight is that this asymmetry shapes all of modern LLM engineering. When practitioners want more capable models, they often increase model depth and dimension, which costs linearly. When they want models that understand longer documents, they face a quadratic penalty. This is why context length has historically lagged so far behind model parameter counts: a 10x increase in context length requires 100x more attention computation, while a 10x increase in parameters (achieved through wider or deeper models) costs only 10x more. The entire field of efficient attention is essentially an attempt to change this ratio.

The theoretical complexity $O(n^2 d)$ tells us about relative growth rates, but it hides constant factors and additive terms that matter enormously in practice. To see the actual numbers, let's build a function that calculates exact FLOPs for self-attention. This exercise converts the abstract complexity formula into concrete operation counts that can be compared against hardware throughput specifications.

The function mirrors our three-stage analysis: score computation contributes $n^2 \cdot d$ operations, softmax contributes $3 \cdot n^2$ operations, and output computation contributes another $n^2 \cdot d$ operations. Notice that the softmax term is negligible relative to the others when $d \gg 3$, which is always true in practice ($d$ is typically in the hundreds). This confirms that the two matrix multiplications, not the softmax, are the computational bottleneck.