Don't Mix Rewards, Mix Policies: Policy Decomposition and Optimization for Multi-Reward RL

1]Fundation Model Research Center, CASIA 2]School of Artificial Intelligence, UCAS 3]Institute for AI Industry Research (AIR), Tsinghua University 4]College of Automotive and Energy Engineering (CAEE), Tongji University \contribution [*]Equal contribution \contribution [ 🖂 ]Corresponding authors \contribution [‡]Work done during internship at THU-AIR

Don’t Mix Rewards, Mix Policies: Policy Decomposition and Optimization for Multi-Reward RL

Ruiming Liang

Yi Zhong

Yizhen Yuan

Yinan Zheng

Tianyi Tan

Tianyue Wang

Haiyun Guo

Jinqiao Wang

Xianyuan Zhan [ [ [ [

Abstract Modern large language models (LLMs) are expected not just to answer correctly, but to adapt their behavior to different human values and use cases. As a result, multi-reward reinforcement learning (RL) has become an increasingly important problem for LLMs, where each reward captures a different aspect of desired behavior. However, optimizing with multiple rewards suffers from a more severe alignment tax issue, where different optimization objectives can trade off or even conflict with each other, leading to unstable and inefficient post-training. In this work, we propose PRISM , a new multi-reward RL framework built upon the idea of policy-space decomposition and composition. Instead of compositing different rewards, PRISM optimizes a set of standalone positive policies and a global negative policy. This alleviates the potential conflict during multi-reward policy optimization, while enabling controllability during inference by flexible policy composition. Experiments on scientific reasoning, tool-use reasoning, and helpfulness–safety alignment show that PRISM consistently outperforms existing multi-reward RL baselines, with extra controllability for inference-time preference control.

\metadata [Emails] liangruiming2024@ia.ac.cn , zhanxianyuan@air.tsinghua.edu.cn

1 Introduction

Large language models [ achiam2023gpt , liu2024deepseek ] are widely deployed as general-purpose assistants, making alignment with human intent a central concern [ askell2021general , ji2023beavertails ] . Reinforcement learning [ sutton1998reinforcement ] is among the most effective tools for this purpose, optimizing LLMs directly against reward signals that encode human preferences [ ouyang2022training , bai2022training ] . Yet human preferences are rarely captured by a single objective: users expect models to be helpful, safe, and honest at the same time [ wang2023aligning , liu2026gdpogrouprewarddecouplednormalization ] , and these goals often conflict, so improving one dimension costs another [ ouyang2022training , yang2024rewards , lin2024mitigating ] . LLM alignment is thus inherently a multi-reward optimization problem [ williams2024multi , liu2026gdpogrouprewarddecouplednormalization ] , substantially harder than single-objective RL.

Existing multi-reward RL methods differ mainly in how they merge reward signals into one training objective. Linear scalarization [ zhou2024beyond , williams2024multi ] sums rewards with manual weights; it is simple but scale-sensitive, and the weights must be retuned per task. Reward shaping and adaptive weighting [ wu2023fine , wang2024interpretable , de2024dynamic , liu2026gdpogrouprewarddecouplednormalization ] mitigate scale mismatch, but still collapse all rewards into one scalar before each update, so conflicting preferences compete within the same gradient step. Constrained or multi-objective RL [ achiam2017constrained , roijers2013survey , li2025gradient ] treats some rewards as Lagrangian constraints in pursuit of Pareto optimality, but is often hard to optimize and unstable in practice.

Despite their differences, all these methods compose preferences in the reward space before policy updates. We argue this principle is fundamentally limiting, for two reasons. First, reward signals differ in scale, distribution, and calibration [ wang2024interpretable , liu2026gdpogrouprewarddecouplednormalization , rame2024warm ] , making reward weights an unreliable interface for controlling trade-offs. Second, collapsing multiple preferences into one optimization signal forces them to compete within the same parameter update, yielding a compromise policy that under-optimizes every preference dimension—an effect we call the multi-reward alignment tax .

We therefore introduce PRISM , a multi-reward RL algorithm built on a different philosophy: do not mix rewards, mix policies . As shown in Figure 1 , PRISM learns one standalone positive policy per reward and a single global negative policy. Desirable behaviors are reward-specific and deserve dedicated optimization directions. In contrast, undesirable behaviors need not be attributed to individual rewards. A single negative policy captures the union of all reward-specific failure modes and penalizes them jointly. Each reward drives updates mainly through its own positive branch. This design alleviates gradient competition among conflicting preferences and mitigates the multi-reward alignment tax.

For efficiency, all sub-policies share a single language model conditioned on different prefix tokens and are composed at sampling time via a weighted sum of logits; the same mixture policy serves both training rollouts and inference. The merge weights thus form an explicit, interpretable interface for preference trade-offs, adjustable at inference time without retraining.

Experiments on three multi-reward settings show that PRISM consistently outperforms reward-space baselines. On GPQA [ rein2023gpqa ] and ScienceQA [ saikh2022scienceqa ] , it achieves the best overall score on every backbone, surpassing the strongest baseline by 17.8, 8.0, and 0.6 points on DeepSeek-R1-1.5B [ shao2024deepseekmath ] , Qwen2.5-1.5B-Instruct, and Qwen2.5-3B-Instruct [ qwen2.5 ] . On the tool-calling benchmark BFCL-v3 [ patil2025bfcl ] and the helpfulness–safety benchmarks Alpaca [ taori2023alpaca ] , HH-RLHF [ ganguli2022red ] , and PKU-SafeRLHF [ dai2024safe ] , it ranks first among all RL methods, and remains more robust and controllable as reward complexity grows.

Our contributions are threefold:

• We identify the multi-reward alignment tax induced by reward-space composition, and advocate composing preferences in policy space instead.

• We propose PRISM , which decomposes multi-reward RL into per-reward positive policies and a global negative policy, realized via prefix-conditioned sub-policies with logit-level composition.

• Experiments on scientific QA, tool-use reasoning, and helpfulness–safety alignment show consistent gains over strong baselines, with training-free inference-time controllability.

2 Preliminaries

2.1 Multi-Reward Policy Optimization

We consider an autoregressive language model parameterized by θ \theta , which defines a conditional distribution π θ ​ ( o ∣ q ) \pi{\theta}(o\mid q) over responses o o given a query q ∼ 𝒟 q\sim\mathcal{D} . Alignment is performed against N N reward functions { R k ​ ( o , q ) } k = 1 N {R{k}(o,q)}{k=1}^{N} , where each R k R{k} encodes a distinct human preference such as correctness, helpfulness, safety, or format compliance. The native multi-reward objective is the system of N N coupled maximization problems

max π ​ { 𝔼 o ∼ π ( ⋅ ∣ q ) ​ [ R 1 ​ ( o , q ) ] , 𝔼 o ∼ π ( ⋅ ∣ q ) ​ [ R 2 ​ ( o , q ) ] , ⋯ 𝔼 o ∼ π ( ⋅ ∣ q ) ​ [ R N ​ ( o , q ) ] \displaystyle\mathrm{max}_{\pi}\ (1)

s . t . \displaystyle s.t. ∫ o π ​ ( o ∣ q ) ​ d o = 1 \displaystyle\int_{o}{\pi(o\mid q)}\mathrm{d}o=1

which, in general, admits no single solution that simultaneously maximizes every R k R{k} [ yang2024rewards ] . The dominant practical workaround is reward-space scalarization [ zhou2024beyond , williams2024multi , liu2026gdpogrouprewarddecouplednormalization ] , which collapses these objectives into one weighted scalar R ​ ( o , q ) = ∑ k = 1 N w k ​ R k ​ ( o , q ) R(o,q)=\sum{k=1}^{N}w{k}R{k}(o,q) and optimizes a single policy against it. As discussed in Section 1 , this composition principle suffers from scale sensitivity, calibration mismatch, and gradient interference between conflicting preferences.

2.2 Group-Relative Advantage Estimation

We adopt the group-relative advantage estimator introduced by GRPO [ shao2024deepseekmath ] as the building block for our policy updates. Given a query q q , a group of G G responses { o i } i = 1 G {o{i}}{i=1}^{G} is sampled from the current policy and scored under a reward R R to obtain { r i } i = 1 G {r{i}}{i=1}^{G} . The group-relative advantage of response i i is

A ^ i = r i − mean ​ ( { r j } j = 1 G ) std ​ ( { r j } j = 1 G ) , \hat{A}{i};=;\frac{r{i}-\mathrm{mean}({r{j}}{j=1}^{G})}{\mathrm{std}({r{j}}{j=1}^{G})}, (2)

which provides a low-variance, group-normalized signal without an explicit value network.

Given the group-relative advantage A ^ i \hat{A}{i} in Eq. 2 computed under a reward function R R , the corresponding GRPO training objective for a policy π θ \pi{\theta} takes the standard clipped surrogate form:

ℒ ​ ( θ ) = − 𝔼 q ∼ 𝒟 , { o i } i = 1 G ∼ π θ old ​ [ min ⁡ ( r i ​ ( θ ) ​ A ^ i , clip ​ ( r i ​ ( θ ) , 1 − ϵ , 1 + ϵ ) ​ A ^ i ) ] , \mathcal{L}(\theta);=;-,\mathbb{E}{q\sim\mathcal{D},,{o{i}}{i=1}^{G}\sim\pi{\theta{\text{old}}}}\left[\min\Big(r{i}(\theta),\hat{A}{i},\ \mathrm{clip}\big(r{i}(\theta),,1{-}\epsilon,,1{+}\epsilon\big),\hat{A}_{i}\Big)\right], (3)

where r i ​ ( θ ) = π θ ​ ( o i ∣ q ) / π θ old ​ ( o i ∣ q ) r{i}(\theta)=\pi{\theta}(o{i}\mid q)/\pi{\theta{\text{old}}}(o{i}\mid q) is the importance ratio and ϵ \epsilon is the clipping threshold. We write ℒ ​ ( θ ; R ) \mathcal{L}(\theta;R) as an explicit function of the scoring reward R R .

In the multi-reward setting, GDPO [ liu2026gdpogrouprewarddecouplednormalization ] mitigates reward-scale mismatch by normalizing each reward-specific advantage within the response group, then summing and re-normalizing the aggregated advantage to preserve relative differences across reward combinations.

A i k \displaystyle A{i}^{k} = r i k − mean ​ ( { r j k } j = 1 G ) std ​ ( { r j k } j = 1 G ) , A i sum = ∑ k = 1 N A i k , \displaystyle=\frac{r{i}^{k}-\mathrm{mean}({r{j}^{k}}{j=1}^{G})}{\mathrm{std}({r{j}^{k}}{j=1}^{G})},\quad A{i}^{\mathrm{sum}}=\sum{k=1}^{N}A_{i}^{k}, (4)

A ^ i \displaystyle\hat{A}{i} = A i sum − mean ​ ( { A j sum } j = 1 G ) std ​ ( { A j sum } j = 1 G ) , \displaystyle=\frac{A{i}^{\mathrm{sum}}-\mathrm{mean}({A{j}^{\mathrm{sum}}}{j=1}^{G})}{\mathrm{std}({A{j}^{\mathrm{sum}}}{j=1}^{G})},

where G G is group size and A ^ i \hat{A}_{i} is GDPO advantage for i i -th response. Nevertheless, this normalization still amounts to composing rewards in advantage space before policy updates, and thus remains a form of implicit reward weighting that is susceptible to the same gradient interference among conflicting preferences.

Figure 1 : Method overview of PRISM

2.3 Positive–Negative Policy Decomposition

Following DIPOLE [ liang2025dichotomous ] and DExperts [ liu2021dexperts ] , the optimal policy π ⋆ \pi^{\star} under a reward function R ​ ( o , q ) R(o,q) can be expressed as a product of expert (PoE) [ hinton2002training ] composition of a positive policy and a negative policy:

π ⋆ ​ ( o ∣ q ) = π + ​ ( o ∣ q ) α ⋅ π − ​ ( o ∣ q ) 1 − α , \displaystyle\pi^{\star}(o\mid q)=;\pi^{+}(o\mid q)^{\alpha}\cdot\pi^{-}(o\mid q)^{1-\alpha}, (5)

where π + \pi^{+} is the positive policy captures behaviors encouraged by the reward, whereas π − \pi^{-} is the negative policy captures behaviors suppressed by the reward. The coefficient α \alpha controls the degree of greediness in the composed policy; increasing α \alpha encourages the final policy to assign higher probability to responses associated with larger rewards.

3 Method

3.1 Multi-Reward Positive-Negative Policy Optimization

Multi-Reward Policy Optimization

Consider a multi-objective optimization problem with N N rewards { R k ​ ( o , q ) } k = 1 N {R{k}(o,q)}{k=1}^{N} as in Eq. 1 . We aim to train once and allow the trade-off to be freely selected at inference time. To achieve it, assuming that there is an arbitrary set of coefficients { α 1 , … , α N } {\alpha{1},\dots,\alpha{N}} , for any such choice, we formulate a KL-regularized one-step local improvement over a reference policy μ \mu [ ziegler2019fine ] , augmented with a shared penalty term:

max π ⁡ 𝔼 o ∼ π ( ⋅ ∣ q ) ​ [ ∑ k = 1 N α k ​ R k ​ ( o , q ) ] − 1 η ​ D KL ​ ( π ∥ μ ) \displaystyle\max\limits{\pi}\ \mathbb{E}{o\sim\pi(\cdot\mid q)}\left[\sum{k=1}^{N}\alpha{k}R{k}(o,q)\right]-\frac{1}{\eta}D{\mathrm{KL}}(\pi|\mu) (6)

− γ ​ 𝔼 o ∼ π ( ⋅ ∣ q ) ​ [ ψ [ R 1 , R 2 , ⋯ , R N ] ​ ( o , q ) ] \displaystyle\qquad-\gamma\mathbb{E}{o\sim\pi(\cdot\mid q)}[\psi{[R{1},R{2},\cdots,R_{N}]}(o,q)]

s . t . \displaystyle\mathrm{s.t.} ∫ o π ​ ( o ∣ q ) ​ d o = 1 , \displaystyle\int_{o}\pi(o\mid q)\mathrm{d}o=1,

∑ k = 1 N α k − γ = 1 . \displaystyle\sum{k=1}^{N}\alpha{k}-\gamma=1.

Here ψ [ R 1 , … , R N ] ​ ( o , q ) \psi{[R{1},\dots,R_{N}]}(o,q) , abbreviated as ψ ​ ( 𝐑 ) \psi(\mathbf{R}) below, is designed to model the union of all reward-specific failure modes: it outputs a large value whenever any one of the individual rewards is low, regardless of how the others behave, so that a single severely violated objective is sufficient to trigger the penalty. Mathematically, ψ \psi is monotone non-increasing in each reward.

Optimal Policy and Policy Decomposition

The optimal solution π ⋆ ​ ( o ∣ q ) \pi^{\star}(o\mid q) of Eq. 6 satisfies [ rafailov2023direct ] π ⋆ ​ ( o ∣ q ) = 1 Z ​ ( q ) ​ μ ​ ( o ∣ q ) ​ exp ⁡ ( η ​ [ ∑ k = 1 N α k ​ R k ​ ( o , q ) − γ ​ ψ ​ ( 𝐑 ) ​ ( o , q ) ] ) \pi^{\star}(o\mid q)=\frac{1}{Z(q)}\mu(o\mid q)\exp\left(\eta\left[\sum{k=1}^{N}\alpha{k}R{k}(o,q)-\gamma\psi(\mathbf{R})(o,q)\right]\right) . Since the constraint in Eq. 6 requires ∑ k = 1 N α k − γ = 1 \sum{k=1}^{N}\alpha{k}-\gamma=1 , we can decompose μ \mu as μ ​ ( o ∣ q ) = ∏ k = 1 N μ ​ ( o ∣ q ) α k μ ​ ( o ∣ q ) γ \mu(o\mid q);=;\frac{\prod{k=1}^{N}\mu(o\mid q)^{\alpha_{k}}}{\mu(o\mid q)^{\gamma}} . Substituting it into the expression for π ⋆ \pi^{\star} , we obtain

π ⋆ ​ ( o ∣ q ) = 1 Z ​ ( q ) ⋅ ∏ k = 1 N μ ​ ( o ∣ q ) α k μ ​ ( o ∣ q ) γ ⋅ ∏ k = 1 N exp ⁡ ( α k ​ η ​ R k ​ ( o , q ) ) exp ⁡ ( γ ​ η ​ ψ ​ ( 𝐑 ) ​ ( o , q ) ) , \pi^{\star}(o\mid q);=;\frac{1}{Z(q)}\cdot\frac{\prod{k=1}^{N}\mu(o\mid q)^{\alpha{k}}}{\mu(o\mid q)^{\gamma}}\cdot\frac{\prod{k=1}^{N}\exp\big(\alpha{k}\eta R_{k}(o,q)\big)}{\exp\big(\gamma\eta\psi(\mathbf{R})(o,q)\big)}, (7)

which can be regrouped term-by-term as

π ⋆ ​ ( o ∣ q ) = 1 Z ​ ( q ) ⋅ ∏ k = 1 N [ μ ​ ( o ∣ q ) ​ exp ⁡ ( η ​ R k ​ ( o , q ) ) ] α k [ μ ​ ( o ∣ q ) ​ exp ⁡ ( η ​ ψ ​ ( 𝐑 ) ​ ( o , q ) ) ] γ . \pi^{\star}(o\mid q);=;\frac{1}{Z(q)}\cdot\frac{\prod{k=1}^{N}\big[\mu(o\mid q)\exp(\eta R{k}(o,q))\big]^{\alpha_{k}}}{\big[\mu(o\mid q)\exp(\eta\psi(\mathbf{R})(o,q))\big]^{\gamma}}. (8)

Note that each bracketed term μ ​ ( o ∣ q ) ​ exp ⁡ ( η ​ R k ​ ( o , q ) ) \mu(o\mid q)\exp(\eta R_{k}(o,q)) and μ ​ ( o ∣ q ) ​ exp ⁡ ( η ​ ψ ​ ( 𝐑 ) ​ ( o , q ) ) \mu(o\mid q)\exp(\eta\psi(\mathbf{R})(o,q)) is, up to normalization, exactly the optimal solution of a single-reward KL-regularized optimal policy.

We define the reward-specific positive policies

π k + ​ ( o ∣ q ) ≜ 1 Z k + ​ ( q ) ​ μ ​ ( o ∣ q ) ​ exp ⁡ ( η ​ R k ​ ( o , q ) ) , k = 1 , … , N , \pi{k}^{+}(o\mid q);\triangleq;\frac{1}{Z{k}^{+}(q)},\mu(o\mid q)\exp\big(\eta R_{k}(o,q)\big),\qquad k=1,\dots,N, (9)

and the global negative policy

π − ​ ( o ∣ q ) ≜ 1 Z − ​ ( q ) ​ μ ​ ( o ∣ q ) ​ exp ⁡ ( η ​ ψ ​ ( 𝐑 ) ​ ( o , q ) ) , \displaystyle\pi^{-}(o\mid q);\triangleq;\frac{1}{Z^{-}(q)},\mu(o\mid q)\exp\big(\eta\psi(\mathbf{R})(o,q)\big), (10)

where r k + = π k + ​ ( o ∣ q ) / μ ​ ( o ∣ q ) r^{+}{k}=\pi^{+}{k}(o\mid q)/\mu(o\mid q) , r − = π − ​ ( o ∣ q ) / μ ​ ( o ∣ q ) r^{-}=\pi^{-}(o\mid q)/\mu(o\mid q) , and Z k + ​ ( q ) Z_{k}^{+}(q) and Z − ​ ( q ) Z^{-}(q) are the corresponding normalizing constants, then π ⋆ \pi^{\star} can be expressed as a composition of the positive and negative policies [ hinton2002training ] :

π ⋆ ​ ( o ∣ q ) ∝ ∏ k = 1 N ( π k + ​ ( o ∣ q ) ) α k ( π − ​ ( o ∣ q ) ) γ . \pi^{\star}(o\mid q);\propto;\frac{\prod{k=1}^{N}\big(\pi{k}^{+}(o\mid q)\big)^{\alpha_{k}}}{\big(\pi^{-}(o\mid q)\big)^{\gamma}}. (11)

The derivation of Eq. 11 shows that, at the optimum, policy composition is decoupled from the optimization of each sub-policy. In particular, the definitions of π k + \pi{k}^{+} and π − \pi^{-} do not involve the coefficients { α k , γ } {\alpha{k},\gamma} used to compose them at inference time. We therefore optimize each sub-policy on its own with the GRPO loss defined in Eq. 3 :

ℒ k + ​ ( π k + ) = − 𝔼 q ∼ 𝒟 , { o i } i = 1 G ∼ μ ​ [ min ⁡ ( r i k ​ A ~ i k , clip ​ ( r i k , 1 − ϵ , 1 + ϵ ) ​ A ~ i k ) ] , \mathcal{L}{k}^{+}(\pi^{+}{k});=;-,\mathbb{E}{q\sim\mathcal{D},,{o{i}}{i=1}^{G}\sim\mu}\left[\min\Big(r{i}^{k},\tilde{A}{i}^{k},\ \mathrm{clip}\big(r{i}^{k},,1{-}\epsilon,,1{+}\epsilon\big),\tilde{A}_{i}^{k}\Big)\right], (12)

ℒ − ​ ( π − ) = − 𝔼 q ∼ 𝒟 , { o i } i = 1 G ∼ μ ​ [ min ⁡ ( r i − ​ A ~ i − , clip ​ ( r i − , 1 − ϵ , 1 + ϵ ) ​ A ~ i − ) ] . \mathcal{L}^{-}(\pi^{-});=;-,\mathbb{E}{q\sim\mathcal{D},,{o{i}}{i=1}^{G}\sim\mu}\left[\min\Big(r{i}^{-}\tilde{A}{i}^{-},\ \mathrm{clip}\big(r{i}^{-},,1{-}\epsilon,,1{+}\epsilon\big),\tilde{A}_{i}^{-}\Big)\right]. (13)

In practice, Eq. 12 and Eq. 13 use monotonically shaped advantages A ~ i k ≜ σ ​ ( A ^ i k ) − c , A ~ i − ≜ c − ( ∏ k = 1 N σ ​ ( A ^ i k ) ) 1 / N \tilde{A}{i}^{k}\triangleq\sigma(\hat{A}{i}^{k})-c,\quad\tilde{A}{i}^{-}\triangleq c-\Big(\prod{k=1}^{N}\sigma(\hat{A}{i}^{k})\Big)^{1/N} where σ \sigma is the sigmoid function and c c is a constant offset. The soft conjunction makes A ~ i − \tilde{A}{i}^{-} large on the union of failure modes, as required of ψ \psi (alternatives are ablated in Appendix LABEL:sec:ablationnegweight ). Moreover, training proceeds iteratively in practice, with the current composed policy π ⋆ \pi^{\star} serving as the reference μ \mu that generates rollouts and anchors the importance ratios at each iteration.

Logit-Level Policy Composition

Since π ⋆ \pi^{\star} , π k + \pi{k}^{+} , and π − \pi^{-} share the same reference policy μ \mu and the same context, taking logarithms converts the product and quotient into a linear combination. The log-partition terms log ⁡ Z ​ ( q ) \log Z(q) , log ⁡ Z k + ​ ( q ) \log Z{k}^{+}(q) , and log ⁡ Z − ​ ( q ) \log Z^{-}(q) collapse into a single token-independent constant. This constant is automatically absorbed by the softmax normalization and therefore does not need to be explicitly computed.

Consequently, composing the N + 1 N+1 sub-policies reduces to a weighted summation at the token-logit level [ o2023contrastive , dekoninck2023controlled ] :

z t ⋆ = ∑ k = 1 N α k ​ z k , t + − γ ​ z t − , z{t}^{\star}=\sum{k=1}^{N}\alpha{k}z{k,t}^{+}-\gamma z_{t}^{-}, (14)

where z t ⋆ z{t}^{\star} , z k , t + z{k,t}^{+} , and z t − z{t}^{-} denote the logits of π ⋆ ( ⋅ ∣ q , o < t ) \pi^{\star}(\cdot\mid q,o{<t}) , π k + ( ⋅ ∣ q , o < t ) \pi{k}^{+}(\cdot\mid q,o{<t}) , and π − ( ⋅ ∣ q , o < t ) \pi^{-}(\cdot\mid q,o_{<t}) , respectively.

Discussion

By decoupling policy optimization from reward trade-off selection, PRISM allows a single training run to yield an entire family of policies spanning all possible coefficient combinations, rather than a single policy tied to one fixed trade-off. This design simultaneously enables inference-time controllability and reduces gradient interference across conflicting objectives, alleviating the multi-reward alignment tax.

3.2 Efficient Implementation

Naively instantiating each sub-policy as an independent model would multiply both parameter count and decoding latency by N + 1 N+1 . We avoid this via two choices: (i) all sub-policies share a single backbone conditioned on lightweight, independently-trainable prefixes, and (ii) rollouts under π ⋆ \pi^{\star} are generated in batched forwards pass rather than N + 1 N+1 sequential ones. Asymmetric update scheme shields the shared backbone from negative branch, making them compatible.

Prefix-conditioned shared backbone.

We realize all sub-policies using one shared language model π θ \pi{\theta} , conditioned on a set of learnable continuous prefix embeddings 𝒫 = { P 1 + , P 2 + , … , P N + , P − } \mathcal{P}={P{1}^{+},P{2}^{+},\dots,P{N}^{+},P^{-}} [ li2021prefix , lester2021power ] , one prefix per sub-policy:

π k + ( ⋅ ∣ q ) ≜ π θ ( ⋅ | [ q , P k + ] ) , π − ( ⋅ ∣ q ) ≜ π θ ( ⋅ | [ q , P − ] ) , k = 1 , … , N . \pi{k}^{+}(\cdot\mid q)\triangleq\pi{\theta}!\big(\cdot,\big|,[q,P{k}^{+}]\big),\quad\pi^{-}(\cdot\mid q)\triangleq\pi{\theta}!\big(\cdot,\big|,[q,P^{-}]\big),\quad k=1,\dots,N. (15)

Backbone-preserving asymmetric training.

The positive branches π k + \pi_{k}^{+} are trained on their respective reward signals, whereas π − \pi^{-} (Eq. 10 ) is trained to maximize ψ ​ ( 𝐑 ) \psi(\mathbf{R}) , thereby capturing the union of reward-specific failure modes, i.e., responses that fail on any one objective. To keep the shared backbone aligned with the positive behaviors, we adopt an asymmetric update scheme that updates θ \theta only through the positive branches [ zhang2024negative ] . Let δ \delta be the learning rate:

θ \displaystyle\theta ← θ − δ ​ ∇ θ ℒ k + , \displaystyle\leftarrow\theta-\delta,\nabla{\theta},\mathcal{L}^{+}{k}, P k \displaystyle P{k} ← P k − δ ​ ∇ P k + ℒ k + , \displaystyle\leftarrow P{k}-\delta,\nabla{P^{+}{k}},\mathcal{L}^{+}_{k}, (16)

θ \displaystyle\theta ← θ , \displaystyle\leftarrow\theta, P − \displaystyle P^{-} ← P − − δ ​ ∇ P − ℒ − . \displaystyle\leftarrow P^{-}-\delta,\nabla_{P^{-}},\mathcal{L}^{-}. (17)

Positive updates (Eq. 16 ) backpropagate through both θ \theta and P k + P_{k}^{+} . Negative updates (Eq. 17 ) apply a stop-gradient to θ \theta , so only P − P^{-} is updated. Consequently, the backbone is trained exclusively on high-quality data, while P − P^{-} still learns a catch-all representation of the union of failure modes that can suppress them under Eq. 14 at inference time.

Figure 2 : Parallel-batch mixture sampling. At each decoding step, the query q q , per-branch prefix ( P k + P{k}^{+} or P − P^{-} ), and the shared context o < t o{<t} are stacked along the batch dimension and processed by the shared backbone π θ \pi{\theta} in a single forward pass. The resulting per-branch logits z k + z{k}^{+} and z − z^{-} are combined with weights α k \alpha_{k} and − γ -\gamma (Eq. 14 ), one token is sampled from π ⋆ \pi^{\star} , and it is broadcast back to all branches for the next step.

Parallel-batch mixture sampling.

Even with a shared backbone, naively sampling from π ⋆ \pi^{\star} would decode the N + 1 N{+}1 branches sequentially at every step. Instead, we place the branches along the batch dimension: as shown in Figure 2 , the decoding state ( q , o < t ) (q,o{<t}) is replicated across all N + 1 N{+}1 prefixes, their logits are computed in a single batched forward pass and composed via Eq. 14 , and the sampled token o t o{t} is broadcast back to all branches so that their KV caches advance in lockstep. Memory and FLOPs still scale linearly with N + 1 N{+}1 , but per-token latency is that of a single forward pass [ pope2022efficiently ] . The same procedure generates rollouts during training and serves inference.

4 Experiments

Table 1: Main results on scientific QA benchmarks : scientific QA results across all evaluated backbones, where small gray values denote absolute changes over the corresponding base model. All numbers are percentages. Rows with “+” denote optimization methods applied to the corresponding backbone. Bold indicates the best result among optimization methods. Fmt: Percentage of rewards that meet the format requirements. Acc: Percentage of correct answers. Joint: Percentage of answers that meet the format requirements and are correct. Avg: Average of format and accuracy.

Results on Scientific QA & GPQA

Model / Method ScienceQA GPQA Overall

Fmt Acc Joint Avg Fmt Acc Joint Avg

DeepSeek-R1-1.5B 0.00 66.64 0.00 22.21 0.00 29.68 0.00 9.89 16.05

• GRPO Sum 74.11 +74.11 67.85 +1.21 54.02 +54.02 65.33 +43.12 16.96 +16.96 29.46 -0.22 6.25 +6.25 17.56 +7.67 41.44 +25.39

• GRPO Product 74.39 +74.39 68.51 +1.87 55.42 +55.42 66.11 +43.90 27.68 +27.68 30.80 +1.12 9.60 +9.60 22.69 +12.80 44.40 +28.35

• GDPO 80.37 +80.37 67.94 +1.30 58.04 +58.04 68.78 +46.57 26.34 +26.34 28.35 -1.33 8.71 +8.71 21.13 +11.24 44.96 +28.91

• PRISM (ours) 95.33 +95.33 69.91 +3.27 68.51 +68.51 77.92 +55.71 81.25 +81.25 33.04 +3.36 28.35 +28.35 47.55 +37.66 62.73 +46.68

  Qwen2.5-1.5B-Instruct 6.82 73.36 5.89 28.69 37.28 27.23 10.49 25.00 26.84

• GRPO Sum 100.00 +93.18 76.26 +2.90 76.26 +70.37 84.17 +55.48 77.46 +40.18 27.90 +0.67 21.21 +10.72 42.19 +17.19 63.18 +36.34

• GRPO Product 98.31 +91.49 75.98 +2.62 74.30 +68.41 82.86 +54.17 37.95 +0.67 27.90 +0.67 11.16 +0.67 25.67 +0.67 54.27 +27.43

• GDPO 99.93 +93.11 76.64 +3.28 76.26 +70.37 84.28 +55.59 76.56 +39.28 29.01 +1.78 21.43 +10.94 42.33 +17.34 63.30 +36.46

• PRISM (ours) 99.53 +92.71 76.72 +3.36 76.45 +70.56 84.23 +55.54 87.95 +50.67 29.24 +2.01 25.67 +15.18 58.45 +33.45 71.34 +44.50

  Qwen2.5-3B-Instruct 44.67 81.78 36.73 54.39 47.99 33.93 16.52 32.81 43.60

• GRPO Sum 99.81 +55.14 83.55 +1.77 83.55 +46.82 88.97 +34.58 87.95 +39.96 29.91 -4.02 27.46 +10.94 48.44 +15.63 68.71 +25.11

• GRPO Product 99.53 +54.86 84.02 +2.24 84.02 +47.29 89.19 +34.80 84.82 +36.83 28.80 -5.13 24.11 +7.59 45.91 +13.10 67.55 +23.95

• GDPO 99.81 +55.14 82.80 +1.02 82.80 +46.07 88.47 +34.08 85.94 +37.95 29.91 -4.02 26.79 +10.27 47.55 +14.74 68.01 +24.41

• PRISM (ours) 98.97 +54.30 84.11 +2.33 84.11 +47.38 89.06 +34.68 88.17 +40.18 31.92 -2.01 28.57 +12.05 49.55 +16.74 69.31 +25.71

  Table 2: Main results on BFCL v3 benchmarks : tool calling results on Qwen2.5-3B-Instruct across all evaluated methods. Acc/R: Percentage of correct answers on RLLA style judgement. Acc/B: Percentage of correct answers on BFCL style judgement. Details of RLLA and BFCL judgement refer to Appendix B.4.2

  Results on BFCL v3

  Model / Method Non-live Live Multi-turn Overall

  Fmt Acc/R Acc/B Fmt Acc/R Acc/B Fmt Acc/R Acc/B Fmt Acc/R Acc/B

  Qwen2.5-3B-Instruct 13.30 8.70 72.17 14.06 5.63 58.30 20.56 7.28 17.34 15.98 7.20 49.27

• GRPO Sum 65.74 40.87 74.78 40.87 31.16 68.02 59.53 17.43 18.28 64.43 27.82 51.27

• GRPO Product 90.17 49.48 74.96 96.97 31.16 60.07 77.51 20.50 17.99 88.22 33.71 51.01

• GDPO 98.87 55.30 76.09 99.33 33.16 62.37 83.22 20.68 17.93 93.81 36.38 52.13

• PRISM (ours) 99.83 56.26 80.35 99.33 31.09 62.22 86.54 21.24 17.81 95.23 36.20 53.46

  Table 3: Main results on Helpfulness-Safety Alignment benchmarks : helpfulness and safety alignment results on Qwen2.5-3B-Instruct across all evaluated optimization methods.

  Results on Helpfulness-Safety Alignment

  Model / Method Alpaca HH-RLHF PKU-SafeRLHF

  Useful Harmless Avg Useful Harmless Avg Useful Harmless Avg

• GRPO Sum 3.03 3.49 3.26 3.03 4.03 3.53 4.79 6.34 5.57

• GRPO Product 3.02 3.48 3.25 3.02 4.03 3.53 4.79 6.34 5.57

• GDPO 2.95 3.44 3.20 3.03 4.02 3.53 4.79 6.33 5.56

• PRISM (ours) 3.17 3.65 3.41 3.14 4.15 3.65 4.86 6.36 5.61

  Figure 3 : Overall BFCL v3 performance under different numbers of training rewards. 1 reward: Train on the correctness reward; 2 reward Trained on correctness and format rewards; 3 reward: Trained on correctness , format , and length rewards. Correctness denotes the RLLA-style accuracy. All scores are normalized with three-reward score of PRISM .

  Figure 4 : Training curves across all methods under the three-reward setting.

  4.1 Experiment Setup

  We evaluate PRISM on three multi-reward alignment scenarios: scientific question answering, tool-use reasoning, and helpfulness–safety alignment.

  Scientific question answering.

  For scientific reasoning, we train models on SciKnowEval [ feng2024sciknoweval ] and evaluate them on GPQA [ rein2023gpqa ] and ScienceQA [ saikh2022scienceqa ] . Further details are in Appendix A.4 .

  Tool-use reasoning.

  For tool-use reasoning, we train models on the dataset collected by ToolRL [ qian2026toolrl ] and evaluate them on BFCL-v3 [ patil2025bfcl ] , a benchmark evaluates the LLM’s ability to call functions. Further details are in Appendix B.4.1 .

  Helpfulness-safety alignment.

  For helpfulness-safety alignment, we use the Alpaca dataset [ taori2023alpaca ] for policy optimization and evaluate the trained models on the corresponding held-out Alpaca test split. We further evaluate models on the HH-RLHF [ ganguli2022red ] and PKU-SafeRLHF [ dai2024safe ] datasets. Further details are in Appendix C.2 .

  Baselines.

  We conduct experiments on DeepSeek-R1-1.5B [ shao2024deepseekmath ] , Qwen2.5-1.5B-Instruct and Qwen2.5-3B-Instruct [ qwen2.5 ] . For each backbone and task, we compare PRISM with representative multi-reward RL baselines, including GRPO [ shao2024deepseekmath ] variants and GDPO [ liu2026gdpogrouprewarddecouplednormalization ] . The GRPO variants include reward weighted summation, denoted as GRPO Sum, and reward product aggregation, denoted as GRPO Prod. Detailed training configurations and hyperparameters are provided in Appendix A.3 , B.3 and C.1 .

  4.2 Main Results

  Tables 3 , 3 , and 3 summarize the main evaluation results across scientific reasoning, tool use, and helpfulness-safety alignment.

  On scientific reasoning (Table 3 ), PRISM attains the best overall score on all three backbones: 62.73 vs. 44.96 for the strongest baseline on DeepSeek-R1-1.5B ( + + 17.77), 71.34 vs. 63.30 on Qwen2.5-1.5B-Instruct ( + + 8.04), and 69.31 vs. 68.71 on Qwen2.5-3B-Instruct. The advantage is most pronounced on the harder GPQA benchmark, where PRISM reaches 47.55 average on DeepSeek-R1-1.5B, more than doubling the best baseline (22.69). Crucially, PRISM achieves the highest Joint score in all six benchmark–backbone settings (e.g., 68.51 vs. 58.04 on ScienceQA and 28.35 vs. 9.60 on GPQA with DeepSeek-R1-1.5B), confirming that it best improves the ability to satisfy all reward criteria simultaneously .

  On tool calling (Table 3 ), PRISM obtains the best overall Fmt (95.23 vs. 93.81 for GDPO) and Acc/B (53.46 vs. 52.13), while its overall Acc/R (36.20) is on par with GDPO (36.38). It sweeps all three metrics in the Non-live setting (99.83 Fmt, 56.26 Acc/R, 80.35 Acc/B) and achieves the best Fmt and Acc/R in Multi-turn, delivering the most balanced profile across settings.

  On helpfulness–safety alignment (Table 3 ), PRISM scores highest on both usefulness and harmlessness across all three evaluation sets, lifting the average to 3.41 on Alpaca ( + + 0.15 over the best baseline), 3.65 on HH-RLHF ( + + 0.12), and 5.61 on PKU-SafeRLHF ( + + 0.04), whereas the baselines are nearly indistinguishable from one another. These consistent gains under continuous reward-model feedback show that policy-space decomposition remains effective beyond rule-based reward settings. Compared with all other multi-reward RL baselines, PRISM demonstrates superior or more stable performance in most settings, further validating its effectiveness in multi-objective optimization scenarios.

  4.3 Alignment Tax under Varying Rewards

  Based on the BFCL-v3 [ patil2025bfcl ] experiments, we further analyze the impact of varying the number of rewards on different algorithms. Specifically, we train the models with one, two and three rewards, respectively.

  All algorithms are trained for 6,000 steps, and the final model performance is recorded. To ensure comparability across different metrics, all metrics are normalized using the corresponding rewards of PRISM under three-reward setting.

  As shown in Figure 3 , PRISM maintains nearly the same correctness performance after adding format and length rewards, while consistently outperforming all baselines. In contrast, GDPO [ liu2026gdpogrouprewarddecouplednormalization ] and GRPO-Prod exhibit a clear drop in correctness as the number of rewards increases; although GRPO-Sum maintains relatively stable correctness, its optimization of format and length rewards is substantially weaker than PRISM . These results suggest that PRISM better mitigates the alignment tax and attains a superior Pareto frontier in multi-reward optimization.

  Shown in Figure 4 , we further visualize the evolution of each reward over training steps under the three-reward setting. Compared with the baselines, PRISM demonstrates substantially higher sample efficiency: it converges to near-optimal performance after roughly 3k steps, while the other methods require more than 6k steps to approach convergence.

  4.4 Inference-Time Controllability

  Figure 5 : Inference-time controllability under different weight configurations on BFCL-Live. Each configuration assigns a dominant weight of 0.55 0.55 to one positive branch while setting the remaining two positive weights to 0.25 0.25 . The negative branch weight is fixed to γ = 0.05 \gamma=0.05 throughout all experiments.

  To investigate inference-time controllability, we evaluate three weight configurations on the BFCL-Live split. Figure 5 shows that increasing the weight α k \alpha{k} consistently improves the corresponding evaluation metric. Specifically, setting α 1 = 0.55 \alpha{1}=0.55 achieves the highest correctness accuracy ( 31.53 % 31.53% ), α 2 = 0.55 \alpha{2}=0.55 yields the best format accuracy ( 98.08 % 98.08% ), and α 3 = 0.55 \alpha{3}=0.55 produces the longest average reasoning length ( 108.0 108.0 ).

  These results suggest that the merge weights provide a simple and effective mechanism for controllable inference-time behavior.

  4.5 Ablation Studies

  Ablations on Policy Composition.

  We conduct component ablations on BFCL-v3 using Qwen2.5-3B-Instruct under the same training and evaluation configuration as the main tool-use experiments. We evaluate three variants: (1) replacing reward-specific positive policies with a shared positive policy, (2) removing the global negative policy, and (3) generating rollouts independently from each branch rather than from the composed policy.

  Table 4: Policy-composition ablations on BFCL-v3 using Qwen2.5-3B-Instruct.

  Method Fmt Acc/R Acc/B

  PRISM 95.23 36.20 53.46

  w/ shared positive policy 94.03 34.38 52.56

  w/o global negative policy 94.86 34.75 52.84

  w/ individual-branch rollouts 94.84 34.59 52.48

  Table 5: Negative-policy weighting ablations on BFCL-v3 using Qwen2.5-3B-Instruct.

  Method Fmt Acc/R Acc/B

  PRISM 95.23 36.20 53.46

  w/ max \max weighting 95.02 33.70 52.43

  w/ LogAvgExp \operatorname{LogAvgExp} weighting 95.06 34.58 52.62

  w/ mean \operatorname{mean} weighting 93.27 33.90 53.12

  Table 5 summarizes the overall results. All three ablations degrade overall performance. Sharing the positive policy causes the largest drops in format and RLLA accuracy, confirming the importance of preserving reward-specific optimization directions. Removing the global negative policy reduces all overall metrics, supporting its role in modeling shared failure modes. Individual-branch rollouts obtain the lowest overall BFCL accuracy, indicating that sampling from the composed policy better aligns the training distribution with composed decoding.

  Ablations on Negative Weighting Function.

  As Table 5 , we further study the weighting function used to train the global negative policy. For the i i -th sampled response, let { A i k } k = 1 N {A{i}^{k}}{k=1}^{N} denote its reward-specific advantages. We compare three alternative weighting functions. The max \max variant uses w i − = − max k = 1 , … , N ⁡ A i k w{i}^{-}=-\max{k=1,\dots,N}A{i}^{k} ; the LogAvgExp \operatorname{LogAvgExp} variant uses w i − = − exp ⁡ ( 1 N ​ ∑ k = 1 N log ⁡ ( A i k ) ) w{i}^{-}=-\exp!\left(\frac{1}{N}\sum{k=1}^{N}\log(A{i}^{k})\right) ; and the mean \operatorname{mean} variant uses w i − = − 1 N ​ ∑ k = 1 N A i k w{i}^{-}=-\frac{1}{N}\sum{k=1}^{N}A_{i}^{k} .

  5 Related Work

  Reinforcement Learning for LLM.

  RL is the dominant tool for aligning LLMs with human preferences. RLHF-PPO [ ouyang2022training , schulman2017proximal ] optimizes a learned reward model with a separate value network, while DPO [ rafailov2023direct ] recasts preference optimization as a stable supervised-style objective. Recent critic-free variants improve scalability and stability: GRPO [ shao2024deepseekmath ] estimates advantages via group-wise normalization, GSPO [ zheng2025group ] reduces importance-ratio variance using sequence-level ratios, and DAPO [ yu2026dapo ] refines the training recipe at scale. All of them, however, optimize a single scalar reward, so multiple preference signals must be collapsed before training.

  Multi-Reward RL.

  Since alignment must satisfy correctness, safety, helpfulness, and format compliance at once, most multi-reward methods scalarize rewards into one weighted objective [ zhou2024beyond , williams2024multi , quan2024dmoerm ] , which is sensitive to reward scales and weights and induces interference among heterogeneous preferences. Refinements include fine-grained rewards [ wu2023fine ] , adaptive weighting [ de2024dynamic ] , reward normalization [ liu2026gdpogrouprewarddecouplednormalization ] , constrained optimization [ achiam2017constrained ] , model merging [ jang2023personalized ] , and preference conditioning [ yang2024rewards ] , yet all of them compose preferences before the policy update. In contrast, PRISM keeps rewards as separate optimization directions and composes the resulting policies in logit space, shifting preference composition from reward space to policy space.

  6 Conclusion

  In this work, we propose PRISM , a policy-space composition framework for multi-reward LLM alignment. Instead of scalarizing heterogeneous rewards into a single optimization signal, PRISM keeps each reward as a separate optimization direction by learning one positive policy per reward. A single global negative policy further captures the union of failure modes. These policies are composed through explicit merge weights, enabling controllable preference trade-offs without retraining. Experiments on scientific reasoning, tool-use reasoning, and helpfulness–safety alignment show that PRISM consistently outperforms reward-space baselines, remains more robust as reward complexity increases, and supports effective inference-time preference control. These results suggest that mixing policies rather than rewards offers a more effective and controllable path toward multi-reward aligned language models.