PaperlayerView exact revision on arXiv

Source-linked AI summary

Subspace Inference Enables Efficient Active Reward Learning from Preferences

Yutai Zhou, Erdem Bıyık

arXiv:2609.04066v1cs.LGcs.AIcs.RO

TL;DR

RLHF preference feedback is sample-inefficient, while uncertainty estimation for active learning is difficult to scale to large neural reward models. PreferenceEKF applies sequential extended-Kalman-filter inference in a low-dimensional parameter subspace, enabling scalable posterior sampling. Across D4RL and V-D4RL, it matches or exceeds Bayesian baselines in sample efficiency and calibration, improves runtime scaling, and yields competitive offline RL policies.

  • Problem

    RLHF preference feedback is sample-inefficient, and scalable uncertainty quantification for active learning remains difficult for large neural-network reward models.

  • Method

    PreferenceEKF performs sequential Bayesian filtering with an extended Kalman filter in a low-dimensional neural-network parameter subspace, enabling posterior sampling for acquisition functions.

  • Results

    PreferenceEKF performs on par with or better than Bayesian deep learning baselines in sample efficiency and calibration, scales better in runtime, and produces competitive offline RL policies.

  • Takeaways & Limitations

    Subspace Bayesian filtering provides a scalable approach to uncertainty-aware preference reward learning for neural reward models.

  • Takeaways & Limitations

    Scaling PreferenceEKF to pixel-based reward models is constrained by EKF updates that scale cubically with observation dimensionality, requiring pretrained image embeddings for high-dimensional inputs.

Abstract

from arXiv · show

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.

1 Introduction

RLHF preference feedback is easy to collect but sample-inefficient, making efficient active querying and uncertainty representation important for scalable reward modeling. PreferenceEKF addresses these challenges by filtering uncertainty in a neural-network subspace and performs competitively across active learning and offline RL evaluations.

  • Binary preference feedback provides at most one bit of information, so learning reward models may require thousands of comparison questions.
  • Active preference learning requires uncertainty estimates for acquisition functions, but Bayesian methods are difficult to scale to large neural networks.
  • PreferenceEKF performs Bayesian filtering in a constructed low-dimensional neural-network subspace using an extended Kalman filter.
  • The subspace posterior enables sampling arbitrary numbers of reward models for uncertainty-based acquisition functions such as expected information gain and disagreement.
  • PreferenceEKF performs on par with or better than Bayesian deep learning baselines in sample efficiency and calibration, with faster and better-scaling runtime.
  • Reward models learned with PreferenceEKF produce offline RL policies with performance competitive with policies optimized using other methods.

2 Related Work

Related work frames active preference learning as sequential uncertainty-guided data collection, while existing Bayesian and approximate methods face scalability or posterior-quality limitations. Subspace inference offers a route to tractable Bayesian neural-network uncertainty estimation without training multiple independent models.

  • Active learning reduces human-feedback collection by iteratively selecting useful samples using the model’s current state, including posterior uncertainty.
  • Bayesian methods support principled uncertainty quantification but are difficult to apply to neural reward models because acquisition functions require high-dimensional posterior samples.
  • Ensembles enable neural active reward learning but require multiple independent models, making training computationally expensive.
  • Dropout avoids ensemble training costs, but studies have reported poor posterior approximation quality.
  • Bayesian filtering provides a principled framework for sequential learning and has only recently been applied to deep networks through subspace methods.
  • Subspace methods exploit the observation that task-relevant neural-network parameters can occupy a lower-dimensional space, improving computational efficiency and Bayesian tractability.

3 Preliminaries

Preference-based reward modeling learns a reward function from binary comparisons between trajectory outcomes rather than direct reward observations. Information-theoretic active learning then selects queries by maximizing preference-label information about the reward-model parameters.

  • Preference-based reward modeling learns an annotator’s reward function from binary labels on pairwise trajectory comparisons without access to the true reward.
  • A parameterized reward model is commonly fit by maximum likelihood using a Bradley-Terry preference model.
  • The Bradley-Terry model uses a temperature parameter to represent noisily optimal annotator behavior and trajectory returns computed from neural-network rewards.
  • InfoGain selects the query maximizing mutual information between its preference label and the reward-model parameters.
  • Sampling approximates InfoGain using models drawn from the current posterior, but this requirement previously restricted the approach to low-dimensional reward models.
  • PreferenceEKF enables sampling high-dimensional neural reward models to compute sampling-based acquisition functions such as InfoGain.

4 Method

PreferenceEKF performs sequential Bayesian inference for neural-network reward parameters in a learned low-dimensional subspace. It maps subspace samples back to the full network for prediction and uses posterior sampling to drive active query selection.

  • Subspace inference: Sampling neural-network posteriors is expensive in the full parameter space, so PreferenceEKF performs inference in a lower-dimensional subspace.
  • EKF inference: The extended Kalman filter updates the posterior sequentially as preference query-label pairs arrive.
  • EKF inference: The dynamics and measurement models use additive Gaussian noise, with neural-network parameters treated as hidden states and preference observations modeled through the Bradley-Terry model.
  • Subspace inference: The method represents full parameters as θ(z) = Az + θ∗, where θ∗ is initialized by SGD and A is obtained from SVD of full-space SGD iterates.
  • Posterior sampling: Subspace samples are projected back to the full parameter space to run neural-network forward passes and estimate predictive distributions.
  • Posterior sampling: The posterior update is deterministic, while posterior sampling supports acquisition-function computation.
  • Active learning procedure: The method can compute disagreement and InfoGain, whereas DeepEnsemble approximates InfoGain by training independent models and Dropout samples dropout masks.
  • Active learning procedure: PreferenceEKF initializes a subspace belief, computes a query from the current belief, obtains its label, and applies an EKF update within the query budget.

5 Experiments

Experiments evaluate PreferenceEKF for preference learning, runtime scaling, calibration, subspace construction, and offline policy optimization across D4RL-based tasks. It generally matches or outperforms Bayesian baselines while offering faster, more scalable inference, with limitations for multimodal preferences and raw pixel inputs.

  • 5.1 Does PreferenceEKF lead to sample-efficient active reward learning?: PreferenceEKF and its random variant match or outperform all baselines in sample efficiency and final test log-likelihood across aggregated D4RL tasks.Active PreferenceEKF outperforms other methods on most individual tasks.
  • 5.1 Does PreferenceEKF lead to sample-efficient active reward learning?: InfoGain is the only tested acquisition function for which all five active methods perform at least as well as their random counterparts.The authors associate this result with InfoGain’s higher sample efficiency in prior work.
  • 5.1 Does PreferenceEKF lead to sample-efficient active reward learning?: PreferenceEKF remains the most successful method in sample efficiency and final log-likelihood when every method uses M = 5 posterior samples.This controls for the default difference between DeepEnsemble using M = 5 and other methods using M = 100.
  • 5.1 Does PreferenceEKF lead to sample-efficient active reward learning?: PreferenceEKF is designed primarily for a single annotator, because EKF’s unimodal Gaussian assumption does not represent multimodal preference distributions well.The authors report poor log-likelihood for all methods on crowd-sourced labels and leave multimodal extensions for future work.
  • 5.2 How does training runtime with PreferenceEKF scale?: PreferenceEKF training is roughly 5× faster than DeepEnsemble and over 40× faster than LLMCMC.The authors attribute this advantage primarily to EKF’s sequential updates, which avoid repeatedly processing previously seen queries.
  • 5.2 How does training runtime with PreferenceEKF scale?: PreferenceEKF scales more gracefully with posterior sample count and neural-network size while retaining favorable final test log-likelihood.Its lower-dimensional subspace update is faster than Dropout’s full-parameter update in the scaling experiments.
  • 5.3 Does PreferenceEKF lead to better model calibration?: PreferenceEKF has the lowest expected calibration error and the second-lowest Brier score, behind active DeepEnsemble.Both calibration metrics are evaluated on the Walker Medium Expert task.
  • 5.4 Ablation study on subspace construction: Random-projection subspaces can provide good reward-learning performance without an initial dataset, decoupling PreferenceEKF from reliance on SGD.The main experiments otherwise use SVD-based subspaces constructed from SGD iterates.

6 Conclusion

PreferenceEKF uses subspace-based extended Kalman filtering to improve active reward learning efficiency, runtime scaling, uncertainty calibration, and downstream policy performance. The approach remains limited by uncertainty about foundation-model-scale applicability and by its unimodal posterior assumption.

  • 6 Conclusion: PreferenceEKF achieved more sample-efficient active reward learning, similarly performant policy optimization, better runtime scaling, and better uncertainty calibration than four Bayesian deep learning methods.The comparison concerns active reward learning and offline RL policy optimization.
  • 6 Conclusion: The method’s applicability to foundation model-scale reward models remains unclear.The authors identify this as a limitation and future-work direction.
  • 6 Conclusion: The EKF’s unimodal Gaussian posterior may require alternatives to represent multimodal preference distributions.The limitation follows from the Gaussian distribution maintained by the extended Kalman filter.
  • 6 Conclusion: Future work will examine whether learned reward-model posteriors support policy exploration and mitigate reward hacking.The paper currently focuses primarily on sample-efficient reward modeling in RLHF.

Broader Impacts

The paper presents PreferenceEKF as an algorithm for active preference-based reward modeling in subjective-evaluation applications. It also warns that subspace reduction and EKF inference may introduce bias amplification or neglect minority preferences.

  • Broader Impacts: PreferenceEKF targets subjective-evaluation applications including natural language processing, personalized recommendations, and human-robot interaction.The authors describe the algorithm as improving the efficiency and accuracy of neural-network training in these applications.
  • Broader Impacts: Reduced-subspace inference and the extended Kalman filter may amplify bias or neglect minority preferences.The authors identify robustness and information loss from subspace reduction as areas for future investigation.

A Technical Appendices and Supplementary Material

The supplementary materials document the implementation stack, compute setup, and aggregation procedure used for experiments and figures.

  • A Technical Appendices and Supplementary Material: Experiments use JAX with Dynamax for EKF, Laplax for Laplace approximation, Blackjax for MCMC, and Unifloral for offline IQL.SciPy is used for statistical tests.
  • A Technical Appendices and Supplementary Material: Unless otherwise stated, experiments run on one node with 8 NVIDIA RTX A6000 GPUs via SLURM sharding.This describes the default computational setup.
  • A Technical Appendices and Supplementary Material: Aggregate figures pool each step across 12 tasks and n seeds per task, reporting the mean over 12n runs with standard-error or 95% bootstrap bounds.The procedure applies to figures such as Figures 1a, 1b, 3a, and A.14.

A.1 EKF with Bradley-Terry Likelihood

The appendix formulates preference learning as sequential Bayesian filtering with Gaussian dynamics and measurement noise, then applies an EKF using an identity parameter-dynamics model and Bradley-Terry preference likelihood. The resulting Gaussian belief is updated through prediction and measurement steps.

  • A.1 EKF with Bradley-Terry Likelihood: Sequential inference combines the previous parameter posterior with dynamics and measurement models after each preference example.The observed data consist of pairwise trajectory queries and binary preference labels.
  • A.1 EKF with Bradley-Terry Likelihood: Additive Gaussian dynamics and measurement noise make the nonlinear inference objective tractable with an EKF.The posterior is represented as a Gaussian with mean and covariance.
  • A.1 EKF with Bradley-Terry Likelihood: Preference learning uses identity parameter dynamics and a Bradley-Terry measurement equal to the model’s probability of preferring one trajectory over another.The probability is computed from the learned reward model’s trajectory returns.
  • A.1 EKF with Bradley-Terry Likelihood: The belief is over subspace coordinates rather than full parameters when subspace inference is used.The appendix uses θ broadly for parameters but specifies z for the subspace belief.
  • A.1 EKF with Bradley-Terry Likelihood: The EKF alternates prediction and update steps to revise the Gaussian belief after each new query and label.With identity dynamics, the dynamics Jacobian is the identity matrix.
  • A.1 EKF with Bradley-Terry Likelihood: The measurement Jacobian contains gradients of trajectory-preference probabilities with respect to reward-model parameters or subspace dimensions.For the Bernoulli Bradley-Terry output, the observation dimension is two.

A.1.1 On linearization of the Bradley-Terry Likelihood:

The Bradley-Terry preference likelihood is linearized around the current EKF parameter estimate, producing a Jacobian that measures parameter sensitivity. Its magnitude is largest for uncertain comparisons and shrinks when the model is confident, reducing subsequent Kalman updates.

  • Measurement function: The EKF measurement function predicts the preference probability for a trajectory pair using the Bradley-Terry sigmoid likelihood.The probability is computed from the sigmoid of the reward difference between the two trajectories.
  • Linearization: The Jacobian Hi is obtained by first-order Taylor expansion and the chain rule around the predicted parameter mean.It captures the sensitivity of the linearized measurement to reward-model parameters.
  • Jacobian structure: The Jacobian weights the difference between trajectory reward gradients by the sigmoid derivative σ′(z).Here z is the reward difference between the compared trajectories.
  • Uncertainty dependence: σ′(z) reaches 0.25 when the compared rewards are equal, where preference uncertainty is highest, and approaches zero as reward differences become large.Thus, uncertain comparisons produce larger measurement sensitivity, while confident comparisons produce smaller sensitivity.
  • Uncertainty dependence: When Hi is small under low preference uncertainty, the Kalman gain tends toward zero and parameter updates become small or absent.This links confidence in the preference likelihood to the size of sequential model updates.

A.1.2 On the locally Gaussian assumption of the Bradley-Terry Likelihood:

The Bradley-Terry likelihood is locally treated as Gaussian for EKF inference by replacing Bernoulli variance with a fixed Gaussian measurement covariance.

  • Local Gaussian assumption: The Bernoulli variance p(1 −p) is maximized at 0.25 when the preference probability is p = 0.5.This corresponds to maximal uncertainty about which trajectory is preferred.
  • Local Gaussian assumption: PreferenceEKF uses a constant Gaussian measurement covariance V = 0.07 · I to represent uncertainty under the locally Gaussian approximation.This covariance replaces the Bernoulli variance in the EKF representation.
  • Local Gaussian assumption: The fixed covariance roughly perturbs the predicted preference probability to account for uncertainty in the measurement model.The approximation is applied under a zero-mean Gaussian noise assumption.

A.1.3 On EKF hyperparameters:

EKF behavior depends on the dynamics, measurement, and initialization covariances, which control regularization and update strength. The selected noise scales must balance overfitting against underfitting, especially with noisy preference labels.

  • Hyperparameters: The EKF hyperparameters are dynamics covariance U, measurement covariance V, and initialization covariance W.The initial belief is specified as b0 = p(θ0) = N(0, W).
  • Hyperparameters: Weak parameter regularization is implemented with dynamics noise U = 0 to support sequential learning without overfitting previously observed data.The dynamics covariance controls how the parameter belief evolves between updates.
  • Measurement noise: Measurement noise addresses flipped preference labels, with synthetic experiments using approximately 5% −10% flipped labels and V = 0.07 · I.The covariance is intended to account for label error in pairwise preferences.
  • Hyperparameter sensitivity: Large W with small V causes overfitting through overly strong posterior updates, whereas small W with large V causes underfitting and little test-likelihood improvement.PreferenceEKF performance was sensitive to all noise hyperparameters.

A.2 Preference-based Reward Learning

The evaluation tests active preference learning, runtime, sampling, acquisition functions, and downstream offline-RL policy performance. PreferenceEKF often improves over baselines, but its advantages depend on statistical significance, acquisition choice, compute budget, and task setting.

  • Statistical testing: Active PreferenceEKF performs on par with its random variant, but does not significantly outperform it.The one-sided bootstrap analysis reports low statistical significance for the average advantage over random querying.
  • Statistical testing: Active PreferenceEKF outperforms active Bayesian baselines in normalized AUC, except for a tie with LLMCMC.Compared with LLMCMC, PreferenceEKF requires less posterior-inference time and does not store all previously observed queries.
  • Acquisition functions: Acquisition choice matters substantially: InfoGain performs best overall, whereas disagreement and entropy do not let active PreferenceEKF or LLMCMC outperform their random variants.The reported ordering is InfoGain, followed by disagreement, then entropy.
  • Model sampling: PreferenceEKF remains effective with only five posterior samples, indicating that its sample-efficiency advantage is not explained solely by using 100 samples.With M = 5 for all methods, PreferenceEKF still outperforms the alternatives in test log-likelihood.
  • Runtime and compute budget: PreferenceEKF retains a runtime lead under matched compute budgets, while SGD-based baselines fail to converge in the reduced-step setting.The original runtime comparison favors PreferenceEKF because it updates only the latest query, whereas baselines repeatedly optimize over all observed data.
  • Policy performance: Offline-RL policies trained from the learned reward models achieve similar rollout performance despite differences in reward-model log-likelihood.The authors report competitive policy performance rather than claiming that PreferenceEKF automatically produces stronger policies.
Loading 2609.04066v1…