A Statistical Framework for Data-dependent Retrieval-Augmented Models
This work addresses a fundamental gap in machine learning for systems that use retrieval augmentation, though it appears incremental as it builds on existing concepts with a new framework.
The authors tackled the problem of understanding and training retrieval-augmented models by proposing a statistical framework with a retriever and predictor, establishing excess risk bounds and validating the approach on open domain question answering.
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.