中文

vAttention:可验证稀疏注意力

机器学习 2026-05-26 v2 人工智能

摘要

用于降低解码延迟的 SOTA 稀疏注意力方法分为两大类:近似 top-kk(及其扩展 top-pp)和最近引入的基于抽样的估计。然而,这些方法在近似全注意力方面存在根本限制:它们无法在各个头和查询向量之间提供一致的近似,且最关键的是缺乏对近似质量的保证,限制了其实际部署。我们观察到 top-kk 和随机抽样是互补的:当注意力分数由少数 token 主导时,top-kk 表现良好;而当注意力分数相对均匀时,随机抽样提供更好的估计。基于这一洞察并利用抽样的统计保证,我们引入了 vAttention,这是第一个在近似精度上提供用户指定 (ϵ,δ)(\epsilon, \delta) 保证的实用稀疏注意力机制(因此是“可验证的”)。这些保证使 vAttention 成为迈向大规模实用、可靠的稀疏注意力部署的引人注目的一步。通过统一 top-kk 和抽样,vAttention 优于两者单独的表现,提供了卓越的质量-效率权衡。我们的实验表明,vAttention 显著提高了稀疏注意力的质量(例如,对于 RULER-HARD 上的 Llama 3.1 8B Instruct 和 DeepSeek-R1-Distill-Llama-8B 提升约 4.5 个百分点),并有效弥合了全注意力与稀疏注意力之间的差距(例如,在各个数据集上,它以高达 20 倍的稀疏度匹配全模型质量)。我们还证明它可以部署在推理场景中以实现快速解码而不影响模型质量(例如,vAttention 在 AIME2024 上以 10 倍稀疏度实现了全模型质量,生成了多达 32K 个 token)。代码:https://github.com/skylight-org/sparse-attention-hub。网页:https://sky-light.eecs.berkeley.edu。

关键词

引用

@article{arxiv.2510.05688,
  title  = {vAttention: Verified Sparse Attention},
  author = {Aditya Desai and Kumar Krishna Agrawal and Shuo Yang and Alejandro Cuadron and Luis Gaspar Schroeder and Matei Zaharia and Joseph E. Gonzalez and Ion Stoica},
  journal= {arXiv preprint arXiv:2510.05688},
  year   = {2026}
}