Subspace Inference Enables Efficient Active Reward Learning from Preferences

AI summary

Overview

Research area: Reinforcement learning from human feedback (RLHF), Bayesian deep learning, and active preference-based reward modeling.

Technical level: Advanced — the paper assumes familiarity with Bayesian filtering (Kalman/EKF), neural network subspaces (SVD of SGD iterates), Bradley-Terry preference models, and offline RL benchmarks.

Scope: The paper introduces PreferenceEKF, a method that performs sequential Bayesian inference over neural network reward model parameters inside a learned low-dimensional subspace, enabling scalable uncertainty-aware active learning from human preference queries.

What This Paper Is About

RLHF learns reward functions by asking humans to compare pairs of behaviors, but each comparison yields only about one bit of information, so thousands of queries may be needed. Active learning can pick the most informative queries, but that requires trustworthy uncertainty estimates over the reward model — and obtaining those estimates for large neural networks is either computationally prohibitive (full Bayesian inference) or unreliable (dropout). This paper's goal is to make principled Bayesian uncertainty quantification cheap enough to drive active query selection for neural network reward models.

Key Contributions

1. PreferenceEKF, a new method that frames active preference-based reward learning as sequential Bayesian filtering and runs an extended Kalman filter (EKF) inside a low-dimensional neural network subspace, rather than over the full parameter space.
2. Scalable acquisition-function computation: because the posterior is a tractable Gaussian in the subspace, an arbitrary number of reward models can be sampled and projected back to full space, making sampling-based acquisition functions such as InfoGain feasible for neural network reward models for the first time.
3. A multi-baseline empirical study across 12 D4RL tasks plus V-D4RL pixel-based control and the real-robot SOAR dataset, comparing against DeepEnsemble, Dropout, Laplace approximation, and last-layer MCMC on sample efficiency, runtime, scaling, and calibration.
4. An ablation on subspace construction, showing that SVD-based and random-projection bases trade off differently with subspace dimension, and that random projection can remove the method's dependence on an initial SGD warm-up dataset.

Main Findings

• Sample efficiency and fit: Across 12 D4RL tasks over 12 seeds, both random-query and active variants of PreferenceEKF match or beat all four baselines in Bradley-Terry log-likelihood on held-out test queries, and active PreferenceEKF leads in most individual tasks.
• Acquisition functions: InfoGain was the only tested acquisition function (versus disagreement and entropy) for which the active variant of every method matched or exceeded its random-query counterpart.
• Speed: PreferenceEKF trains roughly 5x faster than DeepEnsemble and over 40x faster than last-layer MCMC. The advantage comes from the sequential EKF update, which processes only the newest query-label pair, while baselines must retrain on all data seen so far to avoid catastrophic forgetting.
• Scaling: As the number of posterior samples grows, PreferenceEKF's runtime grows far more gracefully than DeepEnsemble's or Dropout's, and it retains the best final log-likelihood. The same favorable trend holds when the reward network itself is enlarged.
• Calibration: PreferenceEKF achieves the lowest expected calibration error among all methods and the second-lowest Brier score (behind active DeepEnsemble), indicating its posterior is genuinely well-calibrated rather than merely cheap.
• Policy optimization: Reward models learned by PreferenceEKF, used to label offline datasets and train policies with Implicit Q-Learning, produce policies competitive with those trained from reward models learned by other methods, and roughly on par with policies trained on ground-truth rewards.
• Fairness of the comparison: When all methods are restricted to the same number of posterior samples (M = 5), PreferenceEKF remains the strongest, showing its gains are not simply from drawing more samples.
• Known limitation: The unimodal Gaussian assumption makes PreferenceEKF suited to a single annotator. On crowd-sourced human labels with multimodal preferences, no method performed well.

Methodology in Plain English

The researchers treat the reward model's parameters as a hidden state that evolves over time and each human preference label as a noisy measurement of that state. This lets them use a classic recursive estimator — the extended Kalman filter — to update a Gaussian belief over the parameters after every query, one query at a time, instead of retraining on the whole history.

The problem is that a full covariance matrix over all network weights is quadratic in the parameter count, which is impossible for a real network. Their fix is to compress: they first run a small number of warm-up SGD steps on an initial preference dataset, then take the singular value decomposition of those parameter iterates to define a low-dimensional basis. All Bayesian inference then happens on a small vector of subspace coordinates (200 dimensions in their experiments), and samples are mapped back to full parameter space through a fixed affine projection.

Because the posterior is a Gaussian in this small space, drawing hundreds of reward models is cheap. Those samples are plugged into the InfoGain acquisition function — a sampling-based approximation of the mutual information between a candidate query's label and the model parameters — to select the next most informative comparison to ask about. The authors also test replacing the SVD basis with random projections, which performs comparably at larger subspace sizes and works even without any warm-up training data.

Why This Matters

This work shows that a classical, deterministic inference algorithm combined with a dimensionality reduction trick can beat modern approximate Bayesian deep learning on the specific task of uncertainty-driven active reward learning — while being dramatically cheaper and better calibrated. It challenges the default assumption that ensembles or dropout are the only practical route to neural reward model uncertainty, and it makes the theoretically superior InfoGain acquisition function usable at neural network scale.

Real-world applications:

• Robotics: Learning reward functions for manipulation and locomotion from human pairwise comparisons, including from sparse success/failure robot data (as demonstrated on the SOAR dataset) and pixel-based observations.
• LLM and agent alignment: Cutting the number of human preference labels needed in RLHF pipelines by selecting maximally informative comparisons.
• Assistive and personalized systems: Adapting to a single user's preferences — vehicle autonomy, exoskeletons, recommendation — where queries are expensive and a unimodal user model is appropriate.
• Offline decision systems: Producing reward models that can label existing logged datasets for policy training without any new environment interaction.

Industry relevance: Annotation is often the dominant cost in RLHF and preference-based training. A method that is 5-40x faster to train, scales better with model size, and needs fewer human queries directly reduces labeling budgets and compute spend, while its sequential one-query-at-a-time update fits naturally into streaming, online alignment pipelines.

Future Directions

• Multimodal and pluralistic preferences: Extending beyond the unimodal Gaussian posterior, for example with particle filters or other nonparametric inference, to handle crowd-sourced annotators with genuinely different reward functions — with the open question of how much inference efficiency this costs.
• Removing the warm-up dependency: The authors show random-projection bases can work without an initial SGD dataset; establishing how reliably this decouples PreferenceEKF from any gradient-based training remains open.
• Truly large reward models: Testing whether subspace filtering scales to language-model-scale reward models and richer input modalities beyond the image-based V-D4RL experiments.
• Reward learning and policy learning jointly: The paper deliberately keeps these separate; the interaction between the learned reward posterior and downstream policy optimization is left for future work.

Target Audience

Researchers and practitioners working on RLHF, preference-based reward learning, active learning, or uncertainty quantification in deep learning; robotics and RL engineers who need scalable Bayesian inference for reward models; and graduate students comfortable with Bayesian filtering and offline RL who want a concrete example of subspace methods making principled inference practical at neural network scale.

Authors’ abstract

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.