HHR：面向高效LLM生成的分层哈希检索

高效的长上下文推理对于大语言模型（LLM）至关重要，但其也带来了严重的计算瓶颈。基于哈希的检索通过将查询和键编码为二进制码，并使用汉明距离选择键，为此提供了一种高效的替代方案。然而，这会导致汉明距离与注意力相关性之间存在严重不匹配。查询-键logit取决于方向相似度和特征幅值，而哈希二值化会丢弃幅值信息，从而造成两类问题：错误检索低logit的键，以及遗漏高logit的键。为解决这些问题，我们提出分层哈希检索（Hierarchical Hash Retrieval，HHR），这是一种通过几何感知键路由（Geometry-Aware Key Routing，GKR）和学习型哈希投影（Learned Hash Projection，LHP）逐步提升检索准确率的由粗到细框架。GKR学习一种按注意力头分别进行的正交变换，以重新分配特征幅值，并推导出更具判别性的页面级logit界限，从而能够有效剪枝低logit的键，同时保留重要候选键。随后，LHP学习一个按注意力头分别构建的投影空间，使汉明距离与真实的查询-键相关性排序相一致，以实现细粒度检索。通过结合GKR和LHP，HHR能够抑制误检并找回漏检，显著提升基于哈希的稀疏注意力的保真度。在多种LLM和基准测试上的大量实验表明，HHR的性能优于现有方法。例如，在LongBench上，HHR将平均得分提升了1.10分；在128K上下文长度下，使用Llama-3.1-8B-Instruct时，其解码速度最高提升3.30倍，端到端速度最高提升2.83倍。代码已公开于https://github.com/lianjunl13-sudo/HHR。