Learn What's Left, Not What's Mastered: Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization

Reinforcement learning (RL) with group-relative advantages has become the de facto standard for post-training language model reasoners. However, when optimizing multiple reward objectives, existing methods typically scalarize the reward vector with a fixed weighted sum before group-wise standardization. We show that this design leads to two fundamental problems: rollouts with distinct reward profiles can receive identical advantages, and all objectives are optimized with fixed relative weights regardless of their current level of saturation. As a result, training continues to allocate gradient budget to already-solved objectives instead of focusing on those with greater remaining headroom. We introduce Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization (SA-MRPO), which standardizes each reward objective independently and adaptively discounts its contribution according to a batch-level estimate of objective saturation. This dynamically reallocates optimization effort toward under-optimized objectives while empirically maintaining performance on those that are already well satisfied. We further show that saturation-aware reweighting can reverse the sign of an update, rather than merely rescale its magnitude. Across mathematical reasoning with two- and three-objective reward combinations, SA-MRPO improves the harder correctness objective over GDPO in 12 of 15 benchmark comparisons, with gains of up to 5% on AIME24. On adaptive reasoning it improves accuracy on all five benchmarks, by 3.8% on average and up to 9.2 % on AMC23, and on coding benchmarks it improves pass rate by up to 2.3%, while in all settings maintaining the easier objectives near their already satisfied levels.

한 줄 요약

SA-MRPO는 다중 보상 최적화에서 포화 상태에 따라 보상 가중치를 동적으로 재조정하여 학습 효율을 향상시킨다.

핵심 기여도

기존 다중 보상 최적화에서 발생하는 **보상 해상도 손실**과 **목표 포화 무시** 문제를 명확히 지적.
**SA-MRPO**를 제안: 각 보상 목표를 독립적으로 정규화하고, **배치 수준의 포화 추정**에 따라 가중치를 적응적으로 조정.
**AIME24**에서 5%, **AMC23**에서 9.2%, **코딩 벤치마크**에서 최대 2.3% 개선.
**GDPO** 대비 15개 중 12개 벤치마크에서 더 높은 정확도 달성.

핵심 아이디어

기존 다중 보상 최적화는 보상 벡터를 고정된 가중치 합으로 스칼라화한 후 그룹별 표준화를 수행한다. 이는 **보상 해상도 손실**과 **목표 포화 무시**라는 두 가지 문제를 유발한다. 예를 들어, 서로 다른 보상 프로파일을 가진 rollout이 동일한 advantage를 받을 수 있으며, 이미 포화된 목표에도 동일한 가중치가 적용되어 학습 효율이 저하된다.

SA-MRPO는 이 문제를 해결하기 위해 **각 보상 목표를 독립적으로 정규화**하고, **목표의 포화 정도에 따라 가중치를 적응적으로 조정**한다. 포화는 배치 평균 보상과 목표의 달성 가능한 범위를 비교하여 추정하며, 목표가 최대치에 가까워질수록 점진적으로 가중치를 낮춘다. 이는 **정책 업데이트의 방향**을 바꾸는 효과도 있다.

기술적 접근법

**GRPO (Group Relative Policy Optimization)** 기반의 다중 보상 최적화 프레임워크 사용.
**GDPO (Group reward-Decoupled Policy Optimization)**는 보상 차원을 독립적으로 정규화하지만, 가중치는 고정.
SA-MRPO는 **정규화된 보상에 포화 상태에 따른 가중치를 적용**.
**가중치 조정 파라미터**는 단일 지수 $\gamma$로 제어.
**목표별 정규화**: $r_k^{(i,j)}$를 각 목표별 평균과 표준편차로 정규화.
**포화 추정**: $r_k^{(i,j)}$의 배치 평균을 목표의 최대 달성 가능 범위와 비교.
**가중치 조정**: $\exp(-\gamma \cdot s_k)$, 여기서 $s_k$는 목표 $k$의 포화 정도.

주요 결과

**수학적 추론**에서 15개 벤치마크 중 12개에서 GDPO 대비 정확도 개선, 최대 5% (AIME24).
**적응 추론**에서 5개 벤치마크 모두 정확도 개선, 평균 3.8%, 최대 9.2% (AMC23).
**코딩**에서 4개 벤치마크 중 3개에서 패스율 개선, 최대 2.3%.
**기존 목표**는 이미 포화 상태를 유지하면서도 **더 어려운 목표**에 대한 성능 향상.

의의 및 한계

SA-MRPO는 다중 보상 최적화에서 **목표별 학습 효율을 동적으로 조정**함으로써, 기존 방법에서 무시되던 **목표 간 불균형**을 해결한다. 특히, **정책 업데이트 방향**을 바꾸는 기능은 기존의 단순 가중치 조정을 넘어선 새로운 접근법이다.

그러나, **포화 추정**은 배치 수준의 평균에 기반하므로, **데이터 배치의 편향**이 결과에 영향을 줄 수 있다. 또한, **가중치 조정 파라미터 $\gamma$**는 실험적으로 조정되어야 하며, **과도한 조정**은 성능 저하를 초래할 수 있다.

실용적 활용

SA-MRPO는 **수학적 추론**, **적응 추론**, **코드 생성** 등 다양한 언어 모델 학습 상황에서 적용 가능하다. 특히, **복수의 보상 목표가 존재하는 상황**에서, **학습 자원을 효율적으로 재할당**하여 어려운 목표를 우선적으로 개선할 수 있다. 이는 **대규모 언어 모델의 학습 효율성**을 높이는 데 기여할 수 있다.