SimKO: 단순 Pass@K 정책 최적화

초록

검증 가능한 보상을 활용한 강화 학습(RLVR)은 대규모 언어 모델(LLM)의 추론 능력을 향상시켜 왔습니다. 그러나 현재의 RLVR 방법론은 탐색보다는 활용에 체계적인 편향을 보이며, 이는 pass@1 성능은 개선되지만 pass@K(K>1) 성능은 감소하는 것으로 나타났습니다. 이 문제를 이해하기 위해, 우리는 RLVR 방법론의 학습 동역학을 어휘 후보군에 대한 토큰 수준의 확률 분포를 추적하여 분석했습니다. 우리의 분석은 상위 1위 후보가 점점 더 많은 확률 질량을 축적하고 다른 후보들의 확률을 억제하는 일관된 확률 집중 효과를 보여주었습니다. 더 중요한 것은, 이러한 과도한 집중이 pass@K 성능 저하와 상관관계가 있다는 점입니다. 이러한 발견에 영감을 받아, 우리는 과도한 집중 문제를 완화하고 탐색을 촉진하기 위한 Simple Pass@K Optimization(SimKO) 방법을 제안합니다. SimKO는 비대칭적인 방식으로 작동합니다. 검증된 정답 응답에 대해서는 상위 K개 후보의 확률을 높이고, 검증된 오답 응답에 대해서는 상위 1위 후보에 더 강한 패널티를 적용합니다. 우리는 이러한 비대칭 설계가 특히 높은 엔트로피를 가진 토큰에 적용될 때 과도한 집중을 완화하는 데 효과적임을 관찰했습니다. 다양한 수학 및 논리 추론 벤치마크에서 SimKO는 넓은 범위의 K에 대해 일관되게 더 높은 pass@K 성능을 제공하며, RLVR의 탐색을 개선하는 간단한 방법을 제시합니다.

English

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.

SimKO: 단순 Pass@K 정책 최적화

SimKO: Simple Pass@K Policy Optimization

초록

Support