Reviving The Classics: Active Reward Modeling in Large Language Model Alignment

TL;DR

This work tackles efficient reward modeling for RLHF by treating active query design through Fisher information, applied to the final linear layer of Bradley-Terry reward models. By employing D-optimal and Past-Aware D-opt designs, the authors balance embedding-space diversity with informative outcomes, achieving superior data efficiency and stability across multiple open-source LLMs. They extend these ideas to nonlinear models via last-layer features and gradient-based approximations, and demonstrate that cross-prompt annotations further boost labeling efficiency. The study provides a practical, scalable framework for interactive annotation in AI alignment, with strong empirical support and a clear link to classical experimental design and Bayesian ideas such as BALD. Overall, the approach offers a principled path to faster, more reliable RLHF alignment in large language models.

Abstract

Building neural reward models from human preferences is a pivotal component in reinforcement learning from human feedback (RLHF) and large language model alignment research. Given the scarcity and high cost of human annotation, how to select the most informative pairs to annotate is an essential yet challenging open problem. In this work, we highlight the insight that an ideal comparison dataset for reward modeling should balance exploration of the representation space and make informative comparisons between pairs with moderate reward differences. Technically, challenges arise in quantifying the two objectives and efficiently prioritizing the comparisons to be annotated. To address this, we propose the Fisher information-based selection strategies, adapt theories from the classical experimental design literature, and apply them to the final linear layer of the deep neural network-based reward modeling tasks. Empirically, our method demonstrates remarkable performance, high computational efficiency, and stability compared to other selection methods from deep learning and classical statistical literature across multiple open-source LLMs and datasets. Further ablation studies reveal that incorporating cross-prompt comparisons in active reward modeling significantly enhances labeling efficiency, shedding light on the potential for improved annotation strategies in RLHF.

Figures (22)

• Figure 1: Comparisons drawn by different strategies to learn a 2D bimodal reward function. The heat map showed the estimated functions. Red dots connected by lines are selected pairs and gray dots on the first column are candidate points to choose from.
• Figure 2: Workflow of active reward modeling and experimental setups. At each round, we start with randomly sampling prompts, generating responses and candidate comparisons, active labeling, and model retraining.
• Figure 3: Comparing annotation efficiency of different methods. (Harmless Dataset, 3 Models, 8 Methods). First row: $1 -$ Spearman's Correlation (lower is better); second row: Best-of-N reward (higher is better). Experiments are repeated with 5 seeds.
• Figure 4: Investigating how annotation batch size choices affect learning performance of our methods. Model: Gemma 2B. The first two columns present results on the Harmless dataset, and the second two columns present results on the Helpful dataset. First row: $1 -$ Spearman's Correlation (lower is better); second row: Best-of-N reward (higher is better). The results presented are from 5 runs with different seeds.
• Figure 5: Comparing annotation efficiency of different methods under the Cross-Prompt annotation setups. (Harmless Dataset, 3 Models, 8 Methods). First row: $1 -$ Spearman's Correlation (lower is better); second row: Best-of-N reward (higher is better). Experiments are repeated with 5 seeds.