🤖 machine learning

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

This paper introduces Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization (SA-MRPO), a novel method that independently standardizes and adaptively reweights multiple reward objectives based on their saturation levels to dynamically shift optimization focus toward under-optimized goals, thereby significantly improving performance on challenging benchmarks while maintaining proficiency on already-solved tasks.

Original authors: Yixuan Wang, Yifei Chen, Haichao Zhang, Haozheng Luo, Xander Wu, Jie Ni, Yun Fu, Nuno Vasconcelos, Yijiang Li

Published 2026-08-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yixuan Wang, Yifei Chen, Haichao Zhang, Haozheng Luo, Xander Wu, Jie Ni, Yun Fu, Nuno Vasconcelos, Yijiang Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the rapidly evolving world of artificial intelligence, researchers are teaching computer programs to think more like humans, particularly when solving complex problems that require step-by-step reasoning. To teach these systems, scientists often use a method called reinforcement learning, where the program tries different approaches and receives feedback in the form of rewards. If the program solves a math problem correctly, it gets a point; if it follows a specific format, it gets another point. The goal is to maximize these rewards so the program learns to produce better answers. However, real-world tasks rarely have just one goal. A helpful assistant must be accurate, but it also needs to be concise, safe, and follow strict formatting rules. The challenge arises when a program tries to juggle all these goals at once. If the program becomes very good at one task, such as keeping answers short, it might stop trying to improve in that area, yet the training system might keep pushing it to be even shorter, wasting valuable learning time that could be spent fixing a different, more difficult problem, like getting the math right.

A team of researchers has identified a flaw in how current systems handle these multiple goals and has proposed a new way to train them that focuses on what still needs to be learned rather than what has already been mastered. In their study, they observed that standard training methods often treat all goals as equally important from start to finish, regardless of how well the program is already performing on them. This leads to a situation where the system keeps polishing a skill it has already perfected, while neglecting a harder skill that still has room for improvement. To fix this, the researchers developed a technique they call Saturation Aware Advantage Reweighting. This method acts like a smart manager that constantly checks the progress of each goal. When a goal is nearly perfect, the system automatically reduces the pressure on it, shifting the focus and energy toward the goals that are still struggling.

The researchers tested this approach on language models tasked with solving difficult mathematical problems and writing computer code. In these tests, the models had to balance getting the answer right with following rules about how long the answer should be or how it should be formatted. Under the old training methods, the models often struggled to improve their accuracy once they had already learned to follow the length rules perfectly. The new method, however, recognized that the length rule was no longer a challenge and stopped wasting effort on it. Instead, it directed the model's attention to the harder task of solving the math correctly. The results were clear: across fifteen different comparisons involving various math competitions and benchmarks, the new method improved the accuracy of the harder tasks in twelve of them. In one specific test involving the American Invitational Mathematics Examination, the improvement was as high as five percent. Crucially, this boost in accuracy did not come at the cost of the easier tasks; the models continued to follow the length and format rules just as well as before.

This approach also worked when the researchers tested it on adaptive reasoning, where the goal was to find the right balance between being correct and being efficient. They created a scenario where the "length" goal had a clear limit: once an answer was short enough, making it shorter provided no extra benefit. The new training method noticed that the model had already reached this limit and stopped trying to shorten the answers further. Instead, it used that saved-up effort to make the answers more accurate. On average, this led to a nearly four percent increase in accuracy across five different benchmarks, with the largest jump being over nine percent on a specific math competition. The models did not become careless with their length; they simply stopped trying to be shorter than necessary and focused on being smarter.

The study also explored how the system decides when to stop pushing on a goal. The researchers introduced a control knob, a single number that determines how aggressively the system should ignore goals that are already satisfied. They found that turning this knob up made the system focus more heavily on the difficult tasks, which improved accuracy but sometimes led to slightly longer answers. Turning it down kept the answers shorter but did not improve the accuracy as much. This showed that the method is flexible and can be tuned to find the best balance for a specific situation. The researchers also tested the method on computer coding tasks, where the model had to write programs that both ran without errors and passed specific tests. Again, the new method helped the model improve its ability to pass the tests while maintaining its ability to write code that runs correctly.

The core discovery is that effective training requires paying attention to the remaining potential for improvement in each area, rather than treating all goals as static targets. By dynamically adjusting how much attention is paid to each goal based on how close it is to being perfect, the system learns more efficiently. The researchers showed that this is not just a theoretical idea but a practical improvement that works across different types of reasoning and different sizes of computer models. They demonstrated that by letting go of the easy wins, the system can achieve much better results on the hard ones, leading to smarter and more capable artificial intelligence without sacrificing the basic rules it must follow. This suggests that the future of training these systems lies in understanding not just what they have learned, but what they still have to learn.

Technical Summary: Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization (SA-MRPO)

1. Problem Statement

Reinforcement Learning with Verifiable Rewards (RLVR) has become the standard for post-training large language models (LLMs) to improve reasoning capabilities. While Group Relative Policy Optimization (GRPO) and its variants have simplified this process by eliminating the need for learned value functions, they typically assume a single scalar reward. In practice, reasoning models must satisfy multiple objectives simultaneously, such as correctness, response length, format compliance, safety, and executability.

The standard approach to multi-objective optimization involves scalarizing these rewards into a fixed weighted sum before performing group-wise standardization. The authors identify two fundamental limitations in this design:

Loss of Reward Resolution: Distinct reward profiles can collapse into the same scalar value. For instance, a rollout with high correctness but poor format might receive the same scalar reward (and thus the same advantage) as a rollout with poor correctness but perfect format, provided the weighted sum is identical.
Ignoring Objective Saturation: Fixed weights treat all objectives as equally important throughout training, regardless of their current state. If an auxiliary objective (e.g., format compliance) is already near-saturation (solved), the optimization process continues to allocate gradient budget to it. This diverts effort away from harder, under-optimized objectives (e.g., mathematical correctness) that still have significant "headroom" for improvement.

Existing methods like Group reward-Decoupled Policy Optimization (GDPO) address the first issue by normalizing rewards independently but fail to address the second, as they still rely on fixed weights that do not adapt to the saturation level of each objective.

2. Methodology: SA-MRPO

The authors propose Saturation Aware Advantage Reweighting for Multi-Reward Policy Optimization (SA-MRPO). This method preserves per-objective normalization but adaptively reweights each objective based on its current degree of saturation.

Core Mechanism

For a batch of rollouts and $n$ reward objectives, SA-MRPO calculates a saturation-aware advantage as follows:

Saturation Estimation: For each objective $k$ , the saturation ratio $s^{(k)}$ is estimated using the batch-mean reward $\bar{r}^{(k)}$ relative to the objective's attainable range $[r^{(k)}_{min}, r^{(k)}_{max}]$ :
$s^{(k)} = \frac{\bar{r}^{(k)} - r^{(k)}_{min}}{r^{(k)}_{max} - r^{(k)}_{min}} \in [0, 1]$
A value of $s^{(k)}$ close to 1 indicates the objective is saturated (near maximum reward), while a value close to 0 indicates significant room for improvement.
Adaptive Reweighting: The prescribed weight $w_k$ for each objective is modulated by a saturation factor controlled by an exponent $\gamma \ge 0$ :
$\tilde{w}_k = w_k (1 - s^{(k)})^\gamma$
As an objective approaches saturation ( $s^{(k)} \to 1$ ), its effective weight $\tilde{w}_k$ decreases. Conversely, objectives with low saturation retain higher weights.
Advantage Construction: The aggregate advantage for a rollout is computed by summing the normalized advantages of individual objectives, weighted by their saturation-adjusted weights:
$\tilde{A}^{(i,j)} = \sum_{k=1}^n \tilde{w}_k A^{(i,j)}_k$
where $A^{(i,j)}_k$ is the group-relative advantage for objective $k$ . This aggregate is then normalized across the batch before being used in the standard clipped surrogate objective.

Theoretical Properties

Generalization: SA-MRPO strictly generalizes GDPO (when $\gamma=0$ ) and GRPO (in the single-objective setting).
Directional Change: Unlike simple rescaling, saturation-aware reweighting can reverse the sign of an update. If a saturated objective conflicts with an under-optimized one, the reduced weight on the saturated objective can cause the aggregate advantage to flip, effectively changing the direction of the policy update to favor the unsaturated objective.
Conflict Handling: The method acknowledges that prioritizing an unsaturated objective may degrade a previously optimized one if their gradients conflict. However, the authors argue this is a necessary trade-off to maximize overall performance, as the gain from improving the under-optimized objective outweighs the loss on the saturated one.

3. Key Contributions

Problem Identification: The paper identifies that scalarized multi-reward optimization suffers from reward-resolution loss and, more critically, ignores objective saturation, leading to inefficient allocation of optimization effort.
Algorithm Proposal: SA-MRPO introduces a mechanism to dynamically reweight normalized reward objectives based on their remaining attainable reward range. It demonstrates that this can alter the direction of policy updates, not just their magnitude.
Empirical Validation: The method is validated across mathematical reasoning, adaptive reasoning, and code generation. It consistently shifts optimization effort from saturated to under-optimized objectives, improving the latter while maintaining performance on the former.

4. Experimental Results

The authors evaluated SA-MRPO against GDPO and baseline models across diverse settings:

Mathematical Reasoning: Using Qwen2.5 models (3B and 7B) with two and three objectives (correctness, length, format), SA-MRPO improved accuracy over GDPO in 12 out of 15 benchmark comparisons. Notable gains include a 5.0% improvement on AIME24 and 3.5% on MATH500 for the 7B model. Crucially, these gains were achieved without significantly degrading performance on the auxiliary length objective.
Adaptive Reasoning: In a setting with an explicitly saturated length objective (where responses under a token limit receive maximum reward), SA-MRPO improved accuracy on all five benchmarks by an average of 3.8%, with a peak gain of 9.2% on AMC23. The method allowed the model to utilize additional tokens for reasoning once the length constraint was satisfied, demonstrating effective reallocation of capacity.
Code Generation: On benchmarks like APPS and Codeforces, optimizing for both executability (easier to satisfy) and pass rate (harder), SA-MRPO improved pass rates on 3 of 4 benchmarks (up to 2.3% gain) while maintaining executability levels comparable to GDPO.
Hyperparameter Sensitivity: The saturation exponent $\gamma$ was shown to directly control the trade-off between improving under-optimized objectives and preserving saturated ones. Moderate values (e.g., $\gamma=0.5$ ) provided the best balance, yielding the highest average accuracy.

5. Significance and Claims

The paper claims that effective multi-reward policy optimization requires more than just balancing reward scales; it must account for the remaining useful improvement in each objective.

Dynamic Allocation: SA-MRPO provides a principled way to dynamically reallocate gradient budget away from "solved" problems toward those that still have headroom, without modifying the underlying GRPO policy update mechanism.
Reversal of Updates: The method is capable of reversing the sign of an update, a capability absent in standard scalarization or fixed-weight decoupling, allowing the policy to actively de-prioritize saturated objectives when they conflict with harder goals.
General Applicability: The approach is shown to be effective across different domains (math, adaptive reasoning, code) and model scales, suggesting that saturation-aware reweighting is a robust strategy for multi-objective RLVR.

The authors conclude that future multi-reward optimization strategies should explicitly consider the saturation state of objectives to avoid wasting optimization capacity on already-satisfied constraints.