SimKO: Simple Pass@K Policy Optimization

Ruotian Peng, Yi Ren, Zhouliang Yu, Weiyang Liu, Yandong Wen

Keywords: large language model, reinforcement learning, learning dynamic

TL;DR: We introduces SimKO, a simple optimization strategy that mitigates over-concentration and improves exploration in reinforcement learning, significantly enhancing pass@K performance by maintaining diversity across candidate outputs.

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. Such bias limits the advancement of LLMs’ reasoning frontier. 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.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 4781