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

Bo-Wen Zhang, Junwei He, Wen Wang, Song-Lin Lv, Wentao Ma, Rongyi Lin, Shuhan Zhong, Lan-Zhe Guo

arXiv:2607.25659 · 2026-07-30 공개 · arXiv · PDF

reinforcement-learning language-models grpo instruction-tuning token-level rubric-guided counterfactual-replay credit-allocation

Abstract

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.

한국어 요약

한 줄 요약

CoRT는 토큰 수준의 크레딧 할당을 위해 카운터패클 레플레이를 활용한 GRPO 개선 방법으로, 평균 4.4%포인트 성능 향상을 보인다.

핵심 기여도

핵심 아이디어

기존 GRPO는 토큰 수준의 크레딧 할당 없이 전체 응답에 대해 동일한 보상을 적용하지만, CoRT는 토큰별로 크레딧을 재분배한다. CoRT는 원본 루브릭 조건이 있는 프롬프트와 조건 없는 프롬프트에서 동일 응답을 재샘플링하여 토큰별 로그-라이클리후드 대비 차이를 계산한다. 이 차이는 루브릭에 의존하는 토큰을 식별하는 데 사용되며, 이를 바탕으로 GRPO의 보상을 토큰 수준으로 재분배한다. 이 방식은 별도의 토큰 관련성 모델 학습 없이도 정책 내부에서 크레딧 할당 신호를 추출할 수 있다는 점에서 혁신적이다.

기술적 접근법

주요 결과

의의 및 한계

CoRT는 토큰 수준 크레딧 할당을 구조화된 루브릭 기반으로 가능하게 하며, 기존 GRPO의 단순성과 안정성을 유지하면서도 정책 업데이트의 효과성을 높인다. 이는 루브릭이 토큰 라이클리후드에 영향을 미치는 경우에만 효과적으로 작동하며, 루브릭이 프롬프트와 중복되거나 응답 보상이 노이즈가 많은 경우 성능이 저하될 수 있다. 또한, CoRT는 토큰 수준 크레딧 할당을 위한 별도 학습 단계를 생략함으로써 훈련 복잡도를 줄인다.

실용적 활용

CoRT는 지시어 기반 대화형 응답 생성, 형식 및 안전성 기준 준수, 다중 기준 평가 등에서 유용하게 활용될 수 있다. 특히, 별도의 토큰 관련성 학습 없이도 정책 내부에서 크레딧 할당을 가능하게 하므로, 대규모 언어 모델의 후처리 훈련 과정을 효율화할 수 있다.