Each Reward to Its Relevant Tokens: Reward-Specific Token-Level Credit Assignment for Multi-Reward RL

AI summary

Overview

Research area: Reinforcement learning with verifiable rewards (RLVR) for large language model post-training, specifically multi-reward optimization and token-level credit assignment.

Technical level: Intermediate (assumes familiarity with GRPO-style policy gradient objectives, advantages, and PPO clipping).

Scope: The paper proposes MRTPO, a method that uses each verifier's own information to redistribute that verifier's reward advantage onto the specific tokens it depends on, evaluated on text-to-SQL, multi-constraint instruction following, and tool calling across six models.

What This Paper Is About

When a language model is trained with several verifiable rewards at once (for example, format and correctness of a tool call, or the many individual constraints of an instruction), existing methods such as GDPO normalize each reward separately but then spread every reward's advantage evenly over every token of the response. Since each verifier typically checks only a small part of the response, a penalty for one mistake dilutes across tokens that had nothing to do with it. MRTPO's goal is to figure out, for each reward, which tokens actually determine it, and to place that reward's advantage mainly there.

Key Contributions

1. A credit-assignment framing of multi-reward RL. The authors show that existing multi-reward methods (GDPO, SAW, DARA, DVAO, SA-MRPO, GD²PO, RVPO) all decide how much each reward contributes at the response level, then give every token the same share, even though only part of the response is relevant to each reward.

2. Reward-specific token-level credit assignment (MRTPO). For each reward, the method replays the already-sampled response under two prompts that differ only in that reward's verifier information, and scores each token by the Jensen–Shannon distance between the two next-token distributions, normalized by ln 2. That score becomes a per-reward, per-token weight, requiring no auxiliary model, no token-level annotations, and no additional generation.

3. Orthogonality to response-level reweighting. MRTPO keeps GDPO's per-reward group advantages and only changes their placement, so it can be stacked on top of reweighting schemes; the paper demonstrates this by combining it with SAW.

4. Empirical validation across three task families and six models, showing better average performance than GRPO and GDPO and faster convergence, plus ablations and diagnostics showing where the replay distance actually lands.

Main Findings

• Best average on all three tasks and all six models. Text-to-SQL average with Qwen2.5-Coder-3B-Instruct: MRTPO 68.60 ±0.15 vs GDPO 67.32 ±1.66 and GRPO 65.37 ±1.18. With Qwen2.5-Coder-7B-Instruct: MRTPO 73.15 ±0.19 vs GDPO 71.55 ±0.27. Adding SAW on top gives the best averages with both models (68.96 ±0.19 and 73.41 ±0.35).

• Instruction following shows the largest gains. On IFBench, whose constraints are not seen in training, MRTPO gains 2.4 and 3.2 points over GDPO on the two models (Qwen3-4B-Instruct-2507: 38.76 prompt-level / 40.92 instruction-level for MRTPO vs 36.39 / 38.84 for GDPO; Qwen2.5-7B-Instruct: 38.87 / 42.00 vs 35.70 / 38.57). Average accuracy is 62.54 ±0.32 and 57.36 ±0.63, best on both models.

• Gain scales with the number of rewards. Text-to-SQL has three rewards and gains 1.3 and 1.6 points in average accuracy; instruction following has five to thirteen constraints per prompt and gains the most; tool calling has two rewards and MRTPO ends close to GDPO, with its advantage showing up as speed rather than final accuracy.

• Faster convergence, most visibly on the format reward. The binary tool-calling format reward, checked on the tags alone, reaches 97% compliance after 10 steps with MRTPO versus 2% for GDPO and 4% for GRPO. The median number of steps to a format reward of 0.8 is 13 for MRTPO, 19 for GDPO and 34 for GRPO. With Coder-7B, MRTPO matches GDPO's final text-to-SQL accuracy at step 60 of 100; with Qwen3-4B-Instruct-2507 it matches every baseline's final IFBench accuracy at step 200 of 300.

• Shorter reasoning. Text-to-SQL responses become 40% shorter than GDPO's with Coder-3B, because reasoning that does not affect the query receives less execution advantage.

• The replay localizes rewards sharply. In text-to-SQL, 80% of a reward's distance falls on 26% to 29% of tokens; the execution reward puts 10% of its distance on the selected columns (twice their token share) and little on the tags, while the format reward puts twice their share on the tags. In tool calling, 80% of a reward's distance falls on 13% to 15% of tokens, and the format reward gives the think tags three times their token share.

• Replay matches known checker locations in instruction following. On 192 HIR-16K responses from each of Qwen2.5-7B-Instruct and Qwen3-4B-Instruct-2507, a satisfied constraint places about 30% of its distance on the part its checker reads, about a tenth of the response; an end-phrase constraint holds 40% of the distance with 4% of the tokens (a ratio of 9.1 on Qwen2.5-7B). A violated constraint moves its distance to where the missing content should be.

• The placement, not just the weighting, causes the gain. Shuffling token weights within each response removes the gain over GDPO (71.03 average vs 73.15 for MRTPO and 71.55 for GDPO). Replaying all rewards jointly with one shared profile reaches 72.56 — better than GDPO but below MRTPO.

• Beats on-policy self-distillation on text-to-SQL. OPSD receives a full reference solution as privileged information yet averages 59.02 ±0.49 (Coder-3B) and 71.49 ±0.54 (Coder-7B), below MRTPO on every text-to-SQL benchmark.

• Modest compute overhead. Only 12% to 21% of response–reward pairs are replayed, and a training step takes 10% to 20% longer than GDPO on the same GPUs. The share of tool-calling groups whose correctness varies falls from 61% in the first ten steps to about 10%.

Methodology in Plain English

The starting point is GDPO, which standardizes each reward within a group of sampled responses so that different reward combinations stay distinguishable. MRTPO takes those same per-reward advantages and changes only where they go.

For each reward, the method assembles two prompts from the original prompt: one containing the text blocks of all verifiers, and one containing the same blocks except the block belonging to that reward. The blocks are things the verifier already uses — a required format description, the SQLite error the model's own query raised, the reference query, or a single instruction constraint. Because the two prompts differ only in one block, any change in the model's prediction between them is attributable to that block.

The model then re-reads its own already-sampled response under both prompts (one forward pass each) and, at each token position, measures the Jensen–Shannon distance between the two next-token distributions, normalized by ln 2 into [0, 1]. A token whose prediction barely moves when the block is removed is irrelevant to that reward; a token whose prediction shifts a lot is where the reward is decided.

These relevance scores are converted into weights that average to one over the response, using w = (1 − β) + β · (q / q̄), so an average-relevance token keeps weight 1 and every token keeps at least 1 − β. The weights are applied to the reward's advantage, the weighted advantages are summed across rewards per token, rescaled by the batch standard deviation, and plugged into the standard clipped surrogate objective. Because the weights average to one, the response's total credit is unchanged from GDPO — only its distribution across tokens moves. During training β is raised from 0 over the first N = 10 steps, and a cap limits how much reallocated advantage a token can receive when several rewards concentrate there.

The authors use β = 0.5 on all three tasks rather than tuning per task. Two ablations test the mechanism: shuffling the weights across positions within a response (same values, wrong locations) removes the gain, and replaying all rewards together with one shared profile falls short of giving each reward its own assignment.

Why This Matters

Impact on research. The paper reframes multi-reward RL as two separable questions — how much each reward contributes (response level) and which tokens receive it (token level) — and shows that existing methods in the DARA/DVAO/GD²PO/RVPO family only answer the first. Because MRTPO is orthogonal to those methods, it provides a drop-in token-level layer that composes with the reweighting literature rather than competing with it. It also offers a route to fine-grained credit without process reward models, Monte-Carlo value estimates, or per-rubric token discriminators, all of which need extra annotations or sampling.

Real-world applications:

• Tool-using agents, where a call must be both well-formed (tags, JSON structure) and semantically correct in function and argument choice — a format penalty should not land on correctly produced arguments.
• Text-to-SQL and database assistants, where the execution error points to specific tables or columns and the credit can follow that pointer instead of covering the entire query.
• Instruction-following systems with many simultaneous constraints (five to thirteen per prompt in this work), where per-constraint credit keeps one violated constraint from being masked by the others.
• Rubric- or verifier-based alignment pipelines in which each criterion comes with checkable information such as a reference answer or a rubric item that can be written as text.

Industry relevance. The method adds only forward passes over already-sampled responses — 10% to 20% longer steps, with replay skipped whenever a reward does not vary within its group — so it does not require a second model, a reward model, or new rollouts. Faster convergence on hard binary rewards (97% versus 2% format compliance at step 10) and 40% shorter reasoning traces translate directly into training cost and inference cost.

Future Directions

• Rewards without checkable text information. Length-based or otherwise unlocalizable rewards cannot be assigned this way and would have to keep GDPO's uniform allocation; how to handle mixed reward sets is left open.
• Adaptive allocation strength. β = 0.5 was used everywhere, but probes at 0.25 and 0.75 show the best β is task-dependent: on text-to-SQL, 0.25 gives a slightly higher average (73.59 vs 73.36 with one seed), while on instruction following 0.5 gives the best average and 0.75 is higher on IFEval and MulDimIF. Per-task or adaptive tuning is identified as a promising but compute-costly direction.
• Scaling to larger models. Experiments cover 1.5B to 7B parameters; whether the effect persists at larger scales is not reported.
• LLM-judged rewards. Extending the approach to rubric items scored by a language model judge, rather than programmatic verifiers, remains to be studied.

Target Audience

Researchers and practitioners working on RL post-training for LLMs who already understand GRPO-style objectives and want finer credit assignment without adding reward models or step-level labels. It is most immediately useful to teams training tool-calling agents, text-to-SQL models, or instruction-following systems with multiple programmatic verifiers, and to anyone building on GDPO-style multi-reward normalization who wants a composable token-level extension.

Authors’ abstract

Language models are increasingly trained against several verifiable rewards at once, such as the correctness and the format of a tool call, or the individual constraints of a complex instruction. Reward-decoupled normalization (GDPO) normalizes each reward separately so that different reward combinations keep distinct advantages, but it still spreads every reward's advantage evenly over the whole response, although each verifier checks only a few positions. We propose MRTPO, which assigns the advantage of each reward to the tokens relevant to it, using the information its verifier checks, without auxiliary models or token-level annotations. For each reward, the policy replays its own response with and without that reward's verifier information, and the per-token Jensen-Shannon distance between the two next-token distributions measures each token's relevance and reassigns the advantage. On text-to-SQL, multi-constraint instruction following and tool calling, MRTPO converges faster and achieves better performance than the compared methods.