자기 증류를 통한 루브릭 기반 강화학습 향상

초록

루브릭 기반 강화학습은 최근 개방형 작업에서 대규모 언어 모델(LLM)을 개선하는 데 유망한 성과를 보이고 있다. 루브릭 기반 강화학습의 널리 알려진 한계는 제한된 탐색이다. 어떤 롤아웃도 충족하지 못하는 기준(Unexplored Criteria, UC)은 최적화 신호를 전혀 받지 못한다. 최근 방법들은 롤아웃 중에 루브릭 정보를 외부 지침으로 통합하여 이 문제를 해결하고자 하지만, 학습-추론 불일치를 초래한다. 즉, 정책은 외부 지침 하에 생성된 롤아웃으로 최적화되는 반면, 추론 시에는 이러한 지침이 없어서 자기회귀 디코딩을 통해 오류가 누적된다. 더욱이 이러한 탐색 중심 접근법은 우리가 억제된 기준(Suppressed Criteria, SC)이라고 명명한 근본적으로 다른 실패 모드를 간과한다. SC는 일부 롤아웃에 의해 충족되지만, 스칼라 보상 집계가 비양수 집계 어드밴티지를 할당함으로써 최적화 중에 학습 신호가 손실되는 기준을 의미한다. 우리의 분석은 SC가 매우 만연함을 보여준다. 훈련 전반에 걸쳐 샘플의 57% 이상이 이 실패 모드를 나타내며, 샘플당 평균 1.8개의 SC가 발생한다. UC와 SC를 동시에 해결하면서 학습-추론 불일치를 도입하지 않기 위해, 우리는 온폴리시 자기 증류를 통해 루브릭 기반 강화학습을 향상시키는 기준 증류 정책 최적화(Criterion-Distilled Policy Optimization, CriPO)를 제안한다. UC의 경우, CriPO는 기준 주입 자기 교사를 구성하고 국소적 순방향 KL 손실을 계산하여 정책에 누락된 행동을 주입한다. SC의 경우, CriPO는 반사실적 자기 교사를 사용하여 음의 어드밴티지를 가진 롤아웃에서 기준 관련 토큰을 찾고 해당 토큰 수준 어드밴티지를 양수 값으로 뒤집어, 그렇지 않으면 억제되었을 유용한 패턴을 보존한다. 의학 및 과학 벤치마크 실험은 CriPO가 루브릭 기반 강화학습을 지속적으로 능가하며, 약 2배 더 적은 최적화 단계로 더 강력한 최종 성능을 달성함을 보여준다.

English

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 2times fewer optimization steps.