Cost-Aware Retrieval-Augmentation Reasoning Models with Adaptive Retrieval Depth
Helia Hashemi, Victor Rühle, Saravan Rajmohan
TL;DR
This work tackles the high computational cost of retrieval-augmented reasoning by introducing Dynamic Search-R1, a framework that enables adaptive retrieval depth and cost-aware reinforcement learning. It formulates an RL objective that incorporates an outcome-based reward and a KL regularization term, and implements both PPO and GRPO training regimes to align the model's reasoning with efficient search interactions. The authors introduce a cost-aware advantage function with two cost models—memory-bound and latency-bound—to jointly optimize accuracy and efficiency, showing that latency-bound penalization reduces end-to-end latency while maintaining competitive accuracy on seven QA datasets. Empirical results demonstrate substantial efficiency gains (latency reductions of ~16–20%) and modest accuracy improvements (~5% EM on average), especially with larger 7B LLMs, and reveal insights into the impact of the cost-weight hyperparameter alpha. The work provides actionable directions for deploying retrieval-augmented reasoning in cost-constrained settings and suggests extensions to agentic AI and multi-modal tasks with energy and environmental considerations.
Abstract
Reasoning models have gained significant attention due to their strong performance, particularly when enhanced with retrieval augmentation. However, these models often incur high computational costs, as both retrieval and reasoning tokens contribute substantially to the overall resource usage. In this work, we make the following contributions: (1) we propose a retrieval-augmented reasoning model that dynamically adjusts the length of the retrieved document list based on the query and retrieval results; (2) we develop a cost-aware advantage function for training of efficient retrieval-augmented reasoning models through reinforcement learning; and (3) we explore both memory- and latency-bound implementations of the proposed cost-aware framework for both proximal and group relative policy optimization algorithms. We evaluate our approach on seven public question answering datasets and demonstrate significant efficiency gains, without compromising effectiveness. In fact, we observed that the model latency decreases by ~16-20% across datasets, while its effectiveness increases by ~5% on average, in terms of exact match.