A Statistical Framework for Data-dependent Retrieval-Augmented Models

TL;DR

This paper introduces a statistical framework for understanding and training retrieval-augmented models in machine learning, focusing on the interaction between retrieval and prediction components and providing theoretical guarantees.

Contribution

It proposes a new end-to-end training method for retrieval-augmented models and offers excess risk bounds that clarify the roles of retriever and predictor.

Findings

Effective end-to-end training method demonstrated on open domain QA

Theoretical excess risk bounds established for retrieval-augmented models

Insights into the contributions of retrieval and prediction components

Abstract

Modern ML systems increasingly augment input instances with additional relevant information to enhance final prediction. Despite growing interest in such retrieval-augmented models, their fundamental properties and training are not well understood. We propose a statistical framework to study such models with two components: 1) a {\em retriever} to identify the relevant information out of a large corpus via a data-dependent metric; and 2) a {\em predictor} that consumes the input instances along with the retrieved information to make the final predictions. We present a principled method for end-to-end training of both components and draw connections with various training approaches in the literature. Furthermore, we establish excess risk bounds for retrieval-augmented models while delineating the contributions of both retriever and predictor towards the model performance. We validate the utility of our proposed training methods along with the key takeaways from our statistical analysis on open domain question answering task where retrieval augmentation is important.