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

TL;DR

SetPO introduces a set-level diversity objective for LLM reasoning, balancing solution diversity and performance by leveraging kernelized similarity and marginal contribution analysis, leading to improved results across benchmarks.

Contribution

This paper proposes a novel set-level diversity objective and a policy optimization method that enhances LLM reasoning diversity without sacrificing accuracy.

Findings

Outperforms strong baselines in Pass@1 and Pass@K metrics

Theoretically proves higher marginal contributions from rarer trajectories

Demonstrates effectiveness across various model scales and benchmarks

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.