💻 computer science

UAMG-RAG: Adaptive Multi-Granularity Retrieval-Augmented Generation with Uncertainty-Guided Re-Retrieval

The paper proposes UAMG-RAG, a framework that dynamically adjusts retrieval depth and document granularity based on question complexity and employs uncertainty-guided re-retrieval to enhance reasoning performance and reduce hallucinations, achieving significant improvements on HotpotQA and WikiHow benchmarks.

Original authors: Shreya B S, Kayarvizhy N

Published 2026-07-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shreya B S, Kayarvizhy N

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart but slightly forgetful assistant (the AI) who is trying to answer your questions. This assistant has read a massive library of books during its training, but it often makes things up (hallucinations) or gives outdated answers because it relies only on its memory.

To fix this, we usually give the assistant a "search engine" (Retrieval-Augmented Generation, or RAG) so it can look up facts in a database before answering. However, the paper argues that most current search engines are like rigid vending machines: they always dispense the exact same amount of snacks (documents) in the exact same size (granularity), no matter if you asked a simple question like "What is 2+2?" or a complex one like "Explain the chain of events that led to the fall of Rome."

This "one-size-fits-all" approach is inefficient. It wastes time looking up too much info for simple questions and not enough for hard ones.

The Solution: UAMG-RAG

The authors propose a new system called UAMG-RAG. Think of this as upgrading the vending machine to a smart, adaptive librarian. This librarian doesn't just fetch books; it figures out how to fetch them based on your specific needs.

Here is how it works, broken down into simple steps:

1. The "Detective" Phase (Understanding the Question)

Before searching, the system acts like a detective. It analyzes your question to guess:

What kind of question is it? Is it a simple fact check, a multi-step puzzle, or a "how-to" guide?
How hard is it? Is it a quick lookup or does it require deep thinking?

2. The "Smart Fetch" Phase (Adaptive Retrieval)

Based on the detective's report, the librarian changes its strategy:

For simple questions: It grabs a tiny, precise snippet (like a single sentence) to get the exact fact quickly.
For complex puzzles: It grabs a whole paragraph or a section of a book to ensure it has all the connecting pieces of the puzzle.
The "Depth" Adjustment: It also decides how many documents to fetch. A simple question gets fewer documents; a complex one gets more.

Analogy: If you ask, "Who is the president?", the librarian hands you a sticky note with the name. If you ask, "How do I bake a sourdough loaf?", the librarian hands you the entire recipe chapter.

3. The "Self-Check" Phase (Uncertainty-Guided Re-Retrieval)

This is the system's "safety net." After the librarian fetches the info and the AI writes an answer, the system asks itself: "Am I 100% sure this answer is backed by the documents I just found?"

The Confidence Meter: It calculates an "uncertainty score."
The Fix: If the score is too high (meaning the AI is guessing or the evidence is weak), the system doesn't just give up. It says, "Wait, I need more proof," and goes back to the library to fetch specific missing pieces of information. It repeats this loop until it feels confident.

Why This Matters (The Results)

The authors tested this "Smart Librarian" against the old "Rigid Vending Machine" using two types of test questions:

HotpotQA: Hard questions that require connecting dots between different facts (Multi-hop reasoning).
WikiHow: Step-by-step instructional questions (Procedural knowledge).

The Findings:

Better Answers: The new system got significantly more correct answers (higher F1 scores) on both types of tests.
Fewer Lies: It made up facts (hallucinations) much less often, especially for the "how-to" questions.
Efficiency: It was actually faster. By not fetching unnecessary books for simple questions, it saved time and computing power.
The Secret Sauce: The authors found that the ability to change the size of the information chunks (from sentences to whole sections) was the most important factor. If they removed this feature, the system's performance dropped drastically.

The Limitations

The authors are honest about where the system still struggles:

Ambiguity: If a question is vague or uses tricky language that doesn't clearly show how hard it is, the "detective" might misjudge the difficulty and fetch the wrong amount of info.
Small Scale: The tests were done on specific datasets. We don't know yet how it would perform in every possible real-world scenario.
Procedural Tasks: For simple "how-to" questions, the improvement was small because the old system was already pretty good at fetching whole sections for those tasks.

Summary

In short, UAMG-RAG is a system that teaches AI to be a flexible researcher rather than a rigid robot. It learns to ask, "How much information do I actually need for this specific question?" and "Am I sure enough to stop searching?" This leads to smarter, more accurate, and more reliable answers.

Technical Summary: UAMG-RAG

Problem Statement

Current Retrieval-Augmented Generation (RAG) systems face significant limitations due to their static retrieval architectures. Existing pipelines typically employ a predetermined, fixed retrieval depth and a single document granularity (e.g., fixed chunk size) regardless of the specific complexity of the user query. This "one-size-fits-all" approach leads to several critical issues:

Inefficiency: Simple factual queries often trigger excessive retrieval (over-retrieval), wasting computational resources, while complex multi-hop or procedural queries suffer from insufficient context (under-retrieval).
Hallucinations: The lack of dynamic adaptation means the retrieved evidence often fails to sufficiently support the generated answer, leading to weak grounding and a higher rate of hallucinations, particularly in reasoning-intensive tasks.
Lack of Validation: Most existing RAG pipelines lack mechanisms to evaluate whether the retrieved facts adequately support the generated output before finalizing the answer.

While recent works like Self-RAG, Corrective RAG, and FLARE have introduced iterative refinement or self-reflection, they largely overlook adaptive granularity selection as a primary optimization variable. Furthermore, many adaptive frameworks rely on computationally expensive graph structures or heavy confidence estimation models that hinder scalability.

Methodology: UAMG-RAG Framework

The authors propose UAMG-RAG (Uncertainty-guided Adaptive Multi-Granularity Retrieval-Augmented Generation), a unified framework designed to dynamically adjust retrieval behavior based on query intent and reasoning complexity. The system operates through a closed-loop pipeline comprising five key modules:

1. Query Intent and Reasoning Complexity Estimation

Upon receiving a query $Q$ , the system first classifies the intent into three categories: factual, multi-hop, or procedural. It simultaneously estimates the reasoning effort ( $E$ ) required, utilizing features such as query length ( $L_q$ ), entity density ( $N_e$ ), and the presence of complex indicators ( $C_q$ ). Queries are categorized into low, medium, or high effort levels.

2. Adaptive Retrieval Policy Selection

Based on the estimated intent ( $I$ ) and effort ( $E$ ), the framework dynamically determines:

Retrieval Depth ( $k$ ): The number of documents to retrieve.
Granularity ( $g$ ): The chunk size level.
The system utilizes a rule-based mapping (derived from pilot validation) where:
Low effort queries trigger shallow retrieval ( $k=7$ ) at sentence-level granularity ( $G1$ ).
Medium effort queries trigger moderate retrieval ( $k=8$ ) at paragraph-level granularity ( $G2$ ).
High effort queries trigger deep retrieval ( $k=10$ ) at section-level granularity ( $G3$ ).

3. Multi-Granularity Indexing

The corpus is indexed at three distinct levels to support the adaptive policy:

G1 (Sentence): ~100 tokens, optimized for precise factual retrieval.
G2 (Paragraph): ~200 tokens, optimized for multi-hop analytical reasoning.
G3 (Section): ~385 tokens, optimized for procedural and long-context reasoning.
Dense retrieval is performed using Sentence-BERT embeddings and FAISS for efficient nearest-neighbor search.

4. Uncertainty-Guided Answer Generation

The generator produces an answer $A$ conditioned on the query and retrieved evidence $D_k$ . Unlike standard RAG, the system does not immediately accept the output. Instead, it calculates a Lightweight Uncertainty Metric ( $U$ ) defined as:
$U = \alpha \cdot H + \beta \cdot (1 - O)$
Where:

$H$ is a proxy for generation confidence (entropy approximation).
$O$ is the Evidence Overlap, measuring the lexical alignment between the generated answer tokens and the retrieved evidence.
$\alpha$ and $\beta$ are weighting parameters.

5. Iterative Selective Re-Retrieval

If the uncertainty metric $U$ exceeds a predefined threshold $\tau$ , the system initiates a refinement loop:

Detection: Identifies weakly grounded answer spans and missing entities.
Reformulation: Constructs a new query $Q'$ targeting the missing information.
Re-Retrieval: Performs adaptive retrieval again with the new query.
This loop continues until the uncertainty falls below the threshold or a maximum iteration limit is reached.

Key Contributions

The paper outlines five primary contributions:

Query-Adaptive Framework: A system that dynamically adjusts both retrieval depth and chunk granularity based on estimated reasoning complexity, moving beyond static $k$ -value policies.
Multi-Granularity Indexing: A unified pipeline supporting sentence, paragraph, and section-level retrieval, allowing the system to select the appropriate context resolution for diverse query types.
Uncertainty-Guided Refinement: A lightweight mechanism that triggers iterative re-retrieval only when the generated response lacks sufficient evidential support, balancing reliability with efficiency.
Comprehensive Evaluation: Extensive experiments demonstrating improvements in factual grounding, hallucination reduction, and retrieval efficiency on HotpotQA and WikiHow datasets.
Ablation Analysis: A detailed study quantifying the relative contribution of each component, highlighting adaptive granularity as the most significant factor in performance gains.

Experimental Results

The framework was evaluated on HotpotQA (multi-hop reasoning) and WikiHow (procedural knowledge) against baselines including BM25+LLM, RAG+Reranker, and Fixed RAG ( $k=5$ ).

Performance Gains: UAMG-RAG achieved an F1 score of 24.9% on HotpotQA and 27.6% on WikiHow, outperforming all baselines.
Hallucination Mitigation: The framework reduced the hallucination rate on procedural questions to 14.2%, significantly lower than the Fixed RAG baseline (16.8%) and BM25+LLM (31.2%).
Efficiency: Despite retrieving a similar volume of documents as standard RAG, UAMG-RAG achieved a 50% reduction in end-to-end latency (3,684ms vs. 7,364ms) by utilizing adaptive early stopping and reducing unnecessary LLM calls for simpler queries.
Ablation Findings: Removing the Adaptive Granularity component resulted in the largest performance drop (-19.5 F1 points), confirming that matching retrieval granularity to query type is the primary driver of success. Removing uncertainty estimation caused a -7.0 point drop.
Query-Type Analysis: The system showed the most significant gains in factual (+5.0 F1) and multi-hop (+4.6 F1) queries. For procedural queries, the gain was marginal (-3.0 F1), as the authors note that static section-level retrieval often already provides sufficient context for instructional tasks.

Significance and Claims

The paper claims that UAMG-RAG addresses the rigidity of current RAG systems by introducing a dynamic, query-aware approach to retrieval. The authors assert that the framework successfully reconciles the trade-off between retrieval accuracy and contextual comprehensiveness.

The significance of the work is framed around three main points:

Reliability: By integrating uncertainty-guided re-retrieval, the system ensures that generated answers are grounded in sufficient evidence, directly mitigating hallucinations.
Efficiency: The adaptive policy prevents the computational waste associated with retrieving excessive context for simple queries or insufficient context for complex ones.
Granularity as a First-Class Optimization: The study posits that adaptive granularity selection is a critical, often overlooked, optimization lever that yields greater performance improvements than iterative refinement alone.

The authors maintain a modest tone regarding their claims, acknowledging limitations such as the reliance on heuristic features for complexity estimation, the use of proxy uncertainty metrics due to API restrictions on model logits, and the fact that the framework's superiority over all existing approaches has not been exhaustively proven. They conclude that while the results are promising, future work is needed to explore transformer-based reasoning models and broader domain generalization.