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

1 Introduction

• Practical multi-reward objectives include tool calls that must be well formed and correct, queries that must execute and return intended results, and instructions with multiple constraints.Each criterion is typically evaluated by a separate verifier that produces its own reward.Source: §1 · p. 1
• GDPO separates reward normalization to preserve distinct reward combinations, but its response-level advantages still give every token the same share of each reward.This can penalize correctly generated content when a verifier checks only a small portion of the response.Source: §1 · p. 1
• A token is relevant to a reward when that reward’s verifier information changes the policy’s prediction for it.Verifier information can include a reference tool call, a query’s execution error, or an individual instruction constraint.Source: §1 · p. 2
• MRTPO replays each response with and without one verifier’s information, then uses token-wise Jensen–Shannon distances between next-token distributions to redistribute that reward’s group-normalized advantage.It assigns token relevance separately for each reward, unlike response-level weighting methods.Source: §1 · p. 2
• Across text-to-SQL, multi-constraint instruction following, and tool calling, MRTPO achieves higher performance than GDPO with faster convergence.The method replays the policy’s own response without auxiliary models, token-level annotations, or additional generation.Source: §1 · p. 2

Multi-reward reinforcement learning

• GDPO preserves distinct reward combinations by normalizing each reward before summation, and concurrent methods adapt reward weights during aggregation.These approaches still operate at the response level; MRTPO can be combined with such reweighting methods, including SAW.Source: §2 · p. 2

Token-level credit assignment and privileged information

• Prior token-level credit methods use process reward models, Monte-Carlo value estimates, or rubric-specific discriminators, which require annotations, additional sampling, or another model.Other approaches use token entropy or remove the entire rubric, but reweight the aggregated advantage rather than each reward separately.Source: §2 · p. 2
• MRTPO uses self-replay and distributional differences under each reward’s information to allocate that reward’s advantage across tokens, while keeping rewards as the training objective.Unlike privileged-information self-distillation, replay determines where each advantage goes rather than replacing reward maximization with teacher-distribution fitting.Source: §2 · pp. 2–3

Multi-reward group optimization

• For each prompt, the behavior policy samples a group of responses, each scored by K verifiers; K can vary by prompt, with each instruction constraint serving as a reward.Group standardization uses each reward’s mean and sample standard deviation over the responses.Source: §3 · p. 3

GRPO and GDPO

• GRPO aggregates rewards before standardization, so responses with different reward combinations but equal sums can receive the same advantage.GDPO instead standardizes each reward before taking a weighted sum; the paper uses equal reward weights.Source: §3 · p. 3

Token-level objective

• In existing methods, every token in a response shares one advantage within the clipped surrogate objective.Figure 2 presents per-token advantages under GDPO and MRTPO using the same per-reward advantages.Source: §3 · p. 3, Figure 2 · p. 3

4.1 Motivation

• GDPO spreads each reward’s advantage across the whole response, so credit from other rewards can outweigh a penalty even when only a few tokens determine that reward.Source: §4.1 · p. 3

4.2 How Do We Find Relevant Tokens?

• For each reward, MRTPO replays the sampled response in contexts differing only in that verifier’s information and scores token relevance by the Jensen–Shannon distance between next-token distributions.For instruction constraints, the full context is the prompt, and the leave-one-out context removes one constraint.Source: §4.2 · pp. 3–4
• In the SQL example, replay associates executability and execution penalties mainly with the nonexistent telefone table and selected columns, while format credit falls on the tags.Source: Figure 3 · p. 4, §4.2 · p. 4

4.3 Redistributing Advantage by Relevance

• MRTPO redistributes each reward’s advantage using token relevance relative to the response average; β controls the relevance-weighted share, and every token retains at least 1−β weight.Weights average to one per response, preserving total credit; training raises β from zero and caps reallocation when rewards concentrate on the same tokens.Source: §4.3 · p. 4
• MRTPO has the best average text-to-SQL performance for both models, with and without SAW, and the best instruction-following average on both models, with its largest gains on IFBench’s unseen constraints.Source: Table 1 · p. 5, Table 2 · p. 5

5.1 Experimental Setup

• The evaluation covers text-to-SQL, instruction following, and tool calling, with GRPO and GDPO baselines and task-specific additional comparisons including SAW, CoRT, and OPSD.Instruction-following prompts carry five to thirteen binary constraint rewards; the experiments train models on each task and average results over independent runs.Source: §5.1 · pp. 5–6

Does token-level assignment improve multi-

• After 10 steps, required format appears on 97% of BFCL cases with MRTPO, versus 2% with GDPO and 4% with GRPO.MRTPO reaches a format reward of 0.8 in a median of 13 steps, versus 19 for GDPO and 34 for GRPO.Source: Figure 4 · p. 6, Does token-level assignment improve… · p. 7
• MRTPO has the best average performance on all three tasks and six models; because it shares per-reward advantages with GDPO, the difference is token-level assignment.MRTPO also combines with SAW to give the best average with both text-to-SQL models, while OPSD remains below MRTPO on every text-to-SQL benchmark.Source: Does token-level assignment improve… · pp. 6–7
• MRTPO passes GDPO’s final text-to-SQL accuracy at step 60 of 100, passes every baseline’s final IFBench accuracy at step 200 of 300, and produces responses 40% shorter than GDPO with Coder-3B.The reported figures are Figure 5a–c; the shorter responses are attributed to less execution advantage for reasoning that does not affect the query.Source: Figure 5 · p. 7, Does token-level assignment improve… · p. 7

Which tokens receive each reward’s advantage?

• Figure 6 shows reward distances concentrate on subsets of tokens: 80% falls on 26%–29% in text-to-SQL and on 13%–15% in tool calling.In text-to-SQL, execution distance emphasizes selected columns and format distance emphasizes tags; in tool calling, format distance emphasizes think tags while correctness distance centers on call-selection reasoning.Source: Figure 6 · p. 7, Which tokens receive each reward’s… · pp. 7–8

Does token-level assignment improve multi-

• MRTPO’s gain grows with reward count: 2.4 and 3.2 points on IFBench, 1.3 and 1.6 points in text-to-SQL average accuracy, and chiefly faster learning in tool calling.Instruction following has five to thirteen constraint rewards per prompt; text-to-SQL has three rewards and tool calling has two.Source: Does token-level assignment improve… · p. 7

Which tokens receive each reward’s advantage?

• For instruction constraints, satisfied checkers place about 30% of distance on the roughly 10% of tokens they read; an end-phrase constraint places 40% on 4% of tokens.For violated constraints, distance shifts toward where missing content should appear; different constraints’ distances on a response are weakly correlated.Source: Table 4 · p. 8, Which tokens receive each reward’s… · p. 8
• Shuffling token weights removes the gain over GDPO, while a shared profile for all rewards improves over GDPO but remains below MRTPO; task-specific β tuning may improve results but was omitted for compute cost.β = 0.5 was used on all tasks; replay makes training steps 10% to 20% longer than GDPO on the same GPUs.Source: Which tokens receive each reward’s… · pp. 8–9

• Across text-to-SQL, instruction following and tool calling, MRTPO outperforms GRPO and GDPO on six models and reaches GDPO’s final accuracy in fewer steps.Source: §6 · p. 9
• MRTPO replays the policy’s response with and without one verifier’s information, then assigns more advantage to tokens whose predictions change most.The method requires no step-level annotations.Source: §6 · p. 9
• MRTPO can localize rewards that come with checkable information, including a reference answer, execution error or rubric item.Source: §6 · p. 9

• MRTPO requires verification information that can be written as text and whose addition or removal changes policy predictions.Source: Limitations · p. 9
• Rewards without localizable verification information, such as response length, retain uniform allocation in GDPO.Source: Limitations · p. 9
• Experiments use 1.5B–7B models and β = 0.5 on all tasks; larger models, LLM-judged rewards and adaptive β remain unstudied.Source: Limitations · p. 9

• The study evaluates credit assignment on text-to-SQL, instruction-following and tool-calling datasets with programmatic verifiers.Source: Ethics Statement · p. 9
• The datasets used are publicly available, and the study involves neither human subjects nor new data collection.Source: Ethics Statement · p. 9
• More efficient optimization of verifiable rewards also makes misspecified rewards easier to optimize.Source: Ethics Statement · p. 9

• Concurrent methods use batch or group statistics, mask or weight responses, or replace summation with a soft minimum; all tokens still receive the same advantage.Source: A Relation to Reweighting Methods · p. 11
• MRTPO changes how each advantage is assigned across tokens, while concurrent methods alter how per-reward advantages combine into a response-level advantage.Source: A Relation to Reweighting Methods · p. 11
• Because MRTPO is orthogonal to these aggregation methods, the paper combines it with SAW as an example.Source: A Relation to Reweighting Methods · p. 11

B Training and Evaluation Parameters

• Results are averaged across independent runs; all tasks use AdamW with β1 = 0.9 and β2 = 0.999, while MRTPO uses β = 0.5, a 10-step smoothstep warm-up, and Eq. 8’s reallocation bound.GRPO standardizes the aggregated reward within each group without batch-level rescaling.Source: B Training and Evaluation Parameters · p. 11
• Text-to-SQL training excludes 3 Spider questions with non-executable references and 82 prompts exceeding 4,096 tokens; OPSD uses released OmniSQL chain-of-thought solutions as reference context.OPSD shares rollout and optimizer settings with the other methods; SAW weights each reward by its batch coefficient of variation, normalized to sum to three.Source: B Training and Evaluation Parameters · p. 11
• Instruction following uses HIR-16K’s first two subsets, comprising 6,998 prompts and 43,496 program-checked constraints; tool-calling numeric or boolean reference strings are converted to schema types.Source: B Training and Evaluation Parameters · p. 11

C Replay contexts

• The SQL replay contexts distinguish format requirements from execution feedback, including the reference query and errors for failing or missing SQL blocks.Source: C Replay contexts · pp. 11–12
• For instruction following, each replay omits one constraint from the HiR concatenation template; Figures 7 and 8 show tool-calling and instruction-following replays.The instruction-following example has 12 constraints, with three replays shown.Source: C Replay contexts · p. 12, Figure 8 · p. 15, Figure 7 · p. 15
• Tool-call replays append format and correctness blocks to the last user message, with correctness first; the format block specifies call structure and the correctness block supplies ordered reference operations and arguments.Source: C Replay contexts · pp. 13–14

D Relevance and Weighting Details

• The token relevance divergence is computed over the full vocabulary; its symmetry and bounded range capture distribution shifts among competing tokens, even when the sampled token’s probability is unchanged.Source: D Relevance and Weighting Details · p. 14
• A smoothstep warm-up raises β from zero over the first N steps, while batch rescaling bounds reallocation and preserves each response’s mean advantage.The rescaling keeps the token-advantage second moment at no more than twice that of the uniform advantage.Source: D Relevance and Weighting Details · p. 14