Retrieval-Enhanced Machine Learning: Synthesis and Opportunities

Selecting $\textrm{C}_{i}$ is not only regarded with deciding what kind of external information should be provided to the $f_{\theta}$ (Pan et al., 2023) but also encompasses the question of when to query (no corpus selection when no retrieval is beneficial to $f_{\theta}$). This can be a critical question as retrieval augmentation can hurt the performance for certain input types (Maekawa et al., 2024; Mallen et al., 2023; Asai et al., 2024). It also can save computational resources by reducing the number of searching (Labruna et al., 2024). There can be several criteria on whether external information can be beneficial to the predictive model. It can be based on term popularity (Mallen et al., 2023), input complexity (Jeong et al., 2024), or a trained model (Asai et al., 2024).

Once one or more $\textrm{C}_{i}$ is selected, a REML system can further benefit itself by choosing the optimal retriever specialized in searching the selected corpora. This is a relatively challenging and new task, and we refer readers to Khramtsova et al. (2023, 2024) for deeper understanding.

In certain scenarios, not all words or components of the input prove relevant for the search objective. Consequently, the common practice of omitting specific segments of the input has been employed in numerous prior studies. In the majority of cases, sequence-to-sequence models are trained to identify and mark the segments that require reduction (Ni et al., 2019; Musa et al., 2019; Yadegari et al., 2022; Khashabi et al., 2017). At times, a straightforward approach such as segmenting the input into distinct chunks and utilizing these segments as queries can be highly effective (Borgeaud et al., 2022). In multi-modal search scenarios, omitting a specific modality from the input and conducting searches across a corpus from different modalities has proven to be valuable (Gui et al., 2022).

In certain scenarios, the input alone may lack essential information required by the search system to yield desired results. In such situations, augmenting the input with additional pertinent data can be done. This process of expansion broadens the context and enhances the search system’s capability to retrieve relevant and meaningful results, thereby improving overall system performance. Typically, expansion is achieved by concatenating the input with previously retrieved results (Zhu et al., 2021; Xiong et al., 2021) or generated text that is conditioned on the input (Wang et al., 2023d; Mao et al., 2021; Liu et al., 2022; Chuang et al., 2023).

For some situations, reshaping the input into a new query based on its inherent structure, instead of mere expansion, is proved to be advantageous. This approach is particularly valuable when crafting structured queries for database (Arcadinho et al., 2022; Dou et al., 2023) and API access (Schick et al., 2023; Qin et al., 2023; Ouyang et al., 2022; Jin et al., 2024).

The conversion operation may results in a transformation of the input space $\mathcal{X}$ into the query space $\mathcal{Q}$. Consequently, it is essential to note that $\mathcal{Q}$ and $\mathcal{X}$ are not necessarily equivalent, signifying that the transformed queries might operate in a distinct space compared to the original input. In complex multi-modal search scenarios, employing the input directly might not be viable. Consequently, transforming the input into a different modality form becomes imperative, ensuring seamless and efficient communication between the predictive model and the search system (Gao et al., 2022; Lin and Byrne, 2022; Wu and Mooney, 2022; Lin et al., 2023; Salemi et al., 2023a).

Moreover, transforming the input from the original input space into the latent space of a language model and retrieving information from the model’s prior interactions with data, like what happens in kNN-LM (Khandelwal et al., 2020), represent additional transformations that alter the query space (Chen et al., 2022; Yogatama et al., 2021; He et al., 2021; Khandelwal et al., 2021; Kassner and Schütze, 2020). Indeed, Neural Turing Machines (Graves et al., 2014; Gulcehre et al., 2017; Rae et al., 2016) and Memory Transformers (Zhong et al., 2022; Wu et al., 2022a; Wan et al., 2022) employ a similar conversion process to translate input into a latent variable. This transformation is essential for enabling effective retrieval from the memory/storage component of these models.

At $t$’th iteration, with slight abuse of notation (simplifying $\textrm{C}_{i}^{t}$ to $\textrm{C}_{t}$), given a query $q_{t}$, a collection $\textrm{C}_{t}$, content-based addressing can be defined as:

Location-based addressing lets the searching system access the corpus purely by storage location such as index without any lexical or contextual comparison between a query and the corpus. Therefore, this is often used for storing (Section 6) or recency-based retrieval. Thus,

With the final address vector $w_{t}$, the retrieval results $r_{t}$ from $g_{\omega}$ can be defined as

where read simply selects the content from corpus in location $w_{t}$ that that is transformed by the $\textit{transform}_{s}$.

Dependent on the task and source of data, the result data will be incomplete prior to consumption. The transformation of data is a general and powerful process that converts data through a separate model depending on the needs of the system. Common reasons requiring further data transformation include decontextualization, translation, and summarization among others, thus $r$ can be transformed by the following equation:

To further optimize the presentation of the result items for size or clarity, multiple items can be combined, e.g., summarizing all items jointly (Wu et al., 2021; Sarthi et al., 2024). Not all REML systems will combine items, as operating over individual items can be efficient and effective. Furthermore, combining items may lead to complications such as miscalibration between the individual scores and scores of a combined result, which is represented by compose:

If not all documents fit into the context for consumption, we discard or truncate documents based on these limits, optimizing for length and potentially other properties such as diversity (Hofstätter et al., 2023; Bahri et al., 2020; Meng et al., 2024), which is represented by truncate:

The full equation for presentation:

Typically, multiple items are retrieved and it is a design choice whether to process retrieved items separately or together. The followings are the main paradigms for consuming retrieval items:

• •

  Single: Only a single item is incorporated in the prediction. This may be sufficient for simple queries, but often information will need to be combined across multiple retrieved items.

• •

  Ensemble: Predictions are made for multiple retrieval items in the single-fashion, then aggregated (Khandelwal et al., 2020; Shi et al., 2024).

• •

  Joint: When the prediction has a sufficiently flexible context representation, then multiple retrieval items can be passed simultaneously in a single inference procedure (Izacard and Grave, 2021b; Lewis et al., 2020a). This is potentially richer than the ensemble approach since each retrieved item is aware of the others. Due to computational limits, only a few retrieval items are able to be processed this way and the ensemble-approach is relatively more scalable.

• •

  Multi-round: A hybrid approach where a subset of retrieved items are processed at a time, and the next subset incorporates information about the retrieved items processed thus far (Jiang et al., 2023). Although more scalable than the joint-approach, this may be slower. That being said, some applications (e.g. dialogue) naturally conform to the multi-round framework.

These paradigms are atomic functions that can be composed to form more complex operations. For instance, given lists of retrieved items, denoted as $r_{0},r_{1},r_{2},r_{3}$, the atomic functions with different granularities can be composed as the following equation:

Independent of the choice of granularity, there are algorithms applicable for consumption. Broadly, they fall into the following categories:

• •

  Extractive: The predictive model is limited to extracting exact information from the retrieved item, e.g., a span of text from a retrieve passage to answer a question (Khandelwal et al., 2020; Lan et al., 2023). This can be achieved through pointer networks (Vinyals et al., 2015), constrained decoding (Hokamp and Liu, 2017; Hu et al., 2019; Post and Vilar, 2018), and other similar techniques.

• •

  Analogical: Learning by example, case based reasoning, and retrieve-and-edit approaches all fall under the category of analogical reasoning. Each of these involves different underlying mechanisms, but essentially the predictive model will be extrapolating from one or more demonstrative examples to make its prediction and not necessarily extracting factual knowledge from the retrieved items (Das et al., 2020).

• •

  Contextual: The predictive model incorporates the retrieved items in its context, but the decoding of the output is not constrained in any way (Shi et al., 2024).

• •

  Latent: A hybrid approach where retrieved items are not incorporated directly into the context, but are instead incorporated in other ways, e.g., by merging hidden states (Yogatama et al., 2021). Similar to the contextual-approach, decoding is not constrained.

Independent of the consumption algorithm, there are different decoding algorithms that can be done to predict a high quality output in REML. Broadly, decoding algorithms explored thus far fall into the following categories:

• •

  Output-only: Search algorithms like beam search (Freitag and Al-Onaizan, 2017) will only consider the model output when scoring candidate predictions (Khandelwal et al., 2020).

• •

  Retrieval-enhanced: Search algorithms like beam search will consider both the model output and the retrieved items when scoring candidate predictions (Lewis et al., 2020a; Asai et al., 2024). This should penalize spurious associations between query and retrieved item.

• •

  Verification-based: The initial output of the predictive model will be scrutinized by a verification module and potentially rejected if a condition is a met, e.g., the output does not entail the retrieved item (Jiang et al., 2023).

Steps can be taken in presentation to speed up inference, such as by compressing passages through summarization or truncating the list of retrieved items (De Jong et al., 2023; Hofstätter et al., 2023). There are other approaches to improve efficiency that are more tightly integrated with consumption, e.g., partially precomputing passage embeddings.

Advanced applications of REML will augment the predictive model output space to incorporate REML-specific information. The most common instance of this is probably to support attribution, so that each part of the output can be traced back to the relevant retrieved item (Gao et al., 2023b; Schuster et al., 2024; Asai et al., 2024; Menick et al., 2022). Other cases involve in-line verification or calls to external tools that would not be easily possible without incorporating the appropriate retrieved item (Asai et al., 2024).

An important aspect of utilizing storage in REML systems is the ability to store and retrieve specific pieces of information. Therefore, the storage system must be capable of generating a specific address for reading from or writing to the storage. Storage location can be divided into abstract location (slots in the storage space) and physical location (where in hardware the storage sits). For abstract location, it boils down to the role of $\textit{address}_{\textit{location}}$, as introduced in Section 4.5. In most cases, it will be a simple rotational function (Graves et al., 2014, 2016) which allows an iteration through a sequence of slots. For better efficiency, the function can store entries into clusters by content (Weston et al., 2015) or layers (Lample et al., 2019), reducing the computation during searching. For physical location, entries that do not need to be in RAM or VRAM can be moved to a disk with memory mapping, while entries that must be in the RAM, such as embedding index to be searched, can be compressed without notable performance degradation (Izacard et al., 2024).

Reading from the storage is an essential part of REML models and closely tied with the search operations. However in Section 6, we focus on how retrievable items are represented in the storage. How they are read are discussed in Section 4.

At time $t$, after the address vector $w_{t}$ is obtained, storage operation can be done by a write function that updates the datastore $\textrm{C}_{t}$ as follows:

With the view of addressing mechanism mentioned in Section 4.5, it can be understood that location-based addressing is used before the write operation, following the framework of Neural Turing Machine (NTM) (Graves et al., 2014, 2016). In most cases, the write operation will simply append the latest entry to the end of a storage (address vector pointing to the next available slot), executed after every new input in sequential order, following the work of neural cache model (Grave et al., 2017). On the other hand, some architectures, such as Memory Neural Network (MemNN) (Weston et al., 2015) or large memory layers (Lample et al., 2019) operate differently, and will generate an address specifying where in memory to write the new entry, potentially overwriting any entry that was there before. Regularization can be applied in memory networks to ensure that a substantial portion of the memory is used. These decisions can be made by concerted effort with the storage management component.

In most cases, a REML system will initialize its storage by processing an entire retrieval corpus. This can be done offline before training, after training before inference, and throughout training as needed. Storage construction is well documented and studied, and the initial storage construction is essentially a series of write and address generation operations. Storage can be constructed as a key-value structure where retrieval space $\mathcal{D}$ can be defined as:

Once the storage is initialized, both for optimal task completion and efficiency, there is a need to schedule the storage operations, and we dub this phase as storage management. Efficient usage of storage can be understood in terms of space and speed which come down to determining three decisions: when, what, and how to store.

Coupled storage is defined as a storage that can be updated online during training and inference of the predictive model and can be jointly optimized. Initial developments of the coupled storage enhancements were primarily led by Neural Turing Machine (NTM) (Graves et al., 2014, 2016) and Memory Network (MemNN) (Weston et al., 2015). They showed abstract operations, such as copying, recall, and sorting, to language reasoning tasks by leveraging external addressable storage, where contents in the storage of these architectures are dense vectors which can be readily used by $f_{\theta}$. We refer the readers to the original papers of the two models for deeper understanding of primary shape of REML models with coupled storage. There are a few notable characteristics about REML systems with coupled memory, including but not limited to:

Decoupled storage is defined as a storage method in REML systems where $g_{\omega}$ operates independent to $f_{\theta}$. In contrast to the coupled storage, where the entries are being updated dynamically online usually influenced by joint optimization of $f_{\theta}$ and $g_{\omega}$, decoupled storage involves offline population of entries, i.e., storage becomes read-only during the online stage. There are a few notable characteristics about REML systems with decoupled memory, including but not limited to:

In a wide range of studies, retrieval models are not optimized. In many of them, queries and documents are in the form of unstructured text. In that case, query and document representations are often computed based on term statistics, such as term frequencies in the documents or document frequencies in the document collection. Using retrieval models such as TF-IDF (Salton and Buckley, 1988), BM25 (Robertson et al., 1995), and query likelihood (Ponte and Croft, 1998), with default parameters, belongs to this category. For example. the Dr.QA model (Chen et al., 2017a) uses the ElasticSearch implementation of TF-IDF for document retrieval from Wikipedia for factoid question answering. The SelfMem model (Cheng et al., 2024) uses BM25 for document retrieval for a number of retrieval-augmented text generation tasks, such as translation and dialogue.

Pre-trained language models (e.g., encoder representations) can be also used to produce latent representations for queries and documents, where simple similarity functions, such as dot product or cosine similarity, are used for computing relevance scores for a query-document pair. Even though these language models went through expensive optimization procedures, their optimizations are not REML-specific. Note that employing language model representations for retrieval with no optimization often do not perform well. For instance, Lien et al. (2024) demonstrated that using plain BERT or RoBERTa representations for zero-shot retrieval is substantially worse than term matching models, like BM25.

More recently, it has been observed that large-scale instruction-tuned language models, such as GPT-3.5, can be carefully instructed to rank a few documents for a given query (Sun et al., 2023). These models can perform effectively, but cannot be applied to large document collections and could be only used as re-ranking models.

Retrieving from databases through structured queries, such as SQL, also belong to this category. A wide range of task-oriented dialogue systems, such as intelligent assistants for travel booking and restaurant reservation, require access to databases for up-to-date availability (Chen et al., 2017b).

REML may take advantage of retrieval models whose optimization is independent of the predictive model’s parameters. Retrieval models are often trained using a set of query-document-relevance triplets. The relevance signal may come from (1) explicit annotations, e.g., from expert assessors or crowdworkers, (2) implicit feedback (Joachims, 2002; Joachims et al., 2017), e.g., user clicks, dwell time, mouse movements, etc., or (3) automatically generated weak signals (also known as distant supervision signals), such as appearance of a phrase, e.g., answer in the context of QA, in the documents, retrieval scores from another retrieval model (Dehghani et al., 2017), annotations produced by large language models (Thomas et al., 2023). Using these training triplets, retrieval models can be optimized using three different approaches: (1) pointwise ranking, (2) pairwise ranking, and (3) listwise ranking. Refer to Liu (2009) for more information on various ranking loss functions.

In the context of REML, a large set of studies, e.g., (Vu et al., 2023; Hashemi et al., 2021; Lyu et al., 2023; Jiang et al., 2023), use commercial search engines, such as Bing or Google, as their retrieval models. These search engines are optimized using a combination of the relevance signals mentioned above. Therefore, they are considered as independent optimization models. A set of approaches, such as (Izacard and Grave, 2021b), train retrieval models on explicitly labeled collections, such as MS MARCO (Campos et al., 2016) or provenance labels in the KILT benchmark (Petroni et al., 2021), and then use the trained models on an often out-of-domain REML scenario. Weak or distant supervision is also used in (Qu et al., 2021, 2020; Salemi et al., 2023b) for open-domain (visual) question answering by assuming that any document that contains the answer phrase is relevant.

In conditional optimization, the retrieval model is optimized, conditioned on the predictive model $f_{\theta}$. A group of conditional optimization approaches use knowledge distillation. For instance, Izacard and Grave (2021a) use the aggregation of cross-attention weights in the fusion-in-decoder architecture as weak signals to train the retrieval model. Here, the decoder that provides the weights plays the role of a teacher model and a dense passage retrieval model plays the role of a student model. Alternatively, Yang and Seo (2020) use the similarity score produced by an answer span selection model, i.e., the reader, as teacher scores and minimize the KL-divergence between them and retrieval scores.

As shown by Izacard and Grave (2021a), knowledge distillation from the downstream ML model to the retrieval model can be done iteratively, as follows:

where the retrieval model in iteration (or epoch) $t+1$ is optimized based on the parameters of the predictive model at iteration $t$, where $\mathcal{L}$ is the downstream loss function.

Similar to retrieval models, predictive models may also be used as a ‘black-box’ systems without REML-specific training. For instance, a wide range of query expansion approaches, such as the Rocchio’s algorithm (Rocchio, 1971), relevance models (Lavrenko and Croft, 2001), and divergence minimization model (Zhai and Lafferty, 2001) expand the queries based on the appearance of terms and concepts in the retrieval results. Using pre-trained large language models in a zero-shot setting is another example that has received considerable attention in recent years (Shi et al., 2024; Salemi et al., 2024b).

Predictive models in REML can be optimized independent of the retrieval model’s parameters. For instance, we can optimize predictive models by assuming that the retrieval model is optimal (i.e., retrieving ground truth relevant documents). In this case, the optimization of predictive model $f_{\theta}$ can be modeled as:

where $\mathcal{L}$ is the loss function for the downstream task. For instance, a number of open-domain QA models are optimized to extract or generate answers given the question and the gold (ground truth) passage (Chen et al., 2017a). Some may relax the optimality assumption of retrieval models and inject non-relevant documents to the ground truth set. These documents can be either sampled randomly or from the output of a retrieval model, but not $g_{\omega}$.

Alternatively, predictive models can be trained conditioned on the retrieval model’s performance. Without loss of generality, this can be seen as an iterative process, where the predictive model is optimized in one iteration and the retrieval model is optionally optimized in the next iteration. With this formulation, the parameters of a predictive model at iteration $t$ can be obtained as follows:

Retrieval and predictive models can be trained jointly. Joint optimization can be modeled end-to-end (explained later in this section) or through multi-task learning. In joint multi-task optimization, for any training instance, both the retrieval results and the predictive model parameters will be updated. For instance, FiD-Light (Hofstätter et al., 2023) generates the documents with positive provenance score in addition to the output text for retrieval-augmented text generation tasks. The generated document IDs are then used for re-ranking the result list. Therefore, this can be seen as a joint optimization of re-ranking and generation.

Following the risk minimization framework, end-to-end optimization in REML can be modeled as follows:

where both parameters sets $\theta$ and $\omega$ get updated simultaneously by optimizing an appropriate loss function for the downstream machine learning task. End-to-end optimization of REML, however, can be challenging. It is mostly due to the top $k$ item selection process of information access models in REML that makes the end-to-end REML model non-differentiable. Existing work make some simplifying assumptions to turn the optimization to a differentiable process. For instance, the RAG model from Lewis et al. (2020b) by marginalizing the retrieved document set to a set of pre-selected documents. A similar approach was later utilized by Sachan et al. (2021) for open-domain question answering. In addition to marginalization, RetGen (Zhang et al., 2022a) and EMDR² (Singh et al., 2021) make a document independence assumption and computes the loss function as a summation over each individual document.

Intrinsic evaluation of a retrieval model focuses on comparing systems according to their isolated retrieval performance. In this case, assuming single-turn retrieval, we can pose two styles of hypothesis. In non-overlapping system comparison, given two model tuples $\mathcal{M}=\langle f_{\theta},g_{\omega}\rangle$ and $\mathcal{M}^{\prime}=\langle f_{\theta^{\prime}},g_{\omega^{\prime}}\rangle$, determine if $\mathop{\mathbb{E}}[\mu(g_{\omega}(f_{\theta}(x)))]>\mathop{\mathbb{E}}[\mu(g_{\omega^{\prime}}(f_{\theta^{\prime}}(x)))]$, where, with some abuse of notation, $f_{\theta}(x)$ and $f_{\theta^{\prime}}(x)$ considers only the query processing for $g_{\omega}$ and $g_{\omega^{\prime}}$. In fixed query processing comparison, given two model tuples $\mathcal{M}=\langle f_{\theta},g_{\omega}\rangle$ and $\mathcal{M}^{\prime}=\langle f_{\theta},g_{\omega^{\prime}}\rangle$, determine if $\mathop{\mathbb{E}}[\mu(g_{\omega}(f_{\theta}(x)))]>\mathop{\mathbb{E}}[\mu(g_{\omega^{\prime}}(f_{\theta}(x)))]$. The choice of metric $\mu$ depends on the task but should be some retrieval metric, unless the retrieval result is not a ranking. Such metrics require some relevance estimate for each item. In the case of REML, this can come from,

1. (1)

   Explicit labels gathered from human raters. This requires instances, targets, and items to be interpretable. ’provenance Labels’ in the KILT benchmark for some tasks such as Natural Questions (Kwiatkowski et al., 2019) and ELI5 (Fan et al., 2019) can be thought as such labels.

2. (2)

   Inferred labels from the target. For example, in QA, we could compute the similarity between a retrieved item and the target. ‘Context Relevance’ from RAGAS (Es et al., 2024) and ARES (Saad-Falcon et al., 2024) can be thought as a variant of this case.

3. (3)

   Attributed labels from the model prediction. For example, in QA, if a model generates an answer correctly, we can try to attribute the answer correctness to each of the retrieved items. This method, drawing inspiration from eRAG (Salemi and Zamani, 2024a), assesses the retrieval model’s performance by quantifying the contribution of each retrieved document towards achieving the correct answer.

Intrinsic evaluation of consumption focuses on comparing systems according to their isolated ability to translate retrieval results into effective predictions. Although extrinsic evaluation measures the effectiveness of the system in general, intrinsic evaluation of consumption focuses on whether a prediction is attributable to retrieval results (e.g., versus information already in the consumption model parameters). In fixed retrieval comparison, given two model tuples $\mathcal{M}=\langle f_{\theta},g_{\omega}\rangle$ and $\mathcal{M}^{\prime}=\langle f_{\theta^{\prime}},g_{\omega}\rangle$, determine if $\mathop{\mathbb{E}}[\mu(f_{\theta}(\mathcal{R}_{\omega,x}),\mathcal{R}_{\omega,x})]>\mathop{\mathbb{E}}[\mu(f_{\theta^{\prime}}(\mathcal{R}_{\omega,x}),\mathcal{R}_{\omega,x})]$. The choice of metric $\mu(\tilde{y},\tilde{\mathcal{R}})$ depends on the task but measures whether the prediction $\tilde{y}$ is related to the retrieval results $\tilde{\mathcal{R}}$. ‘Faithfulness’ from RAGAS (Es et al., 2024) and ARES (Saad-Falcon et al., 2024) can be thought as a variant of this case.