Source-linked AI summary

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.25659v1cs.AI

TL;DR

GRPO applies response-level advantages uniformly across tokens, limiting within-response credit assignment despite rubric-level supervision. CoRT uses counterfactual replay to derive normalized token weights, improving over response-level GRPO in most comparisons while remaining competitive with learned token-relevance methods.

  • Problem

    GRPO aggregates rubric judgments into a response-level advantage applied uniformly to tokens, leaving criterion-specific within-response credit assignment underused.

  • Method

    CoRT rescores fixed sampled responses with and without rubric criteria, converts tokenwise likelihood contrasts into normalized weights, and redistributes the original GRPO advantage.

  • Results

    CoRT improves over response-level GRPO in most comparisons across instruction-following benchmarks and reward granularities, while remaining competitive with learned token-relevance methods.

  • Takeaways & Limitations

    Policy-internal counterfactual likelihood contrasts can guide within-response credit allocation while preserving response-level supervision and GRPO’s simplicity.

  • Takeaways & Limitations

    CoRT uses a frozen scoring policy for the current batch when computing its replay signal.

Abstract

from arXiv · show

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.

Introduction

Rubric-based RL offers structured criteria, but GRPO-style training collapses them into a response-level advantage broadcast uniformly across tokens. CoRT addresses this gap by using fixed-response counterfactual likelihood contrasts to allocate bounded, rubric-sensitive token credit without changing the original reward or requiring a separate relevance model.

  • Motivation: GRPO-style pipelines reduce rubric judgments to a single response-level advantage, applying uniform credit even when tokens satisfy different requirements.This creates a within-response credit-assignment gap despite structured supervision.
  • Motivation: Fixed-response replay compares rubric-conditioned and criteria-free rescoring, revealing which observed tokens’ likelihood depends more on rubric criteria.Holding response tokens fixed isolates policy-likelihood changes rather than introducing relabeled trajectories or new data.
  • Method: CoRT proposes token-level credit weighting for rubric-conditioned GRPO using policy-internal likelihood contrasts from counterfactual replay.The intervention keeps sampled tokens fixed while removing rubric criteria during rescoring.
  • Method: CoRT converts likelihood contrasts into bounded, response-normalized weights that redistribute the signed response-level GRPO advantage toward rubric-dependent tokens.The redistribution preserves the update direction while concentrating credit on tokens more affected by the rubric context.
  • Contribution: Unlike approaches that train a token-level relevance model, CoRT avoids a separate relevance-learning stage while retaining the original response-level reward.The method changes token-level credit allocation rather than the response-level reward itself.

Related Work

Prior work advances instruction following, rubric feedback, policy optimization, and fine-grained credit assignment, but criterion-level rubric information is not directly used to identify which response parts satisfy specific requirements. CoRT addresses this gap by using rubric feedback for within-response, token-level credit assignment.

  • Instruction following and rubric feedback: Instruction-following research improves alignment through curated instruction data, instruction synthesis, and preference optimization, while recent benchmarks emphasize explicit constraints and fine-grained instructions.
  • Instruction following and rubric feedback: Rubric or verifier criterion information is not directly used to determine which response parts satisfy or violate specific requirements.
  • Policy optimization for language models: PPO uses a clipped policy-gradient surrogate, whereas GRPO removes the value model by normalizing rewards within response groups for scalable language-model reinforcement learning.
  • Fine grained credit assignment: Fine-grained credit-assignment methods address sparse outcome-only rewards by scoring intermediate reasoning steps or assigning implicit, hindsight, and hierarchical rewards in long-horizon agentic reinforcement learning.

Rubric Signals in Group-Relative RL

Standard GRPO converts rubric-verifier outcomes into group-relative response advantages but broadcasts each advantage uniformly across all generated tokens. CoRT instead redistributes the unchanged signed advantage using rubric-dependent weights derived from counterfactual replay and response normalization.

  • Rubric Signals in Group-Relative RL: Rubric verifiers assign scalar rewards to sampled responses, which GRPO converts into group-relative advantages using the group reward mean and standard deviation.For each prompt group, responses are sampled under the full rubric prompt x+ = (x, c).
  • Rubric Signals in Group-Relative RL: CoRT replays each sampled response under the criteria-free prompt, contrasts token log probabilities with the full-prompt rollout, and maps the contrasts into bounded token weights.The pipeline derives tokenwise scores from the two prompt conditions rather than training an auxiliary token scorer.
  • Rubric Signals in Group-Relative RL: Standard GRPO applies the same response-level advantage coefficient to every token, leaving token credit agnostic to which tokens express rubric criteria.Although the loss is evaluated tokenwise, the advantage remains uniform across each response.
  • Rubric Signals in Group-Relative RL: CoRT redistributes the signed advantage within each response through response-normalized weights, preserving the response-level reward and group-relative normalization.The shaped token advantage is bA_i,t = sg(w_i,t)A_i.

CoRT

CoRT transforms response-level GRPO advantages into dense, token-level credit by replaying the same sampled responses under criteria-conditioned and criteria-free prompts. It bounds and schedules these replay-derived weights, normalizes them within each response, and preserves the original advantage’s reinforcement direction.

  • Counterfactual replay: CoRT adds counterfactual replay for the same sampled responses, comparing factual and criteria-free token log probabilities to convert sparse response-level advantages into dense token credit.The sampled responses and rewards remain fixed; replay computes per-token log probabilities under the criteria-free prompt.
  • Counterfactual replay: Positive replay contrasts indicate stronger relative dependence on rubric criteria, while small or negative contrasts indicate weak or reversed dependence.The contrast is used as a proxy for rubric dependence, not as a calibrated measure of token importance.
  • Weight construction: CoRT bounds unbounded, potentially heavy-tailed replay contrasts before constructing token credit weights.The bounded score is centered by b, controlled in sharpness by τ, and lies in (−1/2, 1/2).
  • Weight construction: A cubic Hermite SmoothStep schedule activates replay weighting gradually, keeping early updates close to GRPO and introducing token-level weighting after stabilization.The schedule is zero before the ramp, one after it, and has zero slope at both endpoints.
  • Advantage shaping: Response normalization redistributes the signed GRPO advantage across tokens while retaining its original scale and reinforcement or suppression direction.Token weights control update magnitude; the group-relative advantage determines whether tokens are reinforced or suppressed.
  • Computation: CoRT requires one additional criteria-free scoring pass, reusing generated tokens without extra generation, verifier calls, rejection sampling, or token-level relevance labels.Because factual rollout log probabilities are already available, added computation is limited to counterfactual scoring and response normalization.

Experimental Setup

The experiments train Qwen3-4B-Instruct and Qwen2.5-7B-Instruct with CSR and AON rewards, using paired rubric-conditioned and criteria-free contexts for CoRT replay scoring. Validation spans four instruction-following benchmarks under standardized sampling settings, while CoRT redistributes advantages through bounded replay-derived token weights without changing the reward definition.

  • Training data: HIR-16k provides prompts both without instruction lists and with corresponding lists, enabling CoRT’s rubric-conditioned x+ and criteria-free x−only contexts.x+ is used throughout the main reinforcement-learning loop, while x−only is used for counterfactual scoring.
  • Reward modes: CSR rewards the fraction of satisfied positive criteria, whereas AON assigns one only when all positive criteria are satisfied.These reward modes are used for reinforcement-learning training; validation uses each benchmark’s official metrics.
  • Models and optimization: The main matrix trains Qwen3-4B-Instruct and Qwen2.5-7B-Instruct under both CSR and AON, with batch size 64 and 8 sampled responses per prompt.Runs use maximum prompt and response lengths of 2048 and 4096, rollout temperature 1.0, top-p 0.99, top-k 100, actor learning rate 10^-6, warmup ratio 0.03, and weight decay 0.1.
  • CoRT instantiation: CoRT uses a zero-centered replay-margin transform with unit temperature, η = 0.5, a 100-step SmoothStep ramp, and response-mean normalization.Before response normalization, active token multipliers lie in (0.75, 1.25); replay-derived weights apply to both positive and negative group-relative advantages without changing the reward definition.
  • Evaluation: Validation covers IFBench, IFEval, MultiDimIF, and AdvancedIF, evaluating all methods with 5 sampled responses per prompt at temperature 0.7, top-p 0.8, and top-k 20.IFBench and IFEval report strict prompt-level and instruction-level accuracy, MultiDimIF reports multidimensional constraint accuracy, and AdvancedIF uses its judge-based protocol with o3-mini.

Results

CoRT improves over matched response-level GRPO across most reported comparisons and remains competitive with RTT while avoiding a learned relevance discriminator. Its effectiveness extends to Qwen3-14B and alternative policy optimization objectives, with metric-specific variation under sparse rewards.

  • Main comparison: CoRT improves over response-level GRPO in most comparisons across models, reward settings, and benchmarks.The controlled comparisons use the same rollout prompts, scalar rewards, and group-relative normalization.
  • Main comparison: CoRT improves over matched GRPO across the reported metrics on AdvancedIF under the fixed-step judge-based protocol.The gain holds under a fixed-step evaluation setting rather than only the main benchmark matrix.
  • Comparison with learned token relevance: CoRT is competitive with and often stronger than RTT while avoiding a separate relevance-learning stage.RTT trains a token-level relevance discriminator, whereas CoRT derives token weights from policy-internal counterfactual likelihood contrasts.
  • Qwen3-14B results: At Qwen3-14B scale, CoRT improves all reported CSR metrics and improves IFBench and MultiDimIF under AON, but trails GRPO on IFEval.The results indicate effectiveness at larger scale with metric-specific variation under sparse rewards.
  • Compatibility with policy optimization objectives: Adding CoRT improves all reported metrics over DAPO and four of five metrics over GSPO, with only a small decrease on MultiDimIF.The integrations use CSR reward on Qwen3-4B-Instruct-2507 and demonstrate compatibility with different policy optimization objectives.

Training Diagnostics

Training diagnostics show that CoRT’s full integration recipe stabilizes token-weight scale and optimization dynamics better than ablations. Response normalization controls coefficient scale, while scheduled activation reduces abrupt changes in token-weighted updates.

  • Integration ablations: The full recipe keeps mean token weights near one, length clipping near zero, and actor gradient norm and entropy stable.These diagnostics evaluate Qwen2.5-7B under the CSR reward across the full training run.
  • Integration ablations: Removing response normalization drives mean token weight upward, approaching 1.04 near training’s end, alongside sharp increases in length clipping, gradient norm, and entropy.The resulting scale drift suggests unnormalized replay weights can act as a response-level multiplier.
  • Integration ablations: Removing the SmoothStep ramp preserves response-normalized scale control but produces elevated mean weights and higher length clipping than the full recipe.This indicates scheduled activation addresses abrupt changes in the token-weighted update rather than coefficient scale alone.
  • Validation continuation: Around the nominal 500-step horizon, the full recipe retains the strongest validation average and highest 10-step trailing training reward, while ablations move more between checkpoints.IFBench is an exception, with some ablations achieving higher accuracy despite weaker reward, IFEval, or MultiDimIF behavior.
  • Overall diagnosis: Overall, only the full recipe keeps average token coefficients on the GRPO scale while maintaining low length clipping, stable gradient norms, and the strongest average validation and training reward.The diagnostics therefore support distinct roles for response normalization and scheduled activation in this setting.

Conclusion

CoRT addresses GRPO’s within-response credit mismatch by converting counterfactual likelihood contrasts into normalized token weights that redistribute the signed advantage. Across benchmarks and reward granularities, it improves over response-level GRPO in most comparisons and remains compatible with alternative optimizers.

  • These weights redistribute the signed GRPO advantage across tokens instead of assigning the same response advantage to every generated token.This directly addresses the within-response mismatch between tokenwise policy updates and response-level criteria rewards.
  • CoRT uses counterfactual replay without the criteria to derive normalized token credit weights from likelihood contrasts.The weights target tokens that depend more strongly on the criteria context.
  • CoRT improves over response-level GRPO in most comparisons across instruction-following benchmarks and reward granularities.
  • CoRT remains competitive with learned token relevance methods and gains consistently when combined with DAPO and GSPO.These results indicate replay-derived token credit can integrate with policy optimization algorithms beyond standard GRPO.
  • Response normalization and scheduled activation help keep token weights stable during training.

Additional Results

Additional analyses show that CoRT’s replay contrasts align with visible rubric-controlled spans and remain localized under matched-context and criterion-wise controls. Complementary evaluations check whether instruction-following optimization changes general math and knowledge capabilities.

  • Checkpoint analysis: Table 4 reports CoRT at exactly step 500 as a transparent fixed-checkpoint analysis, while the main comparison retains best-validation-checkpoint selection within the first 500 steps.The fixed-checkpoint table does not replace the main selection rule.
  • General capability checks: Table 5 evaluates original instruction-tuned models against step-500 CoRT checkpoints on mean@5 Math500, GPQA-Diamond, and MMLU-Pro results.The CSR and AON comparisons use the same original checkpoint within each model family, while only the CoRT checkpoint varies with reward granularity.
  • Replay examples: CoRT’s replay contrasts highlight visible rubric-controlled tokens, including quotation marks, bullet-count markers, Markdown bullets, postscript prefixes, and fixed-phrase boundaries.These tokens are reported as distinct from generic content words in the balanced-diet HIR example.
  • Replay controls: Matched-context controls show larger average contrasts for criteria-controlled tokens than for other tokens under criteria-free, neutral-filler, and unrelated-criteria replays.The controls test whether the signal is explained by prompt removal, prompt length, or merely the presence of an instruction block.
  • Replay controls: Criterion-wise removal produces localized effects: removing bullet, highlight, or section criteria emphasizes their corresponding markers or words, unlike global-criterion removal.The pandemic-economy response provides the local control case for this analysis.
Loading 2607.25659v1…