Search papers, labs, and topics across Lattice.
This paper addresses the challenge of end-to-end training of early-stage rankers (ESRs) in two-stage ranking systems, where vanilla policy gradient (V-PG) suffers from high variance due to large candidate set sizes. They introduce a "credit-assigned" policy gradient (CA-PG) that computes gradients based on the marginal probability of a target item being chosen in any candidate set. Theoretical analysis and experiments on synthetic and real-world data demonstrate that CA-PG reduces variance, improves convergence speed, and enhances training stability for ESRs, particularly with large candidate sets.
Training early-stage rankers in retrieval systems doesn't have to be a high-variance nightmare: CA-PG offers a stable, faster-converging alternative to vanilla policy gradients.
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.
