用数学方法精确量化稀疏注意力的误差,可实现安全压缩。
A Mathematical Theory of Top-$k$ Sparse Attention via Total Variation Distance
- 基于总变差距离建立统一理论框架,精准刻画截断误差。
- 证明输出误差与头尾向量距离成正比,给出可计算的上界。
- 在BERT等模型上验证,可压缩2-4倍计算量仍保持精度。
我们构建了一个统一的数学框架,用于认证Top-k注意力截断的近似误差,从分布和输出两个层面进行量化。对于单个注意力分布P及其Top-k截断<span style="font-style:italic"></em></span>,我们证明总变差距离等于被舍弃的softmax尾部质量,并满足\mathrm{TV}(P,\hat P)=1-e^{-\mathrm{KL}(<span style="font-style:italic">\hat P\Vert P\end{aligned}
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。