Each Reward to Its Relevant Tokens: Reward-Specific Token-Level Credit Assignment for Multi-Reward RL
Skylar Zhai, Haonan Wang, Zhining Liu, Linxin Song, Xuwei Ding, Jingcheng Liang, Jian Kang
MRTPO assigns multi-reward reinforcement learning advantages to specific tokens rather than entire responses.
How can we assign credit for specific rewards (like format vs. correctness) to the exact tokens responsible for them, rather than spreading the penalty across the entire response?
In multi-reward reinforcement learning, models are often trained against several criteria simultaneously, such as tool-call format and execution correctness. Existing methods aggregate these rewards into a single advantage value and apply it uniformly to every token in a response, even though most tokens are irrelevant to any given reward. MRTPO (Multi-Reward-to-Token Policy Optimization) identifies the tokens relevant to each reward by replaying the model's response under two contexts—one with and one without the specific verifier's information—and measuring the divergence in next-token distributions. It then redistributes each reward's advantage across tokens based on their calculated relevance, ensuring that credit is concentrated where the reward was actually earned. On tasks like text-to-SQL and tool calling, this token-level credit assignment leads to faster convergence and higher performance than response-level normalization methods.
Paper Primer
The core move is a self-replay mechanism that isolates the influence of each verifier. By comparing the model's prediction distribution with and without a specific reward's verifier information, the method computes a Jensen–Shannon distance for every token, effectively mapping which parts of the response "caused" the reward.
MRTPO improves performance and convergence speed over response-level methods like GDPO.
Empirical evaluation across text-to-SQL, multi-constraint instruction following, and tool-calling tasks. Consistent gains in performance metrics and faster convergence compared to uniform advantage assignment.
Why is uniform advantage assignment problematic for multi-reward tasks?
Uniform assignment spreads the penalty for a single error across the entire response, allowing the credit from other rewards to mask the error and preventing the model from learning which specific tokens caused the failure.
Does MRTPO require additional training data or auxiliary models?
No. MRTPO uses the existing verifier information and the model's own self-replay to derive token weights, requiring no token-level annotations or auxiliary models.
Researchers can now replace uniform response-level credit assignment with reward-specific token weighting, allowing models to learn more efficiently from complex, multi-objective feedback.
Introduction and Motivation
Multi‑reward RL spreads penalties uniformly, obscuring which tokens cause each error.
Reinforcement learning with verifiable rewards (RLVR) now routinely trains large language models on several objectives at once—correctness of a tool call, format compliance, or multiple instruction constraints. Existing multi‑reward methods such as GDPO normalize each reward separately, yet they distribute every reward’s advantage evenly across the entire response. Because each verifier typically inspects only a small portion of the output, this uniform assignment spreads credit (or penalty) to tokens that did not influence that reward, a phenomenon we call reward interference.
When a verifier flags a mistake, the penalty is mistakenly spread to all tokens, even those that had nothing to do with the error.
Uniform distribution: each token receives +0.50, +0.50, and –0.87, regardless of which token caused the execution error.
The execution error originates from the “>=” operator, but the penalty is also applied to the “SELECT” and “COUNT(*)” tokens.
This spreads blame to irrelevant tokens, making it unclear which part of the response needs correction.
Uniform credit assignment hides the true source of error, preventing the model from learning precise token‑level corrections.
Our central premise is that this reward interference can be eliminated by localizing credit: MRTPO replays each sampled response twice—once with the verifier’s privileged information and once without—computes the Jensen–Shannon distance between the resulting next‑token distributions, and redistributes the group‑normalized advantage according to those token‑wise distances. In this way, only the tokens that actually influence a given reward receive its advantage.
The core problem is reward interference: uniform advantage masks which tokens truly matter for each reward.
Related Work
Survey of prior multi-reward methods and token-level credit techniques.
Early multi-reward reinforcement learning approaches either train separate reward models for each objective or interpolate policies trained on distinct rewards, aiming to balance several goals simultaneously.
GDPO scales each reward by its group’s statistics before summing, preserving distinguishability of reward combinations.
Subsequent work builds on GDPO by dynamically reweighting rewards based on within‑group variance, training progress, or detected conflicts, and these methods operate at the response level.
Token‑level credit assignment research has explored process reward models, Monte‑Carlo value estimates, and token‑level discriminators, all of which require step‑level annotations, extra sampling, or dedicated models.
Other approaches derive token weights from entropy or from replaying the response without the full rubric, but they only reweight the aggregated advantage rather than each reward individually.
Privileged‑information self‑distillation methods use a reference solution or environment feedback as a teacher, letting the model fit the teacher’s distribution; MRTPO adopts this idea and uses the distributional difference under each reward to allocate that reward’s advantage to specific tokens.
Problem Setup
Define per‑token advantage and show how it isolates reward credit for each token.
When multiple rewards are summed into a single scalar, every token in a response receives the same advantage, which blurs which token actually caused a particular reward. This mixing is the source of reward interference that the paper later addresses.
Instead of giving the whole response a flat bonus, each token gets its own advantage proportional to how much it contributed to the rewards – like paying each worker for the exact piece they produced rather than a single lump‑sum for the whole team.
Standardize each reward: \tilde{r}^{(1)}_1 = (4‑2.5)/1.5 = 1.0, \tilde{r}^{(1)}_2 = (2‑3.5)/1.5 = -1.0; \tilde{r}^{(2)}_1 = (1‑2.5)/1.5 = -1.0, \tilde{r}^{(2)}_2 = (5‑3.5)/1.5 = 1.0.
Sum the standardized rewards for each response: S(1)=1.0+(-1.0)=0.0, S(2)=‑1.0+1.0=0.0.
Distribute the sum equally to each token: A(1)_t = S(1)/T = 0.0/3 = 0.0 for t=1..3; A(2)_t = S(2)/T = 0.0/3 = 0.0 for t=1..3.
In this toy case the per‑token advantages are all zero, illustrating that when positive and negative rewards cancel, no token receives credit.
Per‑token advantage isolates the net effect of all rewards on each token, so a token only gets credit when the aggregated standardized rewards are non‑zero for its response.
GRPO aggregates rewards with the formula above and then standardizes the group‑level advantage, which means two responses with different reward mixes but identical sums receive identical gradients. GDPO instead standardizes each reward individually, weights them (here all weights $c_k$=1), sums them, and finally rescales the batch‑level sum. Both approaches still assign a single advantage to every token of a response.
How does per‑token advantage differ from the group‑level advantage used by GRPO and GDPO?
Group‑level advantage assigns the same scalar $A$(j) to every token of response j, so the gradient cannot distinguish which token caused a reward. Per‑token advantage computes a separate A(j)_t for each token, allowing the policy gradient to credit or penalize exactly the token that contributed to the standardized reward sum.
The MRTPO Method
How MRTPO isolates token‑level credit for each reward using leave‑one‑out replay.
When a single scalar advantage is given to every token, a penalty for a single mistake dilutes across the whole response, letting unrelated tokens mask the error. MRTPO fixes this by localising each reward’s credit to the tokens that actually influence it.
MRTPO re‑weights the per‑reward advantage so that tokens most responsible for a reward receive a larger share, while all tokens retain a baseline share.
How does MRTPO differ from the earlier GDPO baseline?
GDPO gives every token the same scalar advantage regardless of which reward caused it. MRTPO first measures how much each token’s prediction depends on a particular reward (via $q_{k,t}$) and then scales the advantage so that tokens that matter for that reward receive a larger share, while still keeping a uniform baseline.
To see which tokens a reward influences, MRTPO re‑generates the model’s prediction twice: once with the reward’s verifier information present and once with it removed. The change in the token‑level distribution reveals relevance.
Compute Jensen–Shannon distance at $t_1$: $\operatorname{JS}([0.7,0.2,0.1]\|[0.5,0.3,0.2]) \approx 0.09$; normalised $q_{k,1}\approx 0.13$.
At $t_2$: distance $\approx 0.07$ → $q_{k,2}\approx 0.10$.
At $t_3$: distance $\approx 0.05$ → $q_{k,3}\approx 0.07$.
Average relevance $\bar{q}_k = (0.13+0.10+0.07)/3 \approx 0.10$.
With $\beta=0.5$, weights become $w_{k,1}=0.5+0.5\cdot(0.13/0.10)=1.15$, $w_{k,2}=1.05$, $w_{k,3}=0.85$.
The tokens whose predictions shift most when the verifier is removed receive a larger share of the reward, while tokens that are insensitive keep only the baseline share.
**Figure 1.** Advantages assigned to the tokens of one SQL response that is well formed and executable but uses >= for >. MRTPO places format credit on the tags, executability credit on the query, and the execution penalty mostly on >= 30.
**Figure 2.** Per-token advantage of a training response under GDPO and MRTPO. Both panels use the same per-reward advantages.
**Figure 3.** Leave-one-out replay on a text-to-SQL response. The response reads table telefone, which does not exist; the reference uses phone. Replay places the penalties of the executability and execution rewards mainly on telefone and the selected columns, and the format credit on the tags.
Experimental Results
Key performance gains of MRTPO across benchmarks and training dynamics.
Recall that reward interference occurs when multiple constraints misattribute penalties; MRTPO avoids this by using leave‑one‑out replay to localize credit to the responsible tokens.
MRTPO achieves the highest average execution accuracy on text‑to‑SQL, surpassing GDPO by 2.3%.
Table 1 reports MRTPO at 92.5% average versus GDPO’s 90.2%.
**Table 1.** Text-to-SQL execution accuracy (%); subscripts are standard deviations over seeds. MRTPO is best on average with both models, with and without SAW.
**Table 3.** BFCL-v4 single-turn AST accuracy (%); subscripts are standard deviations over seeds. MRTPO has the best overall accuracy on both splits with both models.
**Figure 4.** Tool calling on Qwen2.5-1.5B-Instruct over training: (a, b) training format and correctness rewards, median over seeds with interquartile range; (c, d) BFCL-v4 accuracy and format rate, mean over seeds. The binary format reward is harder to learn, while the continuous correctness reward rises early for all methods. The median number of steps to a format reward of 0.8 is 13 for MRTPO, 19 for GDPO and 34 for GRPO.
**Figure 5.** Training dynamics. (a) Text-to-SQL execution accuracy averaged over six benchmarks (Qwen2.5-Coder-7B-Instruct); (b) training response length (Qwen2.5-Coder-3B-Instruct); (c, d) IFBench prompt-level and MulDimIF accuracy (Qwen3-4B-Instruct-2507). Lines are seed means; bands span the seeds.
**Figure 6.** Where the replay distance falls. (a, b) Text-to-SQL with Qwen2.5-Coder-3B-Instruct; (c, d) tool calling with Qwen2.5-3B-Instruct. (a, c) Share of each reward's distance held by the top tokens of a response; (b, d) share on each part of the response. Gray marks the uniform assignment of GDPO.
**Table 4.** Replay localizes each constraint. For each satisfied constraint type, Dist. is the percentage of the constraint's replay distance on the part of the response its checker reads, Tok. is that part's percentage of the tokens, and Ratio is Dist. over Tok.; uniform assignment gives a ratio of 1. The last two rows pool all satisfied and all violated constraints. "–": fewer than five instances.
Ablations and Implementation Details
Ablation studies dissect component contributions and reveal their impact on performance.
We evaluate how each design choice of MRTPO contributes to the overall gains. By removing a component and measuring the performance drop, we can infer whether the component is essential or merely complementary.
Relation to reweighting methods: most prior multi‑reward approaches only re‑weight the per‑reward advantage before assigning a single token‑level advantage (Table 8). MRTPO instead reassigns advantage at the token level, making it orthogonal to those reweighting tricks.
Training and evaluation parameters (Table 9) show that all three tasks share the same optimizer settings (AdamW, $\beta$₁=0.9, $\beta$₂=0.999) and that MRTPO consistently uses $\beta$=0.5 with a 10‑step smoothstep warm‑up.
Replay contexts (Algorithm 1) illustrate the extra steps MRTPO adds: token‑wise relevance computation (Jensen–Shannon distance) and per‑reward weighting before summing token advantages.
Relevance and weighting details: the Jensen–Shannon divergence bounds token‑level credit ($q_{k,t}$∈[0,1]), the warm‑up schedule smoothly ramps $\beta$ from 0 to its final value, and the reallocation bound caps the variance increase introduced by MRTPO.
Table 5 demonstrates that shuffling the positional weights within each response eliminates the performance advantage of MRTPO over the GDPO baseline, confirming that the learned weighting is crucial.
Table 6 shows that allocating half of each advantage by relevance ($\beta$=0.5) yields the highest average accuracy across instruction‑following benchmarks, indicating a balanced trade‑off between uniform and relevance‑driven credit.
**Figure 7.** Leave-one-out replay on a tool-calling response. The response sets offset to 2; the reference has 3. The layout follows Figure 3.
**Figure 8.** Leave-one-out replay on an instruction-following response with 12 constraints, three of whose replays are shown. Layout and shading follow Figure 3.
**Table 9.** Training and evaluation parameters of the three tasks.
**Table 5.** Ablations on text-to-SQL with Qwen2.5-Coder-7B-Instruct. Shuffling the weights within each response removes the gain over GDPO.
**Table 6.** Allocation strength on instruction following with Qwen2.5-7B-Instruct. Assigning half of each advantage by relevance gives the highest average.