[ICLR 2025] InfoFlow KV: Restoring the "Pulse" of Information Flow in Long-Context RAG

TL;DR

InfoFlow KV is a training-free, plug-and-play method to eliminate the prefilling bottleneck in long-context RAG. By identifying "information-critical" tokens using query-conditioned attention norms and recomputing them under a global RoPE geometry, it achieves a 2.57x speedup over Ring Attention and up to 3.49x over standard baselines while maintaining or even improving answer quality.

Problem & Motivation: The "Causal Disconnect" in RAG

In high-throughput RAG systems, we often precompute and cache KV (Key-Value) states for millions of document chunks to avoid redundant computation. However, when we assemble a prompt with multiple retrieved chunks, a fatal issue arises: Local vs. Global Context.

Each chunk was cached in isolation (Local Causal Mask). When combined, they lack the cross-chunk attention dependencies needed for complex reasoning (Global Causal Mask). Previous attempts to fix this were often "blind":

• EPIC: Recomputes fixed positions (like chunk starts), ignoring the query's actual needs.
• CacheBlend: Looks at representation discrepancies in shallow layers, which might not represent the final information flow to the output.

The authors' insight is simple: We should recompute tokens that the query actually cares about and that are positionally "reachable" within the RoPE geometry.

Methodology: High-Impact Token Selection

The core of InfoFlow KV is treating KV recomputation as an Information Flow problem.

1. The Attention-Norm Criterion

Instead of complex discrepancy measures, the method uses the aggregated attention mass from the prompt to the context tokens: $s_j={\sum }_{i=1}^M{A}_{ij}$ Tokens receiving the most attention are prioritized. This naturally combines semantic relevance (query affinity) and structural importance.

2. Global RoPE Reconstruction

The paper reveals a critical finding: Token selection signals are only reliable if computed under an inference-consistent RoPE geometry. If you compute importance scores using local (0-indexed) positions but then run inference with global positions, the frequency bands of the Rotary Embeddings shift, making the selection suboptimal.

3. Chunk Reordering

To further optimize, InfoFlow KV reorders chunks based on their aggregate importance, moving the most "informative" chunks closer to the prompt to maximize attention effectiveness during decoding.

Experiments & Results: Better, Faster, Stronger

The authors tested InfoFlow KV across three LLM families (Qwen, Llama, ChatGLM) and a multi-modal VLM (Qwen3-VL).

SOTA Comparison

On multi-hop reasoning tasks (2WikiMQA, MuSiQue), InfoFlow KV significantly outperformed previous SOTA methods like CacheBlend and EPIC. For instance, on 2WikiMQA with Qwen, it achieved an F1 of 0.5019, compared to CacheBlend's 0.4330.

The Pareto Frontier

The efficiency gains are striking. As shown in the speed-accuracy trade-off charts, InfoFlow KV stays on the top-left (the "Goldilocks" zone of high accuracy and low latency). At a 15% recomputation ratio for a 32K context, it delivers a 3.49x speedup in Time-To-First-Token (TTFT).

Sequence Parallelism Mastery

In multi-GPU settings, InfoFlow KV avoids the massive communication overhead of All-Gathering full KV caches. By communicating only the recomputed 15% of tokens, it beats Ring Attention's performance and latency consistently at long sequence lengths.

Critical Analysis & Conclusion

Takeaway

The success of InfoFlow KV suggests that the Transformer's attention mechanism is surprisingly robust to "sparse" causal updates, provided those updates are made at the right positional frequencies.

Limitations

• Irregular Kernels: Current hardware-accelerated kernels (like FlashAttention) are optimized for dense/causal masks. InfoFlow KV creates irregular sparse masks during recomputation, which currently limits hardware utilization.
• Layer Sensitivity: Selection works best when using attention norms from intermediate-to-late layers (Layers 22-25), suggesting that "importance" is a high-level semantic feature.

Future Work

The next frontier is likely the development of customized sparse kernels to unlock the full theoretical speedup of these selective recomputation strategies.

Conclusion: InfoFlow KV provides a principled bridge between the efficiency of offline precomputation and the accuracy of full-context inference, making long-context RAG significantly more practical for real-world deployment.