Table of Contents
Fetching ...

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.

Simultaneous Multi-objective Alignment Across Verifiable and Non-verifiable Rewards

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 ulti-ction-ead (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.

Paper Structure

This paper contains 29 sections, 8 equations, 5 figures, 11 tables, 1 algorithm.

Figures (5)

  • Figure 1: Overview of our training-time and test-time alignment framework. Left: we train a single LLM with multiple action heads using head specific DPO losses and a combined loss on the shared backbone. Right: PRM-guided decoding selects the next step among candidates for controllable objective weights.
  • Figure 2: Results with varying action head weights in Math.
  • Figure 3: Results with varying action head weights in Human Values.
  • Figure 4: Alignment performances of a unified PRM trained across 7 dimensions in three domains compared with base model and the specialized PRM trained on each dimension per domain.
  • Figure : PRM-Guided Decoding with Continuing Hidden State