PRISM: Mixing Policies Instead of Rewards for Multi-Objective LLM Alignment

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

Summary

PRISM is a multi-reward reinforcement learning framework from researchers at the Chinese Academy of Sciences, Tsinghua AIR, and Tongji University that shifts multi-objective alignment from reward-space composition to policy-space composition. Instead of collapsing correctness, format, safety, and length rewards into a weighted scalar (the standard GRPO approach that causes whack-a-mole trade-offs), PRISM trains N positive sub-policies each specialized for one reward plus a single global negative policy that captures the union of all failure modes. Sub-policies share a backbone via prefix conditioning and are combined at inference as a linear mixture of logits, yielding a family of policies from one training run. Experiments on ScienceQA, GPQA, BFCL-v3, Alpaca, HH-RLHF, and PKU-SafeRLHF show +17.77 over the strongest baseline on DeepSeek-R1-1.5B and roughly 2x sample efficiency, with inference-time weight tuning enabling real-time behavior control without retraining.

The Multi-Reward Whack-a-Mole Problem

When training LLMs with several reward signals (correctness, format, tool-call compliance, safety), the dominant approach collapses them into a single scalar via linear weighting, adaptive weighting, or constrained optimization, then runs GRPO on that scalar. The result: improving one dimension degrades another, because heterogeneous preference signals compete in the same gradient step. The authors call this multi-reward alignment tax and the broader paradigm reward-space composition.

PRISM's Core Insight: Mix Policies, Not Rewards

PRISM reframes multi-reward alignment as policy-space composition.

1. One positive policy per reward

For N rewards, PRISM trains N independent positive sub-policies $\pi_1^+, \ldots, \pi_N^+$, each specialized for a single reward. Good behaviors are reward-specific and deserve dedicated optimization directions.

2. One global negative policy

A single negative policy $\pi^-$ captures the union of all failure modes: it activates whenever any reward is low. Bad behaviors do not need to be attributed to any single reward, so a shared penalty suffices.

3. Logit-level linear composition at inference

Because all sub-policies share the same reference policy $\mu$ and context, log-ratios decompose into a linear combination of logits:

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

with tunable mixing weights $\alpha_k$ and $\gamma$.

Engineering Tricks That Keep It Practical

Shared backbone + prefix conditioning. All sub-policies share $\pi_\theta$, distinguished by learnable prefix embeddings $\mathcal{P}=\{P_1^+,\ldots,P_N^+,P^-\}$. No N+1 separate models.
Asymmetric training. Positive branches update both the backbone and their own prefixes; the negative branch updates only its prefix with stop-gradient on the backbone, preventing contamination of core capabilities.
Parallel batched sampling. Query, prefixes, and generated tokens are stacked along the batch dimension so all sub-policies run in one forward pass. Per-token latency equals one forward pass; memory and FLOPs grow linearly with reward count.

Experimental Results

Scientific reasoning (correctness + format)

+17.77

+8.04

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

Tool calling (format + correctness + length)

Helpfulness–harmlessness alignment

Alignment tax scaling

Sample efficiency

Inference-Time Control: One Training, a Family of Policies

After training, weights $\alpha_k$ can be adjusted at inference:

$\alpha_1=0.55$ maximizes correctness (31.53%).
$\alpha_2=0.55$ maximizes format accuracy (98.08%).
$\alpha_3=0.55$ maximizes average length (108.0 tokens).

No retraining required; mixing weights yield a continuous family of policies.

Ablations

Sharing positive policies → biggest drop in format and correctness (reward-specific directions matter).
Removing the global negative policy → all metrics drop.
Independent per-branch sampling → lowest overall accuracy (training and inference distributions must match).
Max weighting on the negative policy beats LogAvgExp and mean: "any reward failing badly triggers penalty" aligns with design intent.

Practical Takeaways

1. In multi-reward settings with heterogeneous signals (rule-based + model-based), consider policy-space composition before weighted scalar mixing. 2. Prefix conditioning gives multi-policy behavior without maintaining N models; inference latency stays nearly constant. 3. Stop-gradient on the backbone for any "suppressive" branch protects core capabilities. 4. Train-once-deploy-many becomes feasible: same backbone, different deployments (e.g., medical vs. customer support) via inference-time weights.

Limitations

Validated up to 3 rewards; scaling to more is open.
Mixing weights are currently manual; adaptive per-prompt or per-user adjustment is future work.
Memory and FLOPs grow linearly with reward count, a potential bottleneck for very large backbones.

Reference

Paper: arXiv:2607.29246
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