🤖 machine learning

SOCKET: SOft Collision Kernel EsTimator for Sparse Attention

本文介绍了 SOCKET，一种新颖的稀疏注意力机制，它用概率软碰撞核取代传统的硬局部敏感哈希，以实现高效、低内存的令牌选择，并在长上下文推理期间达到比 FlashAttention 高出 1.5 倍的吞吐量。

原作者： Sahil Joshi, Agniva Chowdhury, Wyatt Bellinger, Amar Kanakamedala, Ekam Singh, Hoang Anh Duy Le, Aditya Desai, Anshumali Shrivastava

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Sahil Joshi, Agniva Chowdhury, Wyatt Bellinger, Amar Kanakamedala, Ekam Singh, Hoang Anh Duy Le, Aditya Desai, Anshumali Shrivastava

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你试图在一本长达数百万页的书中找到最重要的一句话。如果你必须阅读每一页才能找到那句话，这将耗费永恒的时间，并需要海量的内存。这正是大型语言模型（LLM）在尝试理解非常长的对话或文档时所面临的问题。它们会“陷入困境”，因为它们试图关注曾经见过的每一个单词，这既拖慢了速度，又占满了计算机内存。

这篇论文介绍了一种名为SOCKET（SOft Collision Kernel EsTimator，软碰撞核估计器）的新工具来解决这个问题。以下是其工作原理，通过简单的类比进行解释：

问题所在：“硬”搜索与“软”搜索

为了加快速度，以往的方法试图忽略大部分单词，只查看少数几个“重要”的单词。它们使用了一种称为LSH（局部敏感哈希）的技术。

旧方法（硬 LSH）： 想象你正在一个巨大的体育场里寻找一位朋友。旧方法根据一个简单的规则将所有人放入不同的桶中：“如果你穿着红衬衫，就去 A 桶。”
- 如果你的朋友在 A 桶，你就检查他们。
- 如果他们在 B 桶，你就完全忽略他们。
- 缺陷： 这太僵化了。你的朋友可能穿着粉色衬衫（非常接近红色），却被分到了 B 桶。旧方法完全忽略了他们，尽管他们可能是你需要找到的最重要的人。这就像一个“是或否”的开关，却经常错误地翻转。
新方法（SOCKET / 软 LSH）： SOCKET 改变了规则。它不再使用生硬的“是或否”，而是使用一个“概率旋钮”。
- 当你寻找朋友时，系统不会只检查一个桶。它会问：“这个人出现在 A 桶的可能性有多大？B 桶？C 桶？”
- 如果你的朋友穿着粉色衬衫，系统会说：“他们有 70% 的可能性在 A 桶，30% 的可能性在 B 桶。”
- 然后，它将来自许多不同桶的这些“可能性得分”相加，得出一个最终得分。

为何重要：“投票”类比

可以将旧方法比作一种僵化的投票系统，你要么获得一票，要么没有。如果你错过了截止线，即使你非常接近，也会得到零支持。

SOCKET 则像一场加权的人气竞赛。不再是二元化的输赢，每位候选人根据他们在许多不同类别中获得的“票数”（或概率位）获得一个分数。

稳定性： 因为它使用这些平滑的分级得分，关于谁“最重要”的排名要稳定得多。旧方法可能会仅仅因为微小的随机变化，就在最重要的第 1 和第 2 个单词之间互换位置。SOCKET 能保持顺序稳定，因为它看到的是“灰色地带”，而不仅仅是非黑即白。

结果：更快、更智能

通过使用这种“软”评分方法，SOCKET 可以：

更快地找到正确的单词： 它不需要阅读整本书；它只需查看其智能评分系统识别出的顶级候选者。
使用更少的内存： 它不需要存储海量数据来做出这些决定。
更加准确： 在测试中，即使上下文极长（例如 32,000 到 128,000 个单词），它找到正确信息的能力与其他方法一样好（甚至更好）。

核心结论

作者构建了一个自定义的计算机芯片指令（一个"CUDA 内核”），以使这种数学运算变得极快。他们声称，使用 SOCKET，AI 模型阅读和理解长文档的速度比当前标准方法快 1.5 倍，且不会损失准确性。

简而言之：SOCKET 阻止 AI 猜测“是或否”，转而让它询问“可能性有多大？” 这一微小的转变使得 AI 在处理海量文本时能够更加高效、稳定和准确。

技术摘要：SOCKET —— 面向稀疏注意力的软碰撞核估计器

问题陈述

大语言模型（LLM）在长上下文推理过程中面临根本性的可扩展性限制。虽然预填充（prefilling）阶段处理完整的输入序列，但迭代解码阶段要求每一个新生成的 token 都必须关注所有先前缓存的 token。随着上下文长度从 2K 增长至数百万 token，这种稠密注意力机制日益受到内存限制，经常超出 GPU 容量，从而需要代价高昂的 CPU-GPU 数据传输。

现有的稀疏注意力方法试图通过限制计算范围至部分 token 来缓解这一问题。然而，许多方法依赖于数据相关的启发式策略（例如 k-means 聚类），这会导致高昂的首 token 延迟（TTFT）成本，并且缺乏对分布偏移的鲁棒性。其他方法则利用传统的局部敏感哈希（LSH），其分配的是二元碰撞信号。本文指出，这种“硬”形式并不适合排序稳定性；二元碰撞无法表达部分相似性，导致粗糙且嘈杂的排序，从而降低检索质量，尤其是在严格的内存预算下。

方法论：SOCKET

作者提出了SOCKET（SOft Collision Kernel EsTimator，软碰撞核估计器），这是一种与数据无关的稀疏注意力机制，它将 LSH 从二元候选生成器重构为一种 principled（有原则的）、感知相似性的评分核。

核心机制：软 LSH

与传统的 LSH（将查询分配到每个哈希表中的单个离散桶）不同，SOCKET 采用软 LSH：

预计算（Prefill）：键（Keys）使用随机投影被哈希到 $L$ 个独立的哈希表中。每个键在每个表中被分配一个单一的桶 ID，并且值向量的范数被缓存。
软分配（Decoding）：当新查询到达时，它不会被分配到单个桶。相反，它在每个哈希表的所有桶上诱导出一个概率分布。这是通过计算基于查询在随机超平面上投影的软哈希分数，并经过温度控制的 softmax 函数来实现的。
评分：对于每个键，通过聚合所有 $L$ 个表中分配给该键特定桶的概率质量来计算连续分数。该分数由键的值向量范数加权。
选择：基于这些软分数，确定性选择前 $k$ 个键，并在该子集上计算精确注意力。

理论基础

本文提供了理论保证，证明使用软 LSH 分数计算的注意力近似于与 softmax 注意力密切相关的核（具体而言，是角核）。

稳定性：理论分析（定理 3）表明，软桶化引入了一个可控的偏差（ $\epsilon_\tau$ ），随着温度 $\tau \to 0$ 该偏差会消失，同时保持了平滑性。
相关性：引理 4 和实证结果表明，软分数比硬分数更好地保留了投影查询 - 键相似性的方向结构。硬 LSH 依赖于 $\ell_1$ 幅度并丢弃了象限内的方向信息，导致在有限 $L$ 值下排序不稳定。软 LSH 在小信号机制下与投影坐标保持线性关系，从而产生与真实相似性更高的相关性以及更低的估计器方差。

主要贡献

作为稳定排序器的软 LSH：作者提出了一种与数据无关的软 LSH 机制，其作为排序函数的效果显著优于传统的硬 LSH。它用概率性的、感知相似性的聚合取代了二元碰撞信号，从而无需临时投票即可实现稳定的前 $k$ 选择。
通过自定义核实现效率：实现了自定义的 CUDA 评分核和 Flash Decode Triton 后端。这使得 SOCKET 在长上下文解码期间的吞吐量比 FlashAttention 高出1.5 倍，同时保持或超越先前稀疏方法的准确性。
与数据无关的设计：与基于聚类的方法不同，SOCKET 不需要重新训练、校准或离线预处理。这实现了更快的 TTFT 以及对分布偏移的鲁棒性。
理论保证：本文建立了软计数注意力估计器的端到端误差界限，将误差分解为采样方差、有限表近似误差和软桶化偏差。

实验结果

作者在 Llama-3.1-8B-Instruct、Llama-3.2-1B 和 Qwen3 模型（4B、8B、30B）上评估了 SOCKET，测试了两个基准：LongBench（真实世界任务）和 RULER（合成诊断任务）。

准确性：在 LongBench 上，SOCKET 始终优于 Quest、PQCache、Double Sparsity 和 MagicPig 等基线。例如，在 Llama-3.1-8B 模型上，33 倍稀疏度下，SOCKET 的平均得分为 47.83，比最强的基线（Quest）高出约 0.84 分。在 RULER-32K 上，它匹配或超越了先前的方法，特别是在高稀疏度水平下（例如 50 倍）。
效率：SOCKET 在 NVIDIA A100 和 H200 GPU 上展示了显著的吞吐量增益。在 H200 上，随着上下文长度的增加，它超越了 FlashAttention-2 和 FlashAttention-3，在 100K token 时比 FlashAttention-2 实现了1.58 倍的加速，在 140K token 时比 FlashAttention-3 实现了1.58 倍的加速（在 33 倍稀疏度下）。
内存开销：该方法引入了约 15% 的适度内存开销（超出标准 KV 缓存，用于存储哈希表索引和值范数），这显著低于硬 LSH 为实现可比准确性所需的开销（后者需要更多的哈希表）。

意义与主张

本文声称，SOCKET 通过提供一种有原则的、与数据无关的替代方案，解决了长上下文 LLM 推理的关键瓶颈。通过将硬二元哈希转变为软概率评分，SOCKET 实现了更优越的准确性 - 效率权衡。

作者强调，他们的方法：

消除了对数据相关预处理的需求，从而降低了 TTFT 并提高了部署灵活性。
提供了理论稳定性，确保即使哈希表数量有限，token 的排序也能保持一致，这是传统 LSH 的弱点。
带来了实际系统层面的增益，证明了评分核中的算法改进可以直接转化为现代 GPU 架构上的显著吞吐量提升，而不会牺牲模型性能。

这项工作将软 LSH 定位为不仅仅是一个检索工具，而是一个稳健的评分核，可以替代解码阶段的稠密注意力，从而实现 LLM 向数百万 token 上下文的高效扩展。