Stochastic RAG: End-to-End Retrieval-Augmented Generation Through Expected Utility Maximization

This paper introduces Stochastic RAG–a novel approach for end-to-endoptimization of retrieval-augmented generation (RAG) models that relaxes thesimplifying assumptions of marginalization and document independence, made inmost prior work. Stochastic RAG casts the retrieval process in RAG as astochastic sampling without replacement process. Through this formulation, weemploy straight-through Gumbel-top-k that provides a differentiableapproximation for sampling without replacement and enables effective end-to-endoptimization for RAG. We conduct extensive experiments on seven diversedatasets on a wide range of tasks, from open-domain question answering to factverification to slot-filling for relation extraction and to dialogue systems.By applying this optimization method to a recent and effective RAG model, weadvance state-of-the-art results on six out of seven datasets.