Simultaneous Multi-objective Alignment Across Verifiable and Non-verifiable Rewards
Yiran Shen, Yu Xia, Jonathan Chang, Prithviraj Ammanabrolu
TL;DR
The paper tackles multi-objective alignment of large language models across verifiable (e.g., mathematical correctness), non-verifiable (e.g., human values), and interactive tutoring domains. It introduces a unified framework that standardizes Process Reward Model (PRM) training across domains, uses vectorized rewards via Multi-Action-Head DPO (MAH-DPO) to preserve multi-dimensional preferences, and combines this with PRM-guided decoding that continues hidden-state information for fine-grained inference-time control. Empirically, MAH-DPO yields superior joint performance across math accuracy, value alignment, and tutoring engagement, while PRM-guided decoding at test time further improves targeted objectives; the two components also show synergistic gains when used together. A key practical insight is that verifiable rewards benefit most from test-time search on a precise signal, whereas noisier, subjective rewards benefit from training-time representation shaping plus inference-time rebalancing, enabling flexible control without retraining. The approach is demonstrated on three datasets (MATH, UltraFeedback, and Socratic Mind), and a unified PRM across seven dimensions shows cross-domain transfer, indicating potential for scalable, configurable AI assistants that are accurate, safe, and engaging.
Abstract
Aligning large language models to human preferences is inherently multidimensional, yet most pipelines collapse heterogeneous signals into a single optimizeable objective. We seek to answer what it would take to simultaneously align a model across various domains spanning those with: verifiable rewards (mathematical accuracy), non-verifiable subjective preferences (human values), and complex interactive scenarios (multi-turn AI tutoring dialogues). Such multi-objective reinforcement learning setups are often plagued by the individual objectives being at odds with each other, resulting in inefficient training and little user control during inference. We propose a unified framework that: (i) standardizes {process reward model} (PRM) training across both verifiable and non-verifiable settings to better supervise models' chain-of-thought reasoning; (ii) performs {multi-objective alignment} by training the LLM with our $\textbf{M}$ulti-$\textbf{A}$ction-$\textbf{H}$ead $\textbf{DPO}$ (MAH-DPO) and a vectorized reward where the dimensions of the vector correspond to the various objectives instead of a single scalar; and (iii) demonstrates how such a system provides fine-grained inference-time user control. Experiments across math reasoning, value alignment, and multi-turn dialogue show that our framework improves performance across multiple objectives simultaneously, while minimizing cross-objective trade-offs and enabling flexible inference time user control. The code can be found at https://github.com/pearls-lab/multiobj-align.
