UncertaintyRAG: Span-Level Uncertainty Enhanced Long-Context Modeling for Retrieval-Augmented Generation

TL;DR

UncertaintyRAG introduces span-level uncertainty estimation to enhance long-context retrieval-augmented generation, improving robustness, calibration, and performance with less training data under distribution shifts.

Contribution

It proposes a novel span uncertainty method for better calibration and robustness in long-context RAG, along with an efficient unsupervised training technique and flexible retrieval integration.

Findings

Outperforms baselines by 2.03% on LLaMA-2-7B.

Achieves state-of-the-art results using only 4% of training data.

Demonstrates improved robustness and calibration in long-context tasks.

Abstract

We present UncertaintyRAG, a novel approach for long-context Retrieval-Augmented Generation (RAG) that utilizes Signal-to-Noise Ratio (SNR)-based span uncertainty to estimate similarity between text chunks. This span uncertainty enhances model calibration, improving robustness and mitigating semantic inconsistencies introduced by random chunking. Leveraging this insight, we propose an efficient unsupervised learning technique to train the retrieval model, alongside an effective data sampling and scaling strategy. UncertaintyRAG outperforms baselines by 2.03% on LLaMA-2-7B, achieving state-of-the-art results while using only 4% of the training data compared to other advanced open-source retrieval models under distribution shift settings. Our method demonstrates strong calibration through span uncertainty, leading to improved generalization and robustness in long-context RAG tasks. Additionally, UncertaintyRAG provides a lightweight retrieval model that can be integrated into any large language model with varying context window lengths, without the need for fine-tuning, showcasing the flexibility of our approach.