Quadratic Attention Bottleneck › Compute Analysis: The Operations

We have seen that attention requires storing an $n \times n$ matrix of weights. But how much computation does it take to produce those weights and use them? Understanding the computational cost helps us predict training times, estimate inference latency, and appreciate why efficient attention variants are so valuable.

The answer turns out to be $O(n^2 d)$ floating-point operations per attention layer, where $d$ is the model dimension. To see where this comes from, let's trace through exactly what happens when attention runs, quantifying each operation precisely.

Before diving into the math, consider what attention needs to accomplish for each token in the sequence:

1. Compare this token to every other token to determine relevance scores
2. Normalize those scores into a probability distribution
3. Blend information from all tokens according to those probabilities

For a sequence of $n$ tokens, step 1 alone requires $n$ comparisons per token, and we have $n$ tokens, giving us $n \times n = n^2$ comparisons. This is the fundamental source of quadratic scaling: the all-pairs comparison cannot be avoided in standard attention. Steps 2 and 3 must process the results of step 1, so they also touch $n^2$ entries.

Now let's quantify each step precisely and see how the costs add up.

The first operation computes how strongly each query should attend to each key. Mathematically, this is the matrix multiplication $QK^T$, where $Q$ contains all query vectors (one per token) and $K$ contains all key vectors.

When we multiply a matrix of shape $(n \times d_k)$ by a transposed matrix of shape $(d_k \times n)$, we get an $(n \times n)$ output. Each entry in this output matrix is a dot product between one query and one key, requiring $d_k$ multiplications and $d_k - 1$ additions. With $n^2$ such entries, the total operation count for this stage is:

$$\text{Score FLOPs} \approx 2 \cdot n^2 \cdot d_k$$

where:

• $n^2$: the number of query-key pairs (one score per pair)
• $d_k$: the dimension of each query and key vector (typically $d / h$ for $h$ heads)
• $2$: accounts for both multiplications and additions in each dot product (counted as a single fused multiply-add operation on modern hardware)

The $n^2$ appears because every token must be compared against every other token. There is no shortcut here: we do not know which pairs will have high attention until we compute all the scores. This is the irreducible cost of the all-pairs comparison.

Raw dot products can be any real number, positive or negative, and their magnitudes depend on the model's learned parameters. To use them as weights in a weighted average, we need to convert each row of scores into a probability distribution that sums to 1. The softmax function accomplishes this, converting each row of scores into non-negative weights that sum to exactly 1.

For a row of $n$ attention scores $s_1, s_2, \ldots, s_n$, softmax computes:

$$\text{softmax}(s_i) = \frac{e^{s_i / \sqrt{d_k}}}{\sum_{j=1}^{n} e^{s_j / \sqrt{d_k}}}$$

where:

• $s_i$: the raw dot product score for the $i$-th position in the row
• $e^{s_i / \sqrt{d_k}}$: the exponential of the scaled score. This keeps all weights are positive
• $\sum_{j=1}^{n} e^{s_j / \sqrt{d_k}}$: the normalizing constant, summing over all positions in the row
• $\sqrt{d_k}$: the scaling factor that prevents dot products from growing too large in magnitude, which would push softmax into regions of near-zero gradient

The $\sqrt{d_k}$ scaling is a detail from the original Transformer paper, motivated by the observation that without it, dot products grow in magnitude with $d_k$, pushing gradients to near zero during training. Scaling by $\sqrt{d_k}$ keeps the scores in a range where softmax gradients remain healthy.

For each of the $n$ rows in the attention matrix, softmax requires exponentiating $n$ scores, summing those $n$ exponentials, and dividing each of the $n$ exponentials by the sum. This contributes roughly $5n$ operations per row. With $n$ rows:

$$\text{Softmax FLOPs} \approx 5 \cdot n^2$$

While still quadratic in $n$, the softmax cost lacks the $d_k$ factor, making it smaller than score computation when $d_k$ is large. For typical transformers where $d_k = 64$ to $128$, softmax contributes less than 10% of the total attention FLOPs. It is not negligible, but it is not the dominant term.

The final step uses the normalized attention weights to compute a weighted combination of value vectors. Each token's output is a blend of all tokens' values, weighted by attention. In matrix form, this is:

$$\text{Output} = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) \cdot V$$

where the left factor is the $(n \times n)$ attention weight matrix and the right factor $V$ is the $(n \times d_v)$ value matrix.

This is another matrix multiplication: the $(n \times n)$ attention weights multiplied by the $(n \times d_v)$ value matrix produces an $(n \times d_v)$ output. Each of the $n \times d_v$ output elements requires summing across $n$ weighted values:

$$\text{Output FLOPs} \approx 2 \cdot n^2 \cdot d_v$$

where:

• $n^2$: the number of weight-value multiplications (every output position must aggregate from all $n$ positions)
• $d_v$: the dimension of each value vector (typically equal to $d_k$ in standard transformers)
• $2$: multiply-add pairs as before

This matches the cost of score computation exactly when $d_v = d_k$. We have encountered two expensive $O(n^2 d)$ operations, and they dominate the total cost of attention.

The visualization confirms our analysis: score computation and output computation each consume roughly half the total FLOPs (about 49.8% each), while softmax contributes less than 0.4%. The two matrix multiplications involving the $n \times n$ attention matrix dominate entirely. Any optimization that targets the softmax step alone will have negligible impact on the overall compute budget.

Combining all three stages, the total computational cost of self-attention for a single head is:

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

where:

• $n$: sequence length (number of tokens)
• $d_k$: dimension of query and key vectors (typically $d / h$, where $h$ is the number of heads)
• $d_v$: dimension of value vectors (typically equal to $d_k$)
• $d$: model dimension (the full embedding size before splitting into heads)
• $h$: number of attention heads

The simplification to $O(n^2 d)$ holds because $d_k$ and $d_v$ are both proportional to $d$. In standard transformers, each head operates on $d/h$ dimensions, and since we sum across all heads, the total work scales with the full model dimension $d$.

The formula $O(n^2 d)$ reveals an important asymmetry in how the two factors affect complexity:

• Doubling $n$ quadruples the computation because $n$ appears squared
• Doubling $d$ only doubles the computation because $d$ appears linearly

This asymmetry determines which factor becomes the bottleneck. For a model with $d = 768$ processing $n = 4,096$ tokens, the $n^2 = 16{,}777{,}216$ factor dwarfs the $d = 768$ factor by a ratio of over 20,000 to 1. The quadratic term completely dominates, which is why sequence length is the bottleneck rather than model size when processing long documents.

The practical consequence is that scaling up the model (increasing $d$) is far less damaging to inference speed than scaling up the context (increasing $n$). A model twice as wide is only twice as slow. A model with twice the context length is four times as slow. This is why researchers are far more willing to increase model dimension than context length when staying within a compute budget.

Let's make this concrete by computing the actual FLOPs for a GPT-2 style model with $d = 768$ across a range of sequence lengths: