WAS-PO: Using Whitened Advantages to Solve the Correlation Trap in Multi-Reward Reinforcement Learning

Multi-Reward RL is becoming increasingly common in language model alignment—a model may need to simultaneously optimize tool-calling accuracy, response helpfulness, and safety. The standard approach is to normalize each reward and then take a fixed-weight sum. But when these rewards are correlated with each other, this naive approach systematically distorts the weights you are actually applying. A new paper from arXiv proposes WAS-PO (Whitened Advantage Summation Policy Optimization), re-examining this problem from the perspective of covariance structure.

How Correlation Quietly Distorts Your Reward Weights

The core insight of the paper can be summarized in one sentence: when rewards are correlated, weighted summation double-counts their shared variance. Suppose you have two reward signals that are highly correlated—for example, "response accuracy" and "tool-call success rate" often move in the same direction in tool-calling tasks. You assign weight 0.7 to the first reward and 0.3 to the second, and intuitively you think the contribution ratio of the two in the final optimization objective is 7:3.

But it isn't. From the perspective of the ZCA-decorrelated components of the reward signals, the weights actually applied by weighted summation are $\Sigma^{1/2} w$, not the $w$ you selected, where $\Sigma$ is the reward correlation matrix. When $\Sigma$ is not the identity matrix, the components of $\Sigma^{1/2} w$ can diverge significantly from the weight direction you set—some components are amplified, some are compressed, and signs may even flip.

This explains why in practice, tuned reward weights "fail" when you switch to a different task or a different combination of rewards: you are tuning one set of parameters, but the model is optimizing another.

ZCA Whitening is a linear transformation used to convert a correlated random vector into a vector whose components are uncorrelated and have unit variance, while keeping the transformed vector as close as possible to the original vector in Euclidean space. Given a covariance matrix $\Sigma$, the ZCA whitening transformation matrix is $\Sigma^{-1/2}$, and the covariance of the transformed data becomes the identity matrix. In the context of reward aggregation, after the original reward vector is acted on by $\Sigma^{-1/2}$, its components become decorrelated "whitened rewards"; when weights $w$ are then applied, the weights genuinely act on independent components. In contrast, directly taking a weighted sum of the original rewards is equivalent to using weights $\Sigma^{1/2}w$ in the ZCA-whitened space, which is precisely the mathematical root of the distorted weight semantics. The difference between ZCA and PCA whitening is that ZCA does not perform a rotation; the transformed components still preserve the directions of the original coordinate axes as much as possible, which makes it easier to interpret the meaning of the weights.

WAS-PO's Approach: Whiten First, Then Weight

WAS-PO's solution directly targets the above problem: after normalizing advantages but before weighted summation, first whiten the advantages with $\Sigma^{-1/2}$, so that the subsequent weighting acts on decorrelated components. In this way, the weights $w$ you select genuinely act on the ZCA-decorrelated components, and the weight semantics align with the optimization objective.

The paper further derives the scalarized objective that the WAS-PO gradient follows when the normalization statistics are fixed, and presents two noteworthy theoretical results:

• The equivalent weights on original rewards may be negative. Even if all the weights you set for the rewards are positive, the equivalent weights of the whitened aggregate objective on some original rewards may still be negative. This means the model may be "pursuing" the opposite direction of a certain reward—this is not a bug, but a mathematically necessary projection of the correlation structure.
• An exact reversal condition for two rewards. The paper gives an exact two-reward condition under which the reward with the smaller weight gets reversed in the equivalent objective. This provides a testable diagnostic signal for debugging multi-reward setups.

In addition, the paper provides a loose upper bound on the error introduced by estimating $\Sigma$ during training. This bound is rather loose, but at least it shows that estimation error will not amplify without bound.

Batch-level normalization and per-prompt normalization are two different granularities of advantage normalization. Per-prompt normalization standardizes the advantage values separately within each prompt (i.e., each group of responses from the same source), eliminating scale fluctuations caused by difficulty differences across prompts; batch-level normalization computes the mean and standard deviation over all advantage values in the entire training batch and applies a global scaling. The former emphasizes local comparability, while the latter is simpler to compute and can stabilize the global gradient scale. WAS-PO replaces GDPO's per-prompt normalization with batch-level normalization, and this change is not separated from the whitening operation in the experiments, so WAS-PO's shortfall under two rewards may stem from the change in normalization granularity rather than from whitening itself. Understanding this distinction helps determine the true source of the method's improvement.

Experimental Results: Training Reward Improves, but Generalization Does Not Keep Pace

The paper conducted experiments on a tool-calling task using Qwen2.5-1.5B-Instruct, with GDPO as the baseline. The results show a split worth being wary of:

On training reward, WAS-PO consistently wins. In the three-reward setting, all seven random seeds achieved higher average training reward, by 3.8%; in the four-reward setting, all three seeds were superior, by 3.9%. This consistency across all seeds shows that the improvement is not noise.

But the improvement in training reward did not carry over to held-out accuracy. Under three rewards, WAS-PO's held-out accuracy was slightly lower than GDPO's (not significant); under four rewards, it was lower by 0.03 in absolute terms (only two seeds were evaluated). In other words, the model scores higher on the training objective, but does not become better on the task metric we actually care about.

Under two rewards, WAS-PO actually falls behind. In the setting with only two rewards, WAS-PO's training reward is lower than GDPO's. The paper weakly points to the cause with a three-seed controlled experiment: the gap may come from batch-level normalization rather than from whitening itself. It should be noted that there is a confounding factor here—WAS-PO simultaneously replaces GDPO's per-prompt normalization with batch-level normalization, and the paper explicitly states that these two changes were not separated under three or four rewards. So the question of "how much whitening contributes" has not yet received a clean answer from the experimental design.

Verification of the sign prediction. In a mathematical task setting, two rewards are nearly identical by construction. WAS-PO, as the sign condition predicted, suppressed the reward with the smaller weight. This is direct support for the theoretical analysis.

How to Interpret These Results

The value of this paper has two layers. On the theoretical level, it brings to the surface a structural problem in multi-reward aggregation that has long been overlooked: reward correlation distorts weight semantics, and whitening is a principled correction with testable predictions (the sign-reversal condition).

On the experimental level, the results are honest but not particularly impressive. The stable improvement in training reward coexists with unchanged or even slightly decreased held-out accuracy, suggesting that "higher aggregate training reward" does not necessarily equal "a better model"—this is a recurring phenomenon in multi-reward RL, and WAS-PO is no exception. The shortfall under two rewards and the confounding of the normalization change also leave the method's applicable boundaries unclear.

For practitioners working on multi-reward RLHF or multi-objective alignment, the practical takeaway of the paper may be: check your reward correlation matrix. If the correlation between rewards is low, whitening brings limited change; if the correlation is high, then the weights you are currently tuning are very likely not acting in the direction you think—whether or not you adopt WAS-PO, this is worth being vigilant about.

Held-out accuracy refers to the task accuracy of the model on an evaluation dataset that it has not seen during training and that is independent of the training reward signals. In multi-reward reinforcement learning, the training reward is a direct reflection of the optimization objective, but the model may improve the training reward by overfitting to specific patterns in the reward signals (such as exploiting biases in the reward model) rather than genuinely improving task ability. As an external metric, held-out accuracy more truthfully reflects the model's generalization performance. In the WAS-PO experiments, the training reward improved but held-out accuracy did not keep pace, which precisely shows that optimization of the aggregate training reward may be misaligned with the actual task objective—a phenomenon that must be guarded against when evaluating multi-reward RL methods.

One-line Summary

WAS-PO corrects the weight distortion caused by multi-reward correlation through whitened advantages; training reward improves but generalization does not.

Paragraph Summary

The paper points out that in multi-reward reinforcement learning, reward correlation causes the actual weights of naive weighted summation to become Σ^{1/2}w, deviating from the set values. WAS-PO proposes whitening with Σ^{-1/2} after normalizing advantages and before weighting, so that the weights act on decorrelated components. Theoretical derivations show that equivalent weights may be negative, and give an exact condition under which the lighter weight is reversed under two rewards. In tool-calling experiments, training reward under three/four rewards is higher by 3.8%/3.9%, respectively, but held-out accuracy does not improve and may even slightly decrease; under two rewards it falls behind GDPO, and the whitening and normalization changes are not separated. The paper's value lies in revealing a structural problem and providing testable predictions, but the experimental results are limited.