Source-linked AI summary
Interpretable Preferences via Multi-Objective Reward Modeling and Mixture-of-Experts
Haoxiang Wang, Wei Xiong, Tengyang Xie, Han Zhao, Tong Zhang
TL;DR
Conventional RLHF reward models are black boxes that may inadequately represent complex human preferences and can contribute to reward hacking. The paper combines absolute multi-objective reward modeling with prompt-conditioned mixture-of-experts scalarization to make reward preferences more interpretable and steerable. Its Llama-3 8B model achieves state-of-the-art RewardBench performance, surpasses GPT-4-based judging, and approaches the much larger Nemotron-4 340B reward model.
Problem
Conventional Bradley-Terry reward models are black boxes with limited capacity for complex human preferences, making their decisions difficult to interpret.
Method
The paper trains an absolute-rating multi-objective reward model and adds a prompt-conditioned gating network that scalarizes objectives in a mixture-of-experts strategy.
Results
State-of-the-art RewardBench performance is achieved by the Llama-3 8B ArmoRM, which surpasses GPT-4 judges and approaches Nemotron-4 340B.
Takeaways & Limitations
Decomposed and steerable reward scores let humans inspect objective contributions and adjust priorities, including reducing verbosity bias.
Abstract
from arXiv · showhide
Reinforcement learning from human feedback (RLHF) has emerged as the primary method for aligning large language models (LLMs) with human preferences. The RLHF process typically starts by training a reward model (RM) using human preference data. Conventional RMs are trained on pairwise responses to the same user request, with relative ratings indicating which response humans prefer. The trained RM serves as a proxy for human preferences. However, due to the black-box nature of RMs, their outputs lack interpretability, as humans cannot intuitively understand why an RM thinks a response is good or not. As RMs act as human preference proxies, we believe they should be human-interpretable to ensure that their internal decision processes are consistent with human preferences and to prevent reward hacking in LLM alignment. To build RMs with interpretable preferences, we propose a two-stage approach: i) train an Absolute-Rating Multi-Objective Reward Model (ArmoRM) with multi-dimensional absolute-rating data, each dimension corresponding to a human-interpretable objective (e.g., honesty, verbosity, safety); ii) employ a Mixture-of-Experts (MoE) strategy with a gating network that automatically selects the most suitable reward objectives based on the context. We efficiently trained an ArmoRM with Llama-3 8B and a gating network consisting of a shallow MLP on top of the ArmoRM. Our trained model, ArmoRM-Llama3-8B, obtains state-of-the-art performance on RewardBench, a benchmark evaluating RMs for language modeling. Notably, the performance of our model surpasses the LLM-as-a-judge method with GPT-4 judges by a margin, and approaches the performance of the much larger Nemotron-4 340B reward model.
1 Introduction
RLHF reward models guide LLM alignment but conventional models are limited in preference capacity and interpretability. The paper proposes a two-stage multi-objective and mixture-of-experts approach to improve interpretability and RewardBench performance.
- RLHF context: RLHF uses reward models to integrate human preferences into LLM training and guide policy optimization.The standard framework collects preference data, trains a Bradley-Terry reward model, and optimizes the policy with PPO.
- Motivation: Bradley-Terry reward models may inadequately capture complex human preferences and provide no human-interpretable explanation for their scores.Their black-box behavior is associated with reward hacking, including verbosity bias.
- Motivation: Decomposable reward scores could expose objective contributions and allow steering, such as shifting weight from length toward helpfulness.The paper illustrates a score composed of 40% helpfulness and 60% length, then proposes making it 100% helpfulness to mitigate verbosity bias.
- Contribution: The proposed approach first trains a multi-objective reward model, then learns a gating layer that scalarizes objectives in a mixture-of-experts manner.The authors validate the approach with Llama-3 8B and report state-of-the-art RewardBench performance.
2 Related Works
Related RLHF methods use reward models to provide preference signals, while direct preference learning can bypass traditional reward modeling. Existing reward modeling remains constrained by Bradley-Terry assumptions and simplistic integration of multiple objectives.
- RLHF algorithms: PPO-based RLHF and reward-ranked fine-tuning use external reward models to provide supervision or rank sampled responses.PPO-based RLHF supports policy optimization, while RAFT selects high-ranked responses for fine-tuning.
- Direct preference learning: Online iterative direct preference learning methods collect new responses and use external reward models to label them before further fine-tuning.The original DPO is described as offline, whereas subsequent iterative variants explore environments through newly deployed policies.
- Summary: Popular RLHF algorithms require external reward models to provide preference signals for their best performance.
- Reward modeling: Traditional Bradley-Terry reward modeling struggles with complex, intransitive human preferences, and multi-dimensional signals are often combined linearly.The related-work discussion identifies a need for more sophisticated integration of multiple reward objectives.
3 Methodology
The method preserves fine-grained absolute ratings across multiple human-interpretable objectives and uses prompt-conditioned gating to produce a scalar reward. It also corrects verbosity-related correlations before objective scalarization.
- Multi-Objective Reward Modeling: Existing reward models commonly train on binary pairwise Bradley-Terry labels, which can discard fine-grained information from multi-objective absolute ratings.UltraFeedback is described as providing five objectives with ratings from 1 to 5 before binarization.
- Multi-Objective Reward Modeling: Multi-objective regression maps each prompt-response example to a k-dimensional rating vector whose dimensions represent objectives such as helpfulness and truthfulness.
- MoE Scalarization: Fixed linear combinations of objective rewards are rigid for complex application scenarios, motivating context-dependent scalarization.The ArmoRM predicts multiple rewards, but ranking requires reducing them to one scalar.
- MoE Scalarization: A gating layer maps the prompt feature to non-negative coefficients summing to one, which weight the reward objectives into a scalar score.The gating layer can be implemented as a shallow MLP followed by softmax, and prompt features are available during the response forward pass.
- Bias correction: Reward objectives are adjusted with verbosity penalties to reduce bias from correlations between objectives and response length.The adjusted reward vector is denoted r′, and penalty coefficients are selected using a correlation metric over a reference distribution.
- Training: The backbone and regression layer are frozen while the gating layer is trained with Bradley-Terry loss and an additional scaling variable.The implementation uses a ReLU MLP with three hidden layers, Spearman correlation, UltraFeedback as the reference distribution, and ten pairwise datasets.
4 Experiment
The experiments implement ArmoRM and its MoE gating layer, then evaluate the combined reward model on RewardBench against other reward models.
- Implementation: ArmoRM uses a frozen Llama-3 8B backbone with a trained linear layer for regression across 19 objectives from 8 datasets.The objectives include helpfulness, correctness, and verbosity.
- Implementation: The MoE gating layer is a three-hidden-layer ReLU MLP with 1024 hidden units, trained with the backbone frozen on 10 pairwise preference datasets.The gating layer uses Spearman correlation and UltraFeedback as the reference distribution for its correlation metric.
- Implementation: Training uses PyTorch, Transformers, and Scikit-learn, with linear probing on a CPU and gating-layer training on a single NVIDIA A6000 GPU.The gating layer is trained with AdamW for 10,000 steps using a batch size of 1024 and cosine learning-rate decay.
- Evaluation: RewardBench evaluates reward models across Chat, Chat Hard, Safety, Reasoning, and prior-set categories using a weighted-average overall score.The four primary categories have weight 1.0, while prior sets have weight 0.5.
- Results: The combined ArmoRM and MoE model significantly outperforms the Llama-3 8B Bradley-Terry reward model, exceeds GPT-4-based LLM-as-a-judge performance, and approaches Nemotron-4 340B.These comparisons are reported as key observations from Table 1.
5 Conclusion
The paper addresses reward-model interpretability for RLHF by proposing ArmoRM with an MoE strategy and reports state-of-the-art RewardBench performance.
- The proposed approach combines an ArmoRM with an MoE strategy and gating network to improve reward-model interpretability for LLM alignment.
A Experimental Details
The experimental details describe dataset licensing and content considerations, along with the datasets and preprocessing used for multi-objective reward modeling and MoE training.
- Licenses: The model follows the Meta Llama3 license, while the datasets are open-sourced for research use and some also permit commercial use.
- Data considerations: The datasets generally exclude personally identifying information, but several preference datasets deliberately contain offensive content.The datasets named are BeaverTails, PKU-RLHF, and HH-RLHF.
- Multi-objective training data: Multi-objective reward modeling uses datasets providing human-interpretable objectives such as helpfulness, correctness, verbosity, honesty, safety, and readability.The listed sources include HelpSteer, UltraFeedback, BeaverTails-30k, CodeUltraFeedback, Prometheus, and several Argilla datasets.
- Data preprocessing: Ratings from different datasets are normalized to 0–1, with Boolean safety labels mapped to 1 for True and 0 for False.The preprocessing also addresses similar objectives appearing across datasets.
- MoE training data: The MoE gating layer is trained on preference datasets including HelpSteer, UltraFeedback, SHP, HH-RLHF, and several safety, math, code, and reasoning collections.The listed training sources include 10 pairwise preference datasets.
- Preference preprocessing: Datasets without response-pair labels are converted using binarized versions preprocessed by Dong et al.,”