ルーブリックベースのRLの自己蒸留による強化

要旨

ルーブリックベースの強化学習は、近年、自由記述タスクにおけるLLMの改善に有望であることが示されている。ルーブリックベースの強化学習の広く認識されている限界は、探索の乏しさである。すなわち、どのロールアウトも満たすことができない基準（未探索基準、UC）は、最適化信号を受け取らない。最近の手法は、ロールアウト中にルーブリック情報を外部ガイダンスとして組み込むことでこの問題に対処しているが、学習と推論のミスマッチを引き起こす。すなわち、ポリシーは外部ガイダンスの下で生成されたロールアウトに基づいて最適化される一方、このガイダンスは推論時には存在せず、自己回帰デコードを通じて誤差が蓄積される。さらに、これらの探索に焦点を当てた手法は、我々が抑制基準（SC）と呼ぶ根本的に異なる故障モードを見落としている。SCとは、一部のロールアウトによって満たされるものの、スカラー報酬の集約が非正の集約アドバンテージを割り当てるため、その学習信号が最適化中に失われる基準である。我々の分析は、SCが驚くほど一般的であることを明らかにしている。訓練全体を通じて57%以上のサンプルがこの故障モードを示し、サンプルあたり平均1.8個のSCが見られる。UCとSCの両方に同時に対処しつつ、学習と推論のミスマッチを導入しないために、我々は基準蒸留ポリシー最適化（CriPO）を提案する。これはオンポリシー自己蒸留を通じてルーブリックベースの強化学習を強化するものである。UCに対しては、CriPOは基準注入自己教師を構築し、局所的前向きKL損失を計算して欠落した行動をポリシーに注入する。SCに対しては、CriPOは反事実的自己教師を用いて、負のアドバンテージを持つロールアウト内の基準関連トークンを特定し、それらのトークンレベルのアドバンテージを正の値に反転させることで、そうでなければ抑制されてしまう有用なパターンを保持する。医学および科学のベンチマークでの実験は、CriPOがルーブリックベースの強化学習を一貫して上回り、約半分の最適化ステップでより強い最終性能を達成することを実証している。

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.