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

arXiv: https://arxiv.org/abs/2605.26385 · HTML: https://mlsanity.com/paper/2605.26385/

## Key facts

- No citing papers in the map yet (data through 14 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.

## Placed in

- Retrieval-augmented generation / [Agentic RAG, agents and memory](https://mlsanity.com/rag/agentic-rag-agents-and-memory/)
- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

## Builds on (most cited first, 22 of 22)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Reinforcement learning](https://mlsanity.com/paper/2405.10369/) (2024)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Top-K Off-Policy Correction for a REINFORCE Recommender System](https://mlsanity.com/paper/1812.02353/) (2018)
- [Unbiased Learning-to-Rank with Biased Feedback](https://mlsanity.com/paper/1608.04468/) (2016)
- [KuaiRec: A Fully-observed Dataset and Insights for Evaluating Recommender Systems](https://mlsanity.com/paper/2202.10842/) (2022)
- [Computationally Efficient Optimization of Plackett-Luce Ranking Models for Relevance and Fairness](https://mlsanity.com/paper/2105.00855/) (2021)
- [Doubly Robust Off-Policy Evaluation for Ranking Policies under the Cascade Behavior Model](https://mlsanity.com/paper/2202.01562/) (2022)
- [User Fairness, Item Fairness, and Diversity for Rankings in Two-Sided Markets](https://mlsanity.com/paper/2010.01470/) (2020)
- [On component interactions in two-stage recommender systems](https://mlsanity.com/paper/2106.14979/) (2021)
- [Result Diversification in Search and Recommendation: A Survey](https://mlsanity.com/paper/2212.14464/) (2022)
- [The Stereotyping Problem in Collaboratively Filtered Recommender Systems](https://mlsanity.com/paper/2106.12622/) (2021)
- [Learning-to-Rank with Partitioned Preference: Fast Estimation for the Plackett-Luce Model](https://mlsanity.com/paper/2006.05067/) (2020)
- [Policy-Gradient Training of Language Models for Ranking](https://mlsanity.com/paper/2310.04407/) (2023)
- [Hitting the High Notes: Subset Selection for Maximizing Expected Order Statistics](https://mlsanity.com/paper/2012.07935/) (2020)
- [Team Performance with Test Scores](https://mlsanity.com/paper/1506.00147/) (2015)
- [Fast Slate Policy Optimization: Going Beyond Plackett-Luce](https://mlsanity.com/paper/2308.01566/) (2023)

## 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.