CoRT: 토큰 수준의 rubic 기반 정책 최적화를 위한 반사실 재학습 (Counterfactual Replay)

CoRT: Counterfactual Replay for Token-Level Rubric-Guided Policy Optimization

rubric 기반 강화 학습은 언어 모델 학습을 풍부하게 만들 수 있는데, 이는 모델 출력을 명시적인 기준에 따라 평가하기 때문입니다. 그러나 GRPO 스타일 파이프라인에서는 이러한 구조화된 판단들이 스칼라 수준의 응답 보상으로 축소되고, 응답 수준의 이점으로 변환되어 모든 생성된 토큰에 균일하게 전달됩니다. 이는 응답 내에서 명시적인 신용 할당 메커니즘을 제공하지 않으며, 특히 서로 다른 기준이 응답의 다른 부분, 서식 결정 또는 의미적 선택과 관련될 때 문제가 됩니다. 본 논문에서는 rubric 기반 GRPO를 위한 토큰 수준의 신용 가중치 방법인 CoRT를 제안합니다. CoRT는 보조 토큰 점수 모델을 훈련하는 대신, 반사실 재학습을 사용하여 원래의 rubric 기반 프롬프트와 일치된 기준이 없는 프롬프트 모두에서 동일한 샘플 응답을 다시 평가합니다. 결과적으로 얻어지는 토큰별 로그 우도 차이는 rubric 컨텍스트에 대한 의존성의 지표 역할을 합니다. CoRT는 이러한 차이를 경계가 설정되고 응답 수준으로 정규화된 가중치로 변환하고, 이를 사용하여 서명된 GRPO 이점을 토큰 전반에 걸쳐 재분배합니다. 이는 보조 점수기를 도입하거나 응답 수준의 보상을 변경하지 않고 수행됩니다. instruction-tuned 모델 및 다양한 reward granularity에 대한 실험 결과, CoRT는 대부분의 경우 response-level GRPO보다 성능이 향상되었으며, 평균적으로 4.4% 포인트의 이득을 얻었습니다. 또한, 본 방법은 별도의 관련성 학습 단계를 거치지 않고도 학습된 토큰 수준 신용 기준과 경쟁적인 성능을 보입니다. 이러한 결과는 정책 내에서 반사실 로그 우도 차이가 응답 내 신용 할당을 위한 효과적인 훈련 신호가 될 수 있으며, 동시에 GRPO의 단순성과 안정성을 유지할 수 있음을 시사합니다.

Rubric-based reinforcement learning enriches language model training by evaluating model outputs against explicit criteria. Yet in GRPO-style pipelines, these structured judgments are reduced to a scalar response-level reward and converted into a response-level advantage, which is broadcast uniformly to all generated tokens. This leaves no explicit mechanism for allocating credit within a response, even when different criteria are grounded in different spans, formatting decisions, or semantic choices. We propose CoRT, a token-level credit weighting method for rubric-conditioned GRPO. Instead of training an auxiliary token scoring model, CoRT uses counterfactual replay to rescore the same sampled response under the original rubric-conditioned prompt and a matched criteria-free prompt. The resulting tokenwise log-likelihood contrasts serve as a proxy for dependence on the rubric context. CoRT maps these contrasts to bounded, response-normalized weights and uses them to redistribute the signed GRPO advantage across tokens, without introducing an auxiliary scorer or changing the response-level reward. Experiments across instruction-tuned models and reward granularities show that CoRT improves over matched response-level GRPO in the vast majority of comparisons, with an average gain of 4.4 percentage points. The method remains competitive with learned token-level credit baselines while avoiding a separate relevance-learning stage. These results suggest that policy-internal counterfactual likelihood contrasts provide an effective training signal for within-response credit allocation while retaining the simplicity and stability of GRPO.