PaperlayerView exact revision on arXiv

Source-linked AI summary

CoCA: Criterion-Aware Credit Assignment for Rubric-Based Reinforcement Learning

Guangcheng Zhu, Yingfan Ma, Xing Zheng, Haobo Wang, Shenzhi Yang, Kai Tang, Zhengqing Zang, Bowen Song, Weiqiang Wang, Gang Chen

arXiv:2610.17733v1cs.LGcs.AI

TL;DR

Standard rubric-based RL compresses criterion-level feedback into scalar rewards and uniform token advantages, while OPSD’s rubric-conditioned imitation signal is not grounded in criterion outcomes. CoCA derives signed criterion credits from group-relative outcomes and allocates them across tokens using criterion-induced distribution shifts. Across medicine and science benchmarks, it outperforms standard rubric-based RL by 4.2%.

  • Problem

    Standard rubric-based RL reduces criterion-level assessments to scalar rewards and shared token advantages, while OPSD’s finer-grained guidance is not grounded in observed criterion outcomes.

    Source: §1 · p. 1, §3.2 · p. 4

  • Method

    CoCA derives signed criterion credits from group-relative outcomes and uses rubric-conditioned distribution shifts as relevance cues to allocate them across tokens.

    Source: §1 · p. 2

  • Results

    Across medicine and science benchmarks, CoCA outperforms standard rubric-based RL by 4.2%.

    Source: §7 · p. 10

  • Takeaways & Limitations

    CoCA makes more effective use of structured rubric feedback through criterion-aware credit assignment.

    Source: §1 · p. 2

    The medical benchmarks are research testbeds and do not establish clinical safety or suitability for clinical use.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

Rubric-based reinforcement learning (RL) evaluates responses against explicit criteria to construct rewards, offering a viable way to train large language models on open-ended tasks. However, standard methods compress criterion-level feedback into scalar rewards and assign uniform token advantages, leaving structured rubric supervision underutilized. Recent work seeks finer-grained guidance by treating rubrics as privileged information to derive self-distillation targets, but the resulting learning signal is not grounded in criterion outcomes and can lead to reward degradation. This motivates us to reframe the problem as criterion-aware credit assignment, where privileged rubric information serves as a relevance cue for credit allocation rather than a source of imitation targets. Accordingly, we propose CoCA, which derives criterion-level credits from group-relative outcomes and allocates them across tokens based on criterion-induced distribution shifts, yielding both outcome-grounded and fine-grained learning signals. Empirically, CoCA achieves an average gain of +4.2% over standard rubric-based RL across medicine and science benchmarks, demonstrating more effective learning on open-ended tasks.

1 INTRODUCTION

Rubrics structure assessment of open-ended responses, but standard rubric-based RL compresses criterion outcomes into a scalar reward and shares one advantage across all tokens. OPSD adds token-level guidance through rubric-conditioned imitation, yet its updates are not grounded in criterion outcomes and can degrade reward; CoCA instead uses rubric information to allocate outcome-based credit.

1 INTRODUCTION

  • Open-ended medical and scientific tasks admit diverse valid responses and require multidimensional quality assessment, making reliable, scalable reward construction difficult.Rubrics offer explicit, individually assessable criteria for structuring that supervision.Source: §1 · p. 1
  • Standard rubric-based RL aggregates criterion outcomes into one response-level reward, so responses with different criterion profiles can receive the same signal and all tokens share one advantage.This can reinforce criterion-violating content or suppress criterion-satisfying content, depending on the sign of that shared advantage.Source: §1 · p. 1
  • OPSD matches policy predictions to a rubric-conditioned self-teacher, but its updates follow teacher–student prediction differences rather than observed criterion outcomes.The paper reports reward declines when OPSD is optimized alone and degradation when added to rubric-based RL.Source: §1 · p. 2
  • CoCA separates credit determination from allocation: group-relative criterion outcomes set signed credit budgets, while rubric-conditioned distribution shifts guide token relevance.The resulting resolved credits are aggregated into token-level advantages for policy optimization.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries describe rubric-based RL as checklist scoring followed by group-relative policy optimization, while highlighting how aggregation leaves token-level updates untargeted. They also introduce OPSD as a rubric-conditioned self-distillation objective.

2 PRELIMINARIES

  • Rubrics represent response quality as weighted, independently assessable criteria, and the paper uses binary checklist scoring with GRPO as its base RL algorithm.An LLM judge scores each rollout against each criterion; the criterion scores are aggregated into a scalar reward.Source: §2 · p. 2
  • The group-relative advantage is each response’s scalar reward minus the rollout group’s mean reward.The policy is optimized with the GRPO loss, without normalizing by within-group reward standard deviation.Source: §2 · p. 2
  • Because criterion-level evaluations collapse to a response-level advantage shared by every token, the update does not target tokens relevant to individual criteria.The passage identifies this as the limitation of the standard rubric-based signal.Source: §2 · p. 3
  • OPSD uses rubrics as privileged context and trains the policy to match its next-token distribution under ordinary context to the rubric-conditioned distribution.The rubric-conditioned distribution acts as a detached self-teacher, using either the full criterion set or a selected subset.Source: §2 · p. 3

3 CREDIT ASSIGNMENT IN RUBRIC-BASED POLICY OPTIMIZATION

This section frames rubric-based learning as a problem of assigning signed criterion budgets and resolving them across relevant tokens. It analyzes how standard rubric-based RL loses token-level resolution and how OPSD’s fine-grained matching signal can diverge from reward outcomes.

3 CREDIT ASSIGNMENT IN RUBRIC-BASED POLICY OPTIMIZATION

  • Credit assignment has two components: grounding maps observed criterion outcomes to signed budgets, and resolution allocates those budgets across tokens while preserving their values.The allocation is represented by a criterion–token credit matrix whose row mean preserves each criterion’s budget.Source: §3 · p. 3

3.1 RUBRIC-BASED RL: OUTCOME-GROUNDED BUT UNDER-RESOLVED

  • Criterion budgets can be defined by each response’s weighted criterion performance relative to the rollout group, but standard rubric-based RL sums them into one advantage and assigns it uniformly across tokens.That uniform allocation preserves the budgets without representing which tokens are relevant to which criteria.Source: §3.1 · p. 3
  • Uniform token allocation and criterion aggregation obscure strengths and weaknesses: opposing criterion credits can cancel, and tokens can receive the same credit when an oracle resolution would reinforce some and suppress others.The section characterizes the resulting signal as outcome-grounded but under-resolved.Source: §3.1 · pp. 3–4

3.2 OPSD: FINE-GRAINED BUT NOT NECESSARILY OUTCOME-GROUNDED

  • OPSD provides token-specific guidance by matching a rubric-conditioned self-teacher, but the matching updates do not account for which criteria the sampled response satisfies or violates.A teacher’s higher expected reward does not guarantee that its relative token preferences align with reward under the student policy.Source: §3.2 · p. 4
  • A local OPSD matching step can reduce expected reward even when the rubric-conditioned teacher has higher expected reward, if the matching signal’s covariance with reward is negative.The proposition establishes this possibility for finite action spaces with at least three actions and sufficiently small positive step sizes.Source: §3.2 · p. 4
  • The first-order reward change depends on alignment between the matching signal and reward under the student policy; negative alignment can lower reward while reducing distillation loss.The authors offer this local failure mode as a plausible explanation for the reward degradation observed with OPSD.Source: §3.2 · p. 4
  • The proposed alternative is to use observed outcomes to set each criterion’s signed credit budget and privileged rubric information only as a relevance cue for allocating that credit.This refines outcome-grounded credit allocation rather than adding an independent imitation objective.Source: §3.2 · p. 4

4 COCA: CRITERION-AWARE CREDIT ASSIGNMENT

CoCA grounds criterion credits in within-group rubric outcomes, then allocates those credits across tokens using criterion-induced distribution shifts as relevance cues. Its analysis states conditions under which this allocation preserves signed budgets while reducing off-target credit.

4 COCA: CRITERION-AWARE CREDIT ASSIGNMENT

  • CoCA grounds each criterion’s signed contribution through within-group rubric-outcome comparisons, resolves it across relevant tokens, and aggregates the resulting credits into token-level GRPO signals.The overview depicts credit grounding and resolution as the two stages of the framework.Source: §4 · p. 4, Figure 2 · p. 5

4.1 TOKEN-LEVEL ADVANTAGE CONSTRUCTION

  • CoCA’s group-relative grounding rule produces signed criterion budgets whose sum equals the response-level advantage.This keeps criterion credits tied to the observed group-relative outcomes.Source: §4.1 · p. 4
  • For each criterion, CoCA uses Jensen–Shannon divergence between frozen-policy next-token distributions with and without that criterion as an unsigned token-relevance cue.Positive token weights are formed from these scores while preserving the criterion budget’s sign and mean; nonuniform allocation is limited to up to k criteria with the largest credit magnitudes, requiring at most k additional criterion-conditioned forward passes per response.Source: §4.1 · p. 5
  • Aggregating resolved criterion credits can yield nonzero token advantages even when the response-level advantage is zero, if opposing criterion budgets are allocated differently across tokens.CoCA uses these token-level advantages in the GRPO objective without an auxiliary teacher-matching loss.Source: §4.1 · p. 5

4.2 GROUNDING PRESERVATION AND RESOLUTION IMPROVEMENT

  • When relevance scores have a positive mean gap between relevant and irrelevant positions, the analysis shows that controlled allocation concentration can reduce off-target credit while preserving each criterion’s budget and contribution sign.The concentration bound allows stronger concentration for larger mean-score gaps and is more restrictive when the score range widens.Source: §4.2 · pp. 5–6

5 EXPERIMENTS

Across medical and scientific benchmarks, CoCA outperforms the standard rubric-based RL baseline and improves criterion-level correction while reducing regression. Further analyses examine signal conflicts, credit allocation, relevance measures, hyperparameters, and criterion–token relevance.

5.1 EXPERIMENTAL SETUP

  • The evaluation covers in-domain RaR-Medicine and RaR-Science test splits and generalization benchmarks in medicine and science.Medical benchmarks are HealthBench, LLMEval-Med, and MedQA; science benchmarks are ResearchQA and GPQA-Diamond.Source: §5.1 · p. 6

5.2 MAIN RESULTS

  • 64.5% average score; CoCA achieves the best performance on all seven benchmarks and improves over RaR by 4.2%.RaR+OPSD underperforms RaR on every benchmark and trails CoCA by 7.8% on average; CoCA exceeds CoRT by 3.8% on average.Source: §5.2 · pp. 6–7

5.3 FURTHER ANALYSIS

  • Uniform allocation lowers average score by 4.2%, shuffled allocation by 11.4%, and shared criterion profiles by 2.8% relative to CoCA.The comparisons support nonuniform token allocation, alignment of weights with token positions, and distinct allocation profiles for different criteria.Source: Table 2 · p. 7, §5.3 · pp. 7–8

5.2 MAIN RESULTS

  • Correction rises from 36.7% to 46.3% in medicine and 26.5% to 32.0% in science; regression falls from 10.8% to 8.6% and 9.5% to 7.6%, respectively.CoCA has the highest correction and lowest regression rates in both domains; RaR+OPSD shows the opposite pattern to CoCA relative to RaR.Source: Figure 3 · p. 7, §5.2 · p. 7

5.3 FURTHER ANALYSIS

  • Replacing JSD with forward KL and reverse KL reduces average score by 5.3% and 2.8%, while signed log-probability ratio reduces it by 10.3%.The authors suggest JSD’s symmetry and boundedness may explain its advantage, and report that full-distribution changes provide more informative signals than a generated token’s directional probability shift alone.Source: Table 3 · p. 8, §5.3 · p. 8
  • Criterion-credit conflict reaches 54.6% in medicine and 50.4% in science, while distillation misalignment reaches 49.0% and 48.8%, respectively.The conflicts are between response-level advantages and criterion credits, and between OPSD and outcome-grounded gradients.Source: §5.3 · p. 7, Figure 4 · p. 8
  • Even k = 1 improves average scores over uniform allocation by 4.4% in medicine and 1.7% in science; selecting all criteria adds only 0.5% and 0.3%.The authors identify k = 2 as a practical balance, requiring at most two additional forward passes per rollout; performance peaks at β = 5 in medicine and β = 3 in science.Source: Figure 4 · p. 8, §5.3 · p. 8
  • In the medical case study, higher relevance concentrates on the recommendation and its applicable conditions, rather than general care advice.The authors interpret this contrast as criterion-specific semantic relevance rather than general topical relevance.Source: §5.3 · p. 8, Figure 5 · p. 9

6 RELATED WORK

Related rubric-based RL methods refine criteria, reweight them, or localize criterion feedback, but often retain response-level rewards. Rubric-guided generation and self-distillation provide other forms of guidance, while the paper motivates preserving each criterion’s signed contribution during token-level credit assignment.

Rubrics as rewards

Rubrics as guidance

  • Rubric-guided rollout generation promotes exploration, while OPSD uses a rubric-informed teacher to supervise unguided student rollouts [Zhou et al., 2026; Bi et al., 2026; Zhang et al., 2026b; Hou et al., 2026].Guided prefixes can create inference mismatch and exposure bias; OPSD instead introduces a separate imitation objective that may conflict with reward-driven optimization [Gu et al., 2026; Rezaei et al., 2026; Bablani et al., 2026; Ye et al., 2026; Wang et al., 2026].Source: §6 · p. 9

Credit assignment

  • Rubric-based RL must distinguish criteria with opposing contributions, which aggregate-advantage redistribution alone cannot do; CoRT applies positive token weights to one response-level advantage [Zhang et al., 2026a].The related-work discussion frames preserving each criterion’s signed contribution as the key challenge in token-level credit assignment.Source: §6 · p. 9

7 CONCLUSION

CoCA turns criterion-level feedback into fine-grained, outcome-grounded policy updates, achieving a 4.2% average gain over standard rubric-based RL. The appendix proof shows that standard rubric-based RL collapses criterion-level token allocations to uniform row means, potentially losing within-criterion information.

7 CONCLUSION

  • CoCA derives signed criterion credits from observed outcomes and uses privileged rubric information to allocate them across relevant response tokens.This yields fine-grained, outcome-grounded policy updates.Source: §7 · p. 9
  • 4.2% average gain over standard rubric-based RL, with higher correction and lower regression rates across medicine and science.The authors caution that medical benchmarks do not establish clinical safety or suitability for clinical use.Source: §7 · p. 10
  • Medical benchmarks are research testbeds; the results do not establish clinical safety or suitability for clinical use.Source: ETHICS STATEMENT · p. 10

A.1 PROOF OF PROPOSITION 1

  • Standard rubric-based RL projects criterion–token credit matrices onto uniform-across-token rows, replacing each criterion’s allocation with its row mean and discarding within-criterion variation.This projection is independent of the oracle’s token allocation within each row.Source: §A.1 · p. 14
  • The projection can yield zero token credit despite nonzero criterion contributions, or assign credit opposite in sign to the oracle at a token position.Source: §A.1 · p. 14

A.2 PROOF OF PROPOSITION 2

The proof shows that rubric-conditioned matching can lower expected reward even when the teacher has higher expected reward. Its first-order reward effect depends on the covariance between the matching signal and reward under the student policy.

A.2 PROOF OF PROPOSITION 2

  • The local natural-gradient update changes expected reward to first order by η times the student-policy covariance between reward and matching signal.The result assumes a finite action space and a fixed self-teacher during the update.Source: §A.2 · p. 15
  • For |V| ≥ 3, a higher-reward teacher can induce a matching update with negative reward covariance, lowering expected reward for sufficiently small positive steps.The counterexample extends to every action-space size of at least three.Source: §A.2 · pp. 15, 17
  • When the matching signal is nonconstant, a sufficiently small positive step decreases the local matching loss; negative reward covariance means that same step also lowers expected reward.Source: §A.2 · p. 16
  • The proof constructs a teacher with higher expected reward but a matching signal negatively correlated with reward under the student policy.For sufficiently small steps, reward falls even as the distillation objective decreases.Source: §A.2 · pp. 16–17

A.3 PROOF OF PROPOSITION 3

CoCA’s resolution preserves each criterion’s signed credit budget while redistributing credit across tokens. Under a stated concentration bound, it reduces credit magnitude on irrelevant positions relative to uniform allocation.

A.3 PROOF OF PROPOSITION 3

  • CoCA’s allocation preserves each criterion’s budget and sign, so redistribution changes token placement without changing the row’s total credit magnitude.At β = 0, the weights recover uniform allocation.Source: §A.3 · p. 17
  • The improvement condition compares mean relevance scores on relevant and irrelevant token sets, with the positive mean-score gap required to survive exponentiation.Source: §A.3 · pp. 17–18
  • For selected criteria with positive mean-score advantage, nonuniform allocation strictly reduces irrelevant-position credit magnitude versus uniform allocation when 0 < β < 8∆/b^2.The proof applies the comparison to positive and negative criterion budgets.Source: §A.3 · p. 18

B.1 DATASETS, BENCHMARKS, AND EVALUATION

The appendix describes medical and scientific QA training and evaluation datasets, scoring protocols, implementation details, and baseline reproductions. It also specifies how evaluation responses and criterion-conditioned training passes are prepared.

B.1 DATASETS, BENCHMARKS, AND EVALUATION

  • Separate policies are trained on RaR-Medicine and RaR-Science, then evaluated on their test splits and additional medical and scientific benchmarks.Additional benchmarks include HealthBench, LLMEval-Med, MedQA, ResearchQA, and GPQA-D.Source: §B.1 · p. 19
  • Benchmark scores use rubric criteria and specified pass standards; eight GPQA-D responses and two LLMEval-Med responses are averaged, while other benchmarks use one response.The overall score is the unweighted mean of seven benchmark scores.Source: §B.1 · pp. 19–20

B.2 TRAINING AND IMPLEMENTATION DETAILS

  • Training uses verl with FSDP and vLLM; criterion-conditioned passes append a selected criterion’s instruction while retaining prior conversation messages.Response tokens are aligned across conditioned passes, and prompt and padding positions are excluded from normalization and the policy loss.Source: §B.2 · pp. 20–21, Figure 7 · p. 21

B.3 BASELINE REPRODUCTION

  • The authors reproduce five baselines using official implementations or standard procedures, and implement OPSD and RaR+OPSD using the SDPO components and joint policy-distillation losses.For RaR+OPSD, λ = 0.01 balances an auxiliary distillation term against the primary policy loss.Source: §B.3 · p. 21

C.1 RESULTS ACROSS BASE MODELS AND JUDGES

CoCA retains an average-score advantage across tested base models and judge models. Criterion-relevant response positions also show higher mean JSD, supporting its use for criterion-level credit allocation.

C.1 RESULTS ACROSS BASE MODELS AND JUDGES

  • 58.0% and 41.5% average scores for Qwen3-4B and Olmo3-7B-Instruct-SFT exceed RaR by 4.1% and 1.3%, respectively.CoCA also exceeds HeRL by 2.4% and 1.4%, respectively.Source: Table 5 · p. 22, §C.1 · p. 22
  • 64.8% average score with DeepSeek-V4.1-Flash beats RaR by 3.4% and HeRL by 2.5%; relative gaps remain broadly consistent across judges.Absolute scores vary across judges.Source: Table 6 · p. 22, §C.1 · p. 22

C.2 QUANTITATIVE VALIDATION OF CRITERION–TOKEN RELEVANCE

  • The analysis samples 500 criterion–response pairs per domain, annotates criterion-relevant words, manually reviews 20% per domain, and maps spans to response-token positions.Figure 8 shows annotated and remaining words, with subword pieces merged for display.Source: §C.2 · pp. 22–23, Figure 8 · p. 23
  • Figure 9 reports positive mean-JSD gaps in 97.6% of Medicine pairs and 100.0% of Science pairs; 98.8% of all pairs have higher mean JSD at annotated positions.For enrichment ratio thresholds of 1.5, 2.0, and 3.0, the corresponding Medicine and Science fractions are 71.0%/81.2%, 47.0%/61.6%, and 19.2%/31.6%.Source: §C.2 · p. 23

C.3 ANALYSIS OF TRAINING EFFICIENCY

CoCA achieves stronger Medicine validation performance with only modestly higher per-step cost, and its training dynamics, criterion selection, and seed results provide further evidence about its behavior and robustness.

C.3 ANALYSIS OF TRAINING EFFICIENCY

  • 348.94 seconds per step for CoCA versus 336.16 for RaR is a 3.8% increase; credit assignment takes 12.78 seconds, or 3.7% of the total.Judge-based reward evaluation and response generation account for 85.7% of runtime; CoCA reuses criterion judgments and scores selected criteria.Source: §C.3 · p. 24, Figure 10 · p. 24

C.5 EFFECT OF CRITERION SELECTION

  • 61.4% for Random and 63.2% for Top |sj| trail CoCA by 3.1% and 1.3%; selecting criteria with larger |Aij| prioritizes response-specific credit.Both alternatives underperform CoCA on all seven benchmarks; Top |sj| consistently exceeds Random.Source: Table 7 · p. 25, §C.5 · p. 26

C.4 ANALYSIS OF TRAINING DYNAMICS

  • Figure 11 shows CoCA’s reward advantage emerging early in Medicine and later in Science, persisting through the end of training in both domains.In Medicine, both methods approach a plateau; in Science, their initial progress is similar.Source: Figure 11 · p. 25, §C.4 · p. 25
  • Response lengths diverge by domain, while CoCA’s entropy declines more gradually and remains higher late in training in both domains.Medicine responses become longer, but Science responses shorten in the second half even as reward improves further.Source: §C.4 · p. 25, Figure 11 · p. 25

C.6 ROBUSTNESS ACROSS RANDOM SEEDS

  • 0.17% versus 0.24% is CoCA’s versus RaR’s standard deviation of the seven-benchmark average across three seeds; CoCA’s 3.1% lead exceeds the observed variation.Benchmark-level standard deviations range from 0.22% to 0.48% for CoCA and 0.18% to 0.53% for RaR.Source: Table 8 · p. 26, §C.6 · p. 26

D PSEUDOCODE OF COCA

  • Algorithm 1 samples rollouts, obtains criterion outcomes, scores selected criteria via criterion-conditioned token-distribution shifts, then updates the policy using token-level advantages.Source: D PSEUDOCODE OF COCA · p. 27

Related papers

Loading 2610.17733v1…