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

TL;DR

SetPO introduces a set-level diversity objective for large language model reasoning by computing a kernelized semantic diversity measure over grouped rollouts. Each trajectory receives a leave-one-out marginal credit that captures its contribution to the batch's diversity, which is added to the standard group-based policy optimization objective with a tunable weight. The authors prove, via a diversity influence function and a precise s_i decomposition, that rarer trajectories yield larger diversity gains and that the marginal contribution is anti-redundant. Empirically, SetPO consistently improves Pass@1 and Pass@K across model scales from 1.5B to 32B on diverse math benchmarks while increasing diversity scores and maintaining modest compute overhead, demonstrating both effectiveness and robustness.

Abstract

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.

Figures (10)

• Figure 1: Evolution of trajectory embeddings during training for GRPO vs. SetPO (stars denote correct modes; color indicates training steps). Left (GRPO): Although trajectories initially populate multiple modes, training progressively concentrates on a single dominant cluster, exhibiting clear mode collapse. Right (SetPO): Embeddings remain distributed across multiple clusters throughout training, continuing to cover the correct modes at late steps. This indicates that SetPO mitigates mode collapse and preserves rollout diversity.
• Figure 2: Performance of SetPO+GRPO and SetPO+GSPO on the Olympiad benchmark. Shaded regions indicate variance across runs. Both methods improve performance and reduce variability, suggesting that SetPO stabilizes the training dynamics.
• Figure 3: Pass@K performance on the Countdown task for SetPO (SetPO+GRPO) versus the GRPO baseline under varying decoding temperatures, training rollout counts, and KL penalties. SetPO consistently achieves higher Pass@K across these settings.
• Figure 4: Diversity scores across the AIME24 benchmark, computed per problem from 16 Gemini-2.5-Flash generated solutions. Scores range from 1 (lowest diversity) to 5 (highest diversity).
• Figure 5: Pass@k performance of SetPO+GRPO versus the GRPO baseline on Qwen2.5-32B. SetPO+GRPO consistently outperforms GRPO across all k, demonstrating robust gains under varying sampling budgets.

Theorems & Definitions (28)

• Theorem 4.2: Influence function of the diversity functional
• Theorem 4.3: Exact decomposition of the LOO contribution
• Theorem 4.4: Strict anti-redundancy of $s_i$
• Theorem 4.5: Dominance ordering
• Theorem 3.2: Diversity Influence Function
• proof
• Lemma 3.4: Proxy reduction with an explicit remainder bound
• proof
• Theorem 3.5: Local monotonicity via a curvature condition
• proof