Low-Rank Contextual Reinforcement Learning from Heterogeneous Human Feedback

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.

💡 Research Summary

This paper tackles three intertwined challenges that arise when applying reinforcement learning from human feedback (RLHF) to large language models (LLMs): (1) the heterogeneity of human preferences across users, (2) the extremely high dimensionality of the feature space that combines LLM embeddings with user context attributes, and (3) distribution shift between offline collected feedback and the target deployment population. To address these issues, the authors propose a Low‑rank Contextual RLHF framework (LoCo‑RLHF) together with a Pessimism in Reduced Subspace (PRS) algorithm.

The core modeling idea is to replace the homogeneous reward function r(s,a)=θᵀϕ(s,a) with a bilinear contextual reward r(x,s,a)=xᵀΘϕ(s,a), where x∈ℝ^{dₓ} encodes user‑specific context (age, education, demographics, etc.) and ϕ(s,a)∈ℝ^{d_ϕ} is the embedding of a query‑answer pair extracted from a pre‑trained LLM. Θ∈ℝ^{dₓ×d_ϕ} captures the interaction between contexts and LLM features. Directly estimating Θ* would require O(dₓ·d_ϕ) parameters, which is infeasible for typical dimensions (e.g., dₓ≈100, d_ϕ≈4,000). The authors therefore assume that Θ* has a low‑rank structure: Θ*≈UΣVᵀ with rank r≪min(dₓ,d_ϕ). This reduces the effective number of parameters to (dₓ+d_ϕ)·r while preserving the essential cross‑terms.

The PRS algorithm proceeds in three stages. First, a rank‑constrained maximum‑likelihood estimator (MLE) is solved using the Burer‑Monteiro factorization and alternating gradient descent, yielding estimates of the low‑dimensional subspaces U and V. After convergence, a singular‑value decomposition extracts orthogonal bases that define the reduced space. Second, the authors construct confidence bounds that incorporate both the statistical error of the MLE and the additional uncertainty introduced by projecting onto the estimated subspace. This novel analytical framework yields a high‑probability lower bound on the true reward for any context‑state‑action triple within the reduced space. Third, a pessimistic policy is derived by optimizing the worst‑case (lower‑bound) reward, following the pessimism‑principle popular in offline RL. The resulting policy is robust to distribution shift because it deliberately avoids over‑optimistic exploitation of poorly supported regions of the offline dataset.

Theoretical contributions include two main results. Theorem 1 shows that the rank‑constrained MLE attains an estimation error of order Õ(√((dₓ+d_ϕ)·r / n)) in Frobenius norm, where n is the number of pairwise preference labels. Theorem 2 proves that the sub‑optimality gap of the PRS policy scales as Õ(((dₓ+d_ϕ)·r·log(1/δ))/n) with probability at least 1−δ. This bound improves upon the naïve Õ(dₓ·d_ϕ·log(1/δ)/n) bound of existing methods, especially when r is much smaller than the ambient dimensions. In the special case where user preferences follow a pre‑defined group structure, the bound matches that of Zhong et al. (2024), demonstrating tightness.

Empirical evaluation comprises synthetic experiments and a real‑world benchmark called PersonalLLM. In synthetic tests, the authors vary rank, dimensionality, and the distribution of offline feedback. PRS consistently yields smaller sub‑optimality gaps than a greedy policy and a pessimistic policy built on an unconstrained MLE, with gains amplified in low‑rank, high‑dimensional regimes. On the PersonalLLM dataset, which contains millions of human preference judgments together with demographic metadata, LoCo‑RLHF with PRS outperforms a standard homogeneous reward model across several metrics: (i) higher A/B test win rates for generated answers, (ii) improved human‑rated quality scores, and (iii) better alignment with user‑specific satisfaction surveys. Additional ablations that inject noisy, irrelevant features show that PRS remains stable, whereas baseline methods degrade sharply, confirming the robustness of the low‑rank approach.

Related work is surveyed across three domains: (a) heterogeneous‑expert RLHF, where prior methods either train separate reward models per group or aggregate them via welfare functions but lack theoretical guarantees; (b) low‑rank LLM fine‑tuning, which focuses on compressing model weights rather than reward learning; and (c) uncertainty quantification for low‑rank matrix estimation, which traditionally provides asymptotic confidence sets. The present contribution uniquely blends these strands by delivering a provably efficient, personalized RLHF pipeline that is both computationally tractable and theoretically sound.

In conclusion, the paper introduces the first provable low‑rank contextual RLHF framework. By jointly learning a compact representation of user‑context interactions, rigorously quantifying estimation uncertainty, and applying pessimism in the reduced space, it achieves superior personalization, resilience to distribution shift, and scalability to high‑dimensional LLM settings. Future directions include extending the model to dynamic contexts, integrating online updates, and exploring non‑linear low‑rank embeddings via deep networks, thereby moving closer to truly adaptive, safe, and user‑centric LLM deployment.