# Scaling Test-Time Inference with Policy-Optimized, Dynamic Retrieval-Augmented Generation via KV Caching and Decoding

arXiv: https://arxiv.org/abs/2504.01281 · HTML: https://mlsanity.com/paper/2504.01281/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Generation and context integration, one of 15 directions in Retrieval-augmented generation.
- Builds on 56 of its 62 references that are in the map.
- First posted on arXiv 2 Apr 2025.

## Placed in

- Retrieval-augmented generation / [Generation and context integration](https://mlsanity.com/rag/generation-and-context-integration/)

## Builds on (most cited first, 30 of 56)

- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [Gorilla: Large Language Model Connected with Massive APIs](https://mlsanity.com/paper/2305.15334/) (2023)
- [PubMedQA: A Dataset for Biomedical Research Question Answering](https://mlsanity.com/paper/1909.06146/) (2019)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [Accelerating Large Language Model Decoding with Speculative Sampling](https://mlsanity.com/paper/2302.01318/) (2023)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization](https://mlsanity.com/paper/2401.18079/) (2024)
- [Mixture-of-Agents Enhances Large Language Model Capabilities](https://mlsanity.com/paper/2406.04692/) (2024)
- [Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach](https://mlsanity.com/paper/2502.05171/) (2025)
- [Break the Sequential Dependency of LLM Inference Using Lookahead Decoding](https://mlsanity.com/paper/2402.02057/) (2024)
- [Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training](https://mlsanity.com/paper/2309.17179/) (2023)
- [DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads](https://mlsanity.com/paper/2410.10819/) (2024)
- [RAFT: Adapting Language Model to Domain Specific RAG](https://mlsanity.com/paper/2403.10131/) (2024)
- [Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning](https://mlsanity.com/paper/2405.00451/) (2024)
- [Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference](https://mlsanity.com/paper/2407.11550/) (2024)
- [Chain-of-Thought Reasoning Without Prompting](https://mlsanity.com/paper/2402.10200/) (2024)
- [Accessing GPT-4 level Mathematical Olympiad Solutions via Monte Carlo Tree Self-refine with LLaMa-3 8B](https://mlsanity.com/paper/2406.07394/) (2024)
- [Marco-o1: Towards Open Reasoning Models for Open-Ended Solutions](https://mlsanity.com/paper/2411.14405/) (2024)
- [Mutual Reasoning Makes Smaller LLMs Stronger Problem-Solvers](https://mlsanity.com/paper/2408.06195/) (2024)
- [Can 1B LLM Surpass 405B LLM? Rethinking Compute-Optimal Test-Time Scaling](https://mlsanity.com/paper/2502.06703/) (2025)
- [MemoRAG: Boosting Long Context Processing with Global Memory-Enhanced Retrieval Augmentation](https://mlsanity.com/paper/2409.05591/) (2024)
- [Online Speculative Decoding](https://mlsanity.com/paper/2310.07177/) (2023)

## Built on by (most cited first, 5 of 5)

- [Synergizing RAG and Reasoning: A Systematic Review](https://mlsanity.com/paper/2504.15909/) (2025)
- [Gated Differentiable Working Memory for Long-Context Language Modeling](https://mlsanity.com/paper/2601.12906/) (2026)
- [Interaction as Intelligence: Deep Research With Human-AI Partnership](https://mlsanity.com/paper/2507.15759/) (2025)
- [MODE-RAG: Manifold Outlier Diagnosis and Energy-based Retrieval-Augmented Generation Evaluation](https://mlsanity.com/paper/2606.17449/) (2026)
- [Look Before You Leap: Thermodynamic Arbitration of Parametric and Non-Parametric Knowledge in LLM Agents via Self-Regulating Memory Architectures](https://mlsanity.com/paper/2610.05223/) (2026)

## Abstract (arXiv)

We present a comprehensive framework for enhancing Retrieval-Augmented Generation (RAG) systems through dynamic retrieval strategies and reinforcement fine-tuning. This approach significantly improves large language models on knowledge-intensive tasks, including opendomain question answering and complex reasoning. Our framework integrates two complementary techniques: Policy-Optimized RetrievalAugmented Generation (PORAG), which optimizes the use of retrieved information, and Adaptive Token-Layer Attention Scoring (ATLAS), which dynamically determines retrieval timing and content based on contextual needs. Together, these techniques enhance both the utilization and relevance of retrieved content, improving factual accuracy and response quality. Designed as a lightweight solution compatible with any Transformer-based LLM without requiring additional training, our framework excels in knowledge-intensive tasks, boosting output accuracy in RAG settings. We further propose CRITIC, a novel method to selectively compress key-value caches by token importance, mitigating memory bottlenecks in long-context applications. The framework also incorporates test-time scaling techniques to dynamically balance reasoning depth and computational resources, alongside optimized decoding strategies for faster inference. Experiments on benchmark datasets show that our framework reduces hallucinations, strengthens domain-specific reasoning, and achieves significant efficiency and scalability gains over traditional RAG systems. This integrated approach advances the development of robust, efficient, and scalable RAG systems across diverse applications.