Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

InfoFlow KV: Information-Flow-Aware KV Recomputation for Long Context

Xin Teng, Canyu Zhang, Shaoyi Zheng, Danyang Zhuo, Tianyi Zhou, Shengjie Wang

TL;DR

It is shown that a simple attention-norm signal from the query reliably identifies tokens that are both semantically relevant and structurally positioned to propagate information, when computed under an inference-consistent RoPE geometry.

Abstract

Retrieval-augmented generation (RAG) for long-context question answering is bottlenecked by inference-time prefilling over large retrieved contexts. A common strategy is to precompute key-value (KV) caches for individual documents and selectively recompute a small subset of tokens to restore global causal dependencies, but existing methods rely on heuristics or representation discrepancies without modeling whether selected tokens can effectively influence generation. We cast selective KV recomputation as an information flow problem and show that a simple attention-norm signal from the query reliably identifies tokens that are both semantically relevant and structurally positioned to propagate information, when computed under an inference-consistent RoPE geometry. We therefore reconstruct global positional assignments for retrieved chunks and introduce an information-flow-guided chunk reordering strategy. Experiments on LLM and VLM benchmarks demonstrate consistent gains over prior methods under comparable efficiency budgets.

InfoFlow KV: Information-Flow-Aware KV Recomputation for Long Context

TL;DR

It is shown that a simple attention-norm signal from the query reliably identifies tokens that are both semantically relevant and structurally positioned to propagate information, when computed under an inference-consistent RoPE geometry.

Abstract

Retrieval-augmented generation (RAG) for long-context question answering is bottlenecked by inference-time prefilling over large retrieved contexts. A common strategy is to precompute key-value (KV) caches for individual documents and selectively recompute a small subset of tokens to restore global causal dependencies, but existing methods rely on heuristics or representation discrepancies without modeling whether selected tokens can effectively influence generation. We cast selective KV recomputation as an information flow problem and show that a simple attention-norm signal from the query reliably identifies tokens that are both semantically relevant and structurally positioned to propagate information, when computed under an inference-consistent RoPE geometry. We therefore reconstruct global positional assignments for retrieved chunks and introduce an information-flow-guided chunk reordering strategy. Experiments on LLM and VLM benchmarks demonstrate consistent gains over prior methods under comparable efficiency budgets.
Paper Structure (39 sections, 8 equations, 4 figures, 6 tables)

This paper contains 39 sections, 8 equations, 4 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. KV Eviction & Compression
  4. Chunked Input Processing
  5. KV Recomputation
  6. Preliminary
  7. KV Cache
  8. Rotary Positional Embedding (RoPE)
  9. Method
  10. Input Chunking and Prefilling
  11. Token Selection and Recomputation
  12. Selection Method
  13. KV Recomputation
  14. RoPE Geometry
  15. KV reorder
  16. ...and 24 more sections

Figures (4)

  • Figure 1: Context chunks are prefetched independently using chunk-local RoPE. At inference time, retrieved chunks are concatenated with the prompt, global RoPE positions are reconstructed, and prompt-conditioned attention norms are used to select high-impact tokens for full-context KV recomputation. The recomputed KV states are concatenated with cached chunks, restoring cross-chunk interactions. An optional chunk reordering step places more informative chunks closer to the prompt.
  • Figure 2: Speed--accuracy trade-off on LLaMA and Qwen across long-context QA benchmarks. Each curve corresponds to a recomputation budget sweep. Upper-left indicates a better trade-off.
  • Figure 3: Needle-in-a-Haystack accuracy heatmaps on Qwen3-14B under varying context lengths and needle depths.
  • Figure 4: Needle-in-a-Haystack accuracy heatmaps on Qwen3-14B under varying context lengths and needle depths, using attention norms extracted from different Transformer layers.