A Statistical Framework for Data-dependent Retrieval-Augmented Models

Recent advancements in machine learning (ML) have not only led to breakthroughs on long-standing challenging tasks across various fields, but they have also inspired a great deal of interest to develop ML models that can solve even harder tasks (Meinhardt et al., 2022; Lewkowycz et al., 2022; Cramer, 2021) or focus on completely new fields (Austin et al., 2021; OpenAI, 2023; Singhal et al., 2023). While scaling the size of parametric ML models, such as neural networks, is becoming the predominant approach to meet such demands (Brown et al., 2020; Chowdhery et al., 2022; Touvron et al., 2023; Dosovitskiy et al., 2021; Dehghani et al., 2023), the excellent performance realized by this approach is marred by drawbacks such as high computational cost, inefficient storage of world knowledge in parameters, lack of transparency in model behavior, and reduced grounding/factuality of model predictions.

Recognizing these shortcomings, retrieval-augmented models (RAMs) have emerged as a promising alternative. Such models typically employ two components, namely retriever and predictor, during inference on a given input instance: The retriever first identifies instance-specific relevant information from a data-store, and then the predictor jointly processes the retrieved information and the input instance to make a final prediction. In practice, RAMs have enjoyed favorable performance vs. compute trade-off (Borgeaud et al., 2021; Das et al., 2021; Thai et al., 2023) as employing moderate-size parametric models as retriever and predictor in a RAM often matches or exceeds the performance of a much larger standalone ML model that directly maps input instances to predictions. Similarly, conditioning prediction on the retrieved information has shown to exhibit improved grounding (Shuster et al., 2021; Lin et al., 2023; Asai et al., 2023). Furthermore, having access to an external corpus can obviate the need to store task-specific world knowledge in model parameters and enable incorporating dynamically evolving knowledge (Izacard et al., 2022; Liska et al., 2022).

Despite these desirable characteristics, training RAMs presents multiple challenges. The natural approach of independently training retriever and predictor can be sub-optimal (Izacard et al., 2022). Moreover, it requires collecting intermediate supervision on the instance-dependent relevant information to retrieve, which is missing in common datasets and expensive to obtain in general. A common strategy to circumvent the lack of intermediate supervision is to perform end-to-end training which presents its own unique challenges in the context of RAMs. Fundamentally, the retrieval corresponds to the non-differentiable discrete operation of selecting relevant information from a data-store, e.g., via top-k selection based on retriever scores, which prevents direct gradient propagation to the entire retriever. Several clever solutions to above-mentioned issues have been proposed in the literature that focus on different training objectives to propagate learning signal from the predictor into the retriever. However, a formal study that unifies these solutions is missing from the literature.

Another key challenge that prevents the resource-efficient development and deployment of RAMs is the limited understanding of their basic properties such as their generalization behavior and expressive power. For instance, how do the retriever and predictor components interact to ensure good task-specific performance? Are there any principles guiding the selection of the retriever and predictor components? How does (size of) the data-store feature in the final performance of a RAM?

In this paper, we address both aforementioned shortcoming in the literature pertaining RAMs. To unify the training of RAMs, we begin with writing down the natural objective function, which somehow has eluded the literature. This natural objective simply minimizes the expected prediction loss, where the expectation is taken over the distribution induced by the retriever. Empirically, we find this objective to be effective on standard benchmarks: NaturalQuestions (NQ; Kwiatkowski et al., 2019) and TriviaQA (Joshi et al., 2017).

As for the generalization and expressive power, we present an excess risk bound for RAMs that captures the effect of retrieval and prediction function classes. The proposed bound allows us to highlight how retriever and predictor components play complementary roles to reduce approximation error as we increase their respective function class complexity. We also capture the role of data store in improving the model performance by reducing the approximation error. On the generalization front, we carefully decouple the generalization term in the excess risk over the predictor and retriever function classes. This allows us to tightly control the generalization term with only logarithmic dependence on the data store size. As a concrete instantiation for our excess risk bounds, we consider feed-forward neural networks of varying depth for both the retriever and the predictor.

To summarize, our main contributions include:

• •

  We present a principled objective for end-to-end training of RAMs focusing on a classification setting (Sec. 2 3) and draw connections between existing approaches for training RAMs (Sec. 3.6).

• •

  We derive excess risk bound highlighting the role played by retriever and predictor functions classes as well as the data-store towards ensuing improved performance by RAMs (Sec. 3.4); capturing the trade off between model capacities at retriever and predictor (Sec. 3.5).

• •

  We validated the utility of the proposed objective on two standard QA benchmarks: NaturalQuestions (NQ) and TriviaQA (Sec. 4).

In this paper, we focus on developing a systematic understanding of RAMs with learned retrievers in a classification setting where the model has access to a data-store. Towards this, we begin by formally defining the problem setup and providing the necessary background along with the notations used.

Let’s first consider the standard classification setting which requires predicting a class in $\mathscr{Y}$ for a given instance $x\in\mathscr{X}$. Assume that $\mathsf{D}_{XY}$ captures the underlying data distribution and one has access to $n$ training examples $\mathscr{S}_{n}\triangleq\{(x_{i},y_{i})\}_{i\in[n]}$ that are independent and identically distributed (i.i.d.) according to $\mathsf{D}_{XY}$. Given $\mathscr{S}_{n}$, one hopes to learn a classifier $f:\mathscr{X}\to\mathbb{R}^{|\mathscr{Y}|}$ that minimizes the miss-classification error:

$\displaystyle R(f)=\mathbb{P}_{(X,Y)\sim\mathsf{D}_{XY}}\big{[}\operatorname*{arg\,max}_{y\in\mathscr{Y}}f^{y}(X)\neq Y\big{]},$

(1)

where $f^{y}(x)$ denotes the score that $f$ assigns to the $y$-th class, given the input instance $x$. Since directly optimizing the miss-classification error or $0/1$-loss poses computational challenges, one typically selects the classifier that minimizes the empirical risk associated with a well behaved surrogate loss function $\ell:\mathbb{R}^{|\mathscr{Y}|}\times\mathscr{Y}\to\mathbb{R}$ on the training sample $\mathscr{S}_{n}$:

$\displaystyle R_{\ell,n}(f)=\frac{1}{n}\sum_{i\in[n]}\ell\big{(}f(x_{i}),y_{i}\big{)}.$

(2)

The (population) risk associated with the surrogate loss function takes the following form:

$\displaystyle R_{\ell}(f)=\mathbb{E}_{(X,Y)\sim\mathsf{D}_{XY}}\big{[}\ell\big{(}f(X),Y\big{)}\big{]}.$

(3)

Different from the standard classification setup described above, we now consider the classification task with access to a data-store: Given an instance $x$, the classifier can potentially leverage a data-store $\mathscr{I}\subseteq\mathscr{Z}$ – a collection of potentially relevant information or evidences, where $\mathscr{Z}$ denotes the space of all possible evidences. Accordingly, one can define the empirical and population risks of a classifier $f(\cdot,\mathscr{I}):\mathscr{X}\to\mathbb{R}^{|\mathscr{Y}|}$ as follows:

	$\displaystyle R_{\ell,\mathscr{I},n}(f)$	$\displaystyle=\frac{1}{n}\sum_{i\in[n]}\ell\big{(}f(x_{i},\mathscr{I}),y_{i}\big{)},$		(4)
	$\displaystyle R_{\ell,\mathscr{I}}(f)$	$\displaystyle=\mathbb{E}\big{[}\ell\big{(}f(X,\mathscr{I}),Y\big{)}\big{]},$		(5)

where expectation is take over in $(X,Y)\sim\mathsf{D}_{XY}$ as well as the possible randomness in $f(\cdot,\mathscr{I})$. However, due its prohibitive computational cost, such a general classifier that directly processes the entire data-store for each prediction is far from how an additional data-store is utilized by ML models in practice.

This motivates us to study the following explicit retrieval-augmented classification setup to utilize the data-store: Given an input instance $x\in\mathscr{X}$, one first retrieves input-dependent supporting evidences $\mathcal{E}^{x}\subset\mathscr{I}$ with the help of a retriever model which has access to the entire data-store $\mathscr{I}$. Now, given $x$ and $\mathscr{E}^{x}$, one invokes a predictor model to predict the class associated with $x$. Thus, a retriever-augmented classification setup consists of two key components models: 1) retriever model and 2) predictor model, which we formally introduce next.

For ease of exposition, we focus on the setting with $k=|\mathscr{E}^{x}_{\theta}|=1,\forall x\in\mathscr{X},$ in our analysis throughout this paper. This corresponds to retrieving a single supporting evidence for each input instance. Our analysis can be generalized to $k>1$ by working with a $\tilde{\mathscr{I}}\subseteq\mathscr{I}^{k}$ as the new data-store and $\tilde{p}_{\theta,\mathscr{I}}(\cdot|x)$ as a distribution over $\tilde{\mathscr{I}}$ obtained by suitably modifying $p_{\theta,\mathscr{I}}$ in (6). For example, when $k$ supporting evidences are sampled with replacement, then the following holds $\forall(z_{1},\ldots,z_{k})\in\mathscr{I}^{k}$.

$\displaystyle\tilde{p}_{\theta,\mathscr{I}}\big{(}(z_{1},\ldots,z_{k})\big{|}x\big{)}=\prod_{j\in[k]}p_{\theta,\mathscr{I}}(z_{j}|x).$

Empirical risk minimization and excess risk for RAMs.  For a pair of retriever and predictor models parameterized by $\theta$ and $\xi$, respectively, we can define the empirical and population risks associated with a (surrogate) loss function $\ell$ as follows:

	$\displaystyle R_{\ell,\mathscr{I},n}(\xi,\theta)$	$\displaystyle=\frac{1}{n}\sum_{i\in[n]}\sum_{z\in\mathscr{I}}p_{\theta}(z|x)\ell\big{(}h_{\xi}(x_{i},z),y_{i}\big{)},$		(8)
	$\displaystyle R_{\ell,\mathscr{I}}(\xi,\theta)$	$\displaystyle=\mathbb{E}\big{[}\ell\big{(}h_{\xi}(X,\mathcal{E}^{X}_{\theta}),Y\big{)}\big{]}.$		(9)

Note that the expectation in (9) is taken over $(X,Y)\sim\mathsf{D}_{XY}$ as well as the randomness involved in the retrieval stage, e.g., sampling the evidences according to $p_{\theta,\mathscr{I}}(\cdot|x)$ in (6). Given a pair of predictor class $\Xi$ and retriever class $\Theta$, let $(\hat{\xi},\hat{\theta})$ denote the predictor-retriever pair obtained via empirical risk minimization (ERM) as follows:

$\displaystyle(\hat{\xi},\hat{\theta})\in\operatorname*{arg\,min}_{(\xi,\theta)\in\Xi\times\Theta}R_{\ell,\mathscr{I},n}(\xi,\theta).$

(10)

Let $\mathscr{F}_{\rm all}$ denote the set of all measurable functions from $\mathscr{X}\times\mathscr{Z}$ to $\mathbb{R}^{|\mathscr{Y}|}$. The optimal risk for the classification with access to the data-store is achieved by the best possible predictor $f^{\ell}_{\rm opt,\mathscr{I}}\in\mathscr{F}_{\rm all}$ when it has access to the best retrieved evidence in $\mathscr{I}$. In particular, we have

$\displaystyle f^{\ell}_{\mathrm{opt},\mathscr{I}}=\operatorname*{arg\,min}_{f\in\mathscr{F}_{\mathrm{all}}}\mathbb{E}\big{[}\min_{z\in\mathscr{I}}\ell(f(X,z),Y)\big{]}.$

(11)

Given $f^{\ell}_{\mathrm{opt},\mathscr{I}}$, we defined the excess risk of a predictor-retriever pair $(\xi,\theta)$ as follows:

$\displaystyle\Delta_{\ell,\mathscr{I}}(\xi,\theta)=R_{\ell,\mathscr{I}}(\xi,\theta)-R_{\ell,\mathscr{I}}(f^{\ell}_{\mathrm{opt},\mathscr{I}})\triangleq R_{\ell,\mathscr{I}}(\xi,\theta)-\mathbb{E}\big{[}\min_{z\in\mathscr{I}}\ell(f^{\ell}_{\mathrm{opt},\mathscr{I}}(X,z),Y)\big{]}.$

(12)

With the formal definition of the classification setting with access to a data-store and the necessary background in place, we proceed to address the two key objectives of this work: 1) Proposing a natural and efficient joint end-to-end training procedure for the predictor-retriever pair in a RAM; and 2) Developing a rigorous statistical understanding of RAMs focusing on the interaction between predictor and retriever components towards reducing overall excess risk.

Recall that training a RAM involves training both the retriever $r_{\theta}:\mathscr{X}\times\mathscr{Z}\to\mathbb{R}$ and the predictor $h_{\xi}:\mathscr{X}\times\mathscr{I}\to\mathbb{R}^{|\mathscr{Y}|}$ components of the model without access to intermediate supervision on retrieval, which is infeasible to obtain in most practical settings. Thus, it becomes critical to devise methods to jointly train $r_{\theta}$ and $h_{\xi}$ with access to only labeled instances $\mathscr{S}_{n}=\{(x_{i},y_{i})\}_{i\in[n]}\subseteq\mathscr{X}\times\mathscr{Y}$ with the predictor guiding the retriever training based on how valuable the retriever-provided evidences are towards the correct final prediction.

Towards this, we leverage the empirical risk from (8) along with the log-loss $\ell(h_{\xi}(x,z),y)=-\log p_{\xi}(y|x,z)$, where $p_{\xi}(y|x,z)$ is defined in (7). In particular, this leads to the following joint end-to-end training objective:

$\displaystyle\mathscr{L}_{n}(\xi,\theta;\mathscr{I})\triangleq R_{{\rm log},\mathscr{I},n}(\xi,\theta)=-\frac{1}{n}\sum_{i\in[n]}\sum_{z\in\mathscr{I}}p_{\theta,\mathscr{I}}(z|x_{i})\cdot\log p_{\xi}(y_{i}|x_{i},z).$

(13)

Note that the objective in (13) aims to improve the end-to-end performance of a RAM in deployment in the sense that the objective aims to minimize the expected loss given the selected evidences as per the retriever-induced distribution. One can use gradient-based methods to jointly minimize the objective in (13) with respect to $(\xi,\theta)$; however, its efficient implementation is non-trivial due to the sum over entire data-store $\mathscr{I}$. In App. C.1, we discuss some approximate design choices. Lastly, please refer to Sec. 3.6 for connections between our proposed objective in (13) and some of the existing end-to-end training approaches for RAMs.

Next, to study the generalization and expressive power of RAMs, we want to bound the excess risk $\Delta_{\ell,\mathscr{I}}(\hat{\xi},\hat{\theta})$ as defined in (12). We consider $\mathscr{X}$ to be a compact subspace of $\mathbb{R}^{d_{x}}$ and, for simplicity, take $\mathscr{X}\subseteq[-1,1]^{d_{x}}$. Similarly, we consider that each retrieval example $z\in\mathscr{I}$ is embedded in the space $[-1,1]^{d_{z}}$. We consider a data-store that polynomially scales with training data size, i.e., $|\mathscr{I}|={\rm poly}(n)$. For the purpose of analysis, we specialize our log-loss to be bounded by $\ell_{\max}>0$, which is given as

$\displaystyle\ell(h_{\xi}(x,z),y)=\min(\ell_{\max},-\log p_{\xi}(y|x,z))=\min\bigg{(}\ell_{\max},\log\Big{(}\sum_{y^{\prime}\in\mathscr{Y}}\exp(h^{y^{\prime}}_{\xi}(x,z))\Big{)}-h^{y}_{\xi}(x,z)\bigg{)},$

where $p_{\xi}(y|x,z)$ and $h^{y}_{\xi}(x,z)$ are defined in (7).