Shrinking the Variance: Shrinkage Baselines for Reinforcement Learning with Verifiable Rewards
This addresses training stability issues in post-training large reasoning models, though it is incremental as it builds on existing RLVR methods.
The paper tackles the problem of high variance in policy-gradient estimators for Reinforcement Learning with Verifiable Rewards (RLVR) by proposing shrinkage baselines that combine per-prompt and across-prompt means, resulting in lower-variance gradient updates and improved training stability.
Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as a powerful paradigm for post-training large reasoning models (LRMs) using policy-gradient methods such as GRPO. To stabilize training, these methods typically center trajectory rewards by subtracting the empirical mean for each prompt. Statistically, this centering acts as a control variate (or baseline), reducing the variance of the policy-gradient estimator. Typically, the mean reward is estimated using per-prompt empirical averages for each prompt in a batch. Drawing inspiration from Stein's paradox, we propose using shrinkage estimators that combine per-prompt and across-prompt means to improve the overall per-prompt mean estimation accuracy -- particularly in the low-generation regime typical of RLVR. Theoretically, we construct a shrinkage-based baseline that provably yields lower-variance policy-gradient estimators across algorithms. Our proposed baseline serves as a drop-in replacement for existing per-prompt mean baselines, requiring no additional hyper-parameters or computation. Empirically, shrinkage baselines consistently outperform standard empirical-mean baselines, leading to lower-variance gradient updates and improved training stability.