Low-Rank Contextual Reinforcement Learning from Heterogeneous Human Feedback

Our work is closely related to, yet distinct from, recent studies on RLHF with heterogeneous experts and low-rank LLMs.

• •

  RLHF with Heterogeneous Experts: The incorporation of human feedback into the RL framework, introduced by Christiano et al., (2017), has gained significant attention in recent years. Recent studies, such as Ouyang et al., (2022) and Bai et al., (2022), highlight that experts may disagree on optimal responses, and biases in offline datasets can disadvantage underrepresented groups (Prabhakaran et al., , 2021; Feffer et al., , 2023; Kirk et al., , 2023). To address the heterogeneity of human preferences, various approaches have been proposed. One common method is to impose a group structure on individuals and train separate reward models for each group, as seen in Park et al., (2024), and Ramesh et al., (2024). These models are then aggregated using techniques such as social welfare functions, Pareto-optimal solutions (Boldi et al., , 2024), or robust optimization (Chakraborty et al., , 2024). Another line of research focuses on personalized models tailored to individual preferences (Jang et al., , 2023; Kirk et al., , 2023; Poddar et al., , 2024). However, most of these approaches lack theoretical guarantees for their performance. The most closely related work is Zhong et al., (2024), which proposes a multi-party RLHF framework that models and aggregates preferences across multiple individuals. While their method addresses heterogeneity through a pre-defined grouping structure, it requires a predetermined clustering of individual preferences and does not fully account for individual-specific contextual information. In contrast, our work extends this research by introducing a low-rank contextual preference model that incorporates expert characteristics and learns the latent low-rank representation of the reward model.

• •

  Low-rank LLMs: To improve LLM performance, many models leverage high-dimensional parameter spaces (Ouyang et al., , 2022; Bai et al., , 2022). As dimensionality increases, low-rank approximations have been proposed for fine-tuning (Hu et al., , 2021; Ding et al., , 2023; Dettmers et al., , 2024), supported by evidence of intrinsic low-dimensional structures within high-dimensional spaces (Aghajanyan et al., , 2020; Sidahmed et al., , 2024). These models primarily focus on low-rank adoption during pre-training or fine-tuning. In contrast, our work addresses the heterogeneity of human feedback in the reward learning step of RLHF. Additionally, drawing from statistical theory on high-dimensional problems (Wainwright, , 2019), we provide a rigorous theoretical analysis. While low-rank structures have been studied in bandit and RL problems (Jun et al., , 2019; Kang et al., , 2022; Cai et al., , 2023; Stojanovic et al., , 2024; Duan et al., , 2024), their application in RLHF remains underexplored.

In summary, our work presents the first provable low-rank contextual RLHF framework that tackles personalization, distribution shift, and high-dimensionality in heterogeneous human preferences.

We introduce the notations used throughout the paper. Let $[n]$ denote the set $\{1,2,\ldots,n\}$. The Frobenius norm of a matrix is denoted by $\|\cdot\|_{F}$, and the matrix-induced norm $\|v\|_{W}=\sqrt{v^{\top}Wv}$ is denoted by $\|\cdot\|_{W}$. The ball induced by the Frobenius norm, centered at $\Theta_{0}$ with radius $r$ is denoted by $\operatorname*{\mathcal{B}}(\Theta_{0},r):=\{\Theta:\|\Theta-\Theta_{0}\|_{F}\leq r\}$. For a vector $v$, $(v)_{i}$ represents its $i$-th component. We write $\Sigma_{1}\succ\Sigma_{2}$ if $\Sigma_{1}-\Sigma_{2}$ is positive definite. Denote Kronecker product of two matrices as $\otimes$.

To address the heterogeneity in human preferences regarding queries and their corresponding answers, we introduce the contextual preference model. Let $\operatorname*{\mathcal{S}}$ denote the state (query) space, $\operatorname*{\mathcal{A}}$ the action (answer) space, and $\operatorname*{\mathcal{X}}$ the context space. For $t\in[N]$, a human expert with context $x_{t}\in\operatorname*{\mathcal{X}}$ is presented with a state $s_{t}\in\operatorname*{\mathcal{S}}$ and two actions, $a_{t0},a_{t1}\in\operatorname*{\mathcal{A}}$. The expert selects the preferred action, where $y_{t}=1$ indicates a preference for $a_{t1}$ over $a_{t0}$, and $y_{t}=0$ indicates a preference for $a_{t0}$. The offline dataset comprises a collection of (i) contexts, (ii) states, (iii) action pairs, and (iv) expert preferences, represented as ${(x_{t},s_{t},a_{t0},a_{t1},y_{t}):t\in[N]}$ for each individual $t$.

The Bradley-Terry-Luce (BTL) model (Bradley and Terry, , 1952) has been widely used to capture how preferences influence the choices of experts. Let $r_{0}$ and $r_{1}$ denote the preferences associated with actions $a_{0}$ and $a_{1}$, respectively. The BTL model assumes that the probability of selecting an action is proportional to the exponential of its preference. Specifically, the probabilities are given by:

The homogeneous preference model (Ouyang et al., , 2022; Achiam et al., , 2023; Touvron et al., , 2023) assumes that the probability of choosing an action depends solely on the query and the available actions, making it identical across individuals. To account for heterogeneity, we introduce a contextual preference model, which reflects individual-specific preferences. Specifically, we model the reward function using a bilinear form:

where $x\in\operatorname*{\mathcal{X}}$ represents the individual’s context, $\Theta^{*}\in\mathbb{R}^{d_{x}\times d_{\phi}}$ is a parameter matrix, and $\phi(s,a)$ is the feature embedding for the query-action pair. Under this model, the probability of selecting action $a_{1}$ over $a_{0}$, given query $s$ and context $x$, is:

This formulation introduces variability in the choice probabilities by incorporating the context $x$, thereby capturing the heterogeneity among individuals.

The high dimensionality of context and feature embeddings poses significant computational challenges, as estimating the matrix parameter scales with the product of their dimensions. Inspired by the successful application of low-rank approximations in reinforcement learning (Jun et al., , 2019; Hu et al., , 2021), we impose a low-rank structure on the parameter matrix. Let $r$ be the rank of the parameter matrix $\Theta^{*}$, and let $\Theta^{*}=U^{*}D^{*}(V^{*})^{\top}$ represent its singular value decomposition (SVD) of $\Theta^{*}$, where $U^{*}$ and $V^{*}$ are square matrices, and $D^{*}\in\operatorname*{\mathbb{R}}^{d_{x}\times d_{\phi}}$ with $\{d_{i}:i\in[r]\}$ as its $r$ non-zero diagonal entries. The reward function $r_{\Theta^{*}}(x,s,a)$ can then be expressed as:

which simplifies to:

This formulation demonstrates that $U^{*}$ and $V^{*}$ project the original features into a low-dimensional subspace, significantly reducing computational complexity, as illustrated in Figure 3.

Next, we show how the contextual preference model effectively tackles personalization and distribution shift challenges. Let $\rho$ be the joint distribution of context $x$ and state $s$ in the evaluation set. In the personalization problem, the individual feature $x$ is available for the policy, enabling the policy to adapt its responses based on the context of each individual. Formally, the policy is defined as $\pi:\operatorname*{\mathcal{X}}\times\operatorname*{\mathcal{S}}\mapsto\operatorname*{\mathcal{A}}$, allowing it to provide tailored answers for different contexts. The value of the policy is then defined as:

In the personalization problem, the objective is to determine an individualized policy $\pi(x,s)$ that provides personalized responses for each individual with context $x$, even when the the same query $s$ is presented. The goal is to maximize the value of the policy.

The second case is the distribution shift problem. Unlike the personalization problem, the features of individuals are unavailable during evaluation in this second case. In this case, the policy is defined as a function of $s$ alone, i.e., $\pi(s)$. The corresponding value is given by

In the distribution shift problem, the objective is to find a global policy $\pi(s)$ that provides a unified response for all individuals, maximizing the expected value of the policy. It is important to note that the evaluation distribution $\rho$ may differ from the offline training distribution. If the contextual preference model is not considered, a homogeneous reward model would optimize the policy to maximize value based on the offline data distribution, potentially leading to suboptimal performance in the evaluation environment.

For brevity, we use the notation $r_{\theta}(x,s,\pi)$ to represent $r_{\theta}(x,s,\pi(x,s))$ in the personalization problem and $r_{\theta}(x,s,\pi(s))$ in the distribution shift problem. The sub-optimality of a policy is defined as

where $\pi^{*}=\operatorname*{arg\,max}_{\pi}J(\pi)$ denotes the optimal policy. The objective is to determine a policy that minimizes the sub-optimality by leveraging the offline dataset.

We begin by splitting the dataset into two partitions: the first $N_{0}$ observations are used to estimate the low-rank subspace, while the remaining $N-N_{0}$ observations are used to estimate the parameter and quantify the uncertainty in the reduced space. Given the dataset $\{(x_{t},s_{t},a_{t0},a_{t1},y_{t}:t\in[N_{0}]\}$, the contextual preference model in the BTL model (1) implies that the negative log-likelihood is given by:

Based on it, the rank-constrained optimization problem can be formulated as

which can be solved using the Burer-Monteiro formulation (Burer and Monteiro, , 2003). This approach factorizes $\Theta=UV^{\top}$, where $U\in\operatorname*{\mathbb{R}}^{d_{x}\times r}$ and $V\in\operatorname*{\mathbb{R}}^{d_{\phi}\times r}$. To ensure identifiability, we introduce a regularization term following Zheng and Lafferty, (2016). The resulting optimization problem is formulated as:

To solve this non-convex optimization problem, we employ the alternating Factored Gradient Descent (FGD) method. FGD performs gradient descent on each component iteratively:

where $\eta$ denotes the learning rate, and $\nabla_{U},\nabla_{V}$ are the gradients with respect to $U$ and $V$, respectively. By alternately updating each component, the estimated parameter converges to the true parameter, provided that certain conditions on the initial estimate and the learning rate are satisfied (Zheng and Lafferty, , 2016; Zhang et al., , 2023). The full details of the FGD procedure are outlined in Algorithm 1.

In our experiments, we compute the initial estimate $\Theta_{0}$ using the unconstrained maximum likelihood estimator. Convex optimization algorithms, such as L-BFGS (Liu and Nocedal, , 1989), or gradient-based methods can be applied to obtain this initial parameter estimate.

RL algorithms leverage uncertainty quantification to improve performance, such as the “optimism in the face of uncertainty” approach for online RL (Abbasi-Yadkori et al., , 2011) and pessimistic policies for offline RL (Jin et al., , 2021; Rashidinejad et al., , 2021). Existing uncertainty quantification methods are primarily designed for vector parameters, making them less suited to our high-dimensional matrix parameter space $\operatorname*{\mathbb{R}}^{d_{x}\times d_{\phi}}$. The large number of components in this space as well as the low-rank structure of the matrix parameter create significant challenges for effective uncertainty quantification. To address this, we propose the “rotation-truncation-vectorization” (RTV) method, which projects the parameter onto a low-dimensional subspace, reducing its dimensionality and enabling efficient uncertainty quantification.

To illustrate how the RTV method reduces dimensionality, consider the estimate $\hat{\Theta}$ obtained from Algorithm 1. Let $\hat{\Theta}=\hat{U}\hat{D}\hat{V}^{\top}$ be its SVD. Similarly, let $\Theta^{*}=\hat{U}^{*}D^{*}(V^{*})^{\top}$ represent the SVD of the true parameter $\Theta^{*}$. Intuitively, if $\hat{U}$ and $\hat{V}$ are accurate estimates of $U^{*}$ and $V^{*}$, the rotated matrix $\hat{U}^{\top}\Theta^{*}\hat{V}$ is expected to a close approximate of the rank $r$ diagonal matrix $D^{*}$.

For a detailed explanation, let $(U^{*}_{1},U_{2}^{*})$, $(V_{1}^{*},V_{2}^{*})$, $(\hat{U}_{1},\hat{U}_{2})$ and $(\hat{V}_{1},\hat{V}_{2})$ be partitions of $U^{*},V^{*},\hat{U}$ and $\hat{V}$, where $(\cdot)_{1}$ represent the first $r$ columns and $(\cdot)_{2}$ denotes the remaining columns of the matrix. Let $\Theta^{*}_{r}$ be the rotation of $\Theta^{*}$ with respect to the estimated subspace given by:

where $D_{11}^{*}$ is the upper-left $r\times r$ sub matrix of $D^{*}$.

We employ the “subtraction method” commonly used in the low-rank matrix bandit literature (Kveton et al., , 2017; Jun et al., , 2019; Kang et al., , 2022), to reduce the dimension. When $\hat{U},\hat{V}$ are accurate estimates of $U^{*},V^{*}$, the terms $\hat{U}_{2}^{\top}U_{1}^{*}$ and $\hat{V}_{2}^{\top}V_{1}^{*}$ are expected to be negligible, as $U^{*}$ and $V^{*}$ are orthonormal matrices. Therefore, the product $(\hat{U}_{2}^{\top}U_{1}^{*})(\hat{V}_{2}^{\top}V_{1}^{*})^{\top}$ is also expected to be negligible. Based on this intuition, we introduce the concept of “rotation-truncation-vectorization” (RTV) for matrices. For a given matrix $\Theta\in\operatorname*{\mathbb{R}}^{d_{x}\times d_{\phi}}$, let $\Theta_{r}=U^{\top}\Theta V$ be the rotation with respect to orthonormal matrices $U\in\operatorname*{\mathbb{R}}^{d_{x}\times d_{x}}$ and $V\in\operatorname*{\mathbb{R}}^{d_{\phi}\times d_{\phi}}$. Let $(\Theta_{r})_{11}\in\operatorname*{\mathbb{R}}^{r\times r}$, $(\Theta_{r})_{12}\in\operatorname*{\mathbb{R}}^{r\times(d_{\phi}-r)}$, $(\Theta_{r})_{21}\in\operatorname*{\mathbb{R}}^{(d_{x}-r)\times r}$, $(\Theta_{r})_{22}\in\operatorname*{\mathbb{R}}^{(d_{x}-r)\times(d_{\phi}-r)}$ be the sub-block matrices of $\Theta_{r}$. Finally we define the “rotation-truncation-vectorization” of $\Theta$ as

where “$rtv$” is an abbreviation of “rotation-truncation-vectorization”, indicating that the original matrix has been rotated with respect to $U$ and $V$, then truncated and vectorized. Since $(\Theta_{r})^{*}_{22}=(\hat{U}_{2}^{\top}U_{1}^{*})D_{11}^{*}(\hat{V}_{2}^{\top}V_{1}^{*})^{\top}$ is expected to be negligible, the estimation of $\theta^{*}_{rtv}$ should suffice for approximating $\Theta^{*}$, with minimal bias. By reducing the dimension from $d_{x}d_{\phi}$ for $\Theta^{*}$ to $r(d_{x}+d_{\phi})-r^{2}$ of $\theta^{*}_{rtv}$, algorithms operating in the reduced space gain significant computational advantages.

The contextual preference model can be reformulated as:

Here, $(\hat{U}^{\top}x)(\hat{V}^{\top}\phi(s,a))^{\top}=\hat{U}^{\top}(x\phi(s,a)^{\top})\hat{V}$ is the rotation of $x\phi(s,a)^{\top}$ with respect to $\hat{U}$ and $\hat{V}$. Let $Z=x\phi(s,a)^{\top}$ and $z_{rtv}$ denote the rotation-truncation-vectorization of $Z$. Then, the contextual reward function can be expressed as:

where the approximation holds as $(\Theta^{*}_{r})_{22}\approx 0$. This reformulation allows us to focus on estimating the low-dimensional parameter $\theta^{*}_{rtv}\in\operatorname*{\mathbb{R}}^{k}$, where $k=(d_{x}+d_{\phi})r-r^{2}$ rather than the full matrix $\Theta^{*}\in\operatorname*{\mathbb{R}}^{d_{x}\times d_{\phi}}$. This significantly reduces the computational complexity while maintaining the model’s accuracy.

After reducing the problem to a low-dimensional subspace, we re-estimate the parameter in this reduced space. For the remaining $N-N_{0}$ observations, define $Z_{t}=x_{t}(\phi(s_{t},a_{t1})-\phi(s_{t},a_{t0}))^{\top}$ and let $z_{t,rtv}\in\operatorname*{\mathbb{R}}^{k}$ be the rotation-truncation-vectorization of $Z_{t}$ for $t>N_{0}$. The approximate negative log-likelihood function in the reduced space is defined as:

where $\theta_{rtv}\in\operatorname*{\mathbb{R}}^{k}$. As the approximate likelihood function is convex with respect to $\theta_{rtv}$, convex optimization methods can be applied to obtain the estimate $\hat{\theta}_{rtv}$.

Finally, we apply the pessimistic approach, which is widely used in offline reinforcement learning (Li et al., , 2022; Zhu et al., , 2023). Instead of the greedy approach that maximizes the estimated reward function based on the estimated parameter, the pessimistic approach maximizes the pessimistic reward. This reward is defined by penalizing the estimated reward based on the uncertainty quantification of the parameter. Intuitively, actions are penalized if the uncertainty in the estimated reward is high. This approach mitigates the risk of selecting suboptimal actions that may result from inaccurate estimates due to the distribution of the offline dataset.

To quantify uncertainty, we define the empirical second moment in the reduced space: $W=\frac{1}{n_{0}}\sum_{t=1}^{n_{0}}z_{t,rtv}z_{t,rtv}^{\top}$. Using this, we construct the confidence set $\{\theta_{rtv}:\|\theta_{rtv}-\hat{\theta}_{rtv}\|_{W}\leq C\}$ where $\|\cdot\|_{W}$ represents the Mahalanobis norm weighted by $W$, and $C$ is the confidence threshold.

To define the pessimistic reward based on this confidence set, we introduce the approximate reward function with respect to $\theta_{rtv}$:

where $z_{rtv}(x,s,a)$ is the rotation-truncation-vectorization of $x\phi(s,a)^{\top}$. The pessimistic value function is then defined as:

The complete PRS algorithm is summarized in Algorithm 2.

In case where the policy can be parametrized as $\pi_{\theta}$, finding the pessimistic policy can be solved via gradient based methods.

We begin by introducing the boundedness condition and the positive definite second moment condition of the features and the parameter.

The boundedness assumption is common in contextual reinforcement learning problems (Li et al., , 2017; Zhu et al., , 2023; Zhan et al., , 2023; Zhong et al., , 2024). The invertibility of the second moment of the feature is required to guarantee that the parameter is identifiable. Suppose $\operatorname*{\mathbb{E}}_{P}x_{t}x_{t}^{\top}=\Sigma_{x}\succ 0$ and $\operatorname*{\mathbb{E}}_{P}(\phi(s_{t},a_{t1})-\phi(s_{t},a_{t0}))(\phi(s_{t},a_{t1})-\phi(s_{t},a_{t0}))^{\top}=\Sigma_{\phi}\succ 0$, with $x_{t}\perp(s_{t},a_{t0},a_{t1})$. Then $\Sigma=\Sigma\phi\otimes\Sigma_{x}$ and therefore $\Sigma\succ 0$. In other words, if the features from the state-action pairs and contexts both have positive definite second moment and are independent (Zhu et al., , 2023; Zhong et al., , 2024), Assumption 2 is satisfied. When $x_{t}$ is uniformly distributed over standard unit vectors of $\operatorname*{\mathbb{R}}^{M}$, it becomes equivalent with the problem in Zhong et al., (2024), where they assume each query-answer pair is given to each of the $M$ user types.

We first present the estimation bound of the low-rank estimator.

The estimation error bound is proportional to $\sqrt{r}$, whereas without the low-rank assumption, it would scale with $\sqrt{\min\{d_{x},d_{\phi}\}}$. This highlights the critical role of the low-rank assumption in reducing estimation error. This rate aligns with bounds in the low-rank matrix estimation literature, particularly under the assumption of bounded $\ell_{2}$-norms for the context $x$ and feature embedding $\phi(s,a)$. Specifically, if the max-norm of the parameter matrix is bounded, the $\ell_{2}$-norm bound scales with $d_{x}$ and $d_{\phi}$, consistent with results in Negahban and Wainwright, (2011), Xia and Yuan, (2021), and Zhu et al., (2022). Building on this result, we derive the following corollary concerning the error induced by the rotation-truncation-vectorization process.

The details of the proof are provided in Appendices LABEL:sect:EstBound and LABEL:sect:rtv_error. Using these results, we establish the confidence bounds for the estimator $\hat{\theta}_{rtv}$ in the reduced space.

Finally, we derive an upper bound on the sub-optimality gap for the proposed policy.

In low-rank case where $r\ll\min\{d_{x},d_{\phi}\}$, our bound represents a significant improvement over $O\Big{(}C^{*}\cdot\sqrt{\frac{d_{x}d_{\phi}+\log(1/\delta)}{n}}\Big{)}$ achieved by existing methods (Zhu et al., , 2023). Moreover, the setting in Zhong et al., (2024) is a special case of our framework when the true clustering structure of human feedback is known with $M=d_{x}$ and $d=d_{\phi}$. They derived a sub-optimality rate of

where $\tilde{C}^{*}=\max_{m}\|(\Sigma_{m}^{-1/2}\operatorname*{\mathbb{E}}_{s\sim\rho}\phi(s,\pi^{*}(s))\|_{2}$. Assume $r\ll M\ll d$, and ignoring logarithmic factors, their rate simplifies to $M\cdot\sqrt{\frac{(M+d)r^{2}+r\log(\frac{M}{\delta})}{Mn}}\cdot\tilde{C}^{*}$. Noting that their value function is defined as the sum over $M$ types for the utilitarian welfare function, and we can verify that our rate of $\tilde{C}^{*}\cdot\sqrt{\frac{(M+d)r+\log(\frac{1}{\delta})}{n}}$ matches this bound in this special case.