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

Chenyi Li^1,*, Yuan Zhang^2,*,†, Bo Wang³, Guoqing Ma², Wei Tang², Haoyang Huang², Nan Duan²

¹Peking University ²JD Explore Academy, China ³Beijing Institute of Technology

^*Equal contribution ^†Project leader

2026

SetPO augments group-based policy optimization with a leave-one-out diversity credit, rewarding each trajectory according to its marginal contribution to diversity within the rollout set.

Trajectory embeddings showing GRPO collapsing toward fewer modes while SetPO maintains coverage across four correct modes. — **Preserving reasoning modes during training.** GRPO progressively concentrates on a dominant mode, while SetPO keeps multiple correct semantic modes covered.

SetPO Method Overview

SetPO integrates each trajectory’s marginal diversity contribution as a plug-in shaping term for group policy optimization. It uses the standard rollout group and the base advantages produced by GRPO, DAPO, or GSPO.

Start from group policy optimization

Sample a group of outputs, compute task rewards, and obtain the original base advantage.

GRPO / DAPO / GSPO → Ā_i

Build the semantic rollout set

Embed complete trajectories and compute pairwise semantic similarities within the same group.

K_ij = k(o_i, o_j)

Estimate set-level diversity

Aggregate local semantic mass over the sampled trajectories to evaluate the generated batch as a whole.

Ω = {o₁, …, o_G} → D(Ω)

Assign the leave-one-out credit

Contrast the full-set estimate with the subset that excludes one trajectory.

s_i = D(Ω) − D(Ω ∖ {o_i})

Plug SetPO into the base optimizer

Add the weighted credit to the base advantage, then use the augmented advantage in the original surrogate objective.

Â_i = Ā_i + λs_i

Base optimizer signal GRPO / DAPO / GSPO supplies Ā_i.

SetPO diversity credit The same rollout group supplies s_i.

Advantage-only plug-in Ā_i + λs_i replaces Ā_i.

Trajectory-level semantic similarity

Unlike methods that focus on diversity at the token level, SetPO evaluates complete outputs semantically. It uses a bounded, symmetric similarity kernel k(y, y′) ∈ [0, 1] and a pre-defined semantic embedding function e(·).

Local mass

\hat{m} (y; S) = \frac{1}{| S | - 1} \sum_{z \in S ∖ {y}} k (y, z)

A high local mass indicates that y resides in a crowded neighborhood, implying redundancy. A low value places y in a sparse region, signifying relative novelty.

Set-level diversity objective

SetPO aggregates the local semantic mass over the sampled set. Let g : [0, 1] → ℝ be a continuous, non-increasing function. The paper adopts g(x) = −log(1 + x).

Set diversity

D (S) = \frac{1}{| S |} \sum_{y \in S} g (\hat{m} (y; S))

Since g is non-increasing, a smaller local mass yields a larger return. This mechanism rewards semantic novelty and prioritizes the preservation of highly unique trajectories.

Leave-one-out margin for a single trajectory

SetPO contrasts the diversity estimate of the full generated batch Ω with that of the subset Ω ∖ {o_i} to quantify the contribution of a single trajectory.

Marginal credit s_i = D(Ω) − D(Ω ∖ {o_i})

A larger s_i indicates that including o_i yields a more significant gain in the estimated diversity of the set. A small or negative value suggests redundancy.

Set-level policy optimization

Starting from the base advantage Ā_i produced by GRPO, DAPO, or GSPO, SetPO adds the marginal diversity contribution as a plug-in shaping term. The coefficient λ controls the trade-off between task performance and diversity.

SetPO advantage Â_i = Ā_i + λs_i

The augmented advantage Â_i directly replaces the original base advantage in the corresponding GRPO, DAPO, or GSPO objective. SetPO modifies only the advantage term; the remaining design of the chosen base optimizer is unchanged.

Experimental Results

SetPO improves strong group-RL baselines across 1.5B, 7B, and 32B models.

53.8% Best 7B average Pass@1

+5.6 Average gain over GRPO at 7B

<10% Additional wall-clock time

Pass at k curves on four benchmarks showing SetPO plus GRPO above the GRPO baseline for a 32B model. — **Scaling to 32B.** SetPO maintains higher Pass@K on AIME 2024, AIME 2025, AMC 23, and Olympiad.

Average Pass@1 across six math benchmarks

Method	1.5B	7B
GRPO	43.4	47.2
SetPO + GRPO	46.7 +3.3	52.8 +5.6
GSPO	43.2	48.4
SetPO + GSPO	45.2 +2.0	51.1 +2.7
DAPO	44.7	51.7
SetPO + DAPO	46.7 +2.0	53.8 +2.1

Bar chart showing higher diversity scores when SetPO is added to GRPO, GSPO, and DAPO. — **Semantic diversity.** SetPO increases diversity for all three policy optimization baselines.

Countdown Pass at k curves comparing SetPO and GRPO under different rollout and KL settings. — **Countdown.** Higher Pass@K across rollout counts, temperatures, and KL settings.

BibTeX

If you find SetPO useful, please cite our work.

Download paper View code

setpo.bib

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