Summary (Overview)

• CoRT (Counterfactual Replay for Token-Level Rubric-Guided Policy Optimization) addresses the token-level credit assignment problem in rubric-based RL for LLMs, where GRPO uniformly broadcasts response-level advantages to all tokens.
• The method uses a counterfactual replay pass: the same sampled response is rescored under a criteria-free prompt, and the log-likelihood contrast between the two contexts serves as a proxy for rubric dependence.
• This contrast is mapped to bounded, response-normalized token weights that redistribute the signed GRPO advantage toward tokens more dependent on the rubric criteria, without training an auxiliary token scorer or changing the response-level reward.
• Experiments across Qwen3-4B, Qwen2.5-7B, and Qwen3-14B on instruction-following benchmarks (IFBench, IFEval, MultiDimIF, AdvancedIF) show that CoRT improves over response-level GRPO in the vast majority of comparisons, with an average gain of 4.4 percentage points, and remains competitive with learned token-level credit baselines (RTT) while avoiding a separate relevance-learning stage.
• CoRT is also compatible with different policy optimization objectives (DAPO, GSPO) and training diagnostics confirm that response normalization and scheduled activation (SmoothStep ramp) stabilize training dynamics.

Introduction and Theoretical Foundation

Background: Reinforcement learning (RL) is central to LLM post-training, enabling models to follow instructions with complex constraints. Rubric-based RL decomposes response behavior into explicit criteria (e.g., correctness, formatting, safety), providing structured supervision beyond scalar rewards. However, in GRPO-style pipelines, these structured judgments are aggregated into a response-level advantage and applied uniformly to all generated tokens, leaving potential attribution information underused.

Motivation: The paper identifies a mismatch: rubric criteria are evaluated at the criterion level, but the policy update remains uniform at the token level. Previous work like Rubrics-to-Tokens (RTT) learns a separate token-level relevance discriminator, but this requires an additional data-generation pipeline and training stage. The central question: Can token-level credit be derived directly from policy-internal signals, without training a separate relevance estimator?

Key insight: Fixed-response replay under a counterfactual prompt (removing rubric criteria) reveals token-level dependence: tokens whose likelihood depends on the rubric criteria show larger log-probability changes than generic content tokens (see Figure 1 in the paper). This contrast is a natural, policy-internal signal for credit assignment.

Theoretical basis: The paper builds on group-relative policy optimization (GRPO), where for each prompt group, the old policy samples (G) responses and the reward is normalized into a group-relative advantage: [ A_i = \frac{r_i - \mu}{\sigma + \epsilon_r}, ] where (\mu) and (\sigma) are the mean and standard deviation of rewards in the group. The GRPO loss is: [ \mathcal{L}{\text{GRPO}} = -\mathbb{E}{i,t}\left[\min(\rho_{i,t} A_i, \bar{\rho}{i,t} A_i)\right], ] with (\rho{i,t} = \frac{\pi_\theta(y_{i,t}|x^+, y_{i,<t})}{\pi_{\theta_{\text{old}}}(y_{i,t}|x^+, y_{i,<t})}) and (\bar{\rho}{i,t} = \text{clip}(\rho{i,t}, 1-\epsilon_{\text{clip}}, 1+\epsilon_{\text{clip}})). This objective is agnostic to which tokens express the rubric criteria.

Methodology

Core idea: CoRT redistributes the same response-level advantage (A_i) across tokens using token weights (w_{i,t}) that satisfy (\frac{1}{T_i}\sum_{t=1}^{T_i} w_{i,t} = 1), so the average shaped advantage remains (A_i). The weights are derived from counterfactual replay without training an auxiliary model.

Algorithm steps (per prompt group):

1. Standard GRPO: Sample responses (y_i) from prompt (x^+) (with rubric criteria), compute rewards (r_i), and group-relative advantages (A_i).
2. Counterfactual scoring: For each response (y_i), compute log-probabilities under the full prompt (\ell_{i,t}^+ = \log \pi_{\bar{\theta}}(y_{i,t} | x^+, y_{i,<t})) (already available from rollout) and under the criteria-free prompt (\ell_{i,t}^- = \log \pi_{\bar{\theta}}(y_{i,t} | x^-, y_{i,<t})) (by replaying fixed tokens).
3. Contrast computation: (\Delta_{i,t} = \ell_{i,t}^+ - \ell_{i,t}^-).
4. Replay-margin token credit: Map contrasts to bounded scores: [ s_{i,t} = \text{sigmoid}(\tau (\Delta_{i,t} - b)) - \frac{1}{2}, \quad \tilde{w}{i,t} = 1 + \eta \lambda_k s{i,t}, ] where (b) centers the contrast, (\tau) controls sharpness, (\lambda_k \in [0,1]) is a scheduled coefficient, and (\eta) controls weighting strength. Provisional weights lie in ((1 - \frac{\eta\lambda_k}{2}, 1 + \frac{\eta\lambda_k}{2})).
5. Scheduled response normalization: Use a cubic Hermite SmoothStep schedule: [ u_k = \text{clip}\left(\frac{k - k_0}{K}, 0, 1\right), \quad \lambda_k = 3u_k^2 - 2u_k^3, ] where (k_0) is warmup offset and (K) is ramp length. Then normalize within each response: [ \bar{w}i = \frac{1}{T_i}\sum{t=1}^{T_i} \tilde{w}{i,t}, \quad w{i,t} = \frac{\tilde{w}_{i,t}}{\bar{w}i}. ] This ensures (\frac{1}{T_i}\sum_t w{i,t} = 1) and the average shaped advantage remains (A_i).
6. Policy objective: The token-shaped advantage is (\hat{A}{i,t} = \text{sg}(w{i,t}) A_i), and the loss is: [ \mathcal{L}{\text{CoRT}} = -\mathbb{E}{i,t}\left[\min(\rho_{i,t} \hat{A}{i,t}, \bar{\rho}{i,t} \hat{A}{i,t})\right]. ] Gradients are stopped through the weight (w{i,t}).

Computational cost: CoRT adds one additional forward scoring pass per sampled response under the criteria-free prompt, without extra generation, verifier calls, or token-level labels. It is a lightweight modification to GRPO.

Empirical Validation / Results

Experimental setup:

• Training data: HIR-16k, providing prompts with and without instruction lists.
• Reward modes: Constraint Satisfaction Rate (CSR) – fraction of satisfied criteria; All-or-Nothing (AON) – reward 1 only if all criteria satisfied.
• Models: Qwen3-4B-Instruct, Qwen2.5-7B-Instruct, Qwen3-14B-Instruct.
• Evaluation benchmarks: IFBench, IFEval, MultiDimIF, AdvancedIF (with o3-mini judge). All methods evaluated with 5 sampled responses per prompt, temperature 0.7.
• Comparison baselines: Instruct (base), SFT, DPO, GRPO, RTT (learned token relevance). CoRT uses the same rollout and reward as GRPO.

Main results (Table 1, selected):

• CoRT improves over GRPO in all metrics for Qwen3-4B (CSR) and Qwen2.5-7B (CSR), and in most metrics for AON settings.
• CoRT is competitive with RTT (often stronger), while avoiding a separate relevance-learning stage.
• On Qwen3-14B (Table 2), CoRT improves all metrics under CSR and most under AON.
• Compatibility with DAPO and GSPO (Table 3): CoRT improves all metrics over DAPO and four of five over GSPO.

Training diagnostics (Figure 3, Figure 4):

• Ablation study on Qwen2.5-7B CSR: full CoRT (with response normalization and SmoothStep ramp) maintains mean token weight close to 1, low length clipping, stable gradient norm and entropy.
• Removing response normalization causes mean token weight drift and late training spikes.
• Removing the ramp leads to abrupt weighting and instability.
• The full recipe achieves the strongest validation average and training reward around the 500-step horizon.

Theoretical and Practical Implications

Theoretical implications:

• The paper demonstrates that policy-internal counterfactual likelihood contrasts provide a sufficient training signal for within-response credit allocation, without needing an external relevance model. This connects to causal inference (Pearl, 2009) and attribution methods, but uses contrasts for training rather than explanation.
• The method preserves the original reward structure and group-relative normalization, only modifying the distribution of the signed advantage across tokens. This shows that structural criteria can guide not only scoring but also credit assignment.

Practical implications:

• CoRT is a lightweight, plug-and-play addition to GRPO-style pipelines, requiring only a single additional forward pass per response. It avoids the complexity of training a separate token-level relevance model (as in RTT).
• The method is compatible with multiple policy optimization algorithms (GRPO, DAPO, GSPO), making it broadly applicable.
• The scheduled activation and response normalization are crucial for stability, providing practical guidelines for implementing counterfactual-based credit assignment.

Conclusion

Main takeaways:

• CoRT addresses the token-level credit assignment mismatch in rubric-based RL by using counterfactual replay to derive token weights from the policy's own likelihood contrasts.
• It improves over response-level GRPO consistently across models, reward granularities, and benchmarks, with an average gain of 4.4 percentage points.
• It remains competitive with learned token-level methods while being simpler and more lightweight.
• The method is stable in practice due to integration controls (response normalization and SmoothStep ramp).

Future directions:

• Finer-grained interventions (e.g., criterion-specific allocation).
• Longer-horizon credit assignment across turns, actions, and tool calls.
• Extending to settings where criteria have little effect on token likelihoods or rewards are noisy.