Subspace Inference Enables Efficient Active Reward Learning from Preferences

Summary

Abstract Summary

Reinforcement learning from human feedback (RLHF) has emerged as a powerful yet sample-inefficient approach for learning reward models from human preferences, making active learning a critical component in synthesizing informative preference queries. However, effective uncertainty quantification required for active learning remains a key challenge for large neural network reward models. In this paper, we introduce PreferenceEKF, a sample-efficient approach that tracks reward model uncertainty by framing active preference learning as a sequential Bayesian filtering problem. Instead of relying on computationally prohibitive posterior inference over the full neural network parameter space, our method performs sequential inference via an extended Kalman filter within a low-dimensional parameter subspace, continuously updating the reward model posterior as new preference queries arrive. Our approach enables scalable sampling of neural network parameters to efficiently compute acquisition functions for active reward learning. Experiments on the D4RL and V-D4RL benchmarks demonstrate that our approach achieves better sample efficiency, runtime, scalability, and calibration compared to other Bayesian deep learning approaches, and the learned reward models lead to competitive offline reinforcement learning policy performance. This highlights the potential of scalable Bayesian methods for preference-based reward modeling in RLHF. Our code is available at https://github.com/yutaizhou/bnn_pref.

Key Points

Reinforcement learning from human feedback (RLHF) has emerged as a powerful yet sample-inefficient approach for learning reward models from human preferences, making active learning a critical component in synthesizing informative preference queries.
However, effective uncertainty quantification required for active learning remains a key challenge for large neural network reward models.
In this paper, we introduce PreferenceEKF, a sample-efficient approach that tracks reward model uncertainty by framing active preference learning as a sequential Bayesian filtering problem.
Instead of relying on computationally prohibitive posterior inference over the full neural network parameter space, our method performs sequential inference via an extended Kalman filter within a low-dimensional parameter subspace, continuously updating the reward model posterior as new preference queries arrive.
Our approach enables scalable sampling of neural network parameters to efficiently compute acquisition functions for active reward learning.

Preliminaries

Preference-based reward modeling. We consider a Markov decision process (MDP) $\langle\mathcal{S},\mathcal{A},\mathcal{T},r,\gamma\rangle$ with state space $\mathcal{S}$, action space $\mathcal{A}$, transition function $\mathcal{T}$, reward function $r:\mathcal{S}\rightarrow\mathbb{R}$, and discount factor $\gamma\in[0,1)$.

Method

Sampling neural network models to approximate acquisition functions as in Eq. 3 can be expensive due to the high-dimensional parameter space of neural networks ( izmailov2021what ) . We leverage the insight that neural networks are overparameterized and that solutions actually live in a much smaller subspace ( li2018measuring ) , and perform posterior inference within this subspace.

This allows us to sample an arbitrary number of models from a lower-dimensional posterior to approximate Eq. 3 , without, e.g., the overhead of training ensembles.

Experiments

Baselines and Evaluation. We compare PreferenceEKF to four Bayesian deep learning baselines commonly used for reward modeling: DeepEnsemble, Dropout, Laplace, and last-layer Markov chain Monte Carlo (LLMCMC), which we detail in Section A.2.4 .

Sources and Demonstrations

Figures are reproduced from Zhou et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.