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

Each Reward to Its Relevant Tokens. Reward-Specific Token-Level Credit Assignment for Multi-Reward RL. By Skylar Zhai and six others, from the University of Minnesota, the National University of Singapore, and four other institutions. Published on the fifteenth of October. Its code repository is worked out: about three hundred source files, with a readme and a licence.

Language models are increasingly trained against several verifiable rewards at once: the correctness and the format of a tool call, for example, or the individual constraints of a complex instruction. Reward-decoupled normalization, known as GDPO, normalizes each reward separately, so that different combinations of rewards keep distinct advantages. But it still spreads every reward's advantage evenly over the whole response, even though each verifier checks only a few positions in it. A reward about format and a reward about correctness end up crediting the same tokens in the same proportion, whatever those tokens actually did.

MRTPO assigns the advantage of each reward to the tokens relevant to that reward. Relevance is read from the information the reward's own verifier checks, and the method needs no auxiliary models and no token-level annotations.

For each reward, the policy replays its own response twice: once with that reward's verifier information and once without it. At every token, the Jensen-Shannon distance between the two next-token distributions measures how relevant that token is to the reward, and the advantage is reassigned according to that measure. The policy itself, in effect, reports which parts of its answer a given verifier bears on.

The method is tested on text-to-SQL, on multi-constraint instruction following, and on tool calling. Across those settings, MRTPO converges faster and achieves better performance than the compared methods. All three are settings where a response answers to several verifiers at once.

GDPO had already kept rewards apart when computing advantages. MRTPO keeps them apart in position as well, sending each reward's credit to the tokens its verifier actually looks at, and it draws that map from the policy rather than from any extra model or label.