Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Mathematical Theory of Top-$k$ Sparse Attention via Total Variation Distance

Georgios Tzachristas, Lei Deng, Ioannis Tzachristas, Gong Zhang, Renhai Chen

TL;DR

This work provides a rigorous, distribution- and output-level theory for Top-$k$ attention truncation, tying the truncation error to the discarded softmax mass via an exact TV–KL identity. It derives deterministic and probabilistic bounds based on score gaps, blocks, and a head–tail decomposition, and translates them into practical certified selection algorithms (Delta-$k$-Search and MC-Search). The Gaussian score model yields closed-form design rules for the minimal $k$ to achieve a target error, guiding adaptive sparsification in long-context transformers. Empirical validation on BERT and synthetic data confirms scalable certified sparsity and substantial efficiency gains while preserving the prescribed total-variation budget. Overall, the framework offers provable, deployable sparsity with concrete guidelines for real-world Transformer inference and training workflows.

Abstract

We develop a unified mathematical framework for certified Top-$k$ attention truncation that quantifies approximation error at both the distribution and output levels. For a single attention distribution $P$ and its Top-$k$ truncation $\hat P$, we show that the total-variation distance coincides with the discarded softmax tail mass and satisfies $\mathrm{TV}(P,\hat P)=1-e^{-\mathrm{KL}(\hat P\Vert P)}$, yielding sharp Top-$k$-specific bounds in place of generic inequalities. From this we derive non-asymptotic deterministic bounds -- from a single boundary gap through multi-gap and blockwise variants -- that control $\mathrm{TV}(P,\hat P)$ using only the ordered logits. Using an exact head-tail decomposition, we prove that the output error factorizes as $\|\mathrm{Attn}(q,K,V)-\mathrm{Attn}_k(q,K,V)\|_2=τ\|μ_{\mathrm{tail}}-μ_{\mathrm{head}}\|_2$ with $τ=\mathrm{TV}(P,\hat P)$, yielding a new head-tail diameter bound $\|\mathrm{Attn}(q,K,V)-\mathrm{Attn}_k(q,K,V)\|_2\leτ\,\mathrm{diam}_{H,T}$ and refinements linking the error to $\mathrm{Var}_P(V)$. Under an i.i.d. Gaussian score model $s_i\sim\mathcal N(μ,σ^2)$ we derive closed-form tail masses and an asymptotic rule for the minimal $k_\varepsilon$ ensuring $\mathrm{TV}(P,\hat P)\le\varepsilon$, namely $k_\varepsilon/n\approxΦ_c(σ+Φ^{-1}(\varepsilon))$. Experiments on bert-base-uncased and synthetic logits confirm the predicted scaling of $k_\varepsilon/n$ and show that certified Top-$k$ can reduce scored keys by 2-4$\times$ on average while meeting the prescribed total-variation budget.

A Mathematical Theory of Top-$k$ Sparse Attention via Total Variation Distance

TL;DR

This work provides a rigorous, distribution- and output-level theory for Top- attention truncation, tying the truncation error to the discarded softmax mass via an exact TV–KL identity. It derives deterministic and probabilistic bounds based on score gaps, blocks, and a head–tail decomposition, and translates them into practical certified selection algorithms (Delta--Search and MC-Search). The Gaussian score model yields closed-form design rules for the minimal to achieve a target error, guiding adaptive sparsification in long-context transformers. Empirical validation on BERT and synthetic data confirms scalable certified sparsity and substantial efficiency gains while preserving the prescribed total-variation budget. Overall, the framework offers provable, deployable sparsity with concrete guidelines for real-world Transformer inference and training workflows.

Abstract

We develop a unified mathematical framework for certified Top- attention truncation that quantifies approximation error at both the distribution and output levels. For a single attention distribution and its Top- truncation , we show that the total-variation distance coincides with the discarded softmax tail mass and satisfies , yielding sharp Top--specific bounds in place of generic inequalities. From this we derive non-asymptotic deterministic bounds -- from a single boundary gap through multi-gap and blockwise variants -- that control using only the ordered logits. Using an exact head-tail decomposition, we prove that the output error factorizes as with , yielding a new head-tail diameter bound and refinements linking the error to . Under an i.i.d. Gaussian score model we derive closed-form tail masses and an asymptotic rule for the minimal ensuring , namely . Experiments on bert-base-uncased and synthetic logits confirm the predicted scaling of and show that certified Top- can reduce scored keys by 2-4 on average while meeting the prescribed total-variation budget.

Paper Structure

This paper contains 107 sections, 23 theorems, 141 equations, 3 figures, 8 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Contributions
  3. Applications to LLM-based systems
  4. Scope and Structure
  5. Related Work
  6. Fixed and Learned Sparse Patterns.
  7. Low-Rank and Kernel Approximations.
  8. Memory-Efficient Implementations.
  9. Analytical and Theoretical Perspectives.
  10. Position of This Work.
  11. Comparative Guarantees.
  12. Preliminaries and Notation
  13. Attention Scores and Softmax Probabilities
  14. Indexing convention.
  15. Top-$k$ Truncation
  16. ...and 92 more sections

Key Result

Lemma 4.1

For Top-$k$ truncation as above,

Figures (3)

  • Figure 1: Empirical validation of the Gaussian prediction. The theoretical curve $\alpha_{\text{Gauss}}(\sigma)$ (solid) matches the measured averages $\alpha_{\text{emp}}$ (dots) for $\varepsilon\in\{10^{-3},\,5\times10^{-3},\,10^{-2}\}$. Accuracy remains remarkably high up to $\sigma=3$.
  • Figure 2: Certified accuracy--efficiency tradeoff.
  • Figure 3: Diagnostic: mean TV at fixed $k=16$.

Theorems & Definitions (44)

  • Lemma 4.1: Tail--mass identity
  • Remark 4.2: Shift invariance
  • Theorem 4.3: Exact TV--KL identity
  • proof
  • Remark 4.4: Asymmetry of the KL direction
  • Theorem 4.5: Deterministic gap bound
  • proof
  • Remark 4.6
  • Lemma 4.7: Multi-gap tail split
  • proof
  • ...and 34 more