Each Reward to Its Relevant Tokens

🎯 One wrong token, and the penalty lands everywhere

Language models are increasingly trained against several verifiable rewards at once: the correctness and the format of a tool call, or the individual constraints of a complex instruction. GDPO normalizes each reward separately, so different reward combinations keep distinct advantages. But in GRPO, GDPO and SAW, every token of a response shares one advantage.

A verifier checks only a few positions. The penalty for one wrong token is spread over the whole response, where it can be outweighed by the credit other rewards put on the same tokens.

💡 Let each verifier's information point to its tokens

MRTPO assigns the advantage of each reward to the tokens relevant to it, using the information its verifier checks.

1. Write each verifier's information as a text block

   \(Z_k\) is what reward \(k\) checks: the format rule, the response's own SQLite error, the reference query, or one constraint of the instruction. \(P_{\text{all}} = P_0 \oplus\) all blocks; \(P_{\setminus k}\) leaves out \(Z_k\). When the information is already in the prompt (instruction constraints), \(P_{\text{all}}\) is the prompt and \(P_{\setminus k}\) removes one constraint. Sampling still uses \(P_0\).

2. Replay the response under both contexts

   The sampling policy reads its own response once under each context, one forward pass per context, with no new generation. At every token \(t\), the Jensen–Shannon distance \(q_{k,t}\) between the two next-token distributions measures how relevant that token is to reward \(k\).

3. Move each reward's advantage to its tokens

   The distances become per-token weights, and each token's advantage is the weighted sum over rewards. The weights average to 1 per response, so every response keeps its GDPO credit; only the placement changes.

What changes compared with GDPO

The other methods act on the response, so every token of a response gets the same advantage. MRTPO sets a weight per reward and token and can be combined with them.

📊 Best average on all three tasks and all six models

Text-to-SQL

Execution accuracy (%), mean ± standard deviation over 3 seeds.

Instruction following

Accuracy (%), mean ± standard deviation over seeds. The largest gains are on IFBench, whose constraints are not seen in training.

Tool calling

BFCL-v4 single-turn AST accuracy (%), Non-Live and Live overall, mean ± standard deviation over seeds.

⚡ MRTPO learns faster

🔍 Where the credit goes

Each reward's distance concentrates on a small part of the response. In text-to-SQL, 80% of a reward's distance falls on 26–29% of the tokens; in tool calling, on 13–15%. In instruction following, an end-phrase constraint puts 40% of its distance on 4% of the tokens: the phrase itself.

Tool calling

Instruction following

🧪 Using MRTPO on your own task

MRTPO needs two things from a task:

Building the two contexts

For each reward \(k\) the replay compares \(P_{\text{all}}\), which contains the blocks of all rewards, with \(P_{\setminus k}\), the same context with block \(k\) removed.

Writing a block

• One reward, one block. Never merge several constraints, rubric items or verifiers into one.
• Write the information itself, specific to this example. Build every block by code from the example's label or the verifier's own output; a generic requirement ("the query must execute") carries no information.
• Write it as a plain statement, like the rest of the prompt. No instructions to the model, headings, labels, or words about training, replay or scoring.
• Use only this response's own feedback. If there is none, the block is empty; never invent one.
• Write references in the form the policy is asked to produce, with the same tags, code fences and argument types.
• Keep the prompt, the rewards and the blocks consistent, and replay the rollout prompt and the response exactly as sampled.

BibTeX

@article{zhai2026mrtpo,
  title   = {Each Reward to Its Relevant Tokens: Reward-Specific Token-Level
             Credit Assignment for Multi-Reward {RL}},
  author  = {Zhai, Skylar and Wang, Haonan and Liu, Zhining and Song, Linxin and
             Ding, Xuwei and Liang, Jingcheng and Kang, Jian},
  journal = {arXiv preprint arXiv:2610.21439},
  year    = {2026}
}