Quadratic Attention Scaling O(N^2) & In-Context Token Reduction Economics

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}}$$

2. Relative FLOP Reduction Factor Table

Theoretical compute scaling factor relative to baseline sequence length $N$: