A Note on k-NN Gating in RAG
Gérard Biau, Claire Boyer
TL;DR
This note provides a principled statistical framework for balancing a frozen base LM with retrieved evidence in retrieval-augmented generation (RAG). It introduces a k-NN retriever, a gating mechanism $\\lambda(x)$, and a geometry-aware retrieval-trust weight $w_{fact}(x)$ to adaptively fuse evidence per query. The core contributions include a Bayes-optimal hard gate characterization, a discordance-based measure of hallucination, and an asymptotic analysis showing that gating decisions reflect true structural differences between Bayes and the LM under $k\\to\\infty$, $k/n\\to 0$. It further extends to a hybrid geometric-semantic mismatch model to account for distribution shifts between query inputs and memory, providing a foundation for factuality-oriented RAG design in the presence of distributional drift.
Abstract
We develop a statistical proxy framework for retrieval-augmented generation (RAG), designed to formalize how a language model (LM) should balance its own predictions with retrieved evidence. For each query x, the system combines a frozen base model q0 ($\times$ x) with a k-nearest neighbor retriever r (k ) ($\times$ x) through a measurable gate k(x). A retrieval-trust weight wfact (x) quantifies the geometric reliability of the retrieved neighborhood and penalizes retrieval in low-trust regions. We derive the Bayes-optimal per-query gate and analyze its effect on a discordance-based hallucination criterion that captures disagreements between LM predictions and retrieved evidence. We further show that this discordance admits a deterministic asymptotic limit governed solely by the structural agreement (or disagreement) between the Bayes rule and the LM. To account for distribution mismatch between queries and memory, we introduce a hybrid geometric-semantic model combining covariate deformation and label corruption. Overall, this note provides a principled statistical foundation for factuality-oriented RAG systems.