Positional geometry improves sparse key value cache recall in language models
Compression Does Not End at Selection: Post-Selection Positional Geometry for Sparse KV Caches
Computation and Language
Summary
Long-context language models save important past information to help answer questions better, but removing some parts creates gaps that confuse the system. The authors studied how changing the way positions are assigned after selecting what to keep can improve remembering and matching the saved information. They created RoPERelay, a method that rearranges position data to better connect queries with the right memories without extra cost. Tests showed this improves accuracy in most cases with little extra processing time. This work shows that organizing positions is an important step beyond just choosing what to keep.
What this means in practice
- •For language model engineers: Enhance long-term memory recall in transformer models by reassigning position embeddings after sparse key-value selection to improve answer accuracy without extra compute.
- •For natural language processing teams: Deploy RoPERelay to optimize memory caches in long-document question answering systems for higher accuracy with minimal latency increase.
Authors
Wanyang Hong, Zhaoning Zhang
Abstract
Sparse KV compression reduces long-context inference costs by retaining selected historical states, yet these states typically keep their coordinates from the full history. Under Rotary Position Embedding (RoPE), gaps left by evicted tokens continue to affect query-key matching, potentially limiting the use of retained evidence. We study post-selection positional geometry on a fixed selector-produced cache, examining how positional organization affects the use of retained information and which coordinate and continuation rules improve answer quality. To reorganize this geometry, we propose RoPERelay, which assigns compact, order-preserving RoPE coordinates, relocates cached key phases through inverse rotation and re-rotation, and synchronizes the current query and subsequent generation. It reuses existing K/V without renewed selection, historical re-prefill, or additional cache budget. Across 36 comparisons spanning two models, six sparse selectors, and three long-term memory tasks, RoPERelay improves answer quality in 32 cases. For Llama-3.1-8B-Instruct with BM25 on LongMemEval-S, accuracy increases from 23.6% to 35.4%. Measured complete-request time overhead is approximately 1%-3% in the evaluated configurations. These findings establish positional organization as a distinct optimization dimension for sparse KV caches after content selection.