Training a Utility-based Retriever Through Shared Context Attribution for Retrieval-Augmented Language Models

TL;DR

This work tackles the gap in utility-based retrieval for retrieval-augmented language models by prioritizing passages that yield real downstream gains rather than mere semantic relevance. It introduces SCARLet, a three-part framework that builds a shared context for multi-task data synthesis, uses a perturbation-based attribution to estimate passage-level utility and capture inter-passage interactions, and applies one-dimensional clustering to train retrievers. Across ten datasets with in-domain and out-of-domain evaluation, SCARLet consistently improves RALM performance and demonstrates robustness to different generators and corpora. The findings suggest that sharing context and evaluating passage synergy are key to learning task-specific retrieval utility, with practical implications for improving factuality and reasoning in AI systems.

Abstract

Retrieval-Augmented Language Models boost task performance, owing to the retriever that provides external knowledge. Although crucial, the retriever primarily focuses on semantics relevance, which may not always be effective for generation. Thus, utility-based retrieval has emerged as a promising topic, prioritizing passages that provides valid benefits for downstream tasks. However, due to insufficient understanding, capturing passage utility accurately remains unexplored. This work proposes SCARLet, a framework for training utility-based retrievers in RALMs, which incorporates two key factors, multi-task generalization and inter-passage interaction. First, SCARLet constructs shared context on which training data for various tasks is synthesized. This mitigates semantic bias from context differences, allowing retrievers to focus on learning task-specific utility for better task generalization. Next, SCARLet uses a perturbation-based attribution method to estimate passage-level utility for shared context, which reflects interactions between passages and provides more accurate feedback. We evaluate our approach on ten datasets across various tasks, both in-domain and out-of-domain, showing that retrievers trained by SCARLet consistently improve the overall performance of RALMs.

Figures (9)

• Figure 1: The illustration of SCARLet. The upper left part describes the inference process of RALMs. In SCARLet, there are three main stages. First, the shared context is constructed by retrieving external corpus based on the seed data. The synthesizer is instructed with shared context and task information from the task pool, to generate synthetic data. Next, using the shared context as the data source, SCARLet applies perturbation-based utility attribution on the generator, and then, based on the utility scores, samples positive and negative passages for retriever training.
• Figure 2: The performance of the perturbation-based attribution method on the GTI benchmark. The nDCG metrics show that it achieves at least about 80% performance on three datasets, with some exceeding 90%.
• Figure 3: The illustration of the 1D clustering sampling. Based on the utility score, this method clusters the passages into three labels: the high-score passages (green) corresponding to positive samples, the middle-score passages (orange) that will be discarded, and the low-score passages (red) corresponding to negative samples.
• Figure 4: Ablation Study on six in-domain datasets, using BGE as retriever, with two generators. The values in the charts correspond to the metrics of each dataset.
• Figure 5: Case Study on HotpotQA. The passage is from the corpus and has varying recall rankings for different retrievers. Orange text indicates necessary reasoning information for the question.