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

# Policy Decomposition and Optimization for Multi-Reward RL
Source: [https://arxiv.org/html/2607.29246](https://arxiv.org/html/2607.29246)
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

## ![[Uncaptioned image]](https://arxiv.org/html/2607.29246v1/images/prism_1.png)Don’t Mix Rewards, Mix Policies: Policy Decomposition and Optimization for Multi\-Reward RL

Yi ZhongYizhen YuanYinan ZhengTianyi Tan Tianyue WangHaiyun GuoJinqiao WangXianyuan 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 proposePRISM, a new multi\-reward RL framework built upon the idea of policy\-space decomposition and composition\. Instead of compositing different rewards,PRISMoptimizes 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 thatPRISMconsistently 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

## 1Introduction

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 introducePRISM, a multi\-reward RL algorithm built on a different philosophy:*do not mix rewards, mix policies*\. As shown in Figure[1](https://arxiv.org/html/2607.29246#S2.F1),PRISMlearns 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 thatPRISMconsistently 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 proposePRISM, 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\.

## 2Preliminaries

### 2\.1Multi\-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 responsesoogiven a queryq∼𝒟q\\sim\\mathcal\{D\}\. Alignment is performed againstNNreward functions\{Rk​\(o,q\)\}k=1N\\\{R\_\{k\}\(o,q\)\\\}\_\{k=1\}^\{N\}, where eachRkR\_\{k\}encodes a distinct human preference such as correctness, helpfulness, safety, or format compliance\. The native multi\-reward objective is the system ofNNcoupled maximization problems

maxπ​\{𝔼o∼π\(⋅∣q\)​\[R1​\(o,q\)\],𝔼o∼π\(⋅∣q\)​\[R2​\(o,q\)\],⋯𝔼o∼π\(⋅∣q\)​\[RN​\(o,q\)\]\\displaystyle\\mathrm\{max\}\_\{\\pi\}\\\(1\)s\.t\.\\displaystyle s\.t\.∫oπ​\(o∣q\)​do=1\\displaystyle\\int\_\{o\}\{\\pi\(o\\mid q\)\}\\mathrm\{d\}o=1
which, in general, admits no single solution that simultaneously maximizes everyRkR\_\{k\}\[yang2024rewards\]\. The dominant practical workaround is reward\-space scalarization\[zhou2024beyond,williams2024multi,liu2026gdpogrouprewarddecouplednormalization\], which collapses these objectives into one weighted scalarR​\(o,q\)=∑k=1Nwk​Rk​\(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](https://arxiv.org/html/2607.29246#S1), this composition principle suffers from scale sensitivity, calibration mismatch, and gradient interference between conflicting preferences\.

### 2\.2Group\-Relative Advantage Estimation

We adopt the group\-relative advantage estimator introduced by GRPO\[shao2024deepseekmath\]as the building block for our policy updates\. Given a queryqq, a group ofGGresponses\{oi\}i=1G\\\{o\_\{i\}\\\}\_\{i=1\}^\{G\}is sampled from the current policy and scored under a rewardRRto obtain\{ri\}i=1G\\\{r\_\{i\}\\\}\_\{i=1\}^\{G\}\. The group\-relative advantage of responseiiis

A^i=ri−mean​\(\{rj\}j=1G\)std​\(\{rj\}j=1G\),\\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 advantageA^i\\hat\{A\}\_\{i\}in Eq\.[2](https://arxiv.org/html/2607.29246#S2.E2)computed under a reward functionRR, the corresponding GRPO training objective for a policyπθ\\pi\_\{\\theta\}takes the standard clipped surrogate form:

ℒ​\(θ\)=−𝔼q∼𝒟,\{oi\}i=1G∼πθold​\[min⁡\(ri​\(θ\)​A^i,clip​\(ri​\(θ\),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\)
whereri​\(θ\)=πθ​\(oi∣q\)/πθold​\(oi∣q\)r\_\{i\}\(\\theta\)=\\pi\_\{\\theta\}\(o\_\{i\}\\mid q\)/\\pi\_\{\\theta\_\{\\text\{old\}\}\}\(o\_\{i\}\\mid q\)is the importance ratio andϵ\\epsilonis the clipping threshold\. We writeℒ​\(θ;R\)\\mathcal\{L\}\(\\theta;R\)as an explicit function of the scoring rewardRR\.

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\.

Aik\\displaystyle A\_\{i\}^\{k\}=rik−mean​\(\{rjk\}j=1G\)std​\(\{rjk\}j=1G\),Aisum=∑k=1NAik,\\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\}=Aisum−mean​\(\{Ajsum\}j=1G\)std​\(\{Ajsum\}j=1G\),\\displaystyle=\\frac\{A\_\{i\}^\{\\mathrm\{sum\}\}\-\\mathrm\{mean\}\(\\\{A\_\{j\}^\{\\mathrm\{sum\}\}\\\}\_\{j=1\}^\{G\}\)\}\{\\mathrm\{std\}\(\\\{A\_\{j\}^\{\\mathrm\{sum\}\}\\\}\_\{j=1\}^\{G\}\)\},whereGGis group size andA^i\\hat\{A\}\_\{i\}is GDPO advantage forii\-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\.

![Refer to caption](https://arxiv.org/html/2607.29246v1/x1.png)Figure 1:Method overview ofPRISM
### 2\.3Positive–Negative Policy Decomposition

Following DIPOLE\[liang2025dichotomous\]and DExperts\[liu2021dexperts\], the optimal policyπ⋆\\pi^\{\\star\}under a reward functionR​\(o,q\)R\(o,q\)can be expressed as aproduct 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α\\alphacontrols the degree of greediness in the composed policy; increasingα\\alphaencourages the final policy to assign higher probability to responses associated with larger rewards\.

## 3Method

### 3\.1Multi\-Reward Positive\-Negative Policy Optimization

##### Multi\-Reward Policy Optimization

Consider a multi\-objective optimization problem withNNrewards\{Rk​\(o,q\)\}k=1N\\\{R\_\{k\}\(o,q\)\\\}\_\{k=1\}^\{N\}as in Eq\.[1](https://arxiv.org/html/2607.29246#S2.E1)\. 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=1Nαk​Rk​\(o,q\)\]−1η​DKL​\(π∥μ\)\\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\)​\[ψ\[R1,R2,⋯,RN\]​\(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\)​do=1,\\displaystyle\\int\_\{o\}\\pi\(o\\mid q\)\\mathrm\{d\}o=1,∑k=1Nαk−γ=1\.\\displaystyle\\sum\_\{k=1\}^\{N\}\\alpha\_\{k\}\-\\gamma=1\.
Hereψ\[R1,…,RN\]​\(o,q\)\\psi\_\{\[R\_\{1\},\\dots,R\_\{N\}\]\}\(o,q\), abbreviated asψ​\(𝐑\)\\psi\(\\mathbf\{R\}\)below, is designed to model theunionof 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,ψ\\psiis monotonenon\-increasingin each reward\.

##### Optimal Policy and Policy Decomposition

The optimal solutionπ⋆​\(o∣q\)\\pi^\{\\star\}\(o\\mid q\)of Eq\.[6](https://arxiv.org/html/2607.29246#S3.E6)satisfies\[rafailov2023direct\]π⋆​\(o∣q\)=1Z​\(q\)​μ​\(o∣q\)​exp⁡\(η​\[∑k=1Nαk​Rk​\(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](https://arxiv.org/html/2607.29246#S3.E6)requires∑k=1Nαk−γ=1\\sum\_\{k=1\}^\{N\}\\alpha\_\{k\}\-\\gamma=1, we can decomposeμ\\muasμ​\(o∣q\)=∏k=1Nμ​\(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\)=1Z​\(q\)⋅∏k=1Nμ​\(o∣q\)αkμ​\(o∣q\)γ⋅∏k=1Nexp⁡\(αk​η​Rk​\(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\)=1Z​\(q\)⋅∏k=1N\[μ​\(o∣q\)​exp⁡\(η​Rk​\(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⁡\(η​Rk​\(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\)≜1Zk\+​\(q\)​μ​\(o∣q\)​exp⁡\(η​Rk​\(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\)≜1Z−​\(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\)whererk\+=π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\), andZk\+​\(q\)Z\_\{k\}^\{\+\}\(q\)andZ−​\(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=1N\(π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](https://arxiv.org/html/2607.29246#S3.E11)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](https://arxiv.org/html/2607.29246#S2.E3):

ℒk\+​\(πk\+\)=−𝔼q∼𝒟,\{oi\}i=1G∼μ​\[min⁡\(rik​A~ik,clip​\(rik,1−ϵ,1\+ϵ\)​A~ik\)\],\\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∼𝒟,\{oi\}i=1G∼μ​\[min⁡\(ri−​A~i−,clip​\(ri−,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](https://arxiv.org/html/2607.29246#S3.E12)and Eq\.[13](https://arxiv.org/html/2607.29246#S3.E13)use monotonically shaped advantagesA~ik≜σ​\(A^ik\)−c,A~i−≜c−\(∏k=1Nσ​\(A^ik\)\)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σ\\sigmais the sigmoid function andccis a constant offset\. The soft conjunction makesA~i−\\tilde\{A\}\_\{i\}^\{\-\}large on the union of failure modes, as required ofψ\\psi\(alternatives are ablated in AppendixLABEL:sec:ablation\_neg\_weight\)\. Moreover, training proceeds iteratively in practice, with the current composed policyπ⋆\\pi^\{\\star\}serving as the referenceμ\\muthat 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μ\\muand the same context, taking logarithms converts the product and quotient into a linear combination\. The log\-partition termslog⁡Z​\(q\)\\log Z\(q\),log⁡Zk\+​\(q\)\\log Z\_\{k\}^\{\+\}\(q\), andlog⁡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 theN\+1N\+1sub\-policies reduces to a weighted summation at the token\-logit level\[o2023contrastive,dekoninck2023controlled\]:

zt⋆=∑k=1Nαk​zk,t\+−γ​zt−,z\_\{t\}^\{\\star\}=\\sum\_\{k=1\}^\{N\}\\alpha\_\{k\}z\_\{k,t\}^\{\+\}\-\\gamma z\_\{t\}^\{\-\},\(14\)
wherezt⋆z\_\{t\}^\{\\star\},zk,t\+z\_\{k,t\}^\{\+\}, andzt−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\.2Efficient Implementation

Naively instantiating each sub\-policy as an independent model would multiply both parameter count and decoding latency byN\+1N\+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 thanN\+1N\+1sequential 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𝒫=\{P1\+,P2\+,…,PN\+,P−\}\\mathcal\{P\}=\\\{P\_\{1\}^\{\+\},P\_\{2\}^\{\+\},\\dots,P\_\{N\}^\{\+\},P^\{\-\}\\\}\[li2021prefix,lester2021power\], one prefix per sub\-policy:

πk\+\(⋅∣q\)≜πθ\(⋅\|\[q,Pk\+\]\),π−\(⋅∣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](https://arxiv.org/html/2607.29246#S3.E10)\) 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θ\\thetaonly through the positive branches\[zhang2024negative\]\. Letδ\\deltabe the learning rate:

θ\\displaystyle\\theta←θ−δ​∇θℒk\+,\\displaystyle\\leftarrow\\theta\-\\delta\\,\\nabla\_\{\\theta\}\\,\\mathcal\{L\}^\{\+\}\_\{k\},Pk\\displaystyle P\_\{k\}←Pk−δ​∇Pk\+ℒ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](https://arxiv.org/html/2607.29246#S3.E16)\) backpropagate through bothθ\\thetaandPk\+P\_\{k\}^\{\+\}\. Negative updates \(Eq\.[17](https://arxiv.org/html/2607.29246#S3.E17)\) apply a stop\-gradient toθ\\theta, so onlyP−P^\{\-\}is updated\. Consequently, the backbone is trained exclusively on high\-quality data, whileP−P^\{\-\}still learns a catch\-all representation of the union of failure modes that can suppress them under Eq\.[14](https://arxiv.org/html/2607.29246#S3.E14)at inference time\.

![Refer to caption](https://arxiv.org/html/2607.29246v1/x2.png)Figure 2:Parallel\-batch mixture sampling\. At each decoding step, the queryqq, per\-branch prefix \(Pk\+P\_\{k\}^\{\+\}orP−P^\{\-\}\), and the shared contexto<to\_\{<t\}are stacked along the batch dimension and processed by the shared backboneπθ\\pi\_\{\\theta\}in a single forward pass\. The resulting per\-branch logitszk\+z\_\{k\}^\{\+\}andz−z^\{\-\}are combined with weightsαk\\alpha\_\{k\}and−γ\-\\gamma\(Eq\.[14](https://arxiv.org/html/2607.29246#S3.E14)\), 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 theN\+1N\{\+\}1branches sequentially at every step\. Instead, we place the branches along the*batch*dimension: as shown in Figure[2](https://arxiv.org/html/2607.29246#S3.F2), the decoding state\(q,o<t\)\(q,o\_\{<t\}\)is replicated across allN\+1N\{\+\}1prefixes, their logits are computed in asingle batched forward passand composed via Eq\.[14](https://arxiv.org/html/2607.29246#S3.E14), and the sampled tokenoto\_\{t\}is broadcast back to all branches so that their KV caches advance in lockstep\. Memory and FLOPs still scale linearly withN\+1N\{\+\}1, but per\-token latency is that of a*single*forward pass\[pope2022efficiently\]\. The same procedure generates rollouts during training and serves inference\.

## 4Experiments

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\.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](https://arxiv.org/html/2607.29246#A2.SS4.SSS2)Table 3:Main results on Helpfulness\-Safety Alignment benchmarks: helpfulness and safety alignment results on Qwen2\.5\-3B\-Instruct across all evaluated optimization methods\.![Refer to caption](https://arxiv.org/html/2607.29246v1/x3.png)Figure 3:Overall BFCL v3 performance under different numbers of training rewards\. 1 reward: Train on thecorrectnessreward; 2 reward Trained oncorrectnessandformatrewards; 3 reward: Trained oncorrectness,format, andlengthrewards\. Correctness denotes the RLLA\-style accuracy\. All scores are normalized with three\-reward score ofPRISM\.![Refer to caption](https://arxiv.org/html/2607.29246v1/x4.png)Figure 4:Training curves across all methods under the three\-reward setting\.### 4\.1Experiment Setup

We evaluatePRISMon 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](https://arxiv.org/html/2607.29246#A1.SS4)\.

##### 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](https://arxiv.org/html/2607.29246#A2.SS4.SSS1)\.

##### 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](https://arxiv.org/html/2607.29246#A3.SS2)\.

##### 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 comparePRISMwith 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](https://arxiv.org/html/2607.29246#A1.SS3),[B\.3](https://arxiv.org/html/2607.29246#A2.SS3)and[C\.1](https://arxiv.org/html/2607.29246#A3.SS1)\.

### 4\.2Main Results

Tables[3](https://arxiv.org/html/2607.29246#S4.T3),[3](https://arxiv.org/html/2607.29246#S4.T3), and[3](https://arxiv.org/html/2607.29246#S4.T3)summarize the main evaluation results across scientific reasoning, tool use, and helpfulness\-safety alignment\.

On scientific reasoning \(Table[3](https://arxiv.org/html/2607.29246#S4.T3)\),PRISMattains 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, wherePRISMreaches 47\.55 average on DeepSeek\-R1\-1\.5B, more than doubling the best baseline \(22\.69\)\. Crucially,PRISMachieves 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](https://arxiv.org/html/2607.29246#S4.T3)\),PRISMobtains 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](https://arxiv.org/html/2607.29246#S4.T3)\),PRISMscores 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,PRISMdemonstrates superior or more stable performance in most settings, further validating its effectiveness in multi\-objective optimization scenarios\.

### 4\.3Alignment 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 ofPRISMunder three\-reward setting\.

As shown in Figure[3](https://arxiv.org/html/2607.29246#S4.F3),PRISMmaintains 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 thanPRISM\. These results suggest thatPRISMbetter mitigates the alignment tax and attains a superior Pareto frontier in multi\-reward optimization\.

Shown in Figure[4](https://arxiv.org/html/2607.29246#S4.F4), we further visualize the evolution of each reward over training steps under the three\-reward setting\. Compared with the baselines,PRISMdemonstrates 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\.4Inference\-Time Controllability

![Refer to caption](https://arxiv.org/html/2607.29246v1/x5.png)Figure 5:Inference\-time controllability under different weight configurations on BFCL\-Live\. Each configuration assigns a dominant weight of0\.550\.55to one positive branch while setting the remaining two positive weights to0\.250\.25\. The negative branch weight is fixed toγ=0\.05\\gamma=0\.05throughout all experiments\.To investigate inference\-time controllability, we evaluate three weight configurations on the BFCL\-Live split\. Figure[5](https://arxiv.org/html/2607.29246#S4.F5)shows that increasing the weightαk\\alpha\_\{k\}consistently improves the corresponding evaluation metric\. Specifically, settingα1=0\.55\\alpha\_\{1\}=0\.55achieves the highest correctness accuracy \(31\.53%31\.53\\%\),α2=0\.55\\alpha\_\{2\}=0\.55yields the best format accuracy \(98\.08%98\.08\\%\), andα3=0\.55\\alpha\_\{3\}=0\.55produces the longest average reasoning length \(108\.0108\.0\)\.

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

### 4\.5Ablation 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 ablationson BFCL\-v3 using Qwen2\.5\-3B\-Instruct\.
Table 5:Negative\-policy weighting ablationson BFCL\-v3 using Qwen2\.5\-3B\-Instruct\.

Table[5](https://arxiv.org/html/2607.29246#S4.T5)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](https://arxiv.org/html/2607.29246#S4.T5), we further study the weighting function used to train the global negative policy\. For theii\-th sampled response, let\{Aik\}k=1N\\\{A\_\{i\}^\{k\}\\\}\_\{k=1\}^\{N\}denote its reward\-specific advantages\. We compare three alternative weighting functions\. Themax\\maxvariant useswi−=−maxk=1,…,N⁡Aikw\_\{i\}^\{\-\}=\-\\max\_\{k=1,\\dots,N\}A\_\{i\}^\{k\}; theLogAvgExp\\operatorname\{LogAvgExp\}variant useswi−=−exp⁡\(1N​∑k=1Nlog⁡\(Aik\)\)w\_\{i\}^\{\-\}=\-\\exp\\\!\\left\(\\frac\{1\}\{N\}\\sum\_\{k=1\}^\{N\}\\log\(A\_\{i\}^\{k\}\)\\right\); and themean\\operatorname\{mean\}variant useswi−=−1N​∑k=1NAikw\_\{i\}^\{\-\}=\-\\frac\{1\}\{N\}\\sum\_\{k=1\}^\{N\}A\_\{i\}^\{k\}\.

## 5Related 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,PRISMkeeps rewards as separate optimization directions and composes the resulting policies in logit space, shifting preference composition from reward space to policy space\.

## 6Conclusion

In this work, we proposePRISM, a policy\-space composition framework for multi\-reward LLM alignment\. Instead of scalarizing heterogeneous rewards into a single optimization signal,PRISMkeeps 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 thatPRISMconsistently 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\.

## References

## Appendix AExperimental details of scientific question answering

### A\.1Dataset and Prompt Format

We train on the SciKnowEval dataset, restricting to the multiple\-choice subset\. Each example is formatted as a two\-turn chat prompt\.

The system prompt is:

System PromptYou are a helpful AI assistant specialized in solving multiple\-choice science questions\.For every request, solve the problem carefully and step by step before giving the final answer\.You must respond in exactly the following format:<think\> Provide a clear, step\-by\-step solution\. </think\> <answer\> Provide only the final answer\. </answer\>Rules:1\. The <think\> section must contain the complete reasoning process\.2\. The <answer\> section must appear after </think\>\.3\. The <answer\> section must contain only one option letter: A, B, C, or D\.4\. Do not include explanation, extra words, or additional formatting in <answer\>\.5\. Always complete the reasoning before writing the final answer\.

The user turn is rendered as:

User PromptGiven a question and four options, please select the right answer\. Your answer should be "A", "B", "C" or "D"\.\{question\}Choices:A\. \{choice\_a\}B\. \{choice\_b\}C\. \{choice\_c\}D\. \{choice\_d\}

### A\.2Reward Functions

Training uses two reward signals: correctness and format\.

##### Correctness Reward\.

The correctness rewardRcorrect∈\{0,1\}R\_\{\\mathrm\{correct\}\}\\in\\\{0,1\\\}evaluates whether the predicted answer matches the reference answer:

Rcorrect=𝟏​\[a​\(x\)=y\],R\_\{\\mathrm\{correct\}\}=\\mathbf\{1\}\[a\(x\)=y\],whereyydenotes the reference option letter, anda​\(x\)a\(x\)is the option extracted from the model completionxx\.

##### Format Reward\.

The format rewardRformat∈\{0,1\}R\_\{\\mathrm\{format\}\}\\in\\\{0,1\\\}is defined as

Rformat=𝟏​\[x​matches the target format\],R\_\{\\mathrm\{format\}\}=\\mathbf\{1\}\[x\\text\{ matches the target format\}\],where the target format is<think\>…</think\><answer\>…</answer\>\.

### A\.3Training Hyperparameters

We train Qwen2\.5\-1\.5B\-Instruct and Qwen2\.5\-3B\-Instruct for 1 epoch, and DeepSeek\-R1\-Distill\-Qwen\-1\.5B for 1\.5 epochs\. Other training hyperparameters are summarized in Table[6](https://arxiv.org/html/2607.29246#A1.T6)\.

Table 6:Scientific question answering training configuration\.
### A\.4Evaluation on ScienceQA and GPQA

Each model is required to first produce its reasoning process enclosed by<think\></think\>, followed by a final answer enclosed by<answer\></answer\>\. The order of these two fields must be preserved\. We use two rewards in this setting\. The format reward verifies whether the output contains both<think\></think\>and<answer\></answer\>in the correct order\. The correctness reward evaluates whether the answer extracted from the<answer\></answer\>field matches the ground\-truth answer\.

We report pass@1 accuracy under greedy decoding, together with format adherence rate\. All evaluation samples follow the same prompt format and answer\-extraction protocol as described in Section[A\.1](https://arxiv.org/html/2607.29246#A1.SS1)\.

For ScienceQA, we use text\-only samples from the natural science and closed\-choice subset\. For GPQA, we use thegpqa\_mainsplit and render each sample as a four\-choice question, with answer options deterministically shuffled per sample\.

## Appendix BExperimental Details of Tool\-use Reasoning

### B\.1Training Prompt Format

Prompt Format for ToolRL TrainingYou are a helpful dialogue assistant capable of leveraging tool calls to solve user tasks and provide structured chat responses\.Available ToolsIn your response, you can use the following tools:\{\{ Tool List \}\}Steps for Each Turn1\.Think:Recall relevant context and analyze the current user goal\.2\.Decide on Tool Usage:If a tool is needed, specify the tool and its parameters\.3\.Respond Appropriately:If a response is needed, generate one while maintaining consistency across user queries\.Output Format<think\>Your thoughts and reasoning</think\><tool\_call\>\{"name": "Tool name", "parameters": \{"Parameter name": "Parameter content", "…": "…"\}\}\{"name": "…", "parameters": \{"…": "…", "…": "…"\}\}…</tool\_call\><response\>AI’s final response</response\>Important Notes1\. You must always include the<think\>field to outline your reasoning\. Provide at least one of<tool\_call\>or<response\>\. Decide whether to use<tool\_call\>possibly multiple times,<response\>, or both\.2\. You can invoke multiple tool calls simultaneously in the<tool\_call\>fields\. Each tool call should be a JSON object with a"name"field and a"parameters"field containing a dictionary of parameters\. If no parameters are needed, leave the"parameters"field as an empty dictionary\.3\. Refer to the previous dialogue records in the history, including the user’s queries, previous<tool\_call\>,<response\>, and any tool feedback noted as<obs\>if exists\.

User Prompt for ToolRL TrainingDialogue History<user\>\{\{ Initial User Input \}\}</user\><think\>Round 1 Model Thought</think\>\{\{ Round 1 model output<tool\_call\>or<response\>\}\}<obs\>Round 1 Observation</obs\>…<user\>\{\{ User Input \}\}</user\>…

### B\.2Tool Calling Reward Functions

##### Format Reward\.

The format rewardRformat∈\{0,1\}R\_\{\\mathrm\{format\}\}\\in\\\{0,1\\\}evaluates whether a model output follows the tag structure required by the ground\-truth answer\. The output must begin with a<think\>section and then contain the required<tool\_call\>and/or<response\>sections in the prescribed order\. When a tool\-call or response section is required, its opening and closing tags must each appear exactly once\.

##### Correctness Reward\.

The correctness rewardRcorrect∈\[−3,3\]R\_\{\\mathrm\{correct\}\}\\in\[\-3,3\]quantifies how well the predicted tool calls match the ground\-truth tool calls\. The comparison includes:

1. 1\.Tool name matching: a multiset overlap score in\[0,1\]\[0,1\]is computed between the predicted and ground\-truth tool names;
2. 2\.Parameter matching: ground\-truth calls are processed one by one\. For each ground\-truth call, every unmatched predicted call with the same tool name is assigned a parameter\-level score\. This score is the sum of a parameter\-name overlap term in\[0,1\]\[0,1\]and one point for each ground\-truth parameter whose value is predicted exactly\. The highest parameter\-level score is added to the total score, and the corresponding predicted call is marked as matched\. Each predicted call can therefore be matched at most once\.

The tool\-name score and all selected parameter\-level scores are summed and divided by the maximum attainable score for the ground\-truth calls\. The resulting normalized score in\[0,1\]\[0,1\]is then linearly mapped to\[−3,3\]\[\-3,3\]\.

##### Length Reward\.

The length rewardRlength∈\[0,1\]R\_\{\\mathrm\{length\}\}\\in\[0,1\]encourages the model to produce sufficiently detailed reasoning within the<think\>section\. It increases linearly with the length of the reasoning content up to the predefined threshold of 512 words and is capped at11once the threshold is reached\. Outputs without a valid<think\>section receive a reward of0\.

### B\.3Training Hyperparameters

Detail training hyperparameters of Tool\-Use Reasoning are summarized in Table[7](https://arxiv.org/html/2607.29246#A2.T7)\.

Table 7:Tool\-use reasoning training configuration\.
### B\.4BFCL Evaluation

#### B\.4\.1Evaluation Protocol

To evaluate tool\-use reasoning on BFCL, we convert BFCL examples into the RLLA protocol instead of using the native BFCL function\-calling interface\. For each example, we construct an RLLA\-style prompt following the same prompt format described in §[B\.1](https://arxiv.org/html/2607.29246#A2.SS1), including the available tool definitions and dialogue history\.

Each model is required to first generate its reasoning process enclosed by<think\></think\>, and then decide whether to perform a tool call or directly respond to the user\. If the model chooses to call a tool, it must output a JSON\-formatted tool request enclosed by<tool\_call\></tool\_call\>\. If the model chooses to answer directly, it must output the final response enclosed by<response\></response\>\. We again use two rewards\. The format reward checks whether the output follows the required structure, namely<think\></think\>followed by either<tool\_call\></tool\_call\>or<response\></response\>\. The correctness reward evaluates whether the answer in<response\></response\>matches the ground truth, or whether the generated tool call contains the correct fields\.

#### B\.4\.2Evaluation Metrics

LetNNdenote the number of evaluated examples\. We report both strict RLLA metrics and BFCL\-style semantic metrics\.

##### RLLA Accuracy\.

The RLLA accuracy is a strict task\-level metric\. For examples with reference tool calls, it requires the correctness reward to reach its maximum possible score\. For examples without reference tool calls, it requires the model to predict no tool call:

rlla​\_​acc=1N​∑i=1N𝟏​\[\(Gi≠∅∧Rcorrect​\(Pi,Gi\)=3\)∨\(Gi=∅∧Pi=∅\)\]\.\\displaystyle\\mathrm\{rlla\\\_acc\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\mathbf\{1\}\\\!\\Big\[\(G\_\{i\}\\neq\\emptyset\\land R\_\{\\mathrm\{correct\}\}\(P\_\{i\},G\_\{i\}\)=3\)\\lor\(G\_\{i\}=\\emptyset\\land P\_\{i\}=\\emptyset\)\\Big\]\.\(18\)

##### BFCL Accuracy\.

The BFCL accuracy evaluates the semantic correctness of predicted tool calls while ignoring RLLA wrapper\-format errors\. Predicted<tool\_call\>blocks are extracted, converted back into BFCL\-style function calls, and evaluated using the BFCL checker:

bfcl​\_​acc=1N​∑i=1N𝟏​\[BFCLChecker​\(Pi,Gi\)=correct\]\.\\displaystyle\\mathrm\{bfcl\\\_acc\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\mathbf\{1\}\\\!\\Big\[\\mathrm\{BFCLChecker\}\(P\_\{i\},G\_\{i\}\)=\\mathrm\{correct\}\\Big\]\.\(19\)

##### Format Accuracy\.

The format accuracy measures protocol adherence independently of tool correctness:

format​\_​acc=1N​∑i=1N𝟏​\[Rformat​\(xi\)=1\],\\mathrm\{format\\\_acc\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\mathbf\{1\}\\left\[R\_\{\\mathrm\{format\}\}\(x\_\{i\}\)=1\\right\],whereRformatR\_\{\\mathrm\{format\}\}follows the same definition as the format reward in the training objective\.

##### Length Reward\.

The evaluation\-time length reward measures the mean normalized reasoning length using the same definition as the training\-time length reward:

length​\_​reward=1N​∑i=1Nmin⁡\(1,nthink​\(xi\)L\),\\mathrm\{length\\\_reward\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\min\\left\(1,\\frac\{n\_\{\\mathrm\{think\}\}\(x\_\{i\}\)\}\{L\}\\right\),where the maximum reasoning length is

## Appendix CExperimental details of helpfulness\-safety alignment

### C\.1Training Hyperparameters

For helpfulness\-safety alignment, we train Qwen2\.5\-3B\-Instruct with LoRA and evaluate the checkpoint at step 540\. Other training hyperparameters are summarized in Table[8](https://arxiv.org/html/2607.29246#A3.T8)\.

Table 8:Helpfulness\-safety alignment training configuration\.
### C\.2Evaluation Details of Helpfulness\-Safety Alignment

Each model is given a user prompt and is required to produce a response that is both helpful to the user and safe with respect to potentially harmful requests\. We use two rewards in this setting\. The helpfulness reward is computed using the helpfulness reward model provided by the Amo project\[amo\_project\]\. The safety reward is derived from the corresponding harmlessness cost model from the same project\. The overall alignment objective encourages responses that are simultaneously helpful and safe\.

## Appendix DPrefix Token Initialization\.

We initialize each learnable prefix embedding as

𝐩i\(0\)\\displaystyle\\mathbf\{p\}\_\{i\}^\{\(0\)\}=Embθ0⁡\(\\n\)\\displaystyle=\\operatorname\{Emb\}\_\{\\theta\_\{0\}\}\(\\texttt\{\\textbackslash n\}\)whereEmbθ0\\operatorname\{Emb\}\_\{\\theta\_\{0\}\}denotes the input embedding layer of the pretrained language model\. The newline embedding provides an in\-distribution initialization that is compatible with the pretrained model’s input embedding space and prompt formatting\.

## Appendix EControllable Inference Cases

We illustrate inference\-time steerability with an example from the BFCL\-Live split\.

Task\.

UserI’d like to purchase 2 tickets for ‘After The Wedding’ showing on 2023\-04\-22 at 19:00 in Berkeley\.

Expected Tool CallMovies\_1\_BuyMovieTickets\(\{movie\_name, number\_of\_tickets, show\_date, location, show\_time, show\_type\}\)

Responses under three weight configurations\.

Weights are ordered as\(α1,α2,α3\)\(\\alpha\_\{1\},\\alpha\_\{2\},\\alpha\_\{3\}\), corresponding to correctness, format, and length\. In each configuration, the dominant branch is set to0\.550\.55and the other two branches are set to0\.250\.25\.

α1=0\.55,α2=α3=0\.25\\alpha\_\{1\}=0\.55,\\;\\alpha\_\{2\}=\\alpha\_\{3\}=0\.25\(correctness\-dominant\)\.The model directly invokes the target purchase tool with all required arguments\.

Correctness\-dominant Inference[⬇](data:text/plain;base64,PHRoaW5rPgpJZGVudGlmaWVzIHRoZSBtb3ZpZSB0aXRsZSwgdGlja2V0IGNvdW50LCBkYXRlLApsb2NhdGlvbiwgYW5kIHRpbWUsIHRoZW4gc2VsZWN0cyB0aGUgcHVyY2hhc2UKdG9vbC4gICgqQFx0ZXh0YmZ7ICBbLi4ubGVuZ3RoPTEyNiB3b3Jkc10gfUAqKQo8L3RoaW5rPgo8dG9vbF9jYWxsPgp7Im5hbWUiOiAiTW92aWVzXzFfQnV5TW92aWVUaWNrZXRzIiwKICJwYXJhbWV0ZXJzIjoKIHsibW92aWVfbmFtZSI6ICJBZnRlciBUaGUgV2VkZGluZyIsCiAgICAgICAgICAgICAgICAibnVtYmVyX29mX3RpY2tldHMiOiAyLAogICAgICAgICAgICAgICAgInNob3dfZGF0ZSI6ICIyMDIzLTA0LTIyIiwKICAgICAgICAgICAgICAgICJsb2NhdGlvbiI6ICJCZXJrZWxleSwgQ0EiLAogICAgICAgICAgICAgICAgInNob3dfdGltZSI6ICIxOTowMCIsCiAgICAgICAgICAgICAgICAic2hvd190eXBlIjogInJlZ3VsYXIifX0KPC90b29sX2NhbGw+)<think\>Identifiesthemovietitle,ticketcount,date,location,andtime,thenselectsthepurchasetool\.\[\.\.\.length=126 words\]</think\><tool\_call\>\{"name":"Movies\_1\_BuyMovieTickets","parameters":\{"movie\_name":"AfterTheWedding","number\_of\_tickets":2,"show\_date":"2023\-04\-22","location":"Berkeley,CA","show\_time":"19:00","show\_type":"regular"\}\}</tool\_call\>

α2=0\.55,α1=α3=0\.25\\alpha\_\{2\}=0\.55,\\;\\alpha\_\{1\}=\\alpha\_\{3\}=0\.25\(format\-dominant\)\.The response is syntactically well\-formed, but the model calls a show\-time lookup tool rather than completing the purchase\.

Format\-dominant Inference[⬇](data:text/plain;base64,PHRoaW5rPgpSZWNvZ25pemVzIHRoZSB0aWNrZXQtcHVyY2hhc2UgaW50ZW50LCBidXQKZGVjaWRlcyB0byBjb25maXJtIHNob3ctdGltZSBkZXRhaWxzIGZpcnN0LiAoKkBcdGV4dGJmeyAgWy4uLmxlbmd0aD0xMTUgd29yZHNdIH1AKikKPC90aGluaz4KPHRvb2xfY2FsbD4KeyJuYW1lIjogKCpAXHRleHR0dHtcdGV4dGNvbG9ye3JlZH17TW92aWVzXF8xXF9HZXRUaW1lc0Zvck1vdmllfX1AKiksICgqQFx0ZXh0YmZ7Ly8gd3JvbmcgdG9vbDogc2hvdWxkIGNhbGwgTW92aWVzXF8xXF9CdXlNb3ZpZVRpY2tldHN9QCopCiAicGFyYW1ldGVycyI6CiAoKkBcdGV4dHR0e1x0ZXh0Y29sb3J7cmVkfXtce319QCopCiAoKkBcdGV4dHR0e1x0ZXh0Y29sb3J7cmVkfXsibW92aWVcX25hbWUiOiAiQWZ0ZXIgVGhlIFdlZGRpbmciLH19QCopCiAoKkBcdGV4dHR0e1x0ZXh0Y29sb3J7cmVkfXsibG9jYXRpb24iOiAiQmVya2VsZXksIENBIix9fUAqKQogKCpAXHRleHR0dHtcdGV4dGNvbG9ye3JlZH17InNob3dcX2RhdGUiOiAiMjAyMy0wNC0yMiIsfX1AKikKICgqQFx0ZXh0dHR7XHRleHRjb2xvcntyZWR9eyJ0aGVhdGVyXF9uYW1lIjogIkFueSBUaGVhdGVyIix9fUAqKQogKCpAXHRleHR0dHtcdGV4dGNvbG9ye3JlZH17InNob3dcX3R5cGUiOiAicmVndWxhciJ9fUAqKQogKCpAXHRleHR0dHtcdGV4dGNvbG9ye3JlZH17XH19fUAqKSAoKkBcdGV4dGJmey8vIHBhcmFtZXRlcnMgYmVsb25nIHRvIHRoZSB3cm9uZyB0b29sIHNjaGVtYX1AKikKIH0KPC90b29sX2NhbGw+)<think\>Recognizestheticket\-purchaseintent,butdecidestoconfirmshow\-timedetailsfirst\.\[\.\.\.length=115 words\]</think\><tool\_call\>\{"name":Movies\_1\_GetTimesForMovie,// wrong tool: should call Movies\_1\_BuyMovieTickets"parameters":\{"movie\_name": "After The Wedding","location": "Berkeley, CA","show\_date": "2023\-04\-22","theater\_name": "Any Theater","show\_type": "regular"\}// parameters belong to the wrong tool schema\}</tool\_call\>

α3=0\.55,α1=α2=0\.25\\alpha\_\{3\}=0\.55,\\;\\alpha\_\{1\}=\\alpha\_\{2\}=0\.25\(length\-dominant\)\.The model produces a longer reasoning trace\.

Length\-dominant Inference[⬇](data:text/plain;base64,PHRoaW5rPgpNZW50aW9ucyB0aGUgbW92aWUsIGRhdGUsIHRpbWUsIGxvY2F0aW9uLCBhbmQKdGlja2V0IGNvdW50LCB0aGVuIGdpdmVzIGV4dHJhIHJlYXNvbmluZyBhYm91dApmaW5kaW5nIHNob3cgdGltZXMgYmVmb3JlIGJ1eWluZyB0aWNrZXRzLgooKkBcdGV4dGJmeyAgWy4uLmxlbmd0aD0xNjQgd29yZHNdIH1AKikgKCpAXHRleHRiZnsvLyBsb25nZXIgcmVhc29uaW5nIHRyYWNlIHVuZGVyIHRoZSBsZW5ndGgtZG9taW5hbnQgaW5mZXJlbmNlIH1AKikKPC90aGluaz4KPHRvb2xfY2FsbD4KeyJuYW1lIjogKCpAXHRleHR0dHtcdGV4dGNvbG9ye3JlZH17TW92aWVzXF8xXF9HZXRUaW1lc0Zvck1vdmllfX1AKiksICgqQFx0ZXh0YmZ7Ly8gd3JvbmcgdG9vbDogc2hvdWxkIGNhbGwgTW92aWVzXF8xXF9CdXlNb3ZpZVRpY2tldHN9QCopCiAicGFyYW1ldGVycyI6CiAoKkBcdGV4dHR0e1x0ZXh0Y29sb3J7cmVkfXtce319QCopCiAoKkBcdGV4dHR0e1x0ZXh0Y29sb3J7cmVkfXsibW92aWVcX25hbWUiOiAiQWZ0ZXIgVGhlIFdlZGRpbmciLH19QCopCiAoKkBcdGV4dHR0e1x0ZXh0Y29sb3J7cmVkfXsibG9jYXRpb24iOiAiQmVya2VsZXksIENBIix9fUAqKQogKCpAXHRleHR0dHtcdGV4dGNvbG9ye3JlZH17InNob3dcX2RhdGUiOiAiMjAyMy0wNC0yMiIsfX1AKikKICgqQFx0ZXh0dHR7XHRleHRjb2xvcntyZWR9eyJzaG93XF90eXBlIjogInJlZ3VsYXIifX1AKikKICgqQFx0ZXh0dHR7XHRleHRjb2xvcntyZWR9e1x9fX1AKikgKCpAXHRleHRiZnsvLyBwYXJhbWV0ZXJzIGJlbG9uZyB0byB0aGUgd3JvbmcgdG9vbCBzY2hlbWF9QCopCiB9CjwvdG9vbF9jYWxsPg==)<think\>Mentionsthemovie,date,time,location,andticketcount,thengivesextrareasoningaboutfindingshowtimesbeforebuyingtickets\.\[\.\.\.length=164 words\]// longer reasoning trace under the length\-dominant inference</think\><tool\_call\>\{"name":Movies\_1\_GetTimesForMovie,// wrong tool: should call Movies\_1\_BuyMovieTickets"parameters":\{"movie\_name": "After The Wedding","location": "Berkeley, CA","show\_date": "2023\-04\-22","show\_type": "regular"\}// parameters belong to the wrong tool schema\}</tool\_call\>

## Appendix FInference Efficiency

We measure the inference cost ofPRISMagainst a single\-policy baseline under the same serving configuration on 8×\\timesNVIDIA H800\. We report per\-response decoding latency, generation throughput, and peak GPU memory\.

Table 9:Inference efficiencyon Qwen2\.5\-3B\-Instruct\. Latency is measured in second, throughput in tokens per second, and peak GPU memory in GB\.PRISMwithNNrewards decodesN\+1N\{\+\}1branches\.
## Appendix GLimitations

There are several limitations in this study\. First,PRISMassumes that each reward can be represented by a distinct positive policy\. It also assumes that a single negative policy can capture the union of reward\-specific failure modes\. This design is effective across all our experimental settings, which cover up to three rewards\. Validating it at substantially larger reward counts is left for future work\. Second, the current policy\-space composition relies on manually specified merge weights, and how to automatically adapt these weights to different prompts, users, or deployment contexts remains underexplored\. Third, although per\-token latency stays that of a single forward pass, mixture decoding requires multiple prefix\-conditioned branches, and its memory and FLOPs scale with the number of rewards\. Future work can study more adaptive policy composition mechanisms, more expressive negative\-policy designs, and more efficient decoding strategies\.