Know Before You Fetch: Calibrated Retrieval-Budget Allocation for Retrieval-Augmented Generation

2026-06-29

AI Summary

This work addresses the inefficiency of conventional retrieval-augmented generation (RAG) systems, which retrieve a fixed number of passages for all queries, often wasting resources or introducing noise. The authors frame adaptive RAG as a calibrated retrieval budget allocation problem, dynamically deciding—per query—whether to answer closed-book, retrieve a limited or full context, or abstain entirely. The core innovation lies in using a calibrated correctness probability as a unified decision interface, derived by combining sequence log-probabilities with prefix logit uncertainty signals and transforming them into reliable confidence estimates via cross-fold calibration. Experimental results demonstrate substantial improvements in calibration quality—e.g., reducing Expected Calibration Error (ECE) on Natural Questions from 0.643 to 0.009—and achieve Pareto-optimal trade-offs between accuracy and computational cost across model scales on TriviaQA, Natural Questions, and MS MARCO.

Abstract

Retrieval-augmented generation (RAG) typically retrieves a fixed number of passages for every query. This is wasteful when the reader already knows the answer, and it can be harmful when irrelevant or partially relevant passages distract the reader. We formulate adaptive RAG as calibrated retrieval-budget allocation: given a query, decide whether to answer closed-book, retrieve a compact context (k=1), retrieve a full context (k=5), or abstain. The contribution is a probability interface rather than a new raw uncertainty signal. We calibrate sequence log-probability and prefix-logit uncertainty signals into probabilities of correctness, then use these probabilities for graded context selection, selective abstention, and explicit latency/token trade-offs. Across core QA experiments on TriviaQA, Natural Questions, and MS MARCO, with auxiliary PopQA motivation and Qwen/Llama family checks, diagnostic out-of-fold calibration improves probability quality dramatically: for sequence log-probability, ECE drops from 0.275 to 0.062 on TriviaQA, 0.643 to 0.009 on NQ, and 0.711 to 0.031 on MS MARCO. Graded retrieval improves full-context and passage-budget frontiers for both our signal and TARG-style prefix entropy/margin, while retrieval-call AUC remains essentially tied with binary gating because k=1 is still a retrieval call. Held-out train/validation/test threshold experiments report deployable operating points. At matched-accuracy frontier operating points, a measured cost model reveals that gating is not universally faster: it increases latency by about 27% on Qwen3-8B but saves about 8% on Qwen3-32B. These results support a nuanced view of adaptive RAG: calibrated confidence is best understood as a reusable interface for allocating retrieval budget under task and system constraints.

Problem

• Retrieval-Augmented Generation
• adaptive retrieval
• retrieval-budget allocation
• calibrated confidence
• selective abstention

Innovation

• calibrated retrieval
• adaptive RAG
• uncertainty calibration
• retrieval-budget allocation
• selective abstention