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

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 Large Language Model and Vision-Language Model benchmarks demonstrate consistent gains over prior methods under comparable latency.

Many AI systems now answer questions by retrieving long sets of documents, images, or other evidence before generating a response. This makes them more knowledgeable, but it also creates a serious efficiency problem: the model must spend a large amount of time reading all retrieved content before it can begin answering. Existing shortcuts reuse cached computations for each retrieved chunk, but they often miss important interactions between chunks and the user’s question. We tackle this problem by identifying which retrieved tokens are most likely to help the model answer the question and recomputing only those tokens under the correct global context. Our method uses attention signals to find tokens whose cached key–value states can effectively influence downstream generation, while respecting the positional structure used during inference. We also show that reordering independent retrieved chunks can further improve how useful information reaches the final answer. This approach is training-free, so it can be applied to existing language and vision-language models without changing their parameters. The result is faster long-context and retrieval-augmented generation with better accuracy than prior selective recomputation methods under similar efficiency budgets.