DL0042 Attention Computation

Written by

Please break down the computational cost of attention.

Answer

Attention cost splits into a linear-in-n term from projections and a quadratic-in-n term from pairwise interactions. For sequence length $n$ and model dimension $d$ , the total is $O(n d^2 + n^2 d)$ : projections dominate for short sequences, while the $n \times n$ score matrix dominates once $n$ grows past $d$ .

(1) Q/K/V + Output Projections: Four $d \times d$ GEMMs over $n$ tokens: $O(n d^2)$ , linear in sequence length.
(2) Score Matrix $QK^\top$ and Value Mixing $AV$ : Pairwise $n \times n$ work: $O(n^2 d)$ , the quadratic bottleneck.
(3) Softmax: Elementwise over $n^2$ entries: $O(n^2)$ , cheap FLOPs but the $n \times n$ matrix drives memory traffic.

Mathematical Formulation (one attention layer, $h$ heads, $d_k = d_v = d/h$ ):
$\mathrm{Cost}_{\text{proj}} = O(n\, d^2)$
$\mathrm{Cost}_{\text{scores}} = O(n^2\, d_k \cdot h) = O(n^2 d)$
$\mathrm{Cost}_{\text{total}} = O\!\left(n^2 d + n\, d^2\right)$

Where:

$n$ is the sequence length, $d$ the model (hidden) dimension, and $h$ the number of heads.
$Q, K \in \mathbb{R}^{n \times d_k}$ form scores $S = QK^\top \in \mathbb{R}^{n \times n}$ ; attention weights $A = \mathrm{softmax}(S / \sqrt{d_k})$ mix values $V \in \mathbb{R}^{n \times d_v}$ via $O = AV$ .
Multi-head attention costs the same as single-head in big-O: per-head dimensions scale as $d/h$ , so the $h$ heads sum back to $d$ .

Log-log plot of attention and projection FLOPs versus sequence length at hidden dimension 512, with the quadratic attention term overtaking the linear projection term near n equals 512.

Figure 1: Log-log cost curves at $d = 512$ : projection $O(n d^2)$ leads at short lengths, attention $O(n^2 d)$ takes over near $n \approx d$ and dominates at long context.

Regimes: Short sequences ( $n \ll d$ ): the $n d^2$ projection term dominates. Long sequences ( $n \gg d$ ): the $n^2 d$ interaction term dominates; the two balance around $n \approx d$ .