Efficient long‑context inference is a core requirement for large language models (LLMs) but creates a heavy computational bottleneck. Hash‑based retrieval encodes queries and keys into binary codes and selects keys by Hamming distance, offering a fast alternative. The Hamming metric only captures binary similarity and discards feature magnitudes, causing a mismatch with the true query‑key logits $s_{ij}=\frac{q_i\cdot k_j}{\sqrt{d}}$. Consequently, low‑logit keys are falsely retrieved while high‑logit keys may be omitted.
To remedy this, we introduce Hierarchical Hash Retrieval (HHR), a coarse‑to‑fine framework. The first stage, Geometry‑Aware Key Routing (GKR), learns a head‑wise orthogonal transformation $W_h$ that redistributes magnitudes, yielding discriminative page‑level logit bounds $b_{h,p}$ for effective pruning of low‑logit keys while preserving important candidates. The second stage, Learned Hash Projection (LHP), learns a head‑wise projection matrix $P_h$ that aligns Hamming distance with the true relevance ranking, enabling fine‑grained retrieval.
Combining GKR and LHP suppresses false positives and recovers false negatives, substantially improving the fidelity of hash‑based sparse attention. Extensive experiments across diverse LLMs and benchmarks show HHR’s superiority: a 1.10‑point average gain on LongBench, up to 3.30× decoding speedup and 2.83× end‑to‑end acceleration for Llama‑3.1‑8B‑Instruct at a 128K context length. The implementation is publicly available at https://github.com/lianjunl13-sudo/HHR.
Review