{"ID":2878645,"CreatedAt":"2026-06-01T04:54:23.091178241Z","UpdatedAt":"2026-06-01T04:54:23.091178241Z","DeletedAt":null,"paper_url":"https://arxiv.org/abs/2508.18168","arxiv_id":"2508.18168","title":"Improving End-to-End Training of Retrieval-Augmented Generation Models via Joint Stochastic Approximation","abstract":"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.","short_abstract":"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...","url_abs":"https://arxiv.org/abs/2508.18168","url_pdf":"https://arxiv.org/pdf/2508.18168v3","authors":"[\"Hongyu Cao\",\"Yuxuan Wu\",\"Yucheng Cai\",\"Xianyu Zhao\",\"Zhijian Ou\"]","published":"2025-08-25T16:17:16Z","proceeding":"cs.CL","tasks":"[\"cs.CL\"]","methods":"[\"RAG\"]","has_code":false}