AB-RAG: Adaptive Budgeted Retrieval-Augmented Generation for Reliable Question Answering

Abstract: Retrieval-Augmented Generation (RAG) has become the standard way to ground LLMs in external knowledge, yet most systems retrieve a fixed number of passages for every question regardless of its difficulty. This wastes computation on easy questions, starves hard ones, and gives no signal for when a generated answer can be trusted. With a growing share of question answering systems built on top of commercial LLM APIs, a method that can decide how much to retrieve, and how far to trust its own answers, without retraining the underlying model, is of clear practical value. This paper presents AB-RAG (Adaptive Budgeted Retrieval-Augmented Generation), a training-free and backbone-agnostic framework that generates an answer, estimates its confidence from a combination of three signals, and then decides whether to stop or to retrieve more evidence, subject to a fixed retrieval budget. The estimator combines the model's own certainty, the agreement between the answer and the evidence, and the variance of the retrieval scores. For models that expose token probabilities the certainty signal is read directly; for closed APIs it is approximated by self-consistency, so the method works without access to model internals. Across three backbones and two datasets, the central result is that the confidence estimate reliably separates correct from incorrect answers on every backbone, reaching a clean split of 57.6% against 0% Exact Match between high- and low-confidence answers on a factoid dataset. The adaptive policy improves accuracy on capable backbones, and the study reports its negative and nuanced findings honestly, including a confidence signal that proved unsuitable for short answers and a retrieval signal whose sign was found and corrected through measurement. The entire study was carried out on a single consumer laptop with only a few dollars of API spend.

• The paper introduces AB-RAG which dynamically allocates retrieval effort based on query difficulty, optimizing evidence use.
• It combines multi-signal confidence estimation using model certainty, evidence consistency, and retrieval variance to determine answer reliability.
• Empirical results show improved Exact Match scores and cost-accuracy tradeoffs compared to static RAG on datasets like HotpotQA and TriviaQA.

AB-RAG: Adaptive Budgeted Retrieval-Augmented Generation for Reliable Question Answering

Motivation and Problem Setting

Retrieval-Augmented Generation (RAG) has become a dominant paradigm for grounding LLM answers in external knowledge. Conventional RAG pipelines statically retrieve a fixed number of passages per query, regardless of question complexity or model certainty, resulting in inefficient evidence use—over-retrieving for simple queries and under-retrieving for complex or multi-hop queries. Such rigidity leads to unnecessary computational or API costs and absence of a trustworthiness signal on the system’s output.

The paper introduces Adaptive Budgeted RAG (AB-RAG), a training-free, backbone-agnostic system that dynamically adapts retrieval effort to individual query difficulty. Unlike prior work, AB-RAG requires neither model retraining nor access to model internals for closed APIs, while directly integrating trust-calibrated answer selection and explicit control over retrieval budget.

Figure 1: Fixed-depth RAG applies the same retrieval budget to every query, which over-retrieves for easy questions and under-retrieves for hard ones. AB-RAG estimates confidence after each answer and retrieves more only when needed, subject to a budget.

System Architecture and Methodology

AB-RAG formalizes retrieval as an adaptive loop:

1. Retrieve initial evidence using hybrid (sparse + dense) retrieval, rerank the candidates, and generate an initial answer.
2. Estimate the answer’s confidence based on a multi-signal estimator.
3. If confidence is below a tunable threshold and budget remains, retrieve more, update, and repeat; otherwise, stop and output the current answer.

A central property is that AB-RAG works seamlessly with both open-weight models (supporting token-probability confidence calculation) and closed APIs (using self-consistency across multiple samples as a proxy for confidence).

Figure 2: The AB-RAG architecture. Hybrid retrieval and reranking produce an evidence set; the generator answers; the confidence estimator combines three signals; and the decision step either stops or triggers another retrieval round, subject to the budget.

Confidence Signals

The system’s confidence estimator leverages three signals:

• S₁: Model certainty, quantified via (a) mean output token probabilities (if accessible), or (b) self-consistency (majority agreement over independent samples) for black-box models.
• S₂: Evidence consistency, computed as the embedding cosine similarity between the generated answer(s) and retrieved evidence.
• S₃: Variance in retrieval reranker scores—a surrogate for the retriever’s ability to separate relevant from distracting evidence.

Figure 3: The three confidence signals. $S_1$ is the model's own certainty from token probabilities or self-consistency; $S_2$ is the embedding similarity between the answer and the evidence; $S_3$ is the variance of the reranker scores, used as a reward for clean separation.

The adaptive control policy combines these signals into a single confidence value (via weighted sum and clipping), governing whether to halt or expand retrieval based on the user-set threshold and maximal retrieval budget.

Adaptive Retrieval Loop

Upon each iteration:

• Evidence is fetched and reranked.
• An answer and $S_1$ are generated.
• $S_2$ and $S_3$ are computed from current evidence and reranker outputs.
• Combined confidence is compared against threshold; if unmet, evidence set size is increased by a fixed step and the process cycles (until the retrieval budget is exhausted).

A real execution trace is provided to exemplify dynamic evidence augmentation depending on evolving confidence.

Figure 4: A real worked example of the adaptive loop. The first round is below the threshold and triggers more retrieval; the second round crosses the threshold and the loop stops with the correct answer. The confidence values are taken from the actual run.

Experimental Setup and Implementation

Experiments utilize hybrid retrieval (BM25 + BGE embeddings), reciprocal rank fusion, cross-encoder reranking, and three backbone LLMs: a small open-weight model (Qwen2.5-1.5B-Instruct), a mid-sized open model (Llama-3.2-3B), and a closed commercial API (Claude Haiku). Two open-retrieval QA datasets are used: HotpotQA (multi-hop) and TriviaQA (factoid). The evaluation is intentionally kept accessible: all stages are executed on a commodity laptop with modest resource demands.

The evidence pool is pooled from dataset-supplied candidates but remains sufficiently large to demonstrate real retrieval challenges, especially after reranking and evidence size increases.

Results and Analysis

Retrieval Stack Effectiveness

Dense retrievers outperform sparse and hybrid methods for clean text, both in HotpotQA and TriviaQA. Reranking with a cross-encoder further lifts top-$k$ recall, especially at lower $k$, validating the importance of combining semantic and lexical signals before evidence is passed to the generator.

Figure 5: Open-retrieval recall by method on HotpotQA and TriviaQA. The dense retriever is strongest on both datasets, hybrid fusion sits below it on this clean text, and reranking lifts the hard low-$k$ recall.

Static vs. Adaptive RAG

The adaptive policy’s impact is contingent on backbone capability. For strong backbones (e.g., Llama-3.2-3B), AB-RAG achieves higher Exact Match (EM) compared to static RAG, raising EM from 39.5 to 45.0 on HotpotQA, with a significant increase in average retrieval iterations, signifying real utilization of the adaptive loop. On factoid TriviaQA with the Claude API, EM improves from 35.5 (static) to 40.0 (adaptive).

Figure 6: Exact Match for static RAG and AB-RAG across backbones, with 95% bootstrap confidence intervals on the AB-RAG values. AB-RAG improves over static most clearly on the mid-sized model and on the closed model with TriviaQA.

For small models, benefit is minimal; the adaptive mechanism rarely triggers additional retrieval rounds, reflecting the inability of weak models to exploit added evidence.

Confidence as Predictor of Correctness

A central claim is the estimator's ability to separate correct from incorrect outputs. Across architectures and datasets, high-confidence answers achieve much higher EM than low-confidence ones; on TriviaQA with Claude, high-confidence outputs reach 57.6% EM (versus 0% for low confidence), a sharp and operationally significant separation.

Figure 7: High-confidence answers achieve far higher Exact Match than low-confidence answers on every backbone and dataset. The closed model on TriviaQA shows the cleanest separation, 57.6% against zero, with a large low-confidence group.

Cost-Accuracy Tradeoff

The primary architectural knob—the confidence threshold—enables cost-accuracy tradeoff tuning. For capable models, increasing the threshold yields higher retrieval costs and accuracy, supporting practical deployment requirements (latency or API cost constraints vs. required answer reliability).

Figure 8: Cost-accuracy tradeoff as the confidence threshold is swept. The closed model on HotpotQA shows a rising curve, the closed model on TriviaQA holds accuracy while cost rises, and the small Qwen model stays flat.

Signal Ablation and Diagnostics

Signal ablation shows that the model-certainty signal ($S_1$) is the only strong predictor of answer correctness. The evidence consistency signal ($S_2$0), based on answer-passage embedding similarity, provides little to no discriminative value for short-form answers and can be misleading due to dimensionality and context mismatches. The retrieval-variance signal ($S_2$1) is weakly predictive but its sign—whether it is subtracted or added—matters, with higher variance correlating with more reliable retrieval.

Figure 9: Single-signal predictiveness across backbones. $S_2$2, the model-certainty signal, is strongly predictive everywhere; $S_2$3, evidence consistency, is at or below chance for these short answers; $S_2$4, the corrected retrieval-variance reward, is weakly predictive.

Theoretical and Practical Implications

This study establishes that accurate, reliable confidence estimation for answer correctness in RAG can be achieved without model retraining, even under closed-model constraints, by leveraging model uncertainty (via token probabilities or self-consistency). AB-RAG offers robust selective prediction: only answering when confidence is sufficient, an essential property for practical QA deployment. It further allows dynamic cost management via retrieval budget and confidence thresholds—directly translating into monetary savings in API-based production systems.

The negative result that answer-evidence embedding similarity is not useful for short answers guides future work to alternative grounding validation strategies—such as natural language inference or span-level string matching—especially as systems broaden to cover long-form or structured generation.

Future Directions

Potential avenues include:

• Scaling up to web-scale corpora while validating that confidence-based separation holds at larger retrieval depths.
• Designing improved consistency or grounding signals for both short and long answers.
• Automatically learning the combination weights for confidence signals.
• Integrating adaptive retrieval into agentic or multi-agent paradigms, where retrieval, generation, and verification are jointly optimized.
• Addressing answer latency and its tradeoff with reliability in real-world deployments.

Conclusion

AB-RAG delivers a pragmatic, training-free, and backbone-agnostic solution to adaptive, budgeted retrieval-augmented generation, enabling reliable confidence-driven QA even with closed commercial backbones. Its design choices, robust empirical findings—including a high-confidence/zero-confidence split of 57.6%/0% EM on a factoid dataset—and candid reporting of both strengths and limitations offer valuable methodological guidance for further research and deployment in knowledge-intensive NLP systems (2606.29090).