Training Data is More Valuable than You Think: A Simple and Effective Method by Retrieving from Training Data

TL;DR

The paper tackles the high cost of external retrieval in retrieval-augmented NLP by proposing REINA, which retrieves from the training data itself. By indexing labeled input-label pairs and concatenating the top-K retrieved examples with the input across training and inference, REINA yields substantial improvements across 11 datasets spanning summarization, language modeling, machine translation, and question answering, including SOTA results on XSum, BigPatent, and CommonsenseQA. The method demonstrates that training-data retrieval can match or exceed massive pre-trained models in certain settings, offering a scalable and cost-efficient alternative. This approach highlights the continued value of leveraging labeled training data as a knowledge source and provides practical guidance for integrating retrieval into supervised learning pipelines.

Abstract

Retrieval-based methods have been shown to be effective in NLP tasks via introducing external knowledge. However, the indexing and retrieving of large-scale corpora bring considerable computational cost. Surprisingly, we found that REtrieving from the traINing datA (REINA) only can lead to significant gains on multiple NLG and NLU tasks. We retrieve the labeled training instances most similar to the input text and then concatenate them with the input to feed into the model to generate the output. Experimental results show that this simple method can achieve significantly better performance on a variety of NLU and NLG tasks, including summarization, machine translation, language modeling, and question answering tasks. For instance, our proposed method achieved state-of-the-art results on XSum, BigPatent, and CommonsenseQA. Our code is released, https://github.com/microsoft/REINA .

Figures (2)

• Figure 1: REINA pipeline of model training/inference with retrieval from training data. Filter only happens at training, as the same training sample will be retrieved from the index. For each instance, we concatenate the input with the retrieved content, i.e., data and/or labels, for model training and inference.
• Figure 2: Model training with retrieval from the training data ( REINA ). (a) Index on the training data and data retrieval for 4 different tasks. Box in blue is the query or the input sequence to encode. Box in green is the retrieved text. (b-e) Leveraging retrieved data for model training with different structures. For language modeling, we prepend the retrieved data to the query data, and append the retrieved data to the query for all the other tasks. After concatenation, we will directly feed them into Transformers, either Seq2Seq or Encoder-only frameworks, for text generation and answering selection. As we focus on the question answering tasks requiring commonsense reasoning, we have another version of index integrating knowledge graph for more precise retrieval. K: external knowledge from ConceptNet and Wiktionary, src: source language, tgt: target language.