Retrieval-augmented generation (RAG) has become a widely recognized paradigm to combine parametric memory with non-parametric memories. An RAG model consists of two serial connecting components (retriever and generator). A major challenge in end-to-end optimization of the RAG model is that marginalization over relevant passages (modeled as discrete latent variables) from a knowledge base is required. Traditional top-K marginalization and variational RAG (VRAG) suffer from biased or high-variance gradient estimates. In this paper, we propose and develop joint stochastic approximation (JSA) based end-to-end training of RAG, which is referred to as JSA-RAG. The JSA algorithm is a stochastic extension of the EM (expectation-maximization) algorithm and is particularly powerful in estimating discrete latent variable models. Extensive experiments are conducted on five datasets for two tasks (open-domain question answering, knowledge-grounded dialogs) and show that JSA-RAG significantly outperforms both vanilla RAG and VRAG. Further analysis shows the efficacy of JSA-RAG from the perspectives of generation, retrieval, and low-variance gradient estimate.

翻译：检索增强生成（RAG）已成为将参数化记忆与非参数化记忆相结合的广受认可的范式。RAG模型由两个串行连接的组件（检索器与生成器）构成。对RAG模型进行端到端优化的主要挑战在于，需要从知识库中对相关段落（建模为离散隐变量）进行边缘化处理。传统的Top-K边缘化方法与变分RAG（VRAG）存在梯度估计有偏或高方差的问题。本文提出并发展了基于联合随机逼近（JSA）的RAG端到端训练方法，称为JSA-RAG。JSA算法是EM（期望最大化）算法的随机扩展，在估计离散隐变量模型方面具有显著优势。我们在两个任务（开放域问答、知识支撑对话）的五个数据集上进行了大量实验，结果表明JSA-RAG显著优于原始RAG与VRAG。进一步分析从生成质量、检索效果及低方差梯度估计等角度验证了JSA-RAG的有效性。