论文

HHR:分层哈希检索以高效生成 LLM

HHR: Hierarchical Hash Retrieval for Efficient LLM Generation

模型推理KV Cache

摘要

高效的长上下文推理对于大语言模型(LLM)至关重要,但它带来了严重的计算瓶颈。基于哈希的检索通过将查询和密钥编码为二进制代码并使用汉明距离进行密钥选择,提供了一种有效的替代方案。然而,这导致汉明距离和注意力相关性之间严重不匹配。 Query-Key logits 共同依赖于方向相似性和特征量级,而哈希二值化会丢弃量级信息,导致低 logit 键的误报检索和高 logit 键的漏报漏报。为了解决这些失败问题,我们提出了分层哈希检索(HHR),这是一种从粗到细的框架,通过几何感知密钥路由(GKR)和学习哈希投影(LHP)逐步提高检索准确性。 GKR 学习 head-wise 正交变换来重新分配特征量值并导出更具辨别力的页面级 logit 边界,从而能够有效修剪低 logit 键,同时保留重要的候选值。然后,LHP 学习一个头向投影空间,将汉明距离与真实的查询键相关性排名对齐,以进行细粒度检索。通过结合 GKR 和 LHP,HHR 抑制误报并恢复误报,大大提高了基于哈希的稀疏注意力的保真度。跨不同大语言模型和基准的广泛实验表明,HHR 比现有方法实现了卓越的性能。例如,在 LongBench 上,HHR 将平均得分提高了 1.10 分,并且在上下文长度为 128K 的情况下,Llama-3.1-8B-Instruct 的解码加速高达 3.30 倍,端到端加速高达 2.83 倍。该代码可在 https://github.com/lianjunl13-sudo/HHR. 公开获取