Summary (Overview)

• Core innovation: The paper introduces a certified sparse attention method that perturbs each key's score with an independent Exp(1) variable, transforming threshold retrieval into probability-proportional-to-size (PPS) sampling. This enables unbiased Horvitz–Thompson estimation of the softmax output with finite-sample error certificates that hold simultaneously across all threshold levels.
• Key theoretical guarantee: At a suitable threshold, the estimator reads $O(\varepsilon^{-2}\log(1/\delta))$ values regardless of context length, which is proven optimal up to constants via a matching lower bound. The anytime-valid certificates ensure each attention head's error exceeds $\varepsilon$ with probability at most $\delta$ per decoding step.
• Empirical validation: Across three LLMs and ten settings at 16K–128K tokens, no certificate was violated. At equal keys read, the method achieves 4.6–7.9× lower 90th-percentile error than the best of 15 baseline families, while eliminating the heavy error tail (max errors ≤ 0.029 vs. baselines' 0.39–1.62).
• Efficiency: A 4-bit scan with deterministic bounds generates candidates exactly, reading only 16.4% of dense attention's bytes for metadata. On CPU, decoding steps at 32K–128K tokens run 1.1–1.4× faster than dense attention.
• End-to-end quality: The method stays within one point of dense attention on tasks and within 0.1% in PG-19 perplexity, with 79–98% of outputs identical to dense attention's.

Introduction and Theoretical Foundation

Background and Motivation

Sparse attention aims to accelerate long-context decoding by reading only a small fraction of the KV cache. For Llama-3.1-8B with a 128K-token context, each decoding step requires reading ~16 GiB of 16-bit keys and values, making memory bandwidth the bottleneck. Existing heuristic methods—such as attention sinks, heavy hitters, page-level bounds, and approximate nearest neighbors—do not control per-query error, which exhibits heavy tails on real KV caches (90th-percentile errors of 0.051–0.090 but maxima of 0.4–1.6 for the best of 15 baselines).

Theoretical Foundation

The key insight is that perturbing each key's score with an independent exponential variable enables principled sampling. For each key $i$ with perturbation $E_i \sim \mathrm{Exp}(1)$ and unnormalized attention weight $w_i = e^{s_i - c^*}$, the perturbed score threshold condition $s_i - c^* + E_i \geq \log \tau$ holds with probability $\min(1, w_i/\tau)$. This creates threshold (Poisson PPS) sampling, under which Horvitz–Thompson estimates of the softmax numerator and denominator are unbiased. Lowering $\tau$ monotonically adds keys, enabling a geometric grid search where each head stops at the first level whose certificate meets the target error.

Methodology

Problem Setup

For a fixed layer, KV head, and decoding step with $n$ cached keys $k_i$ and values $v_i$:

• Scores: $s_i = q^\top k_i/\sqrt{d}$, maximum $c^* = \max_i s_i$
• Unnormalized weights: $w_i = e^{s_i - c^*} \in (0,1]$
• Denominator: $D = \sum_i w_i$, Numerator: $N = \sum_i w_i v_i$
• Attention output: $o = N/D$, attention weights: $a_i = w_i/D$

Error is measured on three scales: denominator ($|\hat{D} - D| \leq \varepsilon D$), output norm ($\|\hat{o} - o\| \leq \varepsilon\|o\|$), and value scale ($\|\hat{o} - o\| \leq \varepsilon M$, where $M^2 = \sum_i a_i\|v_i\|^2$). The value scale is $\kappa = M/\|o\|$ times the output norm.

Algorithm Components

1. Perturbed threshold sampling. The KV cache is split into a window (recent 64 tokens, always read exactly) and a pool. Each pool key has perturbation $E_i = -\ln u_i$ from a 53-bit uniform hash. The sampled set at threshold $\tau$ is:

$$S(\tau) = \{i \in P: s_i - c^* + E_i \geq \log \tau\}$$

2. Horvitz–Thompson estimation. Light keys have inclusion probability $w_i/\tau$, giving HT weights $\tau$:

$$\hat{D} = D_R + D_B + \tau k, \qquad \hat{N} = N_R + N_B + \tau \sum_{i \in S \cap A} v_i$$

3. Anytime-valid certificates. On grid $\tau_j = 2^{-j/r}$ with failure budgets $\delta_j = \delta/((J+1)\dot{e})$, the algorithm computes:

• Denominator interval: Chernoff bounds on the Poisson-binomial sample count
• Centered numerator: Bernstein inequality in Hilbert space on the centered sum $Z_c = \sum_{i\in A}(\tau X_i - w_i)(v_i - c)$
• Output error bound: $\beta = (\zeta + \|\hat{o} - c\|dD)/(\hat{D} - dD)$
• Stopping rule: Stop at first certified level meeting target error

4. Candidate generation. A 4-bit per-block-per-channel quantization (as in KIVI) with AVX-512 VNNI provides deterministic upper bounds $u_i \geq s_i$ for all scores. Candidates are bucketed by earliest possible inclusion level via counting sort, with 8-bit refinement and exact 16-bit reads only for keys that may enter the sample. Exact maximum computation covers only 0.1–1% of the pool.

Empirical Validation / Results

Error at Matched Keys Read

Table 1. Error/​M at matched keys read (ε = 0.1; baselines read ≤ 1.1× our keys).

Certificate Validity

Across 120 configurations (10 settings × 3 measures × 4 values of ε), no certificate was violated, no layer had a violation, and no head fell back to exact attention. Certificates are conservative: observed errors are about ε/10 at the 90th percentile, with largest errors 0.31ε (value scale) and 0.41ε (output norm).

End-to-End Performance

Table 2. End-to-end score change versus dense attention.

PG-19 perplexity changes by only −0.02% to +0.07% with KL divergence 0.0002–0.0006 from dense attention, whereas oracle top-k reading the same keys raises perplexity by 4.8–11.5%.

CPU Speed

Attention speedups over dense attention: 1.8–3.5× at 64K–128K and 1.0–2.1× at 16K–32K on Llama. Whole fp32 decoding steps are 1.1× faster at 32K and 1.4× at 128K. Fresh perturbations add a median 4% overhead.

Theoretical and Practical Implications

Theoretical Contributions

Theorem 1 (Unbiasedness and variance). $\mathbb{E}\hat{D} = D$, $\mathbb{E}\hat{N} = N$, with $\mathrm{Var}\,\hat{D} \leq \tau D_A$ and $\mathbb{E}\|Z_c\|^2 \leq \tau G_A(c)$.

Theorem 2 (Sample size independent of context). With $\tau = D/m$ and probability at least $1-\delta$:

$$\|\hat{o} - o\| \leq \frac{\sigma_v\sqrt{2L/m} + 2R_o L/(3m)}{1 - \eta}$$

where $L = \ln(4/\delta)$, $\eta = \sqrt{2L/m} + L/(3m)$, $\sigma_v^2 = M^2 - \|o\|^2$. Thus $m = 2L\varepsilon^{-2}(1+o(1))$ values suffice for error $\varepsilon M$, independent of $n$.

Theorem 3 (Lower bound). Any algorithm reading at most $m$ values requires $m \geq (\kappa^2 - 1)\ln(1/(6\delta))/(8.8\varepsilon^2)$ for output-norm error and $m \geq \ln(1/(6\delta))/(17.6\varepsilon^2)$ for value-scale error, proving the PPS estimator is optimal up to constants.

Theorem 4 (Anytime-valid certificates). For each query head and error measure:

$$\Pr[\hat{o} \text{ is certified and its error exceeds } \varepsilon] \leq \delta$$

The proof intersects confidence events across all levels, making the data-dependent stopping level harmless.

Practical Implications

• Reliability: The method eliminates the heavy error tail of heuristic sparse attention, guaranteeing per-call error bounds that scale to full generations via union bounds ($\delta/(HT)$ per call).
• Efficiency floor: The 4-bit scan sets a floor of 16.4% of dense attention's bytes, which caps achievable speedups and allows baselines at equal bytes to match typical (but not worst-case) errors.
• Conservativeness: Certificates read 1.8–3.5× more keys than uncertified central-limit rules; the gap represents the cost of statistical guarantees.

Conclusion

This work demonstrates that exponential perturbations transform threshold retrieval into weight-proportional sampling, enabling a sparse attention kernel that estimates unbiasedly and certifies each call's error. The certificates held in every configuration tested across three LLMs and ten settings, removing the heavy error tail of heuristic approaches. While end-to-end task quality matches the strongest baselines at equal bytes, the method uniquely provides per-call guarantees and closer agreement with dense attention during sampling. The main limitation is the linear candidate-generation scan, which sets a bytes floor; a sublinear index with recall guarantees would remove this bottleneck. The method does not cover closed-loop generation with fixed perturbations, and deployments should not present it as a guarantee of unchanged model behavior—only as per-head, per-step error certificates with probability at least $1-\delta$.