Retrieval-Guided Fine-Tuning as Noisy Estimation: Risk bounds and Architectural Analysis
Authors: Bhargav Lad, Yifan Hao
Organizations: Department of Mathematics Indian Institute of Technology Delhi New Delhi, India · Department of Computer Science University of Illinois Urbana-Champaign Urbana, IL 61801
Retrieval-Guided Fine-Tuning (RAG-FT) incorporates retrieved data directly into the training objective, but the statistical consequences of noisy retrieval during training remain theoretically undercharacterized. We study this question by modeling RAG-FT as an estimation problem in a multi-task linear regression framework, using an OLS proxy for single-layer linear self-attention to obtain finite-sample risk bounds. Under homoscedastic retrieval noise, we show that retrieval failure decays exponentially with task separation relative to noise, and derive explicit finite-sample conditions under which RAG-FT achieves lower risk than both target-only and full-corpus training. We then introduce a Distance-Proportional Noise (DPN) model, in which retrieval quality degrades with rank, and compare two estimators under the same retrieval process: the OLS proxy and the literal, uniform-weight forward pass of linear self-attention. We prove that the attention estimator's bias diverges as Θ(n2q) even under exact retrieval, while OLS risk remains Θ(d/n) for every noise exponent q>0. These results locate the instability not in noisy retrieval itself, but in the fixed, unweighted aggregation of the literal LSA forward pass, which reweighting by reliability empirically removes. We validate the predicted rate separation through direct simulation of the DPN model.
Figures & tables
Figure 1: Risk vs. retrieval volume n under DPN ( d=20 , m=50 ). MC: Monte Carlo; curves: finite- n theory.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 2: (a) LSA risk vs. n across d∈{5,20,50,100} : curves separate where variance scales with d , then collapse as the dimension-independent bias term takes over. (b) Oracle inverse-variance weighting and self-normalization both remain stable where uniform averaging diverges.
Figure
d
m
q
γ
n range
1a
20
50
0.5
0.3
10 – 104
1a
20
50
1
0.2
10 – 103.5
1a
20
50
2
0.1
10 – 102.5
1b
20
50
1
0.2
10 – 103.5
2a
{5,20,50,100}
50
1
0.2
10 – 103.5
2b
20
50
1
0.2
10 – 103.5
Appendix
Table 1: Fixed parameters per figure. σ=1 throughout. Panel 1a sweeps three (q,γ) pairs, each over its own n range (narrower at larger q , since risk grows faster and the largest n values become numerically extreme).
Retrieval-Augmented Generation (RAG) has become the standard way to ground large language models in external knowledge, yet most systems retrieve a fixed number of passages for every question regardless of its difficulty. This wastes computation on easy questions, starves hard ones, and gives no signal for when a generated answer can be trusted. With a growing share of question answering systems built on top of commercial language model APIs, a method that can decide how much to retrieve, and how far to trust its own answers, without retraining the underlying model, is of clear practical value. This paper presents AB-RAG (Adaptive Budgeted Retrieval-Augmented Generation), a training-free and backbone-agnostic framework that generates an answer, estimates its confidence from a combination of three signals, and then decides whether to stop or to retrieve more evidence, subject to a fixed retrieval budget. The estimator combines the model's own certainty, the agreement between the answer and the evidence, and the variance of the retrieval scores. For models that expose token probabilities the certainty signal is read directly; for closed APIs it is approximated by self-consistency, so the method works without access to model internals. Across three backbones and two datasets, the central result is that the confidence estimate reliably separates correct from incorrect answers on every backbone, reaching a clean split of 57.6% against 0% Exact Match between high- and low-confidence answers on a factoid dataset. The adaptive policy improves accuracy on capable backbones, and the study reports its negative and nuanced findings honestly, including a confidence signal that proved unsuitable for short answers and a retrieval signal whose sign was found and corrected through measurement. The entire study was carried out on a single consumer laptop with only a few dollars of API spend.
Ansh Kamthan
Department of Artificial Intelligence and Machine Learning Manipal University Jaipur Jaipur, India
Retrieval-augmented generation (RAG) has become a standard mechanism for grounding language models in external knowledge, yet conventional retrieval based on lexical or semantic similarity is poorly suited for complex reasoning tasks: a semantically similar problem may demand an entirely different solution strategy, while a superficially different problem may share the same underlying reasoning pattern. We propose Retrieval-Augmented Reinforcement Fine-Tuning (RA-RFT), a post-training framework that teaches language models to reason by analogy. RA-RFT uses gold-relevance distillation to train a retriever that ranks contexts by expected reasoning benefit rather than semantic overlap, and then fine-tunes the policy model via reinforcement fine-tuning methods with retrieved analogous demonstrations, so the model learns to leverage reasoning traces under verifiable outcome rewards. We further analyze the diversity of retrieved contexts and find that reasoning-aware retrieval surfaces complementary solution strategies that provide distinct reasoning scaffolds for individual problems. Across challenging mathematical reasoning benchmarks, RA-RFT consistently outperforms standard reinforcement fine-tuning methods. For example, it improves AIME 2025 average@32 accuracy by 7.1 and 2.8 points over GRPO for Qwen3-1.7B and Qwen3-4B respectively -- suggesting that reasoning-aware retrieval is a complementary axis of improvement and orthogonal to advances in reward design or training curricula.
Zilin Xiao, Qi Ma, Chun-cheng Jason Chen +4
Meta Superintelligence Labs · Rice University · Work partially done during internship at Meta
Retrieval-augmented generation (RAG) typically treats retrieval and generation as separate systems. We ask whether an attention-based encoder-decoder can instead retrieve directly from its own internal representations. We introduce INTRA (INTrinsic Retrieval via Attention), a framework where decoder attention queries score pre-encoded evidence chunks that are then directly reused as context for generation. By construction, INTRA unifies retrieval and generation, eliminating the retriever-generator mismatch typical of RAG pipelines. This design also amortizes context encoding by reusing precomputed encoder states across queries. On question-answering benchmarks, INTRA outperforms strong engineered retrieval pipelines on both evidence recall and end-to-end answer quality. Our results demonstrate that attention-based models already possess a retrieval mechanism that can be elicited, rather than added as an external module.
Elad Hoffer, Yochai Blau, Edan Kinderman +3
NVIDIA · Department of Electrical Engineering, Technion, Haifa, Israel