SimKO: Simple Pass@K Policy Optimization

TL;DR

RLVR methods bias toward exploitation, shrinking exploration as token-level distributions concentrate on the top-1; this degrades pass@K. SimKO introduces asymmetric gradient redistribution across top-K tokens, identifies high-entropy forking tokens, and applies top-K label smoothing for positive gradients plus stronger penalties for rank-1 negatives to promote exploration. Empirically, SimKO yields consistent improvements in pass@K (up to $K=256$) while maintaining or improving pass@1 across math and logic benchmarks, and analysis shows reduced concentration and preserved token entropy. This simple, principled approach provides a practical route to balance exploitation and exploration in RLVR and generalizes across backbones and tasks.

Abstract

Reinforcement learning with verifiable rewards (RLVR) has advanced the reasoning capabilities of large language models (LLMs). However, prevailing RLVR methods exhibit a systematic bias toward exploitation over exploration, as evidenced by improved pass@1 but reduced pass@K (K>1) performance. To understand this issue, we analyze training dynamics of RLVR methods by tracking the token-level probability distributions over vocabulary candidates. Our analysis reveals a consistent probability concentration effect where the top-1 candidate increasingly accumulates probability mass and suppresses that of other candidates. More importantly, stronger over-concentration correlates with worse pass@K performance. Inspired by this finding, we propose Simple Pass@K Optimization (SimKO), a method designed to mitigate the over-concentration issue, thereby encouraging exploration. SimKO operates in an asymmetrical manner. For verified-correct responses, it boosts the probabilities of the top-K candidates. For verified-incorrect responses, it applies stronger penalties to the top-1 candidate. We observe that this asymmetric design is particularly effective at mitigating over-concentration when applied at tokens with high entropy. Across various math and logical-reasoning benchmarks, SimKO consistently yields higher pass@K for a wide range of K, providing a simple way to improve RLVR's exploration.

Figures (10)

• Figure 1: SimKO improves pass@K performance on math tasks (AIME24/25, AMC, MATH500, Minerva, Olympiadbench) and logic tasks (Synlogic, BBH) compared to GRPO, as shown in the plots (left and middle). The figure on the right shows the $k$-th highest candidate probabilities averaged over the dataset. The SimKO-trained model exhibits a less concentrated probability distribution compared to GRPO.
• Figure 2: (a) The exploration behavior visualized according to the token-level posterior probability. (b) Comparison of two exploration strategy. (c) An example of two distributions with identical entropy but distinct probability distribution.
• Figure 3: (a)-(c) Training dynamics of average log probability $\Lambda$ and top-K probabilities $\Lambda^{(k)}$ derived by GRPO, NSR, and PSR. (d) The corresponding pass@1 and pass@K results of the RLVL-trained models. Following the setups of zhu2025surprising, we train a Llama3.2-3B-Instruct on a mixture of GSM8K and MATH (Level 1) and train Qwen2.5-Math-7B on the MATH dataset. Complete training dynamics are shown in Figure \ref{['fig:appendix-analysis_onpolicy']}.
• Figure 4: Intuition of the proposed method. (a) We begin by identifying the "forking" tokens, which are high-entropy tokens, and diverge into multiple reasoning paths. (b) For positive samples, we redistribute the probability mass from the top-1 candidate to the top-K candidates, mitigating overconcentration. (c) For negative samples, we apply a strong penalty to the top-1 candidate and a weaker penalty to non-top-1 candidates to prevent the squeezing effect, thereby avoiding sharp distributions and facilitating model exploration.
• Figure 5: The pseudo-code of the policy loss computation with SimKO. SimKO only requires modifying a few lines from a standard policy gradient implementation.