SOCKET：一种稀疏注意力软碰撞核估计器
机器学习
2026-05-11 v2
摘要
在长上下文推理中利用稀疏性是扩大大型语言模型规模的关键因素，因为注意力机制在自回归解码中占据主要成本。稀疏注意力通过限制计算仅针对一部分 token 来降低此开销，但其有效性取决于推理时的高效评分与选择。我们重新审视局部敏感哈希（LSH），并提出 SOCKET（SOft Collision Kernel EsTimator），用概率且基于相似性的聚合取代硬碰撞匹配。传统 LSH 仅产生二元碰撞信号，限制了排序质量且需要大量内存。相比之下，软 LSH 在多个哈希表中累积分级碰撞证据，显著降低内存需求，同时保留前 k 排序。将 LSH 从候选生成器重新框定为稀疏注意力的原则性评分核。基于此属性，SOCKET 实现了无需临时投票的高效 token 选择，跨多个长情境基准与先前稀疏注意力方法持平或超越。配合自定义 CUDA 评分核和 Flash Decode Triton 后端，SOCKET 在吞吐量上实现提升最高可达 1.5 倍于 FlashAttention。
引用
@article{arxiv.2602.06283,
title = {SOCKET: SOft Collision Kernel EsTimator for Sparse Attention},
author = {Sahil Joshi and Agniva Chowdhury and Wyatt Bellinger and Amar Kanakamedala and Ekam Singh and Hoang Anh Duy Le and Aditya Desai and Anshumali Shrivastava},
journal= {arXiv preprint arXiv:2602.06283},
year = {2026}
}
备注
7 figures, 17 tables