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

Overview

SetPO is a set-level policy optimization method for diversity-preserving reasoning in large language models.

In reinforcement learning with verifiable rewards, improved reasoning accuracy is often accompanied by reduced output diversity, which can eventually lead to mode collapse. SetPO addresses this issue by assigning each rollout a reward based not only on its own return, but also on its marginal contribution to the diversity of the sampled rollout set.

For a rollout set $\Omega = {o_1, \ldots, o_G}$, SetPO augments the original advantage as

$$ \hat A_i = A_i + \lambda s_i, $$

where $s_i$ is a leave-one-out diversity marginal that measures how much rollout $o_i$ improves the diversity of the full set.

SetPO is designed as a lightweight plug-in for standard group-based policy optimization pipelines such as GRPO, GSPO, and DAPO.

Toy example of mode collapse vs. diversity preservation

Figure 1: Evolution of trajectory embeddings in a toy multi-modal bandit environment. Each point corresponds to a sampled trajectory embedding, stars denote correct modes, and colors represent different training stages.

Compared with GRPO, which progressively concentrates probability mass onto a single dominant mode, SetPO maintains coverage over multiple semantic modes throughout training. This toy example highlights the central motivation of SetPO: improving reasoning performance while preserving output diversity.

Highlights

• Set-level diversity credit assignment.
  Instead of rewarding trajectories independently, SetPO evaluates each rollout through its contribution to the diversity of the whole sampled set.

• Trajectory-level semantic diversity.
  Diversity is defined over complete reasoning trajectories rather than token entropy or local pairwise penalties.

• Simple integration into existing RL pipelines.
  SetPO only modifies the advantage term, making it easy to combine with existing group-based policy optimization methods.

• Consistent empirical gains.
  SetPO improves strong baselines across multiple model scales and reasoning benchmarks.

Method

SetPO measures diversity over complete semantic trajectories. It defines

$$ m_P(y) = \mathbb{E}[k(y, y')], \qquad \mathcal{F}(P) = \mathbb{E}[g(m_P(y))], \qquad g(x) = -\log(1+x), $$

where:

• $k:\Omega \times \Omega \to [0,1]$ is a bounded semantic similarity kernel over complete trajectories.
• $m_P(y) := \mathbb{E}_{y' \sim P}[k(y,y')]$ is the kernelized local mass around trajectory $y$ under distribution $P$.
• Smaller $m_P(y)$ indicates that $y$ lies in a sparser semantic region and is therefore less redundant.

Given a sampled rollout set $\Omega$, the leave-one-out diversity marginal is

$$ s_i = D(\Omega) - D(\Omega \setminus {o_i}). $$

This gives higher credit to trajectories that expand semantic coverage, while naturally down-weighting redundant ones.

Since SetPO only changes the advantage term, it can be used as a clean drop-in modification to standard group-based policy optimization pipelines.

Main Results

Qwen2.5-Math-7B

Method	GSM8K	MATH500	College Math	AMC23	AIME24	AIME25	Avg
GRPO	89.1	73.5	42.2	53.5	15.4	9.7	47.2
SetPO + GRPO	92.2	80.8	48.3	60.5	21.6	13.6	52.8
GSPO	89.5	73.0	45.0	57.7	17.4	7.9	48.4
SetPO + GSPO	91.6	75.7	47.2	60.9	21.1	10.1	51.1
DAPO	92.2	77.6	47.1	59.6	21.7	11.0	51.7
SetPO + DAPO	93.0	79.2	49.5	62.3	25.3	13.5	53.8

Qwen2.5-Math-1.5B

Method	GSM8K	MATH500	College Math	AMC23	AIME24	AIME25	Avg
GRPO	83.0	67.2	43.1	47.5	11.5	8.2	43.4
SetPO + GRPO	86.9	72.9	45.7	51.3	13.4	9.7	46.7
GSPO	83.5	67.2	43.0	48.8	11.1	5.6	43.2
SetPO + GSPO	85.0	69.4	44.7	50.6	13.2	7.4	45.2
DAPO	85.9	69.5	44.3	50.1	12.0	6.6	44.7
SetPO + DAPO	86.6	71.7	46.5	53.2	13.4	8.9	46.7

Additional Findings

• SetPO improves Pass@K on Countdown across multiple decoding temperatures, rollout counts, and KL settings.
• On AIME24, SetPO remains stronger than the base model even at large sampling budgets.
• SetPO improves diversity scores under LLM-based evaluation.
• Adding SetPO to GRPO, GSPO, and DAPO introduces less than 10% wall-clock overhead in the reported experiments.

Repository Structure

• verl/: vendored verl source tree used by the original training stack
• scripts/: training launcher scripts for SetPO
• assets/: paper PDF and figures used by the project page

Installation

Install dependencies with:

pip install -r requirements.txt

This project builds upon the excellent verl framework. For environment setup, please also refer to the installation strategy and dependency recommendations from the official verl project.

Quick Start

The main training entry released in this repository is:

scripts/run_setpo.sh

The localized script expects the following external assets:

• MODEL_PATH: path to the base model
• EMBEDDING_MODEL_PATH: path to the embedding model
• TRAIN_FILE: path to the training parquet file
• VAL_FILE: path to the validation parquet file
• PYTHON_BIN: Python executable for the training environment

Example:

MODEL_PATH=/path/to/Qwen2.5-Math-7B \
EMBEDDING_MODEL_PATH=/path/to/Qwen3-Embedding-0.6B \
TRAIN_FILE=/path/to/train.parquet \
VAL_FILE=/path/to/test.parquet \
PYTHON_BIN=/path/to/python_env \
bash scripts/run_setpo.sh

Acknowledgements

This project builds upon the excellent verl framework. We thank the verl team for open-sourcing such a practical and powerful RL training system.

Citation

If you find this repository useful, please consider citing our work:

@article{setpo2026,
  title={SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning},
  author={Chenyi Li and Yuan Zhang and Bo Wang and Guoqing Ma and Wei Tang and Haoyang Huang and Nan Duan},
  year={2026},
  eprint={2602.01062},
  archivePrefix={arXiv},
  primaryClass={cs.AI},
  url={https://arxiv.org/abs/2602.01062}
}