PaperlayerView exact revision on arXiv

Source-linked AI summary

Direction-Aware Token Credit Assignment with Reward-Model Gradients

Yikuan Hu, Jifeng Zhu, Minle Su, Jiaqin Yue, Gujie Shao, Yanzhe Chen, Haoxuan Jia, Haoran Xu, Wangzhou Dai

arXiv:2610.32772v1cs.LG

TL;DR

GRPO shares a response-level advantage across tokens, though their updates may differ in local reward utility. DATC maps frozen reward-model gradients into policy logit space and weights token updates by alignment with the signed GRPO direction. It outperforms GRPO across four settings in two model families.

  • Problem

    GRPO applies the same response-level advantage at every token position, although updates at those positions are not necessarily equally beneficial.

    Source: §1 · p. 1

  • Method

    DATC maps a frozen reward-model gradient into policy logit space and allocates token credit by the alignment of the reward direction with GRPO’s signed sampled-action update direction.

    Source: §1 · p. 2

  • Results

    DATC achieves higher mean scores than GRPO across all four evaluation metrics in Base and Instruct settings of Qwen3-14B and Llama-3.1-8B.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The KL-regularized allocation provides a local reward-gain guarantee in the surrogate model.

    Source: §5 · p. 9

    The input gradient is evaluated at the actual sampled response, and the sampled response and its prefixes remain fixed during differentiation.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

GRPO shares a response-level advantage across tokens, leaving unresolved how strongly each token should contribute to learning. Updates induced by the shared advantage reinforce or suppress the sampled token at each position, but these updates do not always point toward increasing reward at the token level. This motivates a work-based view of credit: just as work depends on the component of force along a displacement, token credit should reflect the local reward change along the policy update direction. We introduce Direction-Aware Token Credit Assignment (DATC), which maps frozen reward-model gradients into policy logit space and allocates weights by evaluating the alignment between the reward direction and GRPO's signed sampled-action update direction, while preserving the sign and mean of the response-level advantage. DATC outperforms GRPO in all four settings across two model families.

1 Introduction

GRPO gives every token in a response the same advantage, although token updates can differ in their local reward utility. DATC assigns token credit according to the work of the signed GRPO update along the reward-improving direction.

1 Introduction

  • GRPO shares one response-level advantage across token positions, but those positions’ updates are not necessarily equally beneficial.The paper frames the problem as allocating learning strength according to each position’s reward utility.Source: §1 · p. 1
  • Unlike prefix-score redistribution or token-replacement estimates, DATC evaluates local reward utility along the current policy’s GRPO update direction.The work analogy treats the signed token update as force and the RM reward-improvement direction as displacement.Source: §1 · pp. 1–2
  • The analysis proves an advantage-weighted end-to-end surrogate gradient consistent with REINFORCE, with directional credit aggregating local reward advantages along the sampled-action direction.The reported contribution also includes higher mean scores than GRPO across four metrics.Source: §1 · p. 2
  • DATC outperforms GRPO on all four evaluation metrics across Base and Instruct settings of Qwen3-14B and Llama-3.1-8B.Ablations support the roles of directional information and matching weights to positions.Source: §1 · p. 2

2 Related Work

Prior token-credit work uses fine-grained rewards, model-internal signals, or optimization objectives to evaluate token utility. DATC focuses on local reward effects along signed GRPO update directions.

Fine-grained rewards and value estimation

  • Fine-grained feedback has been obtained through process supervision, implicit process rewards, token-level reward learning, and estimates from sampled continuations.These evaluations motivate allocating group-relative advantages according to local token-update reward effects.Source: §2 · p. 2

Token credit from model-internal signals

  • Token-credit methods use policy entropy, attention, prefix-score differences, reward-model gradients, judge-model gradients, and final-answer sensitivity.Because policy updates redistribute probability among candidate tokens, evaluating sampled-token reinforcement or suppression requires connecting input importance to that adjustment.Source: §2 · p. 2

Token reweighting from optimization objectives

  • DATC’s overview combines the reward-improvement direction with the GRPO update direction to allocate token-level advantages.It assigns larger weights to updates with higher predicted local reward gains.Source: Figure 2 · p. 3
  • Other approaches derive token weights from correct-response likelihood, discriminative gradients, supervised-loss utility, or feedback-conditioned likelihood differences.DATC instead uses surrogate reward gradients to estimate first-order reward changes along signed GRPO update directions.Source: §2 · p. 3

3 Method

DATC maps a frozen reward model’s gradient into policy logit space, then allocates token-level advantages using directional credit and KL-regularized weights. The method analyzes local reward effects, aggregates candidate-token advantages, and preserves the sign and mean scale of the response-level advantage.

3.1 Displacement and Force in Logit Space

  • A straight-through surrogate connects sampled discrete tokens to the frozen reward model and propagates its standardized reward gradient into actor logit space.The sampled response and prefixes are fixed during differentiation.Source: §3.1 · pp. 3–4
  • Directional credit is the inner product of the reward-improvement direction and sampled-action score direction, whose signed effect reflects whether GRPO reinforces or suppresses the token.A positive Ad_t predicts a local reward increase, whereas a negative value predicts a decrease.Source: §3.1 · p. 4

3.2 Aggregating Local Advantages along the Sampled-Action Direction

  • The surrogate gradient has a REINFORCE-like advantage-weighted decomposition, and directional credit aggregates local reward advantages according to their relationship with the sampled-action direction.Candidate-token advantages are defined relative to the current policy’s average reward under the RM’s first-order linear model.Source: §3.2 · p. 5

3.3 Converting Directional Gains into Token-Level Advantages

  • DATC assigns token shares through a KL-regularized allocation, with positive weights of mean 1 that preserve the response-level advantage’s sign and scale.For finite temperature, the optimal allocation has a predicted local-gain guarantee over original GRPO in the first-order independent-logit surrogate; experimental weights are constrained to [1/λ, λ].Source: §3.3 · pp. 5–6

4 Experiments

Across four Base and Instruct settings, DATC outperforms GRPO on all four evaluation metrics. Ablations and token-weight examples examine directional information, weight-to-position matching, allocation concentration, and how DATC assigns token update strengths.

4.1 Experimental Setup

  • In four settings spanning Base and Instruct initializations of Qwen3-14B and Llama-3.1-8B, DATC and GRPO share initialization, rollout budget, reward processing, and LoRA settings.Source: §4 · p. 6, §4.1 · p. 6

4.2 Overall Performance and Training Dynamics

  • Across four Qwen3-14B and Llama-3.1-8B Base/Instruct settings, DATC beats GRPO on all four metrics; Qwen Base AlpacaEval and Arena-Hard win rates rise 12.18 and 10.82 percentage points.Source: Table 1 · p. 7, §4.2 · p. 7
  • For Llama-3.1-8B-Base, DATC versus GRPO yields final-window raw RM rewards of 14.14 versus 7.72 and response lengths of 497.9 versus 605.2 tokens, with GPU time increasing 5.7%–9.1% across settings.Source: §4.2 · p. 7

4.3 Credit Assignment Comparisons and Ablations

  • Removing the angle term lowers RM-Reward to −4.02, AlpacaEval to 1.53%, and Arena-Hard to 1.23%, identifying direction alignment as key to DATC’s gains in this setting.Source: §4.3 · p. 7
  • Permuting weights within responses preserves their set and concentration but lowers AlpacaEval by 16.95 percentage points, Arena-Hard by 12.32 points, and MT-Bench by 0.48 points versus DATC.Source: §4.3 · pp. 7–8
  • On Qwen3-14B-Base, DATC outperforms RED and P2T on all four metrics, using the policy’s candidate-token distribution and sampled-action direction to estimate local reward utility.Source: §4.3 · p. 7, Table 2 · p. 8
  • Among λ values 2, 4, and 8, λ = 4 achieves the highest mean on all four metrics.Source: §4.3 · p. 8

4.4 Token-Level Weight Analysis

  • In example responses, higher weights strengthen reinforcement for positive advantages and suppression for negative ones, with weights varying across specific expressions.Source: §4.4 · p. 8, Figure 4 · p. 9

5 Conclusion

DATC allocates token-level advantages according to alignment between GRPO updates and the reward-model improvement direction. It outperforms GRPO across all four settings of two model families.

5 Conclusion

  • DATC uses a work-based credit allocation aligned with GRPO updates and the reward-model improvement direction, with a local reward-gain guarantee in its surrogate model.Source: §5 · p. 9
  • DATC achieves higher scores than GRPO across all four settings of two model families.Source: §5 · p. 9

A Derivations and Theoretical Properties

The appendix derives DATC’s token credit from a frozen reward model’s end-to-end surrogate gradient, then analyzes its allocation and local-gain properties. It also relates the experimental allocator to this theory.

A.1 End-to-End Surrogate Gradients at the Hard Response

  • The end-to-end surrogate gradient decomposes into candidate-action score gradients weighted by local reward advantages from the reward model’s linearization at the sampled response.The decomposition is an algebraic identity at the fixed expansion point; the reward interpretation comes from the local linear model.Source: §A.1 · pp. 12–13

A.2 From the GRPO Direction to Token Credit

  • DATC’s directional credit combines local reward advantages with the alignment between candidate-action and sampled-action directions to score token updates.The local-gain analysis evaluates these updates in independent logit coordinates; the appendix separately relates that evaluation to shared-parameter updates.Source: §A.2 · pp. 14–15

A.3 KL-Based Weight Allocation and Local Gains

  • KL-regularized allocation preserves the sign and mean of each response-level advantage, and yields nonuniform weights with strictly positive predicted-gain improvement when utilities vary.When utilities are constant, the optimal weights are 1; when the response-level advantage is zero, all token-level advantages are zero.Source: §A.3 · pp. 15–16

A.4 Adaptive Temperature and Weight Range

  • For any amplification ceiling λ > 1, a feasible temperature exists; when the initial maximum weight exceeds λ, a unique temperature above 1 meets the ceiling.The appendix describes finding a feasible temperature by doubling and bisection.Source: §A.4 · p. 17

A.5 Correspondence between the Experimental Allocator and the Theory

  • Within-response standardization and fixed positions preserve the KL-optimal allocation structure, while the two-sided weight constraints admit a finite feasible temperature.The local-gain comparison remains a first-order evaluation at a fixed response in independent logit coordinates.Source: §A.5 · pp. 17–18

B Training Objective and Implementation

The appendix specifies DATC’s policy objective and implementation, from response scoring and reward-model input gradients to standardized credit weights and actor updates.

B.1 From Response Rewards to the Policy Objective

  • The policy objective uses cached token weights with the clipped policy term and reference-policy KL term; GRPO uses the same objective with every weight set to 1.The reward-processing procedure includes length-dependent terms and special handling for empty or degenerate responses.Source: §B.1 · pp. 19–20

B.2 Input Gradients, Vocabulary Reductions, and Position Correspondence

  • A reward-model input backward pass supplies position gradients, which are mapped to sampled tokens using aligned token identities and response positions.Formatting and unmapped positions retain weight 1, leaving content positions to participate in allocation.Source: §B.2 · p. 20

B.4 One Training Iteration

  • Each rollout caches fixed token-level advantages before the policy-update epoch trains the actor’s LoRA parameters.The procedure computes reward-model input gradients, derives directional credit, selects weights, and then updates the actor using the cached coefficients.Source: §B.4 · pp. 21–22

C Experimental Setup and Supplementary Analyses

  • The experiments report that actor, reward model, and sampler each occupy one GPU, with additional computation concentrated in reward-model input backpropagation and vocabulary reductions.The appendix directs readers to its supplementary analysis for end-to-end training overhead.Source: §B.4 · p. 21, C Experimental Setup and Supplementary… · p. 23

C.1 Experimental Setup and Comparison Implementations

The experiments compare DATC and GRPO across two model families, Base and Instruct initializations, and multiple evaluation and ablation settings.

C.1 Experimental Setup and Comparison Implementations

  • The training data comes from UltraFeedback, with approximately 59,000 candidate examples, 2,500 SFT examples, and 2,000 RL prompts.RM-Reward evaluation uses 256 fixed held-out prompts.Source: §C.1 · p. 23
  • The study evaluates Qwen3 and Llama 3.1 in Base and Instruct settings, pairing each family with its corresponding Skywork reward model.Base models receive lightweight SFT before RL; Instruct models proceed directly to RL.Source: §C.1 · p. 23
  • Each RL run processes 2,000 prompts in groups of eight, samples eight responses per prompt, and performs one policy update after each rollout.Main-text evaluations use the end-of-training checkpoint; DATC and GRPO share training initialization and optimization settings.Source: §C.1 · p. 23
  • Evaluation covers reward under the training RM, open-ended preference win rates, multi-turn response quality, and supplementary comparisons against token-credit baselines.The supplementary comparisons include RED and P2T; the main DATC results are reported as means over three random seeds.Source: §C.1 · pp. 23–24
  • The ablations vary directional alignment, within-response position correspondence, and the allowed weight range to examine distinct parts of DATC’s allocation.The weight-range experiment varies λ over 2, 4, and 8; the shuffled-weight variant permutes weights within each response.Source: §C.1 · p. 24

C.2 Reward and Response-Length Dynamics

Across four settings, DATC has higher mean raw reward than GRPO in the final 25 rollouts, while response-length patterns vary by initialization.

C.2 Reward and Response-Length Dynamics

  • 14.14 versus 7.72 mean raw reward and 497.9 versus 605.2 mean response length: DATC versus GRPO in Llama Base’s final 25 rollouts.DATC responses are shorter in both Base settings; lengths are similar in Llama Instruct and longer for DATC in Qwen Instruct.Source: §C.2 · p. 24
  • Higher mean raw reward for DATC than GRPO appears in all four settings over the final 25 rollouts.Figure C1 plots raw reward and mean response length as trailing 25-rollout moving averages.Source: §C.2 · p. 24, Figure C1 · p. 25

C.3 Training Overhead

DATC raises end-to-end training cost by approximately 5.7%–9.1% over GRPO, while the reported evaluations show higher mean scores across four metrics in all four settings.

C.3 Training Overhead

  • DATC increases cumulative GPU time by approximately 5.7%–9.1% over the corresponding GRPO baseline, while achieving higher mean scores across four metrics in all four settings.End-to-end cost includes sampling, scoring, credit computation, and policy updates.Source: §C.3 · p. 26, Figure C2 · p. 26

C.4 Human Evaluation

In human evaluation, DATC wins the largest share of prompts in a four-method comparison. In a separate single-step analysis, independent-logit predicted gains correlate with shared-parameter first-order gains, and DATC's gain is positive in every sampled batch.

C.4 Human Evaluation

  • The human comparison used 120 held-out prompts, split evenly across three training seeds, with final checkpoints and identical decoding settings for each method.The prompts were excluded from training and hyperparameter selection.Source: §C.4 · p. 26
  • 65.00% of prompts were won by DATC, versus 16.67% for GRPO, 15.00% for P2T, and 3.33% for RED.Seven evaluators chose among four anonymized responses per prompt, with a tie-breaking vote when needed.Source: Table C3 · p. 26, §C.4 · p. 26

C.5 Independent-Logit Predictions and Shared-Parameter Single-Step Gains

  • Across 100 batches, independent-logit predicted gain X and shared-parameter first-order gain P had Spearman correlation 0.619, Pearson correlation 0.589, and linear-fit R2 = 0.347; all points had X > 0 and P > 0.The analysis used one AdamW update per method from the same Qwen3-14B-Base SFT checkpoint, with LoRA updates and λ = 4.Source: §C.5 · p. 27, Figure C3 · p. 28

D Limitations & Future Work

DATC is scoped to preference optimization with differentiable reward models and requires reward-model input gradients and aligned actor–reward-model vocabularies. Experiments used LoRA for general preference alignment across Base and Instruct initializations in two model families.

D Limitations & Future Work

  • DATC requires a differentiable reward model, access to its input gradients, and aligned actor–reward-model vocabularies.Token or sequence alignment across tokenizers and gradient estimation for score-only evaluators are proposed as future directions.Source: D Limitations & Future Work · p. 28
  • Experiments cover LoRA-based general preference alignment across Base and Instruct initializations in two model families, not full-parameter training or specialized reward-model tasks.Evaluations across broader model scales and tasks such as mathematical reasoning and code generation are left for future work.Source: D Limitations & Future Work · p. 28

Related papers

Loading 2610.32772v1…