🤖 machine learning

Subspace Inference Enables Efficient Active Reward Learning from Preferences

This paper introduces PreferenceEKF, a sample-efficient active learning method that leverages extended Kalman filtering within a low-dimensional parameter subspace to enable scalable uncertainty quantification for neural network reward models, thereby improving the efficiency and performance of reinforcement learning from human feedback.

Original authors: Yutai Zhou, Erdem Bıyık

Published 2026-09-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yutai Zhou, Erdem Bıyık

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, there is a persistent challenge known as the "sample inefficiency" of learning from human feedback. Imagine teaching a complex computer program to behave in a way that aligns with human values. The most powerful method currently available involves asking people to compare two different outcomes—such as two robot movements or two written responses—and stating which one they prefer. While this feedback is easy for humans to give, it is incredibly sparse; a single preference provides only a tiny sliver of information. To build a reliable model of what humans want, an algorithm must ask thousands of these questions. If the computer asks the wrong questions, it wastes time and money. If it asks the right ones, it learns much faster. The difficulty lies in knowing which questions will be the most informative. To do this, the computer needs to understand what it does not yet know, a concept called uncertainty. However, calculating this uncertainty for massive, modern neural networks is notoriously difficult and computationally expensive, often requiring the training of dozens of separate models just to get a rough estimate.

A team of researchers at the University of Southern California has developed a new approach to solve this bottleneck, allowing computers to learn from human preferences with far greater speed and efficiency. They introduced a method called PreferenceEKF, which treats the process of learning preferences as a continuous, step-by-step filtering problem rather than a massive, one-time calculation. Instead of trying to map out every possible variation of a giant neural network at once, the researchers realized that the network's behavior could be accurately tracked within a much smaller, lower-dimensional space. By focusing their calculations on this compact subspace, they were able to use a classic mathematical tool, the extended Kalman filter, to update the model's understanding in real-time as new answers arrived. This technique allowed them to generate thousands of different versions of the reward model instantly, without the heavy computational cost of training multiple independent networks.

The researchers tested their method against several existing techniques using a variety of standard benchmarks for robotic control and decision-making. They found that their approach was not only significantly faster—running up to forty times quicker than some of the most advanced alternatives—but also more accurate in its predictions. In experiments where the goal was to learn a reward model from a limited number of human comparisons, the new method consistently learned the correct preferences using fewer questions than the other methods. Furthermore, the models it produced were better calibrated, meaning the computer's confidence in its answers matched the actual accuracy of those answers more closely. This precision is vital for active learning, where the system must decide which question to ask next; if the system is unsure, it asks a question to resolve that uncertainty, and if it is confident, it moves on. The new method excelled at this balancing act, leading to reward models that could successfully train robotic policies to perform complex tasks, matching the performance of policies trained with much more expensive and time-consuming methods.

One of the most striking aspects of this work is how it changes the workflow of training these systems. Traditional methods often require the computer to re-train or re-evaluate its entire understanding of the world every time it receives a new piece of feedback, a process that becomes slower as the system grows larger. The new method, by contrast, updates its knowledge sequentially, incorporating only the latest piece of information while maintaining a running estimate of what it has learned so far. This allows the system to scale efficiently, handling larger neural networks and generating more samples of possible reward models without running out of memory or time. The researchers also demonstrated that this approach works even when starting with no initial data, using a random projection technique to build the necessary subspace from scratch, and it showed promise when applied to image-based tasks where the input data is much more complex than simple numbers.

While the method shows great promise, the researchers are careful to note its boundaries. The mathematical framework they used assumes that the preferences being learned come from a single, consistent source. When they tested the system with data from multiple different human annotators who might have conflicting views, the method struggled to capture the full complexity of those differing opinions. This suggests that while the approach is a powerful tool for streamlining the learning process, it is best suited for scenarios where a single, coherent set of preferences is being modeled. Nevertheless, the results indicate a significant step forward in making artificial intelligence more adaptable to human intent. By making the process of learning from feedback faster and more efficient, this work removes a major barrier to deploying intelligent systems in real-world settings, from personalized recommendations to autonomous robots, where the cost of human time is high and the need for rapid, accurate learning is critical.

Technical Summary: Subspace Inference Enables Efficient Active Reward Learning from Preferences

Problem Statement

Reinforcement Learning from Human Feedback (RLHF) is a dominant technique for aligning agents with human intentions, yet it suffers from poor sample efficiency. Since human preference feedback provides at most one bit of information per query, learning a reward model (RM) often requires thousands of comparisons, which is not scalable. Active learning addresses this by judiciously selecting queries to maximize information gain. However, effective active learning requires robust uncertainty quantification (UQ) of the reward model.

While Bayesian methods offer principled uncertainty representation, they struggle to scale to large neural network (NN) reward models due to the computational intractability of performing posterior inference over high-dimensional parameter spaces. Conversely, popular alternatives like ensemble methods and dropout are either computationally expensive (requiring training multiple independent models) or provide poor posterior approximations. The core challenge is to develop a method that enables scalable, sample-efficient active learning for neural network reward models without the overhead of training ensembles.

Methodology: PreferenceEKF

The authors propose PreferenceEKF, a sample-efficient approach that frames active preference learning as a sequential Bayesian filtering problem. The method leverages two key insights:

Sequential Filtering: Instead of batch training, the method updates the reward model posterior sequentially as new preference queries arrive, using an Extended Kalman Filter (EKF).
Subspace Inference: Recognizing that neural networks are overparameterized and that solutions often lie in a low-dimensional subspace, the method performs EKF inference within this subspace rather than the full parameter space.

Technical Workflow

Subspace Construction:
- An initial dataset is used to run Stochastic Gradient Descent (SGD) for a few iterations.
- The resulting parameter iterates are used to construct a low-dimensional subspace. This is done via Singular Value Decomposition (SVD) of the SGD iterates to obtain a projection matrix $A$ , or alternatively via random projections.
- The full parameter space $\theta$ is approximated by an affine mapping $\theta(z) = Az + \theta^*$ , where $z$ represents the low-dimensional subspace parameters ( $|z| \ll |\theta|$ ).
Sequential Bayesian Filtering (EKF):
- The method treats neural network parameters as hidden states evolving over time.
- Dynamics Model: Assumes an identity function with additive Gaussian noise ( $p(\theta_i | \theta_{i-1}) = \mathcal{N}(\theta_i | \theta_{i-1}, U)$ ).
- Measurement Model: Uses the Bradley-Terry (BT) model to predict preference labels given trajectory pairs. The BT likelihood is linearized around the current posterior mean to fit the EKF framework.
- Update Step: Upon receiving a new query-response pair $(Q_i, y_i)$ , the EKF updates the posterior distribution of the subspace parameters $z$ in closed form. The posterior remains Gaussian: $b_i = \mathcal{N}(\mu'_i, \Sigma'_i)$ .
Active Learning Loop:
- The method samples an arbitrary number of models from the learned subspace posterior.
- These samples are projected back to the full space to compute acquisition functions. The authors specifically utilize InfoGain (mutual information between the query label and model parameters) to select the most informative query.
- Unlike ensemble methods that retrain on all data, PreferenceEKF updates only on the most recent query, making it computationally efficient.

Key Contributions

First Subspace Filtering for NN Reward Models: The paper claims to be the first to leverage subspace filtering to train neural network reward models from preference feedback.
Scalable Bayesian Inference: By restricting inference to a low-dimensional subspace, the method makes the Extended Kalman Filter tractable for deep neural networks, avoiding the $O(|\theta|^2)$ complexity of full covariance matrices.
Efficient Acquisition Function Computation: The ability to sample an arbitrary number of models from the subspace posterior allows for the scalable computation of InfoGain, a state-of-the-art acquisition function previously restricted to low-dimensional models.
No Ensemble Training Overhead: Unlike DeepEnsembles, PreferenceEKF does not require training multiple independent models, significantly reducing runtime and memory usage.

Experimental Results

The authors evaluated PreferenceEKF on the D4RL and V-D4RL benchmarks across 12 tasks involving MuJoCo locomotion, Adroit manipulation, and Maze2D navigation. They compared against four baselines: DeepEnsemble, Dropout, Laplace Approximation, and Last-Layer MCMC (LLMCMC).

Sample Efficiency: PreferenceEKF performed on par with or better than all Bayesian deep learning baselines in terms of sample efficiency and final log-likelihood on held-out test queries.
Runtime and Scalability:
- PreferenceEKF achieved roughly 5x speedup compared to DeepEnsemble and over 40x speedup compared to LLMCMC.
- The method scales gracefully with both the number of posterior samples ( $M$ ) and the size of the neural network architecture, whereas baselines (especially ensembles) suffer from memory constraints or linear scaling issues.
Calibration: PreferenceEKF demonstrated superior model calibration, achieving the lowest Expected Calibration Error (ECE) among all methods and the second-lowest Brier score.
Policy Optimization: When the learned reward models were used for offline reinforcement learning (via Implicit Q-Learning), the resulting policies achieved performance competitive with those trained on reward models from other methods, often matching or slightly underperforming the ground-truth reward policy.
Robustness: The method remained effective even without an initial warm-up dataset (using random projections) and showed promise in sparse feedback settings (real robot data) and pixel-based tasks (using pretrained image embeddings).

Significance and Claims

The paper positions PreferenceEKF as a viable solution for scaling Bayesian deep learning to large neural network reward models in RLHF. The authors claim that their approach successfully balances the need for principled uncertainty quantification with computational tractability.

Efficiency: The primary contribution is improving the sample efficiency of preference-based reward learning while drastically reducing the computational cost of training and inference compared to existing Bayesian methods.
Practicality: The method enables the use of high-dimensional acquisition functions like InfoGain for neural networks, which was previously infeasible.
Limitations: The authors modestly acknowledge that the Gaussian assumption of EKF limits the method to single-annotator settings (unimodal posteriors) and that extending it to foundation model-scale reward models or multimodal preference distributions (e.g., diverse human annotators) remains an open challenge. They also note that while their reward models lead to competitive policies, the direct correlation between reward model log-likelihood and final policy performance is complex and not guaranteed to be improved solely by better reward learning.

In conclusion, the paper demonstrates that subspace inference via EKF offers a scalable, efficient, and well-calibrated alternative to ensemble-based and other Bayesian deep learning methods for active reward learning from human preferences.