End-to-End Trainable Retrieval-Augmented Generation for Relation Extraction

TL;DR

ETRAG delivers end-to-end trainable retrieval-augmented generation for relation extraction by replacing the non-differentiable instance selection with differentiable $k$-nearest neighbor (soft) selection and injecting retrieved instances as soft prompts into a text-generation model. Built atop a SuRE-style generation framework, it demonstrates consistent gains on TACRED across data regimes, with notable improvements in low-resource settings and competitive performance at full data. Key contributions include the differentiable $k$NN for instance selection, neural prompting with trainable instance integration, and end-to-end training that aligns the retriever with the relation extraction objective. The work shows retrieved instances tend to share relation labels or entities with the query, validating task-specific instance selection and offering a foundation for broader applications of retrieval-augmented generation in NLP.

Abstract

This paper addresses a crucial challenge in retrieval-augmented generation-based relation extractors; the end-to-end training is not applicable to conventional retrieval-augmented generation due to the non-differentiable nature of instance retrieval. This problem prevents the instance retrievers from being optimized for the relation extraction task, and conventionally it must be trained with an objective different from that for relation extraction. To address this issue, we propose a novel End-to-end Trainable Retrieval-Augmented Generation (ETRAG), which allows end-to-end optimization of the entire model, including the retriever, for the relation extraction objective by utilizing a differentiable selection of the $k$ nearest instances. We evaluate the relation extraction performance of ETRAG on the TACRED dataset, which is a standard benchmark for relation extraction. ETRAG demonstrates consistent improvements against the baseline model as retrieved instances are added. Furthermore, the analysis of instances retrieved by the end-to-end trained retriever confirms that the retrieved instances contain common relation labels or entities with the query and are specialized for the target task. Our findings provide a promising foundation for future research on retrieval-augmented generation and the broader applications of text generation in Natural Language Processing.

Figures (3)

• Figure 1: An overview of ETRAG
• Figure 2: The input template of SuRE. ${$\cdot$} are placeholders to replace with input.
• Figure 3: Change in Extraction Performance with the Number of Instances $k$ Used in Prompts