Shrinking the Variance: Shrinkage Baselines for Reinforcement Learning with Verifiable Rewards

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 cross-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. This baseline serves as a drop-in replacement for existing per-prompt mean baselines, requiring no additional hyperparameters or computation. Empirically, shrinkage baselines consistently outperform standard empirical-mean baselines, leading to lower-variance gradient updates and improved training stability.

Modern reasoning AI systems are trained by having the model practice answering questions and learn from whether its answers are right or wrong. To make this training stable, the system needs an accurate estimate, for each question, of how well it typically performs. The standard approach averages a handful of attempts, but each attempt is expensive, so this average is noisy and training becomes unstable. Our work applies a classical statistics result known as Stein's paradox: when many related quantities must be estimated at once, pooling information across them yields more accurate estimates than treating each independently. Applied to AI training, this means blending a model's success rate on a specific question with its average across the whole batch. The result is provably less noise in the training signal, with no extra cost or tuning knobs. Across math, logic-puzzle, and visual-reasoning experiments, the method yields more stable training and stronger reasoning models, while requiring only a small change to existing pipelines.