규칙 기반 강화 학습의 자기 증류를 통한 성능 향상

Enhancing Rubric-based RL via Self-Distillation

최근 규칙 기반 강화 학습은 개방형 작업에서 LLM(Large Language Models)의 성능을 향상시키는 데 유망한 결과를 보여주고 있습니다. 규칙 기반 강화 학습의 주요 한계점 중 하나는 제한된 탐색 능력입니다. 즉, 어떤 시행에서도 만족하지 못하는 기준(Unexplored Criteria, UC)은 최적화 신호를 받지 못합니다. 최근 연구에서는 이러한 문제를 해결하기 위해 룰 정보를 시뮬레이션 과정에서 외부 지침으로 활용하지만, 이는 학습과 추론 간의 불일치를 야기합니다. 정책은 외부 지침 하에 생성된 시뮬레이션을 기반으로 최적화되지만, 추론 시에는 이러한 지침이 없으므로 자동 회귀 디코딩 과정에서 오류가 누적될 수 있습니다. 또한, 기존 연구들은 탐색 능력 향상에 초점을 맞추고 있지만, '억제된 기준(Suppressed Criteria, SC)'이라는 근본적으로 다른 실패 모드를 간과합니다. 즉, 일부 시뮬레이션에서는 만족되지만, 스칼라 보상 집계 과정에서 긍정적인 강화 신호를 받지 못하여 학습에 반영되지 않는 경우입니다. 분석 결과, SC는 매우 흔하게 발생하며, 전체 샘플의 57% 이상이 이러한 실패 모드를 나타내며, 평균적으로 샘플당 1.8개의 SC가 존재합니다. 본 연구에서는 학습과 추론 간의 불일치를 도입하지 않고 UC와 SC 문제를 동시에 해결하기 위해, 온-정책 자기 증류를 활용한 기준 기반 정책 최적화(Criterion-Distilled Policy Optimization, CriPO) 방법을 제안합니다. CriPO는 UC 문제를 해결하기 위해 기준 주입 방식의 자기 지도자를 구축하고, 누락된 행동을 정책에 주입하기 위한 로컬 Forward-KL 손실을 계산합니다. SC 문제는 반사실적 자기 지도자를 활용하여 부정적인 강화 신호를 받는 시뮬레이션에서 기준과 관련된 토큰을 찾아내고, 해당 토큰의 강화 값을 긍정 값으로 변경하여 유용한 패턴을 보존합니다. 의학 및 과학 분야의 벤치마크 실험 결과, CriPO는 기존 규칙 기반 강화 학습 방법보다 일관되게 우수한 성능을 나타냈으며, 약 2배 적은 최적화 단계를 사용하여 더 높은 최종 성능을 달성했습니다.

Rubric-based RL has recently shown promise in improving LLMs on open-ended tasks. A widely recognized limitation of rubric-based RL is limited exploration: criteria that no rollout manages to satisfy (Unexplored Criteria, UC) receive no optimization signal. Recent methods address this by incorporating rubric information as external guidance during rollout, yet they introduce a train-inference mismatch: the policy is optimized on rollouts produced under external guidance while this guidance is absent at inference time, causing error accumulation through autoregressive decoding. Moreover, these exploration-focused approaches overlook a fundamentally different failure mode that we term Suppressed Criteria (SC) -- criteria that are satisfied by some rollouts yet whose learning signals are lost during optimization because scalar reward aggregation assigns them non-positive aggregate advantages. Our analysis reveals that SC are remarkably prevalent: over 57% of samples exhibit this failure mode throughout training, with an average of 1.8 SC per sample. To simultaneously address both UC and SC without introducing training-inference mismatch, we propose Criterion-Distilled Policy Optimization (CriPO), which enhances rubric-based RL via on-policy self-distillation. For UC, CriPO constructs a criterion-injection self-teacher and computes a localized forward-KL loss to inject missing behaviors into the policy. For SC, CriPO employs a counterfactual self-teacher to locate criterion-relevant tokens in negative-advantage rollouts and flips their token-level advantages to positive values, preserving useful patterns that would otherwise be suppressed. Experiments on medicine and science benchmarks demonstrate that CriPO consistently outperforms rubric-based RL, achieving stronger final performance with approximately $2\times$ fewer optimization steps.