mlsanityData to 7 Oct 2026

Credit-assigned Policy Gradient for Early Stage Retrieval in Two-stage Ranking

arXiv 2605.26385 ↗First posted 25 May 2026cs.IR cs.AI stat.MLcited by 0

On the map: Retrieval-augmented generation → Agentic RAG, agents and memory

Abstract (arXiv)

Large-scale search, recommendation, and retrieval-augmented generation (RAG) systems typically employ a two-stage architecture: an early-stage ranker (ESR) generates a candidate set, which is subsequently re-ranked by a late-stage ranker (LSR). While there are many reinforcement learning (RL) methods for training the LSR, end-to-end training of the ESR has proven challenging. In particular, naive application of "vanilla" policy gradient (V-PG) is not scalable for candidate-set sizes relevant for practical use due to exploding variance. This issue arises because V-PG propagates the gradient to the joint probability of the candidate sets, ignoring the contribution of each specific item in the candidate set to the reward. To mitigate this issue, we propose a novel "credit-assigned" policy gradient (CA-PG), which computes gradients with respect to the probability that the target item is chosen in any candidate set, i.e. marginalizing over all candidate sets that contain it. Our theoretical analysis reveals that CA-PG significantly reduces the variance of V-PG by marginalizing over the specific composition of the candidate set, while preserving the ability to learn the correct ranking of items under a reasonably aligned LSR policy. Experiments on both synthetic and real-world data demonstrate that CA-PG improves the convergence speed and training stability for ESRs utilizing the canonical Plackett-Luce model, especially when the candidate-set size is large.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Agentic RAG, agents and memory, one of 15 directions in Retrieval-augmented generation. Builds on 22 of its 40 references that are in the map. First posted on arXiv 25 May 2026.

Builds on

22 of 40 references.

References · 22

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  2. 2017Proximal Policy Optimization Algorithmscited by 14,441
  3. 2020Retrieval-Augmented Generation for Knowledge-Intensive NLP Taskscited by 9,160
  4. 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 8,851
  5. 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Modelcited by 7,916
  6. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,873
  7. 2024Reinforcement learningcited by 236
  8. 2019Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacementcited by 205
  9. 2018Top-K Off-Policy Correction for a REINFORCE Recommender Systemcited by 181
  10. 2016Unbiased Learning-to-Rank with Biased Feedbackcited by 155
  11. 2022KuaiRec: A Fully-observed Dataset and Insights for Evaluating Recommender Systemscited by 94
  12. 2021Computationally Efficient Optimization of Plackett-Luce Ranking Models for Relevance and Fairnesscited by 35
10 more
  1. 2022Doubly Robust Off-Policy Evaluation for Ranking Policies under the Cascade Behavior Modelcited by 30
  2. 2020User Fairness, Item Fairness, and Diversity for Rankings in Two-Sided Marketscited by 23
  3. 2021On component interactions in two-stage recommender systemscited by 14
  4. 2022Result Diversification in Search and Recommendation: A Surveycited by 11
  5. 2021The Stereotyping Problem in Collaboratively Filtered Recommender Systemscited by 10
  6. 2020Learning-to-Rank with Partitioned Preference: Fast Estimation for the Plackett-Luce Modelcited by 8
  7. 2023Policy-Gradient Training of Language Models for Rankingcited by 5
  8. 2020Hitting the High Notes: Subset Selection for Maximizing Expected Order Statisticscited by 4
  9. 2015Team Performance with Test Scorescited by 3
  10. 2023Fast Slate Policy Optimization: Going Beyond Plackett-Lucecited by 3

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026F-GRPO: Factorized Group-Relative Policy Optimization for Unified Candidate Generation and Ranking
  2. 2025Optimizing Retrieval for RAG via Reinforcement Learning
  3. 2026Optimizing RAG Rerankers with LLM Feedback via Reinforcement Learning
  4. 2026Spend Search Where It Pays: Value-Guided Structured Sampling and Optimization for Generative Recommendation
  5. 2026ProRAG: Process-Supervised Reinforcement Learning for Retrieval-Augmented Generation
  6. 2026On the Policy Gradient Foundations of Group Relative Policy Optimization: Credit Assignment, Gradient Sparsity, and Rank Collapse
  7. 2025Domain-Aware RAG: MoL-Enhanced RL for Efficient Training and Scalable Retrieval
  8. 2026Token-Level Credit Assignment Optimization for Generative Document Retrieval ↗