Top-Theta Attention: Sparsifying Transformers by Compensated Thresholding

It finds and validates that static thresholds can be calibrated per attention head to retain approximately k significant elements per attention row, replacing top-k search with elementwise comparison. Prior content-based sparsification mainly relied on top-k selection, which requires in-row ranking and imposes full-row dependency; this work shows calibrated static thresholds approximate top-k and that thresholds are a model-intrinsic property resilient to data-domain shifts. Evaluated on LLaMA2-7b, LLaMA3-8B, and LLaMA3-70B across ARC-C, ARC-E, Hellaswag, Human-eval, and LongBench; thresholds need only a few hundred calibration samples, and cross-dataset calibration (calibrated on ARC-C, applied to Human-eval/LongBench) retains comparable accuracy.

It proposes two numerical compensation methods, SDC and VMC, that almost fully recover accuracy under aggressive sparsification. SDC estimates the exponent sum of discarded elements and multiplies by a factor R/(R+E) to approximate post-softmax sparsification; VMC compensates missing V-row contributions by adding the mean V row scaled by the discarded attention mass β, with a formal expectation-based justification. On Hellaswag, uncompensated Top-j degrades noticeably, while exp-threshold and exact SDC variants substantially recover accuracy, and combining SDC with VMC nearly closes the gap to baseline; on ARC-C/ARC-E, Top-j alone already outperforms baseline, so compensation offers little benefit.

Top-j reduces attention elements by up to 10× in the prefill phase and V-row reads by 3–10× in generative decoding, with no more than 1% accuracy degradation. Unlike Top-k, Top-j thresholding is an elementwise operation with no row dependency, making it compatible with tiling and distributed inference; it also requires no retraining, only one-time calibration. On Q&A tasks both Top-k and Top-j improve accuracy by 0.2%–1% over baseline while pruning 2–5× of active elements; on Human-eval, post-softmax Top-j keeps pass@1 within 1% of baseline while reducing V rows by 3× for 8B/70B models, and the 70B model trades 1% pass@1 for a 5× reduction; on LongBench, V-row reduction reaches 10× with Rouge-L within 1% of baseline.

Thresholds are resilient to distribution shift, and per-layer, per-head, and per-row calibration are all beneficial. Thresholds are associated with the model rather than the data, so calibration is needed only once per model; experiments show keeping the first two layers at higher k (e.g., 512) is crucial for downstream accuracy, and post-softmax sparsification consistently outperforms pre-softmax. On Human-eval, post-softmax Top-j calibrated on ARC-C is even more accurate than calibration on the same dataset; on MedMCQA, Top-j and Top-k perform equally well; ablations cover layer, head, and row granularity.

The method targets the inference phase of decoder-only Transformers and applies to the GQA and MHSA attention mechanisms covered by the LLaMA family; thresholds require one-time offline calibration per model, needing only a few hundred samples, and can be stored as model-attached parameters (e.g., all thresholds for LLaMA-3-70B total about 11.8 MB). It enables LLM deployment on low-memory-bandwidth systems and is compatible with tiling and distributed inference; the authors explicitly state it is a lossy method requiring guardrails in security-sensitive scenarios.

The authors do not extensively assess wall-clock speedups, noting that meaningful speedups require hardware-aware kernel optimizations; the prototype NPU kernel achieves a 1.17× speedup over IncrementalFlashAttentionV4 in a specific configuration, and sparse kernel efficiency is highly hardware-dependent. Experiments are limited to the LLaMA family, and how model size, sequence length, and attention patterns affect the sparsity-accuracy trade-off remains an open question. Additionally, thresholds calibrated for a target k only guarantee k elements on average, with actual counts varying by input; the authors mitigate this by enlarging the calibration set rather than enforcing a hard cap. Multi-k cumulative calibration (MKC) is memory- and time-intensive, and speeding it up remains a direction for future work.