Aligning Large Language Models (LLMs) to human preferences is inherently multidimensional. It is akin to coaching a decathlete: the goal is to cultivate speed, strength, agility, and mental toughness simultaneously. Traditionally, AI alignment has simplified this challenge by averaging these diverse traits into a single score. This process, known as scalarization, collapses heterogeneous human preferences—such as helpfulness, harmlessness, and factual accuracy—into a single reward metric.
The primary drawback of scalarization is the loss of nuance. When signals are collapsed, a model might become perfectly "safe" by refusing to answer any questions, or highly "helpful" by hallucinating pleasing but incorrect information. To address these compromises, researchers from UC San Diego and Databricks have introduced MAHALO (Multi-Action-Head ALignment with PRM-guided DecOding), a unified framework designed to keep competing objectives distinct, harmonious, and controllable.
The Challenge of Multi-Objective Alignment
Most current alignment pipelines, including standard Reinforcement Learning from Human Feedback (RLHF) and Direct Preference Optimization (DPO), treat human preference as a one-dimensional scale. This architectural choice leads to three critical issues:
- Gradient Interference: During training, conflicting goals (e.g., "be concise" versus "be thorough") produce opposing gradient signals that can cancel each other out, leading to inefficient learning.
- Lack of Inference Control: Once a model is trained on a fixed, weighted average of goals, users cannot easily adjust that balance at runtime.
- Reward Mismatch: Alignment involves both verifiable rewards (e.g., mathematical accuracy) and non-verifiable rewards (e.g., pedagogical engagement or tone). Standard methods struggle to integrate both types of feedback into a single, coherent system.
The MAHALO Framework
To overcome these limitations, MAHALO standardizes step-level supervision and performs vectorized multi-objective alignment using a three-pronged approach:
1. Multi-Action-Head DPO (MAH-DPO)
Instead of relying on a single output head, the MAHALO architecture features a shared transformer backbone equipped with multiple linear "action heads"—one dedicated to each objective. This allows the backbone to learn shared, generalized representations while the individual heads specialize. Crucially, this design provides gradient isolation, preventing different objectives from interfering with one another during model updates.
2. Standardized PRM Training
The framework utilizes Process Reward Models (PRMs) to provide granular feedback on every step of a response, rather than just evaluating the final output. The authors adapted PRM training for both verifiable domains (using hindsight relabeling for math problems) and subjective domains (using an "LLM-as-Judge" approach to evaluate traits like engagement).
3. PRM-Guided Decoding with KV-Cache Continuity
At inference time, the model generates multiple candidate trajectories for the next step, and the PRM selects the optimal path. To make this computationally feasible and prevent the degradation associated with re-encoding text, MAHALO carries the internal hidden states (KV-cache) forward. This innovation yields a 4x speedup over standard search methods.
Experimental Scope and Key Findings
The researchers stress-tested MAHALO across three distinct domains:
- Math Reasoning (MATH): Balancing raw problem-solving accuracy with an engaging pedagogical tone.
- Human Values (UltraFeedback): Navigating the complex trade-offs between helpfulness, honesty, and truthfulness.
- Interactive Tutoring (Socratic Mind): Managing multi-turn dialogues where the objective is to guide a student toward an answer without explicitly revealing it.
Key Insights
- The "Verifiability" Rule of Thumb: The nature of the reward dictates the most effective alignment strategy. The authors found that verifiable rewards (like math accuracy) benefit most from test-time search via PRM-guided decoding. Conversely, subjective rewards (like human values) benefit more from the multi-head training of MAH-DPO.
- Gradient Isolation Prevents Collapse: By separating the action heads, the model avoids the performance collapse often seen in single-head models. In a stress test balancing "Helpfulness" and "Honesty," MAHALO maintained high scores in both metrics, whereas standard DPO saw one metric severely degrade the other.
- Smooth Inference Control: Because the objective heads are distinct, users can dynamically adjust the weights of different traits at inference time. This allows for slider-like control—making a tutor more "engaging" or "concise" on the fly without requiring model retraining.
- Unified PRMs act as Generalists: A single PRM trained on a mixture of all seven objectives across the three domains performed exceptionally well, indicating that traits like "good reasoning" and "good style" share transferable features.
Limitations and Future Directions
While MAHALO represents a significant advancement in multi-objective alignment, the authors identify a few areas for future exploration:
- Scaling Objectives: Experiments successfully balanced up to five simultaneous objectives. Future work must evaluate how the architecture scales when managing dozens of fine-grained preferences.
- LLM-as-Judge Noise: For non-verifiable rewards, the system relies on a judge model (e.g., GPT-4o) for labeling. Biases or inconsistencies in the judge can propagate into the PRM.
- Computational Overhead: Although highly optimized via KV-cache continuity, maintaining multiple heads and performing candidate search at inference still incurs a slight latency increase (~13%) compared to a vanilla single-head model.
Paper Citation: Shen, Y., Xia, Y., Chang, J. D., & Ammanabrolu, P. (2026). Simultaneous Multi-objective Alignment Across Verifiable and Non-verifiable Rewards. Proceedings of the 43rd International Conference on Machine Learning (ICML).