Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Low-Rank Contextual Reinforcement Learning from Heterogeneous Human Feedback

Seong Jin Lee, Will Wei Sun, Yufeng Liu

TL;DR

This work addresses heterogeneity in human feedback for RLHF by introducing LoCo-RLHF, a framework that models personalized preferences with a contextual bilinear reward r(x,s,a) = x^T Θ^* φ(s,a) and imposes a low-rank structure Θ^* = U^* D^* (V^*)^T to reduce dimensionality. It develops the Pessimism in Reduced Subspace (PRS) algorithm, which first estimates a low-rank subspace via rank-constrained MLE (solved by Burer–Monteiro factorization), then projects onto a reduced space using rotation-truncation-vectorization (RTV), and finally optimizes a pessimistic policy by constructing a confidence set and minimizing over it. Theoretical guarantees show a sub-optimality gap bound of the form J(π^*) − J(π̂) = O( C^* sqrt( ((d_x + d_φ) r + log(1/δ)) / N ) ) with high probability, with notably tighter rates in low-rank settings compared to full-rank baselines; extensions recover known results in special grouping structures. Empirical results demonstrate that PRS achieves superior performance and robustness to distribution shifts in personalized RLHF tasks, especially when the underlying true rank is low, highlighting the framework’s scalability and practical impact for heterogeneous human feedback systems.

Abstract

Reinforcement learning from human feedback (RLHF) has become a cornerstone for aligning large language models with human preferences. However, the heterogeneity of human feedback, driven by diverse individual contexts and preferences, poses significant challenges for reward learning. To address this, we propose a Low-rank Contextual RLHF (LoCo-RLHF) framework that integrates contextual information to better model heterogeneous feedback while maintaining computational efficiency. Our approach builds on a contextual preference model, leveraging the intrinsic low-rank structure of the interaction between user contexts and query-answer pairs to mitigate the high dimensionality of feature representations. Furthermore, we address the challenge of distributional shifts in feedback through our Pessimism in Reduced Subspace (PRS) policy, inspired by pessimistic offline reinforcement learning techniques. We theoretically demonstrate that our policy achieves a tighter sub-optimality gap compared to existing methods. Extensive experiments validate the effectiveness of LoCo-RLHF, showcasing its superior performance in personalized RLHF settings and its robustness to distribution shifts.

Low-Rank Contextual Reinforcement Learning from Heterogeneous Human Feedback

TL;DR

This work addresses heterogeneity in human feedback for RLHF by introducing LoCo-RLHF, a framework that models personalized preferences with a contextual bilinear reward r(x,s,a) = x^T Θ^* φ(s,a) and imposes a low-rank structure Θ^* = U^* D^* (V^*)^T to reduce dimensionality. It develops the Pessimism in Reduced Subspace (PRS) algorithm, which first estimates a low-rank subspace via rank-constrained MLE (solved by Burer–Monteiro factorization), then projects onto a reduced space using rotation-truncation-vectorization (RTV), and finally optimizes a pessimistic policy by constructing a confidence set and minimizing over it. Theoretical guarantees show a sub-optimality gap bound of the form J(π^*) − J(π̂) = O( C^* sqrt( ((d_x + d_φ) r + log(1/δ)) / N ) ) with high probability, with notably tighter rates in low-rank settings compared to full-rank baselines; extensions recover known results in special grouping structures. Empirical results demonstrate that PRS achieves superior performance and robustness to distribution shifts in personalized RLHF tasks, especially when the underlying true rank is low, highlighting the framework’s scalability and practical impact for heterogeneous human feedback systems.

Abstract

Reinforcement learning from human feedback (RLHF) has become a cornerstone for aligning large language models with human preferences. However, the heterogeneity of human feedback, driven by diverse individual contexts and preferences, poses significant challenges for reward learning. To address this, we propose a Low-rank Contextual RLHF (LoCo-RLHF) framework that integrates contextual information to better model heterogeneous feedback while maintaining computational efficiency. Our approach builds on a contextual preference model, leveraging the intrinsic low-rank structure of the interaction between user contexts and query-answer pairs to mitigate the high dimensionality of feature representations. Furthermore, we address the challenge of distributional shifts in feedback through our Pessimism in Reduced Subspace (PRS) policy, inspired by pessimistic offline reinforcement learning techniques. We theoretically demonstrate that our policy achieves a tighter sub-optimality gap compared to existing methods. Extensive experiments validate the effectiveness of LoCo-RLHF, showcasing its superior performance in personalized RLHF settings and its robustness to distribution shifts.
Paper Structure (15 sections, 4 theorems, 26 equations, 6 figures, 2 algorithms)

This paper contains 15 sections, 4 theorems, 26 equations, 6 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Notations
  4. Problem Formulation
  5. Contextual Preference Model
  6. Value of Policy
  7. Algorithm
  8. Estimation of the Low-rank Subspace
  9. Reduction to the Low-rank Subspace
  10. Pessimism in the Reduced Low-rank Space
  11. Theory
  12. Assumptions
  13. Main Results
  14. Numerical Studies
  15. Conclusion

Key Result

Theorem 1

Suppose Assumptions assm:bounded_features and assm:cov hold. Let $r$ denote the rank of $\Theta^*$, with $\sigma_1 \ge \sigma_2 \ge \cdots \sigma_r > 0$ as its singular values. Assume we have an initial estimate $\Theta_0 \in \mathop{\mathrm{\mathbb{R}}}\limits^{d_x\times d_\phi}$ such that $\|\Thet with probability at least $1-\delta$ for some constant $C_1 > 0$.

Figures (6)

  • Figure 1: Illustration of personalization in RLHF.
  • Figure 2: Illustration of distribution shift in RLHF.
  • Figure 3: Illustration of the Low-rank Contextual Preference Model.
  • Figure 4: Sub optimality gap of our proposed PRS policy compared with MLE-Greedy and MLE-Pessimistic in the random action distribution environment and the imbalance action distribution environment.
  • Figure 5: Sub optimality gap of our proposed PRS policy compared with MLE-Greedy and MLE-Pessimistic when the imbalance differs in datasets.
  • ...and 1 more figures

Theorems & Definitions (8)

  • Remark 1: Flexibility of the bilinear model
  • Remark 2: Connection with zhong2024provable
  • Theorem 1: Informal Estimation Bound
  • Remark 3: Initial Estimate $\Theta_0$
  • Corollary 1: Error induced by Rotation-Truncation-Vectorization
  • Lemma 1: Confidence Bound on $\theta_{rtv}$
  • Theorem 2: Upper Bound of Sub-optimality
  • Remark 4: Discussion on Concentratability Coefficient