Sketching Attention Sinks in LLMs

Direct implementations of transformer attention blocks have computation requirements proportional to $N^2 d$ and memory requirements proportional to $N^2$, where $N$ is the context length and $d$ is the attention head dimension. However, this quadratic-in-$N$ scaling does not exploit the structure inherent in many trained LLMs where attention sinks cause a few tokens to receive most of the attention mass. In particular, the attention matrix $\mathbf{A}$ is extremely concentrated, with a low stable rank, while the value matrix $\mathbf{V}$ has near-uniform row norms and a high stable rank. We exploit this asymmetry with FullSketch, a drop-in attention variant for pre-trained LLMs that identifies the dominant sink columns of $\mathbf{A}$ before their full calculation, thereby approximating $\mathbf{A}\mathbf{V}$ in sub-quadratic $\mathcal{O} \left(N^{1+\gamma}d\right)$ compute and $\mathcal{O} \left(N^{1+\gamma} + Nd\right)$ memory, where $\gamma\in (0,1)$ is a tunable parameter to balance computational cost and efficacy. We benchmark our method on the LLM Qwen2-0.5B: holding token perplexity within $10$% of full attention reduces attention FLOPs by a factor of $1.43$, and at a fixed perplexity target the measured saving is a constant factor across context lengths rather than an asymptotic one.