BalanceRAG: Joint Risk Calibration for Cascaded Retrieval-Augmented Generation Zijun Jia1 , Yuanchang Ye3 , Sen Jia2 , Yiyao Qian4 , Haoning Wang1 , Baojie Chen1 , Diyin Tang1 , Jinsong Yu1,* , Zhiyuan Wang4,* 1
Beihang University, Shenzhen Institute of Advanced Technology, 3 Zhejiang University of Finance & Economics, 4 University of Electronic Science and Technology of China 2
Correspondence: [email protected], [email protected]
Abstract
400
1
Number of Examples
arXiv:2605.20084v1 [cs.CL] 19 May 2026
Large language models (LLMs) can enhance factuality via retrieval-augmented generation (RAG), but applying RAG to every query is unnecessary when the model-only answer is reliable. This motivates cascaded RAG: each query is first handled by an LLM-only branch, escalated to a RAG fallback only if the primary branch is uncertain, and abstained from when neither branch is sufficiently trustworthy. However, calibrating such cascades stage by stage may be conservative, since the final utility depends on joint uncertainty thresholding of LLM-only and RAG. In this work, we develop BalanceRAG to certify threshold pairs at a target risk level. Given uncertainty scores from the two branches, BalanceRAG frames each threshold pair as an operating point on a two-dimensional lattice and identifies safe operating points using sequential graphical testing. This enables risk-adaptive threshold calibration, controlling the system-level error rate among accepted points, while retaining more examples. Furthermore, BalanceRAG extends to multi-risk calibration, allowing retrieval usage to be bounded together with the selectionconditioned risk. Experiments on three opendomain question answering (QA) benchmarks across multiple LLM backbones demonstrate that BalanceRAG meets prescribed risk levels, preserves higher coverage and more accepted correct examples, and reduces unnecessary retrieval calls compared with always-on RAG.
Both Correct Both Wrong Only LLM-only Correct Only LLM-RAG Correct
300 200 100 0
−60
−40
−20 0 20 SLLM−RAG − SLLM−only
40
60
Figure 1: Distribution of the per-example score differences between RAG and LLM-only. SLLM-RAG and SLLM-only are the similarity scores between each path’s prediction and the ground-truth answer. The x-axis reports SLLM-RAG − SLLM-only , with positive values favoring RAG and negative values favoring LLM-only, while the y-axis reports the number of examples. Colors distinguish whether both branches are correct, both are wrong, or only one branch is correct.
2020; Karpukhin et al., 2020; Asai et al., 2024; Jiang et al., 2023). However, always-on retrieval can be unnecessary and costly, since retrieval and long-context generation increase latency, memory use, and token consumption relative to direct LLM inference (Wang et al., 2024a; Li et al., 2024). Figure 1 shows that LLM-only and RAG have similar correctness scores for many queries, suggesting that retrieval mainly adds overhead when the model-only answer is already reliable. At the same time, a non-negligible fraction of examples are answered correctly only with retrieval. Thus, retrieval is neither uniformly necessary nor safely removable (Mallen et al., 2023). A desirable system should keep easy examples on the LLM-only path, escalate uncertain queries to RAG, and abstain when neither path is reliable enough. Prior adaptive RAG frameworks follow this principle by routing queries with uncertainty or complexity signals (Wang et al., 2025a; Jeong et al., 2024), but they are mostly heuristic policies optimized for
Introduction
Large language models (LLMs) have been widely deployed into real-world open-ended question answering (QA) scenarios (Pandit et al., 2025; Duan et al., 2025), but they still often produce fluent yet factually incorrect responses (Wang et al., 2025d; Huang et al., 2024). Retrieval-augmented generation (RAG) can mitigate this issue by grounding model output in external evidence, improving factuality on knowledge-intensive tasks (Lewis et al., 1
A. BalanceRAG: Risk-Controlled Cascading Inference 2 Risk-Controlled Router 𝝀! = 𝒕$𝟏, 𝒕$𝟐
1 LLM-only (Primary Path) Uncertainty Score 𝒖𝟏
LLM
Input Question Q: Who directed the movie Inception?
Answer: I’m not completely sure, but I think it was Christopher Nolan.
Accept LLM-only Output Final Answer
Is
Yes
𝒖 𝟏 ≤ 𝒕.𝟏 ?
(from LLM-only) No
Final Answer Yes Is 𝒖𝟐 ≤ t 𝟐 ?
Uncertainty Score 𝒖𝟐
LLM (with retrieved context)
Answer: According to the retrieved evidence, Inception was directed by Christopher Nolan.
Abstentions (not output)
Accept RAG Output
3 RAG Fallback (Secondary Path)
Retriever (Top-k docs)
Accepted answers (LLM-only or RAG)
I’m not completely sure, but I think it was Christopher Nolan.
1
0
Outputs to User (Under Target Risk Control)
0
No
1
Guarantee
According to the retrieved evidence, Inception was directed by Christopher Nolan.
(from RAG)
The population selectionconditioned risk is bounded by 𝛼 with high probability
Abstain/No Output (Both paths are uncertain)
ℙ 𝑅 𝜆% ≤ 𝛼 ≥ 1 − 𝛿
B. Risk-Controlled Joint Calibration (on Calibration Set) 1 Candidate Threshold Lattice 𝒕𝟐
2 Nodewise Evaluation For each candidate 𝜆 = 𝑡% , 𝑡&
Discretize thresholds
Acceptance 𝑀 𝜆
…
Routed Errors 𝐾 𝜆
1 …
…
Hypothesis Test 𝐻! : 𝑅 𝜆 > 𝛼
0.5 …
…
Obtain p-value
3 Sequential Graphical Testing (SGT)
4 Select Certified Threshold
Propagate significance budget 𝛿# on the lattice
𝑝# ≤ 𝛿#
UGD (Uniform Grid Diffusion)
DWD (Diamond-Wavefront Diffusion)
Yes No
𝒑𝝀 = ℙ Bin 𝑴 𝝀 , 𝜶 ≤ 𝑲 𝝀
…
0 0
… 0.5 … 1
Safe (reject 𝐻! )
𝒑𝝀
𝒕𝟏 0 1 2 3 4 5 6 7 8 9
Legend 𝒖𝟏 , 𝒖𝟐
Budget (more intense color = larger budget) Initialize starting node
Not Safe
Choose 𝜆& that maximizes acceptance
Uncertainty scores (lower is better) Pass the threshold Fail the threshold Data flow Budget flow (SGT) Budget amount
Figure 2: Overview of BalanceRAG: risk-controlled cascading inference (A) and joint threshold calibration via Sequential Graphical Testing (B). α is the user-specified risk level (upper bound).
average accuracy without finite-sample guarantees at a user-specified level. This leaves a deploymentcritical question: after adaptively combining LLMonly and RAG outputs, how can we control the error rate of the final accepted answers? To address this gap, we propose BalanceRAG, a training-free framework for risk-controlled cascaded RAG. As shown in Figure 2, given uncertainty scores from the LLM-only and RAG paths, BalanceRAG first attempts the LLM-only answer, invokes the RAG fallback only when the LLMonly is uncertain, and abstains if the fallback is also unreliable. The router is governed by two uncertainty thresholds, and BalanceRAG calibrates their joint operating point rather than tuning them stage by stage. Stage-wise calibration can be valid in cascaded systems (Jung et al., 2024), but it may be overly conservative as the reliability and utility of the final routing cascade depend on the systemlevel threshold choices for both branches. Specifically, BalanceRAG formulates joint calibration as a multiple testing problem over a twodimensional threshold lattice, where each threshold pair defines a candidate operating point, for which we test the null hypothesis that the systemlevel selection-conditioned risk exceeds α. Rejecting this null certifies the corresponding pair as safe. Although a Bonferroni correction provides valid error control (Bauer, 1991), it is overly conservative for this structured lattice: neighboring threshold pairs typically induce highly correlated
routing decisions, and safe operating points often form contiguous regions. To exploit this structure, BalanceRAG performs Sequential Graphical Testing (SGT) (Bretz et al., 2009), which sequentially propagates significance budget across the lattice and reallocates unused budget from certified nodes to their successors. The final router then selects, among all certified threshold pairs, the one with the largest acceptance rate, while preserving highprobability control of the system-level selectionconditioned error rate (Wang et al., 2026, 2025c). Furthermore, we extend BalanceRAG to multi-risk calibration, enabling simultaneous control of answer error and retrieval usage. Our contributions are threefold. First, to the best of our knowledge, BalanceRAG is the first framework to study finite-sample risk control for cascaded LLM/RAG routing, shifting adaptive retrieval from heuristic confidence estimation to statistically certified reliability. Second, we formulate joint threshold calibration as multiple testing over a two-dimensional lattice and adopt SGT to certify high-utility operating points. Third, we establish high-probability control of the systemlevel selection-conditioned error rate and extend the framework to multi-risk calibration.
2
Related Work
Adaptive Retrieval. Adaptive RAG aims to avoid unnecessary retrieval by routing queries based on 2
empirical signals. Previous studies explore when RAG is needed by comparing parametric and nonparametric memory (Mallen et al., 2023), or route queries using complexity estimation and learned policies (Jeong et al., 2024; Tang et al., 2025). Other methods trigger retrieval during generation based on token-level confidence, self-reflection, or information needs (Jiang et al., 2023; Asai et al., 2024; Su et al., 2024), while recent approaches further exploit retrieval-score statistics and agentic search (Wang et al., 2025a; Li et al., 2025). These methods reduce retrieval usage while preserving average answer quality, but they are mainly heuristic routing policies without statistical guarantees. BalanceRAG is orthogonal to these signals: they can serve as uncertainty estimates, while our calibration layer certifies the final LLM-RAG cascade at a prescribed selection-conditioned risk level. Risk Control in Selective Prediction. Conformal prediction (Wang et al., 2024b, 2025e,b; Tan et al., 2025; Jia et al., 2025) and risk control (Angelopoulos et al., 2024; Li et al., 2026) provide finite-sample guarantees for task-specific risk targets. Recent work extends risk control to selected foundation-model outputs (Gui et al., 2024), AIassigned labeling (Huang et al., 2025), cascaded LLM judges (Jung et al., 2024), and selective QA or routing systems (Wang et al., 2026, 2025c). These works provide the statistical foundation for reliable selection, but they do not investigate joint risk calibration for cascaded LLM-only/RAG routing. BalanceRAG targets this setting by certifying threshold pairs on a two-dimensional routing lattice, reducing conservativeness compared with stage-wise or Bonferroni-style calibration while preserving finite-sample risk control.
3
Methodology
3.1
Problem Formulation
Let A(y ⋆ , ŷ) ∈ {0, 1} indicate whether a generated answer ŷ is admissible with respect to the ground truth y ⋆ . We define the acceptance indicator as Sλ (x) = 1[ŷλ (x) ̸= ∅]. (1) The accepted-error indicator is Zλ (x, y ⋆ ) = Sλ (x) (1 − A(y ⋆ , ŷλ (x))) .
(2)
Here, Zλ (x, y ⋆ ) = 1 means that the router outputs an incorrect answer rather than abstaining. The population risk of interest is the selectionconditioned error rate: R(λ) = P (Zλ (X, Y ⋆ ) = 1 | Sλ (X) = 1) . (3) Given a target risk level α and error level δ, our goal is to select a threshold pair λ̂ with high acceptance while satisfying P R(λ̂) ≤ α ≥ 1 − δ. (4) 3.2
Joint Calibration with SGT
BalanceRAG calibrates the composed LLMonly/RAG router directly. Let Λ = {λj }N j=1 denote the candidate threshold pairs on a twodimensional lattice. For each candidate λj , we test Hj : R(λj ) > α. (5) Rejecting Hj certifies λj as safe with respect to the target risk level. Since the final threshold pair will be selected adaptively from all certified candidates, we require family-wise error rate (FWER) (Holm, b denote the certi1979) control at level δ. Letting Λ fied set, FWER control ensures ! P sup R(λ) ≤ α
≥ 1 − δ.
(6)
b λ∈Λ
b including the Therefore, any later choice from Λ, acceptance-maximizing one, remains valid. A Bonferroni correction also provides this guarantee, but it is often overly conservative on dense threshold grids. BalanceRAG instead uses SGT (Bretz et al., 2009), which preserves FWER control while exploiting the lattice structure to recycle significance budget to more useful operating points. Step 1: p-value computation. On the calibration set Dcal = {(xℓ , yℓ⋆ )}nℓ=1 , for each threshold pair (i) (j) λi,j = (t1 , t2 ), we compute the accepted count and accepted-error count: n n X X Mi,j = Sλi,j (xℓ ), Ki,j = Zλi,j (xℓ , yℓ⋆ ).
Given a prompt x, BalanceRAG considers two inference branches: an LLM-only branch and an RAG fallback branch. Branch b ∈ {1, 2} returns an answer ŷb (x) together with an uncertainty score ub (x), where smaller values indicate higher confidence. Given a threshold pair λ = (t1 , t2 ), the router first accepts the LLM-only answer if u1 (x) ≤ t1 . Otherwise, it invokes the RAG fallback and accepts ŷ2 (x) if u2 (x) ≤ t2 . If both checks fail, the system abstains by returning ∅. Thus, the final output ŷλ (x) ∈ {ŷ1 (x), ŷ2 (x), ∅} is determined by the two thresholds jointly.
ℓ=1
3
ℓ=1
(7)
t1(0) t2(0)
δ
t2(1)
0
1 2
1 2
⋯
0 1 2
t2(N)
0
t1(1) 1 2
1 2
1 2
1
0 1 2
0 1 2
0 1 2
0
t1(M)
⋯ 1 2
1 2
1 2
1
1 2
0 1 2
0 1 2
0 1 2
0
1 2
1 2
1
0
t1(0) t2(0)
δ
t2(1)
0
1 2
1
0
2 3
1
⋯
0
0 3 4
1
0
t2(N)
0
t1(1) 1 2
1 3
1 4
1
0 1 3
0 2 4
0 3 5
0
t1(M)
⋯ 2 3
2 4
2 5
1
3 4
0 1 4
0 2 5
0 3 6
0
3 5
3 6
1
DWD δi,j =
0 1
1
0 1
0
Figure 3: Budget diffusion on the threshold lattice. UGD splits budget evenly across successors, while DWD gives equal accumulated budget to nodes on the same diagonal.
pi,j ≤ δi,j .
δu ← δu + δi,j g(i,j)→u .
(8)
gk,u ← gk,u + gk,(i,j) g(i,j),u ,
(9)
λ̂ = arg max M (λ).
(10)
(17)
b λ∈Λ
As the safe region varies across datasets and models, we split off a small portion of the calibration set as an initialization split. On this split, we choose the highest-acceptance node among those satisfying p ≤ δ; if no such node exists, we choose the node with the smallest p-value. The selected node receives the initial budget δ, while all formal SGT p-values are computed on the remaining calibration data. As shown in Figure 3, we consider two diffusion rules over successors u ∈ N (i, j). Uniform Grid Diffusion (UGD) assigns 1 UGD g(i,j)→u = . (11) |N (i, j)|
Since this selection is restricted to the certified set, it inherits the simultaneous risk guarantee in Eq. (6). Theorem 1 (Risk control of calibrated routing). Under i.i.d. calibration data, for any target risk α and error level δ, any threshold pair λ̂ selected from the SGT-certified set satisfies P R(λ̂) ≤ α ≥ 1 − δ. (18) The proof is given in Appendix A.3. 3.3
Multi-Risk Control
The same certification procedure can control multiple system-level risks. Besides the selectionconditioned answer error,
Diamond-Wavefront Diffusion (DWD) assigns (i + 1)1[u = (i + 1, j)] i+j+2 (j + 1)1[u = (i, j + 1)] + . i+j+2
(16)
where k is a predecessor and u is a successor. This budget recycling preserves validity while giving later, more permissive candidates greater power to be certified. Step 4: final threshold selection. After SGT terminates, BalanceRAG selects the certified threshold pair with the largest accepted count:
which is super-uniform under Hλi,j . Step 2: graph prior and budget diffusion. The candidate thresholds form a natural partial order: increasing either threshold makes the router more permissive and usually increases acceptance. We therefore direct edges from each node to its immediate relaxed neighbors:
DWD g(i,j)→u =
(15)
Because the graph is acyclic, SGT can apply the corresponding rewiring update
We therefore use the one-sided exact binomial pvalue
N (i, j) ⊆ {(i + 1, j), (i, j + 1)}.
(14)
Once node (i, j) is certified, its budget is transferred to uncertified successors:
For a fixed candidate λi,j , conditioned on Mi,j ,
pi,j = P(Bin(Mi,j , α) ≤ Ki,j ) ,
(13)
so threshold pairs with the same total relaxation receive the same testing opportunity. Appendix A.2 proves this property. Step 3: SGT budget propagation. SGT maintains a local testing budget δi,j at each active node and certifies the node whenever
0
Ki,j | Mi,j ∼ Bin(Mi,j , R(λi,j )) .
δ , d+1
R(1) (λ) = P (Zλ = 1 | Sλ = 1) , (12)
(19)
we also control the fallback invocation rate, R(2) (λ) = P(Bλ = 1),
Under DWD, each node on diagonal i + j = d receives accumulated budget
(20)
where Bλ indicates whether the input is routed to the RAG fallback. Given targets α1 and α2 , each 4
node receives one p-value per risk, denoted by (1) (2) pi,j and pi,j . The intersection-union rule (Berger, 1982) combines them as (1) (2) pi,j = max pi,j , pi,j . (21)
Evaluation Metrics. We consider four dimensions. Reliability is measured by the empirical selectionconditioned error rate (Err.) under different target risk levels α. Coverage is measured by the acceptance rate (Cov.). Utility is measured by the number of accepted correct samples (Corr.), which captures how many useful answers the system produces. Cost is measured by the average number of tokens and the average memory increase per sample after uncertainty scores are obtained, so it reflects the execution cost induced by the calibrated routing policy rather than the overhead of any specific uncertainty estimator. Correctness Criterion and Uncertainty Estimator. By default, we use semantic similarity (Reimers and Gurevych, 2019) with a threshold of 0.7 to determine whether a model answer aligns with the ground truth in the admission function. In the robustness analysis, we further consider bidirectional entailment (Kuhn et al., 2023) and LLM-as-a-Judge (Zhang et al., 2024) as alternative correctness criteria. For uncertainty estimation, we consider five estimators: semantic entropy (SE) (Farquhar et al., 2024), the sum of graph Laplacian eigenvalues (EIGV), eccentricity (ECC), the degree matrix statistic (DEG) (Lin et al., 2023), and length-normalized predictive entropy from the model outputs (SELF) (Malinin and Gales, 2020). Unless otherwise stated, we use SE as the default uncertainty estimator. Hyperparameters. Following standard splitcalibration protocols (Angelopoulos et al., 2024), we use a calibration-test split ratio of 0.5 and repeat each experiment over 100 random calibration-test splits. In each calibration set, we reserve 40% of the examples for initial-node selection and use the remaining data for formal SGT calibration.
Feeding this node-level p-value into the same SGT procedure yields P R(1) (λ̂) ≤ α1 , R(2) (λ̂) ≤ α2 ≥ 1−δ. (22) Thus, multi-risk BalanceRAG jointly certifies answer reliability and fallback usage, allowing the final router to trade off coverage and retrieval frequency under explicit risk constraints.
4
Experiment
4.1
Experimental Settings
Datasets and Base LLMs. We evaluate BalanceRAG on three open-ended QA datasets: TriviaQA (Joshi et al., 2017), SQuAD v2 (Rajpurkar et al., 2018), and Natural Questions (NQ) (Kwiatkowski et al., 2019). To assess robustness across model families and scales, we employ eight backbone LLMs from Qwen (Yang et al., 2025, 2024), LLaMA (Grattafiori et al., 2024), and Vicuna (Zheng et al., 2023). Baselines. We compare BalanceRAG with three groups of baselines. First, we include fixed-path systems, LLM-only and LLM-RAG, both calibrated by fixed sequence testing (Bauer, 1991) as single-branch risk-controlled baselines. Second, we also evaluate two BalanceRAG variants: BalanceRAG-UGD, which uniformly diffuses the significance budget to neighboring nodes, and BalanceRAG-DWD, which uses diamondwavefront diffusion to balance testing opportunities among threshold pairs with the same total relaxation. We employ DWD as the default variant. Third, we compare with UCB-based methods for selection-conditioned risk control. UCB-CascadedCP adapts the cascaded judge framework of Jung et al. (2024) to our LLM-only/RAG routing setting with the Clopper-Pearson bound, while UCBCascaded-HFD replaces it with a distribution-free Hoeffding bound. For completeness, Appendix 6 reports additional baselines, including AdaptiveRAG (Jeong et al., 2024), Self-Route (Li et al., 2024), empirical risk selection, and Bonferroni correction, covering uncalibrated adaptive routing, simple empirical selection, and conservative multiple-testing correction under the same target risk levels.
4.2
Utility and Coverage under Risk Control
Reliability under Target Risk Control. We first examine whether BalanceRAG achieves systemlevel risk control at user-specified target levels. Figure 4 reports the empirical test-set selectionconditioned error rate across backbones and target risk levels, where solid curves denote the mean over repeated calibration/test splits and shaded regions show one standard deviation. BalanceRAG provides a finite-sample, high-probability guarantee that the population selection-conditioned risk of the calibrated router is at most the target level α. The empirical test-set Err. closely tracks the target line, which is consistent with this guarantee. 5
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.150
90 Cov.
0.125
75 Err.
Cov.
Err.
0.150
0.175
100
0.125
75
50 0.100
0.100 0.12 0.14 0.16 Target Risk Level α
25
0.18
60 0.10
(a) Qwen2.5-3B
0.10
Err.
Cov.
90
0.125
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(b) Qwen2.5-7B
0.175
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.150
0.18
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.150 Err.
0.175
0.12 0.14 0.16 Target Risk Level α
75
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(c) Qwen3-4B
0.125
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.18
0.10
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.18
0.12 0.14 0.16 Target Risk Level α
0.18
Cov.
100
0.12
80
60
40
0.100
0.09 0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(e) Llama3.2-3B
0.18
0.10
0.18
Err.
Cov.
75
50
100
0.12 0.09
0.12 0.14 0.16 Target Risk Level α
25
0.18
80
60
0.08 0.10
0.10
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.15
0.12
0.12 0.14 0.16 Target Risk Level α
(f) Llama3.1-8B
100
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.16
0.12 0.14 0.16 Target Risk Level α
Cov.
0.10
Err.
0.12 0.14 0.16 Target Risk Level α
(d) Qwen3-14B
60
0.12 0.14 0.16 Target Risk Level α
70 0.10
0.15
0.125
0.18
80
80 Cov.
Err.
0.150
0.18
Err.
0.175
0.12 0.14 0.16 Target Risk Level α
0.12 0.14 0.16 Target Risk Level α
90
60 0.10
0.18
100
0.100 0.100
0.12 0.14 0.16 Target Risk Level α
Cov.
0.10
0.06 0.10
(g) Vicuna-v1.5-7B
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(h) Vicuna-v1.5-13B
Figure 4: Err. (left in each panel pair) and Cov. (right) under different target risk levels α across eight LLMs spanning four model families on TriviaQA. 90% quantile
15 0 20 10 0
Target Risk Level α
Density
Density
Target Risk Level α
30 15 0 30
It outperforms non-RAG baselines and matches or exceeds RAG in many settings. Notably, BalanceRAG accepts more test samples and retains more correctly accepted predictions (Corr.), while maintaining empirical Err. near or below the risk level. This advantage reflects the heterogeneous value of retrieval: retrieval can help knowledgeintensive cases, but may also introduce noisy or distracting context for easier queries. By jointly calibrating the LLM-only and fallback RAG branches, BalanceRAG keeps cases where the LLM-only answer is sufficient on the primary path and routes harder cases to LLM-RAG when the fallback is sufficiently reliable. Appendix 6 further compares BalanceRAG with additional baselines under the same target risk levels. BalanceRAG is more reliable than Adaptive-RAG, Self-Route, and empirical risk selection, while achieving higher coverage than Bonferroni. Overall, these results show that joint calibration improves the reliability-coverageutility trade-off under fixed risk control.
90% quantile
30 15 0 30 15 0 30 15
0.08
0.10
0.12 E RR .
0.14
0.16
0
0.10
0.12
0.14 E RR .
0.16
0.18
Figure 5: Err. distributions for Qwen2.5-3B (Left) and Llama-3.2-3B (Right) on TriviaQA.
Figure 5 further shows the distribution of test-set Err. over repeated random splits. When δ = 0.1, the fraction of splits with test-set Err. exceeding α remains below 10% at every target level. Moreover, joint calibration keeps BalanceRAG feasible under strict risk targets, whereas branch-wise baselines such as LLM-only and UCB-Cascaded often fail to find a feasible operating point due to branch-level constraints and stage-wise δ splitting. Coverage and Utility under Risk Control. We next evaluate how much useful output is retained at the same risk level α. Figure 4 and Table 1 show that BalanceRAG consistently achieves the best or near-best coverage and utility among the baselines. 6
Table 1: Results on TriviaQA, SQuAD v2, and NQ under different target risk levels. Bold numbers indicate the best Corr. under the same dataset, model, and risk level. “-” denotes infeasibility, i.e., no non-empty certified operating point is found under the target risk level. Method
Metric
TriviaQA
LLM-RAG UCB-Cascaded BalanceRAG
0.11 988.0 0.0955 1390.0 0.0949 1028.0 0.0941 1643.0 0.099
0.12 1116.0 0.1070 1532.0 0.1075 1488.0 0.1037 1774.0 0.1091
0.13 1214.0 0.1175 1646.0 0.1192 1749.0 0.1113 1873.0 0.1201
0.14 1283.0 0.1274 1760.0 0.1302 1850.0 0.1214 1973.0 0.1287
0.15 1371.0 0.1373 1876.0 0.1407 1917.0 0.1320 2073.0 0.1384
0.11 1780.0 0.1004 2036.0 0.1006 1849.0 0.0978 2288.0 0.0991
0.12 1928.0 0.1108 2198.0 0.1110 2084.0 0.1072 2430.0 0.1084
0.13 2070.0 0.1215 2382.0 0.1214 2264.0 0.1178 2584.0 0.1206
0.14 2165.0 0.1318 2529.0 0.1314 2360.0 0.1272 2691.0 0.1316
0.15 2218.0 0.1415 2652.0 0.1413 2434.0 0.1367 2727.0 0.1396
Corr. Err. Corr. Err. Corr. Err. Corr. Err.
0.11 1546.0 0.1008 1484.0 0.0975
0.12 1617.0 0.1112 1591.0 0.1066
0.13 1666.0 0.1216 1654.0 0.1164
0.14 1714.0 0.1316 1699.0 0.1271
0.15 137.0 0.1345 1755.0 0.1416 1744.0 0.1354 1745.0 0.1375
0.11 1803.0 0.1000 1784.0 0.1010
0.12 1831.0 0.1095 1819.0 0.1118
0.13 1850.0 0.1175 1836.0 0.1210
0.14 194.0 0.1319 1853.0 0.1193 1835.0 0.1263 1842.0 0.1241
0.15 211.0 0.1324 1853.0 0.1193 1837.0 0.1269 1843.0 0.1240
Corr. Err. Corr. Err. Corr. Err. Corr. Err.
0.15 0.16 0.17 119.0 0.1568 489.0 532.0 562.0 0.1283 0.1390 0.1497 447.0 492.0 531.0 0.122 0.1312 0.1404
0.18 125.0 0.1552 590.0 0.1601 563.0 0.1520 575.0 0.1506
0.19 134.0 0.1572 615.0 0.1697 590.0 0.1582 615.0 0.1662
0.15 158.0 0.1055 734.0 0.1331 698.0 0.1214 709.0 0.1315
0.16 191.0 0.1184 752.0 0.1422 718.0 0.1307 741.0 0.1420
0.17 240.0 0.1362 760.0 0.1466 736.0 0.1408 753.0 0.1511
0.18 283.0 0.1493 763.0 0.1478 743.0 0.1492 757.0 0.1523
0.19 320.0 0.1630 763.0 0.1478 747.0 0.1581 759.0 0.1517
α
SQuAD v2
LLM-only LLM-RAG UCB-Cascaded BalanceRAG α LLM-only
NQ
LLM-RAG UCB-Cascaded BalanceRAG
4.3
Llama3.1-8B
Corr. Err. Corr. Err. Corr. Err. Corr. Err.
α LLM-only
Llama3.2-3B
brated routing policy rather than the overhead of a particular uncertainty-estimation method. Under the same backbone and decoding setup, LLMonly and RAG perform fixed single-path inference, whereas BalanceRAG first attempts the LLM-only branch and invokes the RAG fallback only when the primary answer does not pass its calibrated threshold. We report the average memory increase and average token usage per sample. As shown in Figure 7, BalanceRAG lies between the two fixedpath endpoints: it incurs more cost than LLM-only, but substantially lower cost than always-on RAG, because many queries can be accepted directly by the primary branch without invoking retrieval.
Routing Allocation and Cost Efficiency
Routing Allocation across Models. Figure 6 shows the routing allocation of BalanceRAG across eight backbones. BalanceRAG consistently keeps a substantial fraction of queries on the LLMonly branch, thereby reducing unnecessary fallback usage while maintaining the prescribed risk control. The LLM-only share also reflects the strength of the backbone: stronger models can safely answer more queries without retrieval, whereas weaker models require more frequent fallback to satisfy the same risk target. In our results, stronger backbones keep more than 60% of accepted queries on the LLM-only branch, medium-strength backbones stay around 40-50%, and weaker backbones remain around 25-30%.
4.4
Multi-Risk Control
In deployment, practitioners may want to control not only the selection-conditioned answer risk but also the frequency of fallback retrieval. We therefore extend BalanceRAG to the multi-risk setting by adding a constraint α2 on the fallback invocation rate in addition to the answer-risk target α1 . Table 2 shows that stricter α2 values lead Bal-
Routed Execution Cost. We further evaluate the routed execution cost of BalanceRAG at calibrated operating points. Since BalanceRAG does not depend on a specific uncertainty estimator, we measure cost after uncertainty scores are obtained; this isolates the execution cost induced by the cali7
LLM-only
80 60 40 20 0.12
0.14 0.16 Target Risk α
60 40 20 0
0.18
0.12
LLM-only
80 60 40 20 0.12
0.14 0.16 Target Risk α
LLM-RAG
40 20 0.12
LLM-only
60 40 20 0.12
(e) Llama3.2-3B
0.14 0.16 Target Risk α
0.14 0.16 Target Risk α
60 40 20
(f) Llama3.1-8B
0.14 0.16 Target Risk α
0.18
(d) Qwen3-14B LLM-only
80 60 40 20 0.12
0.12
Rejected
100
0
0.18
LLM-RAG
Rejected
80
0
0.18
LLM-RAG
100
(c) Qwen3-4B
80
0
LLM-only
60
Rejected
100
0.18
Rejected
80
0
0.18
Sample Allocation (%)
100
0
0.14 0.16 Target Risk α
LLM-RAG
100
(b) Qwen2.5-7B
Rejected
Sample Allocation (%)
Sample Allocation (%)
LLM-RAG
LLM-only
80
(a) Qwen2.5-3B LLM-only
Rejected
Sample Allocation (%)
0
LLM-RAG
100
Sample Allocation (%)
Rejected
Sample Allocation (%)
LLM-RAG
Sample Allocation (%)
Sample Allocation (%)
LLM-only
100
0.14 0.16 Target Risk α
0.18
(g) Vicuna-v1.5-7B
LLM-RAG
Rejected
100 80 60 40 20 0
0.12
0.14 0.16 Target Risk α
0.18
(h) Vicuna-v1.5-13B
LLM-only
LLM-RAG
BalanceRAG
Tokens per Sample
Avg. Increase Memory (MB)
Figure 6: Routing allocation of test samples in BalanceRAG on TriviaQA (mean), including LLM-only acceptance, LLM-RAG fallback acceptance, and rejection. 600 400 200 0 0.10
0.11
0.12
0.13
0.14
0.15
0.16
£10
LLM-only
0 0.10
0.17
0.11
0.12
3
LLM-only
LLM-RAG
0.5 0.0 0.11
0.12
0.13
0.14
0.14
0.15
0.16
0.17
(b) Qwen2.5-7B Tokens
BalanceRAG
1.0
0.10
0.13
Risk Level ®
Tokens per Sample
£10
Table 2: Single-risk vs. multi-risk routing at selectionconditioned risk target α1 = 0.15. α2 caps the fallback rate to LLM-RAG; “Route” shows the LLM-only/LLMRAG split.
BalanceRAG
1
Risk Level ®
1.5
LLM-RAG
2
(a) Qwen2.5-7B Memory Avg. Increase Memory (MB)
3
0.15
0.16
0.17
Risk Level ®
(c) Llama-3.1-8B Memory
£10
3
LLM-only
LLM-RAG
BalanceRAG
Method
α2
Route (Only/RAG)
Cov.
Corr.
Err.
Qwen2.5-3B
LLM-only Single-risk Multi-risk Multi-risk
0.6 0.5
1.0000 / 0.0000 0.2995 / 0.6505 0.3464 / 0.5620 0.3859 / 0.4709
0.4432 0.9499 0.9084 0.8568
1137 2421 2332 2206
0.1341 0.1400 0.1338 0.1313
Qwen2.5-7B
LLM-only Single-risk Multi-risk Multi-risk
0.3 0.2
1.0000 / 0.0000 0.6414 / 0.3582 0.7559 / 0.2438 0.8257 / 0.1729
0.9105 0.9996 0.9996 0.9986
2422 2738 2719 2692
0.1412 0.1159 0.1221 0.1298
Llama3.1-8B
LLM-only Single-risk Multi-risk Multi-risk
0.3 0.2
1.0000 / 0.0000 0.5935 / 0.4024 0.7200 / 0.2593 0.7748 / 0.1702
0.8118 0.9959 0.9793 0.9450
2218 2727 2689 2599
0.1415 0.1396 0.1372 0.1357
Llama3.2-3B
LLM-only Single-risk Multi-risk Multi-risk
0.3 0.2
1.0000 / 0.0000 0.4162 / 0.3539 0.4872 / 0.2504 0.5320 / 0.1279
0.5092 0.7701 0.7376 0.6599
1371 2073 1989 1781
0.1373 0.1384 0.1366 0.1355
2
1
0 0.10
0.11
0.12
0.13
0.14
0.15
0.16
0.17
Risk Level ®
(d) Llama-3.1-8B Tokens
Figure 7: Memory and token cost on TriviaQA across risk levels for Qwen2.5-7B and Llama-3.1-8B
anceRAG to certify operating points with lower fallback rates, shifting more queries to the cheaper LLM-only branch. This reduces retrieval usage, with only a moderate decrease in coverage and correctly accepted samples. At the same time, both risks remain controlled in all reported configurations: the empirical selection-conditioned risk stays below α1 = 0.15, and the fallback invocation rate stays below α2 . These results show that multi-risk BalanceRAG provides an explicit mechanism for navigating the reliability-utility-retrieval trade-off.
5
Model
ing finite-sample, high-probability control of the selection-conditioned error rate. The framework also supports multi-risk calibration to bound fallback retrieval usage together with answer risk. Across open-domain QA benchmarks and diverse backbones, BalanceRAG meets prescribed risk levels, preserves high coverage and accepted correct answers, and reduces unnecessary retrieval calls. These results suggest a principled path from heuristic routing to reliable cascaded LLM-RAG.
Conclusion
Limitations
In this paper, we develop BalanceRAG, a trainingfree framework for joint threshold calibration in risk-controlled cascaded LLM-RAG systems. By certifying candidate threshold pairs on a twodimensional lattice with SGT, BalanceRAG selects high-utility operating points while provid-
BalanceRAG relies on the standard splitcalibration assumption that the calibration and deployment examples are drawn from the same distribution. If the query distribution, retrieval corpus, or backbone model changes substantially, 8
the calibrated thresholds may need to be updated. This is a common limitation of finite-sample calibration methods and can be addressed in practice through periodic recalibration. The guarantee is also defined with respect to the chosen correctness criterion. Although we evaluate semantic similarity, bidirectional entailment, and LLM-as-a-Judge variants, open-domain QA correctness can still be ambiguous, especially for partially correct or underspecified answers. This reflects a broader challenge in evaluating free-form generation rather than a limitation specific to BalanceRAG. Finally, our experiments focus on two-branch LLM-RAG cascades. Extending the calibration principle to multi-stage retrieval, tool-augmented systems, or agentic workflows is a promising direction, but may require more scalable search and calibration over larger routing spaces.
language models using semantic entropy. Nature, 630(8017):625–630. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad AlDahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, and 1 others. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783. Yu Gui, Ying Jin, and Zhimei Ren. 2024. Conformal alignment: Knowing when to trust foundation models with guarantees. Advances in Neural Information Processing Systems, 37:73884–73919. Sture Holm. 1979. A simple sequentially rejective multiple test procedure. Scandinavian journal of statistics, pages 65–70. Huipeng Huang, Wenbo Liao, Huajun Xi, Hao Zeng, Mengchen Zhao, and Hongxin Wei. 2025. Selective labeling with false discovery rate control. arXiv preprint arXiv:2510.14581. Yue Huang, Lichao Sun, Haoran Wang, Siyuan Wu, Qihui Zhang, Yuan Li, Chujie Gao, Yixin Huang, Wenhan Lyu, Yixuan Zhang, Xiner Li, Hanchi Sun, Zhengliang Liu, Yixin Liu, Yijue Wang, Zhikun Zhang, Bertie Vidgen, Bhavya Kailkhura, Caiming Xiong, and 52 others. 2024. Position: TrustLLM: Trustworthiness in large language models. In Proceedings of the 41st International Conference on Machine Learning, pages 20166–20270.
References Anastasios Angelopoulos, Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster. 2024. Conformal risk control. In International Conference on Learning Representations, volume 2024, pages 55198–55218.
Soyeong Jeong, Jinheon Baek, Sukmin Cho, Sung Ju Hwang, and Jong C Park. 2024. Adaptive-rag: Learning to adapt retrieval-augmented large language models through question complexity. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 7036–7050.
Akari Asai, Zeqiu Wu, Yizhong Wang, Avi Sil, and Hannaneh Hajishirzi. 2024. Self-rag: Learning to retrieve, generate, and critique through self-reflection. In International Conference on Learning Representations, volume 2024, pages 9112–9141. Peter Bauer. 1991. Multiple testing in clinical trials. Statistics in Medicine, 10(6):871–890.
Zijun Jia, Diyin Tang, Hongyu Long, and Jinsong Yu. 2025. Coverage-guaranteed speech emotion recognition via calibrated uncertainty-adaptive prediction sets. Engineering Applications of Artificial Intelligence, 159:111721.
Roger L. Berger. 1982. Multiparameter hypothesis testing and acceptance sampling. Technometrics, 24(4):295–300. Frank Bretz, Willi Maurer, Werner Brannath, and Martin Posch. 2009. A graphical approach to sequentially rejective multiple test procedures. Statistics in Medicine, 28(4):586–604.
Zhengbao Jiang, Frank F Xu, Luyu Gao, Zhiqing Sun, Qian Liu, Jane Dwivedi-Yu, Yiming Yang, Jamie Callan, and Graham Neubig. 2023. Active retrieval augmented generation. In Proceedings of the 2023 conference on empirical methods in natural language processing, pages 7969–7992.
Jinhao Duan, Xinyu Zhao, Zhuoxuan Zhang, Eunhye Grace Ko, Lily Boddy, Chenan Wang, Tianhao Li, Alexander Rasgon, Junyuan Hong, Min Kyung Lee, Chenxi Yuan, Qi Long, Ying Ding, Tianlong Chen, and Kaidi Xu. 2025. GuideLLM: Exploring LLM-guided conversation with applications in autobiography interviewing. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 5558–5588.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. 2017. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1601–1611. Jaehun Jung, Faeze Brahman, and Yejin Choi. 2024. Trust or escalate: Llm judges with provable guarantees for human agreement. arXiv preprint arXiv:2407.18370.
Sebastian Farquhar, Jannik Kossen, Lorenz Kuhn, and Yarin Gal. 2024. Detecting hallucinations in large
9
Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020. Dense passage retrieval for opendomain question answering. In Proceedings of the 2020 conference on empirical methods in natural language processing (EMNLP), pages 6769–6781.
Shrey Pandit, Jiawei Xu, Junyuan Hong, Zhangyang Wang, Tianlong Chen, Kaidi Xu, and Ying Ding. 2025. Medhallu: A comprehensive benchmark for detecting medical hallucinations in large language models. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages 2858–2873.
Lorenz Kuhn, Yarin Gal, and Sebastian Farquhar. 2023. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. arXiv preprint arXiv:2302.09664.
Pranav Rajpurkar, Robin Jia, and Percy Liang. 2018. Know what you don’t know: Unanswerable questions for squad. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 784–789.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, and 1 others. 2019. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453–466.
Nils Reimers and Iryna Gurevych. 2019. Sentence-bert: Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing (EMNLP-IJCNLP), pages 3982–3992.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, and 1 others. 2020. Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in neural information processing systems, 33:9459–9474.
Weihang Su, Yichen Tang, Qingyao Ai, Zhijing Wu, and Yiqun Liu. 2024. Dragin: Dynamic retrieval augmented generation based on the real-time information needs of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 12991–13013.
Xiaoxi Li, Guanting Dong, Jiajie Jin, Yuyao Zhang, Yujia Zhou, Yutao Zhu, Peitian Zhang, and Zhicheng Dou. 2025. Search-o1: Agentic search-enhanced large reasoning models. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages 5420–5438.
Binyu Tan, Zhiyuan Wang, Jinhao Duan, Kaidi Xu, Heng Tao Shen, Xiaoshuang Shi, and Fumin Shen. 2025. Conformal lesion segmentation for 3d medical images. arXiv preprint arXiv:2510.17897. Xiaqiang Tang, Qiang Gao, Jian Li, Nan Du, Qi Li, and Sihong Xie. 2025. Mba-rag: a bandit approach for adaptive retrieval-augmented generation through question complexity. In Proceedings of the 31st International Conference on Computational Linguistics, pages 3248–3254.
Ye Li, Anqi Hu, Yuanchang Ye, Shiyan Tong, Zhiyuan Wang, and Bo Fu. 2026. Set-valued prediction for large language models with feasibility-aware coverage guarantees. arXiv preprint arXiv:2603.22966. Zhuowan Li, Cheng Li, Mingyang Zhang, Qiaozhu Mei, and Michael Bendersky. 2024. Retrieval augmented generation or long-context llms? a comprehensive study and hybrid approach. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing: Industry Track, pages 881– 893.
Hairu Wang, Yuan Feng, Yukun Cao, Xike Xie, and S Kevin Zhou. 2025a. Skewroute: Training-free llm routing for knowledge graph retrieval-augmented generation via score skewness of retrieved context. In Findings of the Association for Computational Linguistics: EMNLP 2025, pages 11324–11340.
Zhen Lin, Shubhendu Trivedi, and Jimeng Sun. 2023. Generating with confidence: Uncertainty quantification for black-box large language models. arXiv preprint arXiv:2305.19187.
Qingni Wang, Tiantian Geng, Zhiyuan Wang, Teng Wang, Bo Fu, and Feng Zheng. 2025b. Sample then identify: A general framework for risk control and assessment in multimodal large language models. In International Conference on Learning Representations, volume 2025, pages 64280–64297.
Andrey Malinin and Mark Gales. 2020. Uncertainty estimation in autoregressive structured prediction. arXiv preprint arXiv:2002.07650.
Xiaohua Wang, Zhenghua Wang, Xuan Gao, Feiran Zhang, Yixin Wu, Zhibo Xu, Tianyuan Shi, Zhengyuan Wang, Shizheng Li, Qi Qian, and 1 others. 2024a. Searching for best practices in retrievalaugmented generation. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 17716–17736.
Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Hajishirzi. 2023. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. In Proceedings of the 61st annual meeting of the association for computational linguistics (volume 1: Long papers), pages 9802–9822.
Zhiyuan Wang, Tianlong Chen, Yue Zhang, Heng Tao Shen, Xiaoshuang Shi, Kaidi Xu, and 1 others.
10
2025c. Lec: Linear expectation constraints for falsediscovery control in selective prediction and routing systems. arXiv preprint arXiv:2512.01556. Zhiyuan Wang, Jinhao Duan, Lu Cheng, Yue Zhang, Qingni Wang, Xiaoshuang Shi, Kaidi Xu, Heng Tao Shen, and Xiaofeng Zhu. 2024b. ConU: Conformal uncertainty in large language models with correctness coverage guarantees. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 6886–6898. Zhiyuan Wang, Jinhao Duan, Qingni Wang, Xiaofeng Zhu, Tianlong Chen, Xiaoshuang Shi, and Kaidi Xu. 2026. Coin: Uncertainty-guarding selective question answering for foundation models with provable risk guarantees. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, pages 33764– 33772. Zhiyuan Wang, Jinhao Duan, Chenxi Yuan, Qingyu Chen, Tianlong Chen, Yue Zhang, Ren Wang, Xiaoshuang Shi, and Kaidi Xu. 2025d. Word-sequence entropy: Towards uncertainty estimation in free-form medical question answering applications and beyond. Engineering Applications of Artificial Intelligence, 139:109553. Zhiyuan Wang, Qingni Wang, Yue Zhang, Tianlong Chen, Xiaofeng Zhu, Xiaoshuang Shi, and Kaidi Xu. 2025e. SConU: Selective conformal uncertainty in large language models. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 19052– 19075. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and 1 others. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388. An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, and 1 others. 2024. Qwen2. 5 technical report. arXiv preprint arXiv:2412.15115. Ruiyang Zhang, Hu Zhang, and Zhedong Zheng. 2024. Vl-uncertainty: Detecting hallucination in large vision-language model via uncertainty estimation. arXiv preprint arXiv:2411.11919. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, and 1 others. 2023. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances in neural information processing systems, 36:46595–46623.
11
A
Proofs
This appendix provides the proofs deferred from Section 3. Throughout, {(xℓ , yℓ⋆ }nℓ=1 is the i.i.d. calibration set, and for any threshold pair λ = (t1 , t2 ) we use the notation M (λ), K(λ), R(λ) from Section 3.1. We also write Fm,q for the CDF of the binomial distribution Bin(m, q). A.1
Super-uniformity of the binomial p-value
Lemma 1 (Super-uniformity). For each lattice node λi,j , let (23)
pi,j = FM (λi,j ), α K(λi,j ) . Under the null hypothesis Hλi,j : R(λi,j ) ≥ α, pi,j is super-uniform, i.e. for every t ∈ [0, 1], PHλi,j pi,j ≤ t ≤ t.
(24)
Proof. Fix λ = λi,j and abbreviate M = M (λ), K = K(λ), R = R(λ). Conditionally on M , the samples accepted by λ are i.i.d. Bernoulli with error probability R, hence (25)
K | M ∼ Bin(M, R).
Under Hλ we have R ≥ α, which implies the stochastic ordering Bin(M, R) ⪰st Bin(M, α); equivalently, for every integer k, FM,R (k) ≤ FM,α (k). (26) By the generalized probability integral transform applied to the (possibly discrete) CDF FM,R , for every t ∈ [0, 1], P FM,R (K) ≤ t | M ≤ t. (27) Combining Eq. 26 and Eq. 27 gives P(pi,j ≤ t | M ) = P FM,α (K) ≤ t | M
≤ P FM,R (K) ≤ t | M
(28)
≤ t. Marginalizing over M yields PHλi,j (pi,j ≤ t) ≤ t. A.2
Shell-equal budget property of DWD
We record here a structural property of DWD that motivates its design (cf. Eq. 13 in the main text). The result is self-contained and does not depend on any hypothesis-testing content; while it is not used in the proof of Theorem 1, it provides useful intuition for the behavior of DWD on the lattice. Lemma 2 (Shell-equal budget). Let the lattice Λ have root (0, 0) with initial budget δ0,0 = δ and all other budgets set to zero. If the edge weights follow the DWD rule in Eq. 12 and each node forwards its budget to neighbors once reached, then every node (i, j) with i + j = d accumulates DWD δi,j =
δ . d+1
(29)
Proof. We prove Eq. 29 by induction on d. DWD = δ = δ/(0 + 1). Base case (d = 0). The only node is (0, 0) itself, which trivially receives δ0,0 Inductive step. Assume the claim holds on diagonal d. Let (a, b) be any node on diagonal d + 1, so a + b = d + 1. Budget reaches (a, b) only from its in-lattice parents (a − 1, b) and (a, b − 1). Using Eq. 12, a a = , (a − 1) + b + 2 d+2 b b DWD g(a,b−1)→(a,b) = = . a + (b − 1) + 2 d+2 DWD g(a−1,b)→(a,b) =
12
(30) (31)
By the inductive hypothesis, both parents carry budget δ/(d + 1). Summing the incoming contributions, δ a δ b · + · d+1 d+2 d+1 d+2 a+b δ · = d+1 d+2 δ d+1 δ = · = , d+1 d+2 d+2
DWD δa,b =
(32)
where we used a + b = d + 1 in the penultimate step. For boundary nodes with only one in-lattice parent (a = 0 or b = 0), the missing term in Eq. 32 has numerator zero, and the “sole neighbor receives weight 1” convention makes the lone parent contribute δ/(d + 1) · 1 = δ/(d + 2) after the same algebra. Hence every node on diagonal d + 1 accumulates δ/(d + 2), completing the induction. Remark 1. Lemma 2 implies that the total budget on each diagonal is conserved: summing δ/(d + 1) over the d + 1 nodes of diagonal d yields δ. A.3
Proof of Theorem 1
Proof of Theorem 1. The proof combines two ingredients: the validity of the binomial p-value (Lemma 1) and the FWER control of SGT. (i) Valid p-values. By Lemma 1, the p-value pi,j defined in Eq. 9 is super-uniform under the null Hλi,j for every lattice node. (ii) FWER control by SGT. Both the UGD weights (Eq. 11) and the DWD weights (Eq. 12) are P non-negative, satisfy g(i,j)→(i,j) = 0, and obey u∈N (i,j) g(i,j)→u ≤ 1 at every node. Hence, the proposed graph P satisfies the regularity conditions of Bretz et al. (2009). Combined with the initial budget allocation (i,j) δi,j ≤ δ and the super-uniform p-values from (i), Algorithm 1 of Bretz et al. (2009) yields a sequentially rejective multiple testing procedure that strongly controls the family-wise error rate b returned by at level δ. Therefore, with the null hypotheses defined as Hλ : R(λ) > α, the certified set Λ the procedure satisfies P sup R(λ) ≤ α ≥ 1 − δ, (33) b λ∈Λ
which recovers Eq. 6. b to λ̂. Let E := {sup b R(λ) ≤ α}; by (ii), P(E) ≥ 1 − δ. On E, the selection rule in (iii) From Λ λ∈Λ b so Eq. 17 picks λ̂ ∈ Λ, R(λ̂) ≤ sup R(λ) ≤ α. (34) b λ∈Λ
Hence E ⊆ {R(λ̂) ≤ α}, and P R(λ̂) ≤ α
≥ P(E) ≥ 1 − δ.
(35)
Remark 2. The proof does not rely on the specific form of the selection rule in Eq. 17; any measurable b preserves the guarantee. rule that picks λ̂ ∈ Λ
13
B
Details of Experimental Setup
Details of Datasets and Models We evaluate BalanceRAG on three open-domain QA datasets, TriviaQA, SQuAD v2, and Natural Questions (NQ). All datasets are processed into a unified question-answer evaluation format, where each example contains a question, an answer, an id, and the corresponding prompt. The preprocessing stage filters out examples with empty questions, empty answers, non-ASCII questions or answers, or inputs that exceed the maximum context length of the model. The default maximum numbers of retained examples are 7000 for TriviaQA, 5000 for SQuAD v2, and 2000 for NQ. Since the cleaning results after generation may differ across models, the actual number of examples used in the experiments can vary slightly. We therefore report the final number of evaluated examples as N in Table 3. We use eight backbone LLMs, including Qwen2.5-3B-Instruct, Qwen2.5-7B-Instruct, Qwen3-4B, Qwen3-14B, Llama3.2-3B-Instruct, Llama3.1-8B-Instruct, Vicuna-v1.5-7B, and Vicuna-v1.5-13B. Table 3 reports the base error rates of the two branches before calibration, where Only denotes the error rate of the LLM-only branch, RAG denotes the error rate of the LLM-RAG branch, and ∆ = Only − RAG denotes the change in error rate brought by retrieval. Table 3: Branch-level base error rates. N denotes the number of examples. Only and RAG are the pre-calibration error rates of the LLM-only and LLM-RAG branches, and ∆ = Only − RAG indicates the retrieval gain. Blue cells highlight the largest retrieval gain within each dataset. Model
N
Only ↓
RAG ↓
∆↑
Model
N
Only ↓
RAG ↓
∆↑
TriviaQA Qwen2.5-3B Qwen2.5-7B Qwen3-4B Qwen3-14B
5929 6196 6510 6063
0.2739 0.1699 0.2167 0.1921
0.1558 0.1298 0.1246 0.1070
0.1181 0.0401 0.0921 0.0851
Llama3.2-3B Llama3.1-8B Vicuna-7B Vicuna-13B
6248 6365 4183 3036
0.2769 0.2016 0.2417 0.1989
0.2130 0.1475 0.1831 0.1607
0.0639 0.0541 0.0586 0.0382
SQuAD Qwen2.5-3B Qwen2.5-7B Qwen3-4B Qwen3-14B
4475 4375 4557 4360
0.2291 0.2027 0.2203 0.2122
0.1321 0.1191 0.0898 0.0954
0.0970 0.0836 0.1305 0.1168
Llama3.2-3B Llama3.1-8B Vicuna-7B Vicuna-13B
4482 4210 3461 3031
0.3135 0.3192 0.3389 0.2748
0.1772 0.1202 0.1410 0.1488
0.1363 0.1990 0.1979 0.1260
NQ Qwen2.5-3B Qwen2.5-7B Qwen3-4B Qwen3-14B
1732 1755 1858 1807
0.3816 0.3419 0.3617 0.3298
0.1819 0.1744 0.1335 0.1384
0.1997 0.1675 0.2282 0.1914
Llama3.2-3B Llama3.1-8B Vicuna-7B Vicuna-13B
1808 1792 1599 1196
0.3711 0.3203 0.4953 0.3972
0.2163 0.1479 0.1932 0.1898
0.1548 0.1724 0.3021 0.2074
Details of Corpus Construction For each dataset, we build a dataset-level retrieval corpus from the raw evidence fields. Each document consists of a title and a passage. TriviaQA uses evidence from Wikipedia entity pages and web search results, SQuAD v2 uses the article title and context passage, and Natural Questions uses the preprocessed long-answer passage or the majority-voted long-answer candidate. We remove documents with missing titles or passages, strip extra whitespace, and deduplicate exact title–passage pairs across the corpus. The resulting corpus is cached and reused for dense retrieval. Details of Answer Generation and Retrieval All datasets use a one-shot prompt. The LLM-only branch provides only the question itself and one example question-answer pair. The LLM-RAG branch prepends retrieved evidence to the same question-answer format and asks the model to answer based on the evidence. The most likely answer is generated with beam search using a beam size of 5, without sampling. The sampled generations for uncertainty estimation are generated with sampling, where each input has 20 sampled answers, the temperature is 1.0, and top-p is 0.9. The default maximum generation length is 30 tokens. The RAG branch uses dense retrieval with all-MiniLM-L6-v2 as the retriever encoder. During retrieval, the document title and passage are concatenated and encoded, the embeddings are L2-normalized, and the documents are ranked by dot product. For each question, we use the top-4 retrieved documents by default, and each passage is truncated to at most 1200 characters. Details of Correctness Criteria By default, we use semantic similarity as the correctness criterion. Specifically, we compute the similarity between the ground-truth answer and the most likely generation 14
using a sentence-transformers cross-encoder, and set 0.7 as the default threshold. An example is treated as correct when its similarity score is no lower than this threshold. This correctness label is used to construct the error indicator in the calibration risk. In addition to the default setting, we also use entailment and LLM-as-a-Judge in the robustness experiments. The entailment setting uses DeBERTa-v3-large-mnlifever-anli-ling-wanli for NLI-based judgment. The LLM-as-a-Judge setting uses Qwen3-14B as the judge and formulates the judgment as a three-way classification problem, where both correct and partial are treated as acceptable correct answers. Algorithm 1 Single-branch calibration for LLM-only and LLM-RAG Require: Dcal = {(xℓ , ub (xℓ ), cb (xℓ ))}n ℓ=1 , where ub is uncertainty and cb = 1 denotes correctness; branch b ∈ {LLM, RAG}; target risk α; error level δ Ensure: threshold t̂ 1: T ← sorted unique values of {ub (xℓ )}n ℓ=1 2: t̂ ← ∅ 3: for t ∈ T in ascending order do 4: m ← |{ℓ : ub (xℓ ) ≤ t}|; w ← |{ℓ : ub (xℓ ) ≤ t, cb (xℓ ) = 0}| 5: if m = 0 then continue 6: end if 7: p(t) ← P Bin(m, α) ≤ w 8: if p(t) ≤ δ then 9: t̂ ← t ▷ keep scanning for a looser threshold 10: else 11: break ▷ monotone stopping 12: end if 13: end for 14: return t̂
Algorithm 2 UCB-Cascaded calibration Require: Dcal = {(xℓ , u1 (xℓ ), u2 (xℓ ), c1 (xℓ ), c2 (xℓ ))}n ℓ=1 ; target risk α; error level δ; bound type B ∈ {CP, HFD} Ensure: thresholds (t̂1 , t̂2 ) 1: Split the confidence budget: δstage ← δ/2 # Stage 1: calibrate t1 on the full calibration set 2: T1 ← sorted unique values of {u1 (xℓ )}n ℓ=1 3: t̂1 ← −∞ 4: for t ∈ T1 in ascending order do 5: m ← |{ℓ : u1 (xℓ ) ≤ t}|; w ← |{ℓ : u1 (xℓ ) ≤ t, c1 (xℓ ) = 0}| 6: R̄(t) ← U PPER B OUNDB (m, w; δstage ) 7: if R̄(t) ≤ α then 8: t̂1 ← t 9: else 10: break 11: end if 12: end for # Stage 2: calibrate t2 on the residual set 13: R ← {ℓ : u1 (xℓ ) > t̂1 } 14: T2 ← sorted unique values of {u2 (xℓ ) : ℓ ∈ R} 15: t̂2 ← −∞ 16: for t ∈ T2 in ascending order do 17: m ← |{ℓ ∈ R : u2 (xℓ ) ≤ t}|; w ← |{ℓ ∈ R : u2 (xℓ ) ≤ t, c2 (xℓ ) = 0}| 18: R̄(t) ← U PPER B OUNDB (m, w; δstage ) 19: if R̄(t) ≤ α then 20: t̂2 ← t 21: else 22: break 23: end if 24: end for 25: return (t̂1 , t̂2 )
▷ largest feasible so far
▷ deferred samples
Details of Uncertainty Estimators The default uncertainty estimator is semantic entropy based on sampled answer clusters, denoted as SE. For each input, we sample 20 answers to compute the uncertainty score. In addition to SE, we also evaluate EigV, Deg, Ecc, and SELF. SE computes the entropy of the semantic-cluster distribution over sampled generations. EigV, Deg, and Ecc construct a graph from the 15
pairwise similarity matrix of sampled generations and derive uncertainty scores from graph-structure statistics. SELF uses the length-normalized sentence entropy of the most likely generation. Details of Baselines We provide three algorithms to clarify the calibration procedures of different baselines. Algorithm 1 describes single-branch calibration, which is used for LLM-only and LLM-RAG. These methods use the uncertainty score from a single branch and search over a one-dimensional threshold space for a feasible point that satisfies the target risk. Algorithm 2 describes UCB-Cascaded. This method first calibrates the LLM-only threshold on the full calibration set, then calibrates the RAG threshold on the residual subset not accepted by the LLM-only branch. It splits δ equally between the two stages, with each stage using δ/2. UCB-Cascaded includes two variants based on the Clopper–Pearson upper bound and the Hoeffding upper bound. Algorithm 3 describes BalanceRAG, which jointly calibrates the LLM-only and RAG branches on a two-dimensional threshold lattice and uses sequential graphical testing for high-probability risk control. We report two graph diffusion priors, UGD and DWD, where DWD is the default main method. Additional baselines in the appendix include Adaptive RAG, Self-Route, Empirical Risk < α, and Bonferroni correction. Adaptive RAG routes queries using question-complexity labels. Self-Route routes queries using route-judge labels. Empirical Risk < α selects the threshold pair with the largest number of accepted examples within the empirically feasible region. Bonferroni correction applies a conservative adjustment by dividing δ by the number of tested nodes. Algorithm 3 BalanceRAG calibration via SGT Require: Dcal ; target risk α; error level δ; graph prior P ∈ {UGD, DWD} Ensure: thresholds (t̂1 , t̂2 ) # Step 1: choose the initial node 1: Discretize u1 , u2 into grids of size M × N , yielding lattice Λ = {λi,j } 2: Split Dcal into a seed split Dseed and an SGT split Dsgt 3: for each λi,j ∈ Λ do seed seed 4: Mij ← MDseed (λi,j ); Kij ← KDseed (λi,j ) 5: end for seed seed seed 6: Sseed ← {λi,j : Mij > 0, Kij /Mij ≤ α} 7: if Sseed ̸= ∅ then seed 8: s ← arg maxλi,j ∈Sseed Mij 9: else seed seed , α) ≤ Kij 10: s ← arg minλi,j :M seed >0 P Bin(Mij ij 11: end if # Step 2: compute node-wise p-values 12: for each λi,j ∈ Λ do 13: Mij ← MDsgt (λi,j ); Kij ← K Dsgt (λi,j ) 14: pij ← P Bin(Mij , α) ≤ Kij 15: end for
▷ Eq. 9
# Step 3: initialize budget and graph prior 16: δs ← δ; δi,j ← 0 for all λi,j ̸= s b←∅ 17: Λ 18: Set edge weights g(i,j)→u from P with root node s (Eq. 11 or 12) # Step 4: sequential graphical testing b such that pij ≤ δi,j do 19: while ∃ λi,j ∈ Λ \ Λ 20: Choose any such λi,j maximizing Mij b←Λ b ∪ {λi,j } 21: Λ b do 22: for each successor u ∈ N (i, j) with λu ∈ /Λ 23: δu ← δu + δi,j · g(i,j)→u 24: end for 25: δi,j ← 0 26: end while
▷ set of certified thresholds
▷ certify λi,j ▷ propagate budget
# Step 5: select the best certified threshold pair 27: λ̂ ← arg maxλi,j ∈Λb Mij
▷ Eq. 17
28: return λ̂ = (t̂1 , t̂2 )
16
Prompt for LLM-Only QA ### System: This is a bot that correctly answers questions. ### User: {few-shot question} ### Assistant: {few-shot answer} ### User: {question} ### Assistant:
Prompt for RAG QA ### System: This is a bot that correctly answers questions. Answer the question based on the provided evidence. ### Evidence: [Doc 1] Title: {title_1} {passage_1} [Doc 2] Title: {title_2} {passage_2} ... ### User: {few-shot question} ### Assistant: {few-shot answer} ### Evidence: {retrieved evidence} ### User: {question} ### Assistant:
Prompt for LLM-as-a-Judge Correctness Evaluation You are an expert evaluator for open-ended QA correctness. Given a question, a ground-truth answer, and a model’s answer, decide which option best describes the model’s answer: A. correct – semantically equivalent to the ground-truth answer. B. partial – related and contains some correct information but is incomplete or partially wrong. C. incorrect – not compatible with the ground-truth answer. Respond by selecting exactly one of A, B, or C. ### Question: {question} ### Ground truth answer: {ground-truth answer} ### Model answer: {model answer} ### Answer:
Figure 8: Prompt templates used for LLM-only QA, RAG QA, and LLM-as-a-Judge correctness evaluation.
17
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.150
0.12
50
75
0.125
50 0.100
0.10
25 0.12 0.14 0.16 Target Risk Level α
25
0.18
0.10
(a) Qwen2.5-3B
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.18
0.10
0.18
0.10
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
100 0.150 80 Err.
0.125
0.12 0.14 0.16 Target Risk Level α
(b) Qwen2.5-7B
Cov.
0.150
0.12 0.14 0.16 Target Risk Level α
0.12 0.14 0.16 Target Risk Level α
0.18
0.12 0.14 0.16 Target Risk Level α
0.18
0.12 0.14 0.16 Target Risk Level α
0.18
0.12 0.14 0.16 Target Risk Level α
0.18
90 Cov.
0.10
Err.
100
Cov.
75 Cov.
Err.
0.14
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
100
Err.
0.16
0.125
75
60 0.100
0.100
60
40 0.18
0.10
(c) Qwen3-4B
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.10
0.125
60
30
0.12 0.14 0.16 Target Risk Level α
0
0.18
0.10
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.18
0.10
60
30
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(f) Llama3.1-8B
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.150
Err.
0.12
0.12
0.10
90
Cov.
Err.
0.12 0.14 0.16 Target Risk Level α
0.10
90
0.10
(e) Llama3.2-3B
0.14
0.18
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.14
0.100 0.10
0.12 0.14 0.16 Target Risk Level α
(d) Qwen3-14B
90
Cov.
Err.
0.150
0.18
Err.
0.175
0.12 0.14 0.16 Target Risk Level α
Cov.
0.12 0.14 0.16 Target Risk Level α
60
0.125
90
Cov.
0.10
0.100
30
60
30
0.075
0.08 0.10
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(g) Vicuna-7B
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(h) Vicuna-13B
Figure 9: Err. (left in each panel pair) and Cov. (right) under different target risk levels α across eight LLMs spanning four model families on SQuAD v2.
C
Additional Experimental Results
Risk Control Holds Across Datasets Figures 9 and 10 show the results of BalanceRAG on SQuAD v2 and NQ. Consistent with the TriviaQA results in the main text, BalanceRAG achieves stable risk control across different models and target risk levels, while its coverage is usually close to or even higher than that of LLM-RAG. More importantly, under stricter values of α, the UCB-Cascaded baselines often fail to find a feasible point, which suggests that calibrating the two branches separately can be overly restrictive under a tight risk budget. In contrast, BalanceRAG still returns feasible solutions, showing that system-level calibration preserves reliability across datasets with different retrieval difficulty. Correct Acceptances Remain Competitive Under Risk Control Tables 7, 8, and 9 compare the number of correctly accepted examples under the same target risk levels. The results show that the risk control of BalanceRAG is not obtained by simply rejecting a large number of examples. Instead, BalanceRAG preserves high utility while satisfying the risk constraint. On TriviaQA, BalanceRAG is often close to or better than RAG, and clearly outperforms UCB-Cascaded. On SQuAD v2 and NQ, RAG usually retains more correctly accepted examples, but BalanceRAG still maintains comparable utility and is more likely to remain feasible than the UCB-Cascaded baselines under strict values of α. These results show that BalanceRAG maintains utility close to RAG while providing feasible risk control. Risk Control Is Stable Across Uncertainty Estimators Figure 11 and Table 4 examine whether BalanceRAG depends on a particular uncertainty estimator. Deg, SELF, SE, EigV, and Ecc provide different forms of uncertainty signals, and therefore affect sample ranking, calibrated thresholds, and final coverage. The results show that, although different estimators lead to different utility levels, BalanceRAG 18
0.200
75
0.175
50
0.150
0.125
0.17 0.19 0.21 Target Risk Level α
0.23
0.17 0.19 0.21 Target Risk Level α
0.125
0.23
0.15
99 Err.
0.14
0.13
97
0.23
0.15
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.200
98
0.17 0.19 0.21 Target Risk Level α
(b) Qwen2.5-7B
Cov.
Err.
0.15
60
100
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
80
40 0.15
(a) Qwen2.5-3B
0.16
100
0.150
25 0.15
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
Err.
Err.
Cov.
0.175
100
Cov.
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.17 0.19 0.21 Target Risk Level α
0.23
0.17 0.19 0.21 Target Risk Level α
0.23
0.17 0.19 0.21 Target Risk Level α
0.23
0.17 0.19 0.21 Target Risk Level α
0.23
90 Cov.
0.200
0.175
75
0.150 60 0.17 0.19 0.21 Target Risk Level α
0.23
0.15
(c) Qwen3-4B
Err.
Cov.
0.18
0.23
0.15
0.17 0.19 0.21 Target Risk Level α
0.23
0.15
(d) Qwen3-14B
100
0.21
75
0.18 Err.
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.21
0.17 0.19 0.21 Target Risk Level α
50
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
100
75 Cov.
0.15
0.15
0.15 0.12
25
25
0.12 0.17 0.19 0.21 Target Risk Level α
0.23
0.15
(e) Llama3.2-3B
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.15
0.17 0.19 0.21 Target Risk Level α
0.23
0.15
(f) Llama3.1-8B
0.21
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.18
75 Err.
0.18
0.23
100
Cov.
Err.
0.21
0.17 0.19 0.21 Target Risk Level α
50
0.15
90 Cov.
0.15
50
0.15
75 25
0.12
0.12 0.15
0.17 0.19 0.21 Target Risk Level α
0.23
0.15
(g) Vicuna-7B
0.17 0.19 0.21 Target Risk Level α
0.23
0.15
0.17 0.19 0.21 Target Risk Level α
0.23
0.15
(h) Vicuna-13B
Figure 10: Err. (left in each panel pair) and Cov. (right) under different target risk levels α across eight LLMs spanning four model families on NQ. Table 4: Effect of uncertainty quantification methods on BalanceRAG on TriviaQA. AUROC results reflect the performance of different uncertainty quantification methods on the LLM-only and RAG branches. Coverage is reported under different target risk levels α after joint risk calibration. Uncertainty Deg SELF SE EigV Ecc
AUROC Only/RAG ↑
Cov. @ α = 0.10 ↑
Cov. @ α = 0.14 ↑
Cov. @ α = 0.18 ↑
0.7476 / 0.6978 0.5839 / 0.6800 0.7248 / 0.6544 0.7488 / 0.6987 0.6900 / 0.6706
81.97 70.25 84.11 82.87 72.89
99.95 99.43 99.96 99.95 99.49
99.96 99.93 99.96 99.95 99.95
still maintains risk control near the target levels across these settings. This finding suggests that the reliability of the method mainly comes from the calibration procedure itself, rather than from any specific uncertainty score. In other words, the uncertainty estimator affects utility and acceptance size, but does not change the basic ability of BalanceRAG to achieve risk control. Risk Control Is Stable Across Correctness Criteria Figures 12 and 13 evaluate BalanceRAG under LLM-as-a-Judge and entailment criteria. These two correctness criteria change which examples are treated as correct, and therefore affect the error signals observed during calibration and the final accepted set. The results show that BalanceRAG still maintains risk control under different correctness criteria, indicating that it does not rely on a single correctness definition. More importantly, BalanceRAG does not assume that one branch is always more reliable. Instead, it adjusts the accepted set and routing decisions according to the calibration results under the current evaluation criterion. Therefore, when the 19
0.120
0.12
Err.
90 Cov.
Err.
0.15
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
75
90 Cov.
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.105
75 0.090
0.09
60 0.18
0.10
0.14 Target Risk Level α
(a) EigV
0.18
0.10
0.175
Cov.
90
0.12
0.18
0.10
0.14 Target Risk Level α
0.18
0.10
0.14 Target Risk Level α
0.18
(b) SELF
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.150
75
0.14 Target Risk Level α
100
90
Cov.
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.15 Err.
0.14 Target Risk Level α
Err.
0.10
0.125
80 0.09
0.100
60 0.10
0.14 Target Risk Level α
0.18
0.10
0.14 Target Risk Level α
(c) Deg
0.18
0.10
0.14 Target Risk Level α
0.18
(d) Ecc
Figure 11: Err. and Cov. on TriviaQA with Qwen2.5-7B using different UQ methods. Table 5: Branch-level error rates under different correctness evaluation methods on TriviaQA. Model
LLM-as-a-Judge
N
Entailment
Only ↓
RAG ↓
Only ↓ RAG ↓
Llama3.2-3B 6248 0.2623 Llama3.1-8B 6365 0.1277
0.2308 0.2687
0.3782 0.2591 0.2353 0.1983
Model
N
LLM-as-a-Judge
Entailment
Only ↓
RAG ↓
Only ↓ RAG ↓
Qwen2.5-3B 5929 0.2889 Qwen2.5-7B 6196 0.1921
0.2663 0.2515
0.3719 0.1881 0.2352 0.1451
correctness criterion changes, the system can adapt through recalibration rather than being fixed toward either LLM-only or RAG.
0.125
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.16 Err.
Cov.
Err.
60
45
0.100
90
Cov.
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.150
0.14
0.12 60
30 0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(a) Qwen2.5-3B
0.18
0.10
0.18
0.10
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.15 Err.
Cov.
75
0.12
0.12 0.14 0.16 Target Risk Level α
(b) Qwen2.5-7B
90
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.16
0.12 0.14 0.16 Target Risk Level α
0.12 0.14 0.16 Target Risk Level α
0.18
0.12 0.14 0.16 Target Risk Level α
0.18
90 Cov.
0.10
Err.
75
0.12
75
60 0.08
0.09 0.10
0.12 0.14 0.16 Target Risk Level α
0.18
45
0.10
(c) Llama3.2-3B
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(d) Llama3.1-8B
Figure 12: Err. and Cov. on TriviaQA with LLM-as-a-Judge for correctness evaluation.
Calibration Size Mainly Affects Stability Figures 14 analyze the effect of the calibration-test split ratio. With fewer calibration examples, threshold estimation is more sensitive to random splits, and the results may fluctuate more. With more calibration examples, the estimates are generally more stable. In the experiments, the empirical risk of BalanceRAG consistently satisfies risk control as the split ratio changes. Moreover, under some settings of α, the variance of Err. decreases as the calibration size increases. These results show that the calibration size mainly affects stability, rather than changing the basic mechanism by which the method achieves risk control. Routing Allocation Reflects Branch Complementarity Table 10 shows how accepted examples are allocated between the LLM-only and RAG branches. The allocation ratios vary across models and datasets, indicating that BalanceRAG neither always selects LLM-RAG nor always preserves LLM-only. 20
In some settings, stronger models allow more examples to remain in the LLM-only branch, thereby reducing retrieval calls. In other settings, more examples are routed to RAG to use the additional information provided by retrieval. This variation shows that the two branches are complementary, and that BalanceRAG uses calibrated risk signals to decide when to trust each branch. Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
80
0.12
60
90 Cov.
0.14
Err.
0.150 Cov.
Err.
0.175
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.16
0.125
40
75
60 0.100
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(a) Qwen2.5-3B
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.18
0.10
0.18
0.10
Target Risk Level LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG-DWD BalanceRAG-UGD
0.16
60
Err.
0.12
0.12 0.14 0.16 Target Risk Level α
(b) Qwen2.5-7B
75
Cov.
Err.
0.15
0.12 0.14 0.16 Target Risk Level α
0.12 0.14 0.16 Target Risk Level α
0.18
0.12 0.14 0.16 Target Risk Level α
0.18
90 Cov.
0.10
0.12
75
45 60
0.08
0.09 30 0.10
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(c) Llama3.2-3B
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
0.12 0.14 0.16 Target Risk Level α
0.18
0.10
(d) Llama3.1-8B
Figure 13: Err. and Cov. on TriviaQA with entailment for correctness evaluation.
Multi-Risk Calibration Provides a Cost Control Knob Table 11 evaluates the multi-risk variant, where the second risk constraint limits the fallback invocation rate of RAG. Compared with single-risk routing, the multi-risk setting can explicitly control the frequency of RAG calls by adjusting α2 . When this constraint becomes stricter, the system reduces retrieval usage and may also lower overall coverage. This result should be understood as a controllable cost-coverage trade-off, rather than a cost-free performance gain. In practical deployment, this setting allows users to adjust system behavior according to the retrieval budget while maintaining calibrated answer reliability. Table 6: Reliability–coverage comparison on Llama3.1-8B with δ = 0.1 over 500 data splits. We report mean accepted-answer risk, coverage, and guarantee success rate under target risk levels α ∈ {0.10, 0.11, 0.12}. Success rate denotes the fraction of calibration/test splits where the empirical test risk is no larger than α. Adaptive-RAG and Self-Route are uncalibrated adaptive routing baselines that always commit to an answer. Method
α = 0.10 Err. ↓
α = 0.11
Cov. ↑ Succ. ↑
Err. ↓
α = 0.12
Cov. ↑ Succ. ↑
Err. ↓
Cov. ↑ Succ. ↑
Adaptive-RAG (uncal.) 0.1516 100.0 Self-Route (uncal.) 0.1389 100.0 Empirical Risk < α 0.1017 81.2
0.0 0.0 40.2
0.1516 100.0 0.1389 100.0 0.1118 87.9
0.0 0.0 41.6
0.1516 100.0 0.1389 100.0 0.1227 93.5
0.0 0.0 34.8
Bonferroni BalanceRAG
100.0 94.2
0.0780 0.0997
100.0 95.4
0.0884 0.1093
100.0 95.6
0.0783 0.0894
51.9 71.8
57.6 79.7
72.0 86.0
Additional Baselines Table 6 compares BalanceRAG with Adaptive RAG, the Empirical Risk < α baseline, and Bonferroni correction. Adaptive RAG and the Empirical Risk < α baseline have mean risks above the target α and lower guarantee success rates, showing that simple adaptive retrieval or empirical-risk selection does not provide the desired risk control. Bonferroni correction is more reliable, but its coverage is much lower because the correction is more conservative. In contrast, BalanceRAG achieves mean risk below the target level, obtains a high guarantee success rate, and preserves higher coverage than Bonferroni correction. These results show that BalanceRAG strikes a better balance between high-probability reliability and coverage. Qualitative Cases Explain Adaptive Routing Qualitative examples illustrate the routing behavior behind the aggregate results. In some cases, the LLM-only branch already produces the correct answer, while the retrieved content introduces misleading information. In such cases, BalanceRAG preserves the 21
Err.
Split Ratios 0.15 (3:17)
0.16 0.14 0.12 0.1
0.20 (4:16) 0.25 (5:15) 0.30 (6:14) 0.35 (7:13) 0.40 (8:12) 0.45 (9:11)
0.1
0.12 0.14 Target Risk Level α
0.16 Split Ratios
0.16
0.15 (3:17)
Err.
0.14
0.20 (4:16)
0.12
0.25 (5:15)
0.1
0.30 (6:14) 0.35 (7:13) 0.40 (8:12) 0.45 (9:11)
Err.
0.1
0.12 0.14 Target Risk Level α
0.16 Split Ratios
0.21 0.19 0.17 0.15
0.15 (3:17) 0.20 (4:16) 0.25 (5:15) 0.30 (6:14) 0.35 (7:13) 0.40 (8:12) 0.45 (9:11)
0.15
0.17 0.19 0.21 Target Risk Level α Figure 14: Risk control across various calibration-test split ratios. From top to bottom, the results are reported on TriviaQA with Vicuna-v1.5-7B, SQuAD v2 with Qwen2.5-7B, and NQ with LLaMA-3.1-8B.
cheaper direct answer. In other cases, the LLM-only branch makes an error, while LLM-RAG corrects the answer by using external evidence, making routing to retrieval more appropriate. There are also cases where both branches are correct, suggesting that always invoking retrieval is not necessary. These cases do not serve as standalone statistical evidence, but they provide an intuitive explanation of the routing policy. The goal of BalanceRAG is not to favor a fixed branch, but to decide which output can be safely accepted under calibrated risk.
22
Table 7: Correctly accepted samples comparison on TriviaQA. We report the number of correctly accepted samples (Corr.) under different target risk levels α. Bold numbers indicate the best Corr. under the same LLM and risk level. LLMs
Methods / α
0.1
0.11
0.12
0.13
0.14
0.15
0.16
0.17
0.18
Qwen2.5-3B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1473 – – 1569
798 1680 1291 – 1733
885 1907 1633 – 1907
971 2105 1834 – 2110
1049 2280 1952 1555 2279
1137 2400 2049 1799 2421
1227 2475 2143 1955 2498
1316 2500 2237 2051 2512
1397 2501 2304 2152 2514
Qwen2.5-7B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
1856 2118 1820 – 2366
1988 2320 1992 1716 2573
2148 2490 2284 1880 2702
2255 2633 2377 2035 2737
2347 2691 2442 2205 2738
2422 2694 2499 2301 2738
2473 2694 2531 2392 2738
2522 2694 2534 2455 2738
2564 2694 2557 2502 2738
Llama3.2-3B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
870 1367 832 – 1347
988 1390 1028 – 1643
1116 1532 1488 820 1774
1214 1646 1749 964 1873
1283 1760 1850 1328 1973
1371 1876 1917 1661 2073
1443 1949 1991 1799 2159
1513 2036 2065 1899 2229
1589 2117 2142 1990 2290
Llama3.1-8B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
1587 1856 1553 – 2073
1780 2036 1849 1395 2288
1928 2198 2084 1607 2430
2070 2382 2264 1810 2584
2165 2529 2360 1992 2691
2218 2652 2434 2148 2727
2277 2711 2486 2248 2734
2345 2713 2505 2351 2736
2402 2713 2482 2400 2737
Vicuna-7B-v1.5
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
792 1256 981 542 1297
868 1330 1102 603 1367
935 1401 1236 776 1430
992 1455 1306 876 1487
1049 1512 1354 1011 1533
1115 1564 1395 1147 1580
1189 1611 1426 1266 1629
1250 1648 1469 1333 1667
1294 1675 1504 1355 1694
Vicuna-13B-v1.5
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
775 875 750 – 983
861 939 868 669 1041
919 987 925 694 1091
961 1041 962 806 1124
1002 1101 998 888 1168
1043 1163 1031 946 1203
1081 1207 1070 989 1237
1110 1250 1100 1035 1262
1138 1270 1127 1078 1274
Qwen3-4B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 2505 – – –
1764 2627 1746 – 2549
1871 2759 1834 1726 2718
1995 2835 2016 1781 2807
2093 2849 2126 1902 2839
2159 2849 2240 2048 2843
2233 2849 2463 2133 2843
2301 2849 2533 2224 2843
2357 2849 2587 2338 2843
Qwen3-14B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 2521 – – 2499
– 2660 – – 2634
1836 2706 2310 – 2629
1964 2707 2362 – 2629
2073 2707 2467 2058 2660
2176 2707 2538 2215 2661
2259 2707 2533 2359 2661
2314 2707 2525 2410 2661
2363 2707 2497 2420 2661
23
Table 8: Correctly accepted samples comparison on SQuAD. We report the number of correctly accepted samples (Corr.) under different target risk levels α. Bold numbers indicate the best Corr. under the same LLM and risk level; underlined numbers indicate the second best. LLMs
Methods / α
0.1
0.11
0.12
0.13
0.14
0.15
0.16
0.17
0.18
Qwen2.5-3B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1450 – – 1357
– 1653 – – 1464
– 1785 – – 1649
– 1873 – – 1798
– 1929 – – 1887
333 1942 1916 – 1922
362 1942 1927 – 1929
402 1942 1928 – 1930
452 1942 1925 – 1930
Qwen2.5-7B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1723 – – 1611
– 1836 – – 1698
– 1896 – – 1821
– 1924 – – 1895
443 1929 1913 – 1915
468 1929 1915 – 1916
551 1929 1913 – 1916
682 1929 1906 1914 1916
878 1929 1890 1914 1916
Llama3.2-3B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1438 – – 1377
– 1546 – – 1484
– 1617 – – 1591
– 1666 – – 1654
– 1714 – – 1699
137 1755 1744 – 1745
144 1784 1775 – 1781
164 1809 1798 – 1805
211 1829 1818 – 1828
Llama3.1-8B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1763 – – 1723
– 1803 – – 1784
– 1831 – – 1819
– 1850 – – 1836
194 1853 1835 – 1842
211 1853 1837 – 1843
227 1853 1835 – 1843
257 1853 1834 – 1843
288 1853 1830 – 1843
Vicuna-7B-v1.5
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1330 – – 1258
– 1376 – – 1327
– 1411 – – 1374
– 1443 – – 1415
– 1464 – – 1447
– 1479 – – 1468
130 1484 1467 – 1479
138 1484 1473 – 1482
150 1484 1473 – 1482
Vicuna-13B-v1.5
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1027 – – 950
– 1094 – – 1018
– 1143 – – 1097
– 1183 – – 1149
– 1224 – – 1189
– 1256 – – 1224
128 1280 1245 – 1256
136 1291 1266 – 1276
152 1292 1279 – 1281
Qwen3-4B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 2069 – – 2053
– 2075 – – 2050
– 2075 – – 2031
– 2075 – – 2032
– 2075 – – 2032
781 2075 2019 – 2032
838 2075 2018 – 2032
948 2075 2008 2019 2032
1120 2075 1983 2015 2032
Qwen3-14B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 1953 – – 1942
– 1972 – – 1948
– 1973 – – 1905
– 1973 – – 1914
– 1973 – – 1923
974 1973 1912 – 1924
1094 1973 1903 – 1924
1251 1973 1882 1910 1924
1386 1973 1853 1895 1924
24
Table 9: Correctly accepted samples comparison on NQ. We report the number of correctly accepted samples (Corr.) under different target risk levels α. Bold numbers indicate the best Corr. under the same LLM and risk level; underlined numbers indicate the second best. LLMs
Methods / α
0.15
0.16
0.17
0.18
0.19
0.20
0.21
0.22
0.23
Qwen2.5-3B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 461 – – 435
– 508 – – 470
– 565 – – 510
– 622 – – 554
– 667 – – 600
168 696 607 – 646
180 707 645 – 683
193 708 669 – 697
209 708 685 – 701
Qwen2.5-7B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 535 – – 544
– 577 – – 562
– 623 – – 595
– 674 – – 628
– 707 – – 672
264 720 682 – 697
287 724 706 – 715
309 725 712 – 716
341 725 712 703 718
Llama3.2-3B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 489 – – 447
– 532 – – 492
119 562 – – 531
125 590 563 – 575
134 615 590 – 615
146 640 615 – 634
159 661 635 – 657
174 680 653 – 674
189 694 669 633 694
Llama3.1-8B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
158 734 698 – 709
191 752 718 – 741
240 760 736 – 753
283 763 743 717 757
320 763 747 735 759
345 763 744 749 759
365 763 740 753 759
382 763 735 750 759
401 763 730 744 759
Vicuna-7B-v1.5
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 541 – – 478
– 564 – – 508
– 584 – – 536
– 600 – – 556
– 615 – – 582
– 629 – – 601
– 639 – – 617
88 643 – – 628
90 644 629 – 635
Vicuna-13B-v1.5
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 387 – – 373
– 410 – – 400
– 430 – – 413
– 448 – – 418
– 462 – – 416
– 471 – – 427
– 478 – – 439
– 482 – – 449
– 483 – – 460
Qwen3-4B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 799 – – 792
– 804 – – 799
– 806 – – 802
– 806 – – 797
– 806 – – 785
– 806 – – 780
– 806 – – 779
– 806 – – 778
– 806 – – 777
Qwen3-14B
LLM-only LLM-RAG UCB-Cascaded-CP UCB-Cascaded-HFD BalanceRAG
– 747 – – 727
– 772 – – 761
– 777 – – 769
– 778 – – 763
– 778 – – 757
– 778 – – 752
383 778 740 – 752
397 778 741 – 753
421 778 738 – 752
25
Table 10: Route allocation (%) of BalanceRAG. We report the allocation ratio of accepted samples routed to the LLM-only branch (Only) and the RAG branch. Cov. denotes the total accepted ratio, i.e., Only+RAG . Dataset
LLMs Only
TriviaQA
Dataset
Qwen2.5-3B Qwen2.5-7B LLaMA-3.2-3B Vicuna-7B-v1.5 Vicuna-13B-v1.5 LLaMA-3.1-8B Qwen3-4B Qwen3-14B
27.7 62.2 37.1 33.5 51.9 51.2 32.4 29.8
LLMs Only
SQuAD
Dataset
Qwen2.5-3B Qwen2.5-7B LLaMA-3.2-3B Vicuna-7B-v1.5 Vicuna-13B-v1.5 LLaMA-3.1-8B Qwen3-4B Qwen3-14B
15.8 21.7 7.4 6.8 9.7 7.7 29.2 28.3
LLMs Only
NQ
Qwen2.5-3B Qwen2.5-7B LLaMA-3.2-3B Vicuna-7B-v1.5 Vicuna-13B-v1.5 LLaMA-3.1-8B Qwen3-4B Qwen3-14B
13.8 17.8 15.1 12.4 1.9 22.4 1.9 5.1
α = 0.12 RAG Cov. 44.5 35.8 26.7 43.2 28.2 34.5 61.7 67.7
Only
72.2 98.0 63.8 76.7 80.1 85.7 94.1 97.5
29.9 64.1 41.1 33.2 54.5 53.3 58.7 64.0
α = 0.12 RAG Cov.
Only
66.9 71.3 72.1 81.9 71.0 89.6 70.2 69.2
82.7 93.0 79.5 88.7 80.7 97.3 99.4 97.5
17.3 24.0 8.4 6.4 9.7 8.0 38.6 45.6
α = 0.17 RAG Cov.
Only
55.4 62.2 53.4 66.0 78.8 76.6 97.7 93.9
69.2 80.0 68.5 78.4 80.7 99.0 99.6 99.0
20.0 25.0 14.3 12.2 13.2 24.5 26.4 25.0
α = 0.14 RAG Cov. 58.3 35.8 31.4 50.6 33.2 44.1 41.2 35.9
Only
88.2 99.9 72.5 83.8 87.7 97.4 99.9 99.9
30.7 64.1 41.6 35.0 54.7 62.1 58.8 64.1
α = 0.14 RAG Cov.
Only
79.5 75.9 78.5 89.0 79.9 91.9 61.4 54.3
96.8 99.9 86.9 95.4 89.6 99.9 100.0 99.9
18.3 25.1 7.8 6.4 10.2 8.4 38.6 45.7
α = 0.19 RAG Cov.
Only
63.5 67.2 67.4 75.2 70.3 75.3 73.3 74.1
26
83.5 92.2 81.7 87.4 83.5 99.8 99.7 99.1
20.9 34.9 17.5 12.4 19.6 25.1 39.8 46.1
α = 0.16 RAG Cov. 68.4 35.8 39.7 56.2 40.1 37.8 41.2 35.8
Only
99.1 99.9 81.3 91.2 94.8 99.9 100.0 99.9
31.1 64.1 41.5 35.0 54.7 63.1 58.8 64.1
α = 0.16 RAG Cov.
Only
81.6 74.8 85.3 93.2 86.6 91.6 61.4 54.3
99.9 99.9 93.1 99.6 96.8 100.0 100.0 100.0
18.8 25.1 8.0 6.5 10.0 8.4 38.6 45.7
α = 0.21 RAG Cov.
Only
76.0 64.4 72.0 82.6 70.3 74.7 60.0 53.7
96.9 99.3 89.5 95.0 89.9 99.8 99.8 99.8
20.7 36.6 18.0 11.9 23.5 25.1 43.2 48.7
α = 0.18 RAG Cov. 68.9 35.8 46.7 62.2 45.0 36.8 41.2 35.8
100.0 99.9 88.2 97.2 99.7 99.9 100.0 99.9
α = 0.18 RAG Cov. 81.1 74.8 90.1 93.4 89.9 91.6 61.4 54.3
99.9 99.9 98.1 99.9 99.9 100.0 100.0 100.0
α = 0.23 RAG Cov. 79.1 63.2 78.8 87.5 73.3 74.7 56.7 51.1
99.8 99.8 96.8 99.4 96.8 99.8 99.9 99.8
Table 11: Route allocation (%) of multi-risk BalanceRAG. We report the allocation ratio of accepted samples routed to the LLM-only branch (Only) and the RAG branch. α2 denotes the fallback invocation rate cap, chosen as the nearest lower multiple of 10% below the corresponding risk-only RAG allocation. Cov. denotes the total accepted ratio. Dataset
Qwen2.5-3B Qwen2.5-7B LLaMA-3.2-3B TriviaQA Vicuna-7B-v1.5 Vicuna-13B-v1.5 LLaMA-3.1-8B Qwen3-4B Qwen3-14B Dataset
SQuAD
Dataset
NQ
α2
α = 0.12 Only RAG Cov. α2
α = 0.14 Only RAG Cov. α2
α = 0.16 Only RAG
Cov.
α2
α = 0.18 Only RAG
Cov.
40 30 20 40 20 30 60 60
34.2 70.1 38.0 40.6 62.9 59.4 59.1 63.9
35.3 74.8 45.3 37.9 63.8 61.7 62.4 71.8
37.1 75.7 51.9 42.9 63.3 74.0 63.7 73.0
57.7 24.2 25.7 45.9 30.5 25.8 36.3 27.0
94.8 100.0 77.6 88.8 93.8 99.8 100.0 100.0
60 30 40 60 40 30 40 30
42.7 75.8 51.1 40.5 65.4 74.6 63.8 73.2
57.1 24.1 35.1 55.8 33.8 25.4 36.2 26.8
99.7 100.0 86.2 96.3 99.2 100.0 100.0 100.0
α2
α = 0.12 Only RAG Cov. α2
α = 0.14 Only RAG Cov. α2
α = 0.16 Only RAG
Cov.
α2
α = 0.18 Only RAG
Cov.
60 70 70 80 70 80 70 60
19.4 25.6 10.4 9.3 11.1 16.4 38.3 44.4
25.1 33.1 18.8 12.9 18.0 12.3 44.1 52.8
24.9 35.1 17.4 9.9 17.2 13.9 44.5 54.5
74.9 64.8 72.6 88.2 75.9 86.0 55.4 45.4
99.8 99.9 90.0 98.1 93.0 99.9 100.0 99.9
80 70 90 90 80 90 60 50
25.4 35.6 10.3 12.7 21.6 14.1 44.7 54.7
74.5 64.4 86.1 87.2 76.0 85.8 55.2 45.2
99.9 100.0 96.4 99.9 97.5 99.9 100.0 100.0
α2
α = 0.17 Only RAG Cov. α2
α = 0.19 Only RAG Cov. α2
α = 0.21 Only RAG
Cov.
α2
α = 0.23 Only RAG
Cov.
50 60 50 60 70 70 90 90
21.3 35.0 20.1 15.7 3.3 33.0 9.2 11.2
24.2 37.3 20.5 16.4 15.9 37.3 43.3 50.0
26.0 43.1 21.9 15.3 23.4 38.8 52.5 54.9
89.8 95.5 84.9 91.0 84.4 99.8 97.0 98.9
70 60 70 80 70 70 50 50
32.3 48.7 27.1 18.5 29.1 38.9 55.2 56.1
96.5 99.2 90.4 94.8 91.4 99.8 99.3 99.7
LLMs
LLMs Qwen2.5-3B Qwen2.5-7B LLaMA-3.2-3B Vicuna-7B-v1.5 Vicuna-13B-v1.5 LLaMA-3.1-8B Qwen3-4B Qwen3-14B LLMs Qwen2.5-3B Qwen2.5-7B LLaMA-3.2-3B Vicuna-7B-v1.5 Vicuna-13B-v1.5 LLaMA-3.1-8B Qwen3-4B Qwen3-14B
34.0 25.0 0.4 34.2 13.1 24.1 31.2 30.0
55.5 65.3 65.9 77.1 64.9 76.7 60.2 50.5
39.7 42.9 41.6 53.9 63.2 61.6 80.1 78.0
68.2 95.0 38.4 74.8 76.0 83.5 90.3 93.9
74.9 90.9 76.3 86.4 76.0 93.0 98.5 94.9
61.0 77.9 61.6 69.5 66.6 94.5 89.3 89.2
50 30 30 50 30 40 40 30
70 70 70 80 70 90 60 50
60 60 60 70 70 70 70 70
27
46.8 25.1 25.5 45.2 21.3 33.4 36.7 27.2
66.4 65.9 64.6 77.2 65.4 87.4 55.8 46.6
50.8 49.6 53.7 64.4 57.7 61.9 55.4 47.4
82.2 99.9 70.8 83.1 85.0 95.1 99.2 99.0
91.5 99.0 83.5 90.1 83.4 99.7 99.9 99.4
75.0 87.0 74.2 80.8 73.6 99.2 98.7 97.5
60 30 30 50 40 30 40 30
80 70 80 90 80 90 60 50
70 60 70 80 70 70 50 50
63.8 52.4 63.0 75.7 60.9 61.0 44.5 44.0
64.2 50.5 63.3 76.3 62.3 60.9 44.1 43.6
Case 1: LLM-only is correct while RAG is misled
ID: sfq_8482_5657 α = 0.10 t1 = 0.286 t2 = 0.748 uOnly = 0.286 uRAG = 2.466 Route: Only ✓ ................................................................................... Question Charleston is the capital of which US state? ................................................................................... Gold Answer West Virginia ................................................................................... LLM-only Answer ✓ Charleston is the capital of West Virginia. ................................................................................... RAG Answer Charleston is not the capital of any US state. Case 2: LLM-only gives the exact answer while RAG rejects it
ID: qz_453_1207 α = 0.10 t1 = 0.748 t2 = 0.469 uOnly = 0.469 uRAG = 2.020 Route:Only ✓ ................................................................................... Question For what novel did J. K. Rowling win the 1999 Whitbread Children’s book of the year award? ................................................................................... Gold Answer Harry Potter and the Prisoner of Azkaban ................................................................................... LLM-only Answer ✓ Harry Potter and the Prisoner of Azkaban. ................................................................................... RAG Answer She did not win the 1999 Whitbread Children’s book of the year award. Figure 15: Qualitative examples where BalanceRAG preserves the LLM-only branch. In both examples, the LLM-only answer is correct and passes the first-stage uncertainty threshold, while retrieval introduces misleading evidence and produces an incorrect RAG answer.
28
Case 3: RAG corrects a wrong entity generated by LLM-only
ID: bt_4365_238 α = 0.10 t1 = 0.286 t2 = 0.748 uOnly = 4.222 uRAG = 0.286 Route: RAG ✓ ................................................................................... Question Who in verse sailed to the land where the bong tree grows? ................................................................................... Gold Answer The owl and the pussycat ................................................................................... LLM-only Answer Alfred Noyes in his poem “The Highwayman” sailed to the land where the bong tree grows. ................................................................................... RAG Answer ✓ The Owl and the Pussy Cat sailed to the land where the bong tree grows. Case 4: RAG fixes an incorrect surname answer
ID: dpql_4593_2651 α = 0.10 t1 = 0.286 t2 = 0.748 uOnly = 3.304 uRAG = 0.286 Route: RAG ✓ ................................................................................... Question What is the surname of the title character of Anne of Green Gables? ................................................................................... Gold Answer Shirley ................................................................................... LLM-only Answer Cuthbert is not the surname of the title character of Anne of Green Gables. ................................................................................... RAG Answer ✓ The surname of the title character of Anne of Green Gables is Shirley. Figure 16: Qualitative examples where BalanceRAG routes to the RAG branch. The LLM-only branch has high uncertainty and gives an incorrect answer, while the RAG branch has lower uncertainty and produces the correct answer.
29
Case 5: Both branches are correct, but LLM-only is selected earlier
ID: dpql_6316_2767 α = 0.10 t1 = 0.286 t2 = 0.748 uOnly = 0.286 uRAG = 2.139 Route: Only ✓ ................................................................................... Question Who co-starred with George Clooney in the 2013 movie Gravity? ................................................................................... Gold Answer Sandra Bullock ................................................................................... LLM-only Answer ✓ Sandra Bullock co-starred with George Clooney in the 2013 movie Gravity. ................................................................................... RAG Answer Sandra Bullock co-starred with George Clooney in the 2013 movie Gravity. Case 6: Both branches answer correctly, so the cheaper LLM-only route is sufficient
ID: qf_3204_522 α = 0.10 t1 = 0.286 t2 = 0.748 uOnly = 0.286 uRAG = 2.020 Route: Only ✓ ................................................................................... Question If you were eating a Knickerbocker Glory, what would you be eating? ................................................................................... Gold Answer An ice-cream sundae ................................................................................... LLM-only Answer ✓ A Knickerbocker Glory is a type of ice cream sundae. ................................................................................... RAG Answer A Knickerbocker Glory is a type of ice cream sundae. Figure 17: Qualitative examples where both branches produce correct answers, but BalanceRAG selects the cheaper LLM-only route. These cases illustrate that always invoking RAG is unnecessary when the LLM-only branch is already sufficiently confident.
30