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

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.

한 줄 요약

MRTPO는 여러 보상의 영향을 각각의 관련 토큰에 할당함으로써 다중 보상 강화 학습의 성능을 향상시킨다.

핵심 기여도

MRTPO는 GDPO와 달리 각 보상의 영향을 관련 토큰에만 할당함으로써 보상의 분산을 개선한다.
토큰별 Jensen-Shannon 거리를 사용해 보상별 관련성을 측정하고, 이에 따라 보상을 재할당한다.
텍스트-투-SQL, 다중 제약 지시 준수, 툴 호출에서 MRTPO는 GDPO보다 더 빠르게 수렴하며 더 높은 성능을 보인다.
추가적인 보조 모델, 토큰 레벨 어노테이션 없이도 보상 정보만으로 토큰 관련성을 파악한다.

핵심 아이디어

기존의 GRPO와 GDPO는 여러 보상을 합산하거나 정규화하여 전체 응답에 동일한 보상 값을 할당한다. 그러나 실제 보상은 보통 응답의 일부 토큰만을 기준으로 계산되므로, 전체 토큰에 동일한 영향을 주는 것은 비효율적이다. MRTPO는 각 보상의 검증 정보를 기반으로 정책을 재실행하고, 이때의 토큰 분포 변화를 통해 해당 보상에 영향을 주는 토큰을 식별한다. 이 과정에서 **Jensen-Shannon 거리**를 사용해 토큰별 관련성을 측정하고, 보상을 해당 토큰에만 할당한다. 이는 보상별로 **context pair**를 구성하고, **self-replay**를 통해 정책을 조정하는 방식이다.

기술적 접근법

**MRTPO (Multi-Reward-to-Token Policy Optimization)**: 각 보상에 대해, 정책이 해당 보상의 검증 정보가 있는 상태와 없는 상태에서 응답을 재생성한다.
**Jensen-Shannon distance**: 두 상태에서의 토큰 분포 간의 차이를 측정하여, 해당 보상에 영향을 주는 토큰을 식별한다.
**Group standardization**: 보상별로 정규화한 후, **token-level advantage**를 재할당한다.
**Clipped surrogate objective**: 기존 GRPO와 동일한 손실 함수를 사용하지만, **A_t^(j)**만 수정하여 토큰별 보상을 반영한다.
**No auxiliary models or annotations**: 보상 정보만으로 토큰 관련성을 파악한다.

주요 결과

**Text-to-SQL**, **multi-constraint instruction following**, **tool calling** 데이터셋에서 MRTPO는 GDPO보다 더 빠르게 수렴하며, 최종 정확도도 더 높았다.
MRTPO는 6개의 모델에서 GDPO의 최종 정확도를 **더 적은 학습 단계**로 달성했다.
**Jensen-Shannon distance** 기반의 토큰 관련성 측정이 보상 분산을 개선하는 데 기여했다.

의의 및 한계

MRTPO는 다중 보상 강화 학습에서 보상의 분산을 개선함으로써 정책의 학습 효율성을 높인다. 특히, **보상별 토큰 관련성**을 정량적으로 측정하고 이를 기반으로 보상을 재할당하는 방식은 기존 연구와 차별화된다. 그러나 MRTPO는 **검증 정보가 필요한 보상**에만 적용 가능하며, 검증 정보가 없는 보상에는 활용이 제한된다. 또한, **Jensen-Shannon distance 계산**은 계산 비용이 발생할 수 있어 대규모 모델에서는 성능 저하가 우려된다.

실용적 활용

MRTPO는 텍스트-투-SQL, 툴 호출, 복합 지시 준수 등 여러 자연어 처리 작업에 적용 가능하다. 특히, **각 보상이 특정 토큰에만 영향을 주는 경우**에 효과적이다. 예를 들어, SQL 쿼리의 실행 성공 여부나 툴 호출의 형식 정확도와 같은 **검증 가능한 보상**이 있는 상황에서 활용할 수 있다.