PRISM: Don't Mix Rewards, Mix Policies — A Framework for Multi-Objective LLM Alignment

Forum topic · ✨步子哥 · 2026-08-04

Summary

A July 2026 paper from researchers at the Chinese Academy of Sciences, Tsinghua AIR, and Tongji University introduces PRISM, a multi-reward reinforcement learning framework that abandons scalar reward aggregation in favor of policy-space composition. The authors argue that combining heterogeneous reward signals (correctness, format, safety, length) into one scalar gradient update causes a structural 'multi-reward alignment tax,' where each objective is sub-optimally served. PRISM instead trains N reward-specific positive policies sharing a backbone plus one global negative policy capturing all failure modes, conditioned through learnable prefix embeddings. At inference, sub-policies are linearly combined at the logit level in a single forward pass, enabling runtime weight tuning without retraining. Experiments on ScienceQA, GPQA, BFCL-v3, Alpaca, HH-RLHF, and PKU-SafeRLHF show PRISM matches or beats the strongest baselines, with gains up to +17.77 points on DeepSeek-R1-1.5B, while remaining robust as reward count grows.

What PRISM Solves

A common pain point in multi-objective LLM post-training: combining rewards for correctness, output format, and tool-use into one scalar and optimizing via GRPO inevitably leaves every dimension under-optimized. Tuning weights fixes one objective while collapsing another — a structural defect the paper calls the multi-reward alignment tax.

PRISM's thesis is provocative: "Don't Mix Rewards, Mix Policies." Instead of composing in reward space, it composes in policy space.

How PRISM Works

1. One positive policy per reward

Given N rewards $R_1, R_2, \ldots, R_N$, PRISM trains N independent positive policies $\pi_1^+, \pi_2^+, \ldots, \pi_N^+$, each specialized for one reward — like giving each chef an independent pan instead of throwing all ingredients into one pot.

2. One global negative policy

A single $\pi^-$ captures the union of all failure modes: whenever any reward score is low, this branch is activated. The insight: ideal behavior is reward-specific and deserves a dedicated optimization direction; non-ideal behavior does not need attribution to a single reward, so one shared negative policy suffices.

3. Logit-level linear composition

At inference, because all sub-policies share the same reference $\mu$ and context, log-policy ratios combine linearly into logits:

$$z_t^\star = \sum_{k=1}^{N} \alpha_k z_{k,t}^+ - \gamma z_t^-$$

$\alpha_k$ and $\gamma$ are tunable mixing weights. Adjusting them at runtime shifts the model's behavior — no retraining required.

Three Engineering Tricks

Shared backbone + prefix conditioning. All N+1 sub-policies share $\pi_\theta$ and differ only via learnable prefix embeddings $\mathcal{P} = \{P_1^+, \ldots, P_N^+, P^-\}$.
Backbone-preserving asymmetric training. Positive branches update both $\theta$ and their prefix; the negative branch updates only $P^-$ with stop-gradient on the backbone. This prevents the shared trunk from being polluted by negative samples.
Parallel batched sampling. At each token, queries, prefixes, and history are stacked along the batch dimension so one forward pass produces all N+1 logits, which are then mixed. Per-token latency equals one forward pass; memory and FLOPs scale linearly with reward count.

Experimental Results

Scientific reasoning (correctness + format)

| Backbone | PRISM | Strongest baseline | Gain | |---|---|---|---| | DeepSeek-R1-1.5B | 62.73 | 44.96 | +17.77 | | Qwen2.5-1.5B-Instruct | 71.34 | 63.30 | +8.04 | | Qwen2.5-3B-Instruct | 69.31 | 68.71 | +0.60 |

On GPQA, DeepSeek-R1-1.5B with PRISM reaches 47.55 versus 22.69 for the strongest baseline — roughly a 2× improvement.

Tool use (format + correctness + length) — BFCL-v3

Helpfulness vs. harmlessness

Robustness to more rewards

PRISM keeps correctness nearly constant.
GDPO and GRPO-Prod degrade noticeably.
GRPO-Sum holds accuracy but weakly optimizes format and length.

Sample efficiency

Runtime Controllability

A single trained PRISM model behaves as a family of policies:

$\alpha_1 = 0.55$ → highest correctness (31.53%)
$\alpha_2 = 0.55$ → highest format accuracy (98.08%)
$\alpha_3 = 0.55$ → longest inference length (108.0 tokens)

This is the equivalent of a mixing console: one training run yields a tunable policy family.

Ablations

1. Shared positive policy (merging all positives) → biggest drop in format and correctness, confirming reward-specific optimization directions are essential. 2. Removing the global negative policy → all metrics drop, confirming its role in capturing shared failure modes. 3. Independent branch sampling instead of sampling from the composed $\pi^\star$ → lowest overall accuracy, confirming train/inference distribution alignment matters.

A max-weighted negative activation beats LogAvgExp and mean, matching the design intent: any severely failing reward should trigger joint penalization.

Engineering Principles

Match optimization granularity to the granularity of what is being optimized.
Division of labor beats forced unification: positives are specialists, the negative is a generalist.
Compose at the lowest representation layer (logits, not outputs).
Align training and inference distributions by sampling from the composed policy.

Limitations and Open Questions

Validated only up to 3 rewards; scaling to more is untested.
Mixing weights are still manually set; adaptive, prompt- or user-aware weighting is open.
Memory and FLOPs scale linearly with reward count, which may bottleneck large-model + many-reward settings.

Reference

Paper: arXiv:2607.29246 (publication date 2026-07-31)
Authors: Ruiming Liang, Yi Zhong, Yizhen Yuan, Yinan Zheng, Tianyi Tan, Tianyue Wang, Haiyun Guo, Jinqiao Wang, Xianyuan Zhan
Affiliations: Institute of Automation CAS / University of CAS / Tsinghua AIR / Tongji University
Code: not yet open-sourced at publication; watch the authors' GitHub for updates.

Key Points

Mixing heterogeneous rewards into a single scalar is a structural limitation, not a tuning problem — the paper labels this the multi-reward alignment tax.
PRISM trains N reward-specific positive sub-policies plus one global negative sub-policy, all sharing a backbone via learnable prefixes.
Inference-time composition is a linear logit combination $\sum_k \alpha_k z_k^+ - \gamma z^-$, so behavior is controllable at runtime without retraining.
Asymmetric training (stop-gradient on the backbone for the negative branch) prevents the shared trunk from being contaminated by negative samples.
Per-token inference latency equals a single forward pass regardless of reward count; memory and compute grow linearly.
Empirical gains include +17.77 on DeepSeek-R1-1.5B in scientific reasoning, ~2× improvement on GPQA, and simultaneous gains in helpfulness and harmlessness where baselines trade them off.
PRISM converges in roughly half the training steps of reward-mixing baselines because it eliminates gradient conflict between competing objectives.
Validated only up to 3 rewards; adaptive weights, larger reward sets, and efficient multi-branch decoding remain open problems.