Quadratic Attention Scaling O(N^2) & In-Context Token Reduction Economics
Deriving Computation Savings in Transformer Self-Attention Layers
Mathematical intuition
Self-attention matrix multiplication QK^T requires O(N^2 d) operations for sequence length N, yielding quadratic FLOP reductions when prompt sequence length is compressed.
- 01Self-attention compute complexity scales quadratically O(N^2) with prompt length N.
- 02Compressing prompt length by retention factor gamma reduces query-key matrix multiplication FLOPs to gamma^2 N^2.
- 03API provider billing scales linearly with billed tokens while serving latency decreases during the prefill phase.
1. Mathematical Derivation of Attention FLOPs
Standard Scaled Dot-Product Attention computes:
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
For sequence length $N$, batch size $B=1$, number of heads $H$, and head dimension $d_k$:
1. $Q K^T$: Matrix multiplication between $(N \times d_k)$ and $(d_k \times N)$ yields $(N \times N)$. This operation requires $2 N^2 d_k$ FLOPs per head.
2. Softmax multiplication with $V$: Matrix multiplication between $(N \times N)$ and $(N \times d_k)$ requires $2 N^2 d_k$ FLOPs per head.
Total multi-head self-attention prefill FLOPs:
$$\text{FLOPs}_{\text{Attn}} = 4 H N^2 d_k = 4 N^2 d_{\text{model}}$$
When sequence length $N$ is compressed to $N' = \gamma N$ (where $\gamma \in (0, 1)$):
$$\text{FLOPs}_{\text{Attn}}' = 4 (\gamma N)^2 d_{\text{model}} = \gamma^2 \cdot \text{FLOPs}_{\text{Attn}}$$
Ratio_FLOPs = (FLOPs') / (FLOPs) = \gamma^2
2. Relative FLOP Reduction Factor Table
Theoretical compute scaling factor relative to baseline sequence length $N$:
| Retention Factor (gamma) | Token Reduction (1 - gamma) | Attention Compute Factor (gamma^2) | FLOP Savings |
|---|---|---|---|
| 1.0 (Baseline) | 0% | 1.00 | 0% |
| 0.8 | 20% | 0.64 | 36% |
| 0.6 | 40% | 0.36 | 64% |
| 0.5 | 50% | 0.25 | 75% |
| 0.3 | 70% | 0.09 | 91% |
- [Vaswani et al. (2017)]Attention Is All You Need (Advances in Neural Information Processing Systems)
- [Dao et al. (2022)]FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness