Improving End-to-End Training of Retrieval-Augmented Generation Models via Joint Stochastic Approximation

• URL: http://arxiv.org/abs/2508.18168v1
• Date: Mon, 25 Aug 2025 16:17:16 GMT
• Title: Improving End-to-End Training of Retrieval-Augmented Generation Models via Joint Stochastic Approximation
• Authors: Hongyu Cao, Yuxuan Wu, Yucheng Cai, Xianyu Zhao, Zhijian Ou,
• Abstract summary: Retrieval-augmented generation (RAG) has become a widely recognized paradigm to combine parametric memory with non-parametric memories.<n>A major challenge in end-to-end optimization of the RAG model is that marginalization over relevant passages is required.<n>In this paper, we propose and develop joint approximation (JSA) based end-to-end training of RAG.<n>The JSA algorithm is an extension of the EM (expectation-maximization) algorithm and is particularly powerful in estimating latent variable models.
• Score: 9.493788719707835
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Predict the Retrieval! Test time adaptation for Retrieval Augmented Generation [66.36556189794526]
  TTARAG is a test-time adaptation method that dynamically updates the language model's parameters during inference to improve RAG system performance in specialized domains.<n>Our method introduces a simple yet effective approach where the model learns to predict retrieved content, enabling automatic parameter adjustment to the target domain.
  arXiv  Detail & Related papers  (2026-01-16T17:07:01Z)
• Attributing Response to Context: A Jensen-Shannon Divergence Driven Mechanistic Study of Context Attribution in Retrieval-Augmented Generation [52.3707788779464]
  We introduce a novel Jensen-Shannon Divergence driven method to Attribute Response to Context (ARC-JSD)<n>ARC-JSD enables efficient and accurate identification of essential context sentences without additional fine-tuning, gradient-calculation or surrogate modelling.<n> Evaluations on a wide range of RAG benchmarks, such as TyDi QA, Hotpot QA, and Musique, using instruction-tuned LLMs in different scales demonstrate superior accuracy and significant computational efficiency improvements.
  arXiv  Detail & Related papers  (2025-05-22T09:04:03Z)
• Accelerating Adaptive Retrieval Augmented Generation via Instruction-Driven Representation Reduction of Retrieval Overlaps [16.84310001807895]
  This paper introduces a model-agnostic approach that can be applied to A-RAG methods.<n>Specifically, we use cache access and parallel generation to speed up the prefilling and decoding stages respectively.
  arXiv  Detail & Related papers  (2025-05-19T05:39:38Z)
• Direct Retrieval-augmented Optimization: Synergizing Knowledge Selection and Language Models [83.8639566087953]
  We propose a direct retrieval-augmented optimization framework, named DRO, that enables end-to-end training of two key components.<n>DRO alternates between two phases: (i) document permutation estimation and (ii) re-weighted, progressively improving RAG components.<n>Our theoretical analysis reveals that DRO is analogous to policy-gradient methods in reinforcement learning.
  arXiv  Detail & Related papers  (2025-05-05T23:54:53Z)
• RAGO: Systematic Performance Optimization for Retrieval-Augmented Generation Serving [9.962031642362813]
  Retrieval-augmented generation (RAG) is emerging as a popular approach for reliable LLM serving.<n>RAG is a structured abstraction that captures the wide range of RAG algorithms.<n> RAGO is a system optimization framework for efficient RAG serving.
  arXiv  Detail & Related papers  (2025-03-18T18:58:13Z)
• Chain-of-Retrieval Augmented Generation [72.06205327186069]
  This paper introduces an approach for training o1-like RAG models that retrieve and reason over relevant information step by step before generating the final answer.<n>Our proposed method, CoRAG, allows the model to dynamically reformulate the query based on the evolving state.
  arXiv  Detail & Related papers  (2025-01-24T09:12:52Z)
• Retrieval-Augmented Generation for AI-Generated Content: A Survey [38.50754568320154]
  Retrieval-Augmented Generation (RAG) has emerged as a paradigm to address such challenges. RAG introduces the information retrieval process, which enhances the generation process by retrieving relevant objects from available data stores. In this paper, we comprehensively review existing efforts that integrate RAG technique into AIGC scenarios.
  arXiv  Detail & Related papers  (2024-02-29T18:59:01Z)
• Understanding Augmentation-based Self-Supervised Representation Learning via RKHS Approximation and Regression [53.15502562048627]
  Recent work has built the connection between self-supervised learning and the approximation of the top eigenspace of a graph Laplacian operator. This work delves into a statistical analysis of augmentation-based pretraining.
  arXiv  Detail & Related papers  (2023-06-01T15:18:55Z)
• Advancing Semi-Supervised Task Oriented Dialog Systems by JSA Learning of Discrete Latent Variable Models [22.249113574918034]
  JSA-TOD represents the first work in developing JSA based semi-supervised learning of discrete latent variable conditional models. Experiments show that JSA-TOD significantly outperforms its variational learning counterpart.
  arXiv  Detail & Related papers  (2022-07-25T14:36:10Z)
• Understanding Overparameterization in Generative Adversarial Networks [56.57403335510056]
  Generative Adversarial Networks (GANs) are used to train non- concave mini-max optimization problems. A theory has shown the importance of the gradient descent (GD) to globally optimal solutions. We show that in an overized GAN with a $1$-layer neural network generator and a linear discriminator, the GDA converges to a global saddle point of the underlying non- concave min-max problem.
  arXiv  Detail & Related papers  (2021-04-12T16:23:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.