Reproduce “Stochastic RAG: End-to-End Retrieval-Augmented Generation through Expected Utility Maximization” on OpenResearch

About this paper

Hamed Zamani, Michael Bendersky

Stochastic RAG introduces a novel framework for end-to-end optimization of Retrieval-Augmented Generation (RAG) models, reformulating retrieval as a stochastic process to enable direct maximization of the final output's utility. This approach achieves state-of-the-art results on six out of seven knowledge-intensive KILT benchmark datasets.

Key insights

• End-to-end RAG optimization can be achieved by framing retrieval as a stochastic process and utilizing differentiable approximations like the straight-through Gumbel-top-k method.
• Directly optimizing the entire RAG pipeline for the utility of the final generated output, instead of relying on proxy metrics or separate component optimization, yields superior performance.
• Relaxing common simplifying assumptions regarding document independence and marginalization allows for more realistic and effective learning in RAG models that consume multiple retrieved passages.

Summary by alphaXiv.