SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning

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.

Large language models can solve complex problems, such as math questions, more effectively when trained with reinforcement learning using verifiable feedback. However, this often leads to mode collapse, where the model repeatedly produces similar solutions. In this work, we introduce SetPO, an algorithm that encourages more diverse outputs. SetPO evaluates how much each generated solution contributes to overall diversity within a batch of sampled solutions using a leave-one-out marginal contribution: rarer solutions have a greater impact. This information is then used to guide training, helping the model balance correctness and diversity. Experiments across multiple model sizes show that SetPO improves both accuracy and the variety of solutions, consistently outperforming strong baselines on standard benchmarks.