SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning
This addresses the issue of narrow solution sets in LLM reasoning for tasks like mathematics, though it is an incremental improvement over existing methods.
The paper tackles the problem of reduced solution diversity in LLM reasoning when using reinforcement learning with verifiable rewards, by introducing a set-level diversity objective that improves both Pass@1 and Pass@K metrics across benchmarks.
Reinforcement learning with verifiable rewards has shown notable effectiveness in enhancing large language models (LLMs) reasoning performance, especially in mathematics tasks. However, such improvements often come with reduced outcome diversity, where the model concentrates probability mass on a narrow set of solutions. Motivated by diminishing-returns principles, we introduce a set level diversity objective defined over sampled trajectories using kernelized similarity. Our approach derives a leave-one-out marginal contribution for each sampled trajectory and integrates this objective as a plug-in advantage shaping term for policy optimization. We further investigate the contribution of a single trajectory to language model diversity within a distribution perturbation framework. This analysis theoretically confirms a monotonicity property, proving that rarer trajectories yield consistently higher marginal contributions to the global diversity. Extensive experiments across a range of model scales demonstrate the effectiveness of our proposed algorithm, consistently outperforming strong baselines in both Pass@1 and Pass@K across various benchmarks.