AIAny
Icon for item

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

Allocates token-level credit in rubric-conditioned GRPO by counterfactually replaying the same response under rubric and criteria-free prompts, using tokenwise log-likelihood contrasts to compute bounded, response-normalized weights that redistribute GRPO advantages without training an auxiliary scorer.

Introduction

Most rubric-based GRPO pipelines collapse structured, multi-criteria judgments into a single scalar response reward and then broadcast the resulting advantage uniformly across all tokens. The core insight behind this work is that internal policy likelihoods can be replayed under a matched criteria-free prompt to reveal which tokens depend on rubric context, enabling tokenwise credit assignment without a separate relevance model.

Key Findings
  • Counterfactual replay: Rescoring the same sampled response under the original rubric-conditioned prompt and a matched criteria-free prompt yields tokenwise log-likelihood contrasts that correlate with rubric dependence; those contrasts are mapped to bounded, response-normalized weights.
  • Performance gains: Applying these token weights to redistribute signed GRPO advantages improves downstream metrics over matched response-level GRPO in the majority of experiments, with an average improvement reported as 4.4 percentage points across tested settings.
  • Practical trade: The method matches or approaches learned token-level credit baselines while avoiding an explicit relevance-learning stage or an auxiliary token scorer, preserving GRPO's simplicity and stability.
  • Generality: Evaluated across instruction-tuned models and different reward granularities, the approach is broadly applicable where rubric-conditioned prompts are available and policy likelihoods can be cheaply recomputed.
Who it's for and tradeoffs

Great fit if you train LLMs with rubric-based reward signals and need within-response credit assignment but want to avoid a separate token-relevance model. The method is appealing when you can afford replaying model log-likelihoods for sampled responses and when rubric cues are meaningfully localized to token spans.

Look elsewhere if your training budget cannot tolerate extra likelihood evaluations (counterfactual replay increases compute and I/O), if rubric effects are highly diffuse across tokens so likelihood contrasts are weak, or if you require explicit, human-interpretable token-level explanations beyond relative likelihood-based weights.

Where it fits

CoRT sits between coarse response-level advantage broadcasting (simple but blind to within-response structure) and learned token-level scorers (explicit but costly and potentially unstable). It offers a middle ground that leverages the policy itself for attributions, making it appealing for GRPO-style pipelines that value stability and minimal extra modeling.

How it works (brief)
  1. Sample a response under the rubric-conditioned prompt as usual. 2) Recompute token log-probabilities for the same response under a matched criteria-free prompt (counterfactual). 3) Form tokenwise contrasts (rubric vs criteria-free) and map them to bounded, response-normalized weights. 4) Multiply the signed GRPO advantage by these weights to produce a redistributed token-level advantage used for policy updates.

Overall, CoRT offers a lightweight, policy-internal signal for token-level credit allocation that improves training outcomes in rubric-guided RLHF/GRPO scenarios while avoiding the complexity of learning a separate token relevance model.

Information

  • Websitearxiv.org
  • AuthorsBo-Wen Zhang, Junwei He, Wen Wang, Song-Lin Lv, Wentao Ma, Rongyi Lin, Shuhan Zhong, Lan-Zhe Guo
  • Published date2026/07/28

More Items

Rewinds world-model trajectories to test alternative robot actions and turn outcome comparisons into action-level training signals. Real-robot data continually corrects the model; across three physical tasks, success rates beat the strongest baseline by 10.67–11.85 percentage points.

Trains an 8B model to search, synthesize evidence, and write cited research reports. Its reinforcement-learning rewards use evolving rubrics grounded in retrieved evidence and contrasting model responses.

Trains long-horizon web researchers by combining agentic mid-training, synthetic trajectories, and reinforcement learning. A 30.5B-parameter model activates 3.3B parameters per token and learns to plan, search, and synthesize evidence.