Search papers, labs, and topics across Lattice.
This paper addresses the challenges of sparse rewards in reinforcement learning by optimizing inference-time objectives like pass@K and max@K through improved policy-gradient estimators. The authors introduce a Leave-Two-Out (L2O) baseline that maintains policy-gradient unbiasedness while centering realized batch advantages, resulting in a new method called MaxPO, which operates efficiently in quadratic time. Empirical results demonstrate that the L2O baseline significantly reduces gradient variance and outperforms existing non-centered methods, enhancing the effectiveness of group-based RL for large language model post-training.
Centering advantages in policy gradients can drastically reduce variance and improve performance in reinforcement learning tasks.
Reinforcement learning with verifiable rewards is widely used for post-training reasoning models, but sparse outcome rewards make exploration difficult. A complementary approach is to optimize inference-time objectives such as pass@K and max@K directly, yet existing policy-gradient estimators for these objectives use different signals, baselines, and normalizations, making their relationships unclear. We study this issue through baseline design and advantage centering. Starting from the advantage estimator of a leading method in the field, we show that it is policy-gradient unbiased but yields a non-centered advantage. We then introduce a Leave-Two-Out baseline that preserves policy-gradient unbiasedness while making realized batch advantages exactly centered. The resulting method, MaxPO, has an efficient quadratic-time implementation and integrates naturally into group-based RL for LLM post-training. We further derive the canonical finite-batch advantage for max@K, providing a unified view of existing advantage estimators. Empirically, we verify that the L2O baseline reduces gradient variance and outperforms non-centered alternatives.
