Iterative Preference Learning from Human Feedback: Bridging Theory and Practice for RLHF under KL-Constraint

Introduction

Reinforcement Learning from Human Feedback (RLHF) (Christiano et al., 2017; Ziegler et al., 2019) has emerged as a powerful paradigm to align modern generative models like Large Language Models (LLMs) and diffusion models with human values and preferences. This approach has shown significant effectiveness in applications such as ChatGPT (OpenAI, 2023), Claude (Anthropic, 2023), Bard (Google, 2023), and LLaMA2 (Touvron et al., 2023), by making the built AI system helpful, harmless, honest and controllable (Ouyang et al., 2022; Bai et al., 2022a).

Despite its effectiveness, RLHF’s implementation often involves ad-hoc practices and extensive algorithmic tuning in the entire pipeline, including preference data collection (it is hard to select representative humans (Bai et al., 2022a), larger language models (Wang et al., 2024) or program compiler (Wang et al., 2023b) to provide feedback), preference/reward modeling (reward misspecification and misgeneralization (Hong et al., 2022; Gao et al., 2023)), and model optimization (instability of training (Choshen et al., 2019) and distribution shift issue (Michaud et al., 2020; Tien et al., 2022)). Meanwhile, the resulting models of RLHF typically suffer from issues like performance degeneration if we impose strong optimization pressure toward an imperfect reward function (Michaud et al., 2020; Tien et al., 2022; Gao et al., 2023), which contains bias and approximation error from the data collection and preference modeling (Gao et al., 2023; Wang et al., 2023d). Casper et al. (2023) also discussed many other challenges of RLHF. Thus, it is important to understand the mathematical principle of the RLHF process, as well as the connections among its different steps, which should be able to motivate future algorithmic design in principle.

In current RLHF theory, the agent’s objective is to maximize an observed reward function, with the optimal policy typically being deterministic and reward-greedy (Agarwal et al., 2019). However, in practical RLHF applications, merely maximizing the reward function is often insufficient and probably results in overfitting, as the generative model must simultaneously ensure both diversity and high fidelity in its outputs. A deterministic maximizer of the reward tends to compromise on these aspects significantly. For example, the maximizer of the “safety reward” tends to avoid providing answers all the time, which contradicts the LLM’s training objective. The situation worsens due to bias and approximation errors in reward modeling, leading to the critical problem of reward hacking, where the model often repeats superfluous, pleasing yet irrelevant words to appease the reward model (Michaud et al., 2020; Tien et al., 2022; Casper et al., 2023). Thus, it is important to model diversity and high fidelity in the theoretical framework beyond the reward. Notably, the most widely used mathematical objective function for this goal can be regarded as a reverse-KL regularized contextual bandit problem (Ziegler et al., 2019; Wu et al., 2021a; Ouyang et al., 2022; Rafailov et al., 2023; Liu et al., 2023a). The KL regularized contextual bandit additionally imposes a constraint that the optimal policy cannot move too far away from the original policy (i.e. the starting checkpoint of the LLM). A major difference between this objective function from traditional contextual bandit (Langford & Zhang, 2007) is that the optimal policy is stochastic, which is closer to the practical generative models. See an intuitive illustration why such a target is appealing in Figure 1. Despite numerous proposed procedures for this formulation, a rigorous theoretical analysis remains open. This paper provides a theoretical analysis of the regularized contextual bandit problem in both offline and online settings, aiming to inform and motivate practical algorithmic designs. Our contributions are summarized as follows:

We formally formulate the RLHF process as a reverse-KL regularized contextual bandit problem in RLHF theory, which more accurately reflects real-world alignment practices (Ouyang et al., 2022; Bai et al., 2022a; Rafailov et al., 2023) compared to existing theoretical frameworks. Meanwhile, we deliver a comprehensive theoretical analysis in offline, online, and hybrid settings for the formulated framework, where the three settings are complementary to each other and hold their own values in practical applications;

We introduce algorithms designed to address this formulated problem, which incorporate new uncertainty estimation or version space construction, and different non-symmetric exploration structures to handle the introduced KL penalty, as well as the challenges of preference learning;

Moving towards practical applications, we demonstrate that the proposed algorithms can be practically implemented and empirically outperform existing strong baselines like DPO (Rafailov et al., 2023) and RSO (Liu et al., 2023a) in real-world LLM experiments.

There is a rich literature in RLHF and we refer the interested readers to the survey papers like Casper et al. (2023) for a more comprehensive review. We focus on the papers that are most related to our work here.

RLHF has attracted considerable attention in the past few years, especially after its tremendous success in ChatGPT (OpenAI, 2023). We refer interested readers to Wirth et al. (2017); Casper et al. (2023) for a detailed survey but focus on the most related works here. The standard RLHF was popularized by Christiano et al. (2017), which served to direct the attention of the RL community to the preference-based feedback. The most popular and standard RLHF framework is outlined in the InstructGPT paper (Ouyang et al., 2022), Claude (Bai et al., 2022a) and the LLaMA2 report (Touvron et al., 2023) in detail, which typically consists of three steps starting from the pretrained model: supervised finetuning, reward modeling, and reward optimization. The effectiveness of this framework has been showcased by many recent generative models, like ChatGPT (OpenAI, 2023), Bard (Google, 2023), Claude (Anthropic, 2023), and LLaMA2 (Touvron et al., 2023). However, it is also noteworthy to indicate that the RLHF process often leads to degeneration in the performance of generation, commonly referred to as the “alignment tax” in the literature (Askell et al., 2021). This is usually because of the imperfection of the reward model and the model can make use of these imperfections to chase for a high reward. This phenomenon is referred to as the reward hacking (Michaud et al., 2020; Tien et al., 2022). It is also possible to apply RLHF to general generative models, like the diffusion model (Hao et al., 2022; Wu et al., 2023; Lee et al., 2023; Dong et al., 2023). In this work, we use the terminology and analysis of LLMs for better illustration, and defer the study of general generative models to future work.

RLHF algorithms. Proximal Policy Optimization (PPO) (Schulman et al., 2017) is the most well-known algorithm in LLM alignment literature. However, its instability, inefficiency, and sensitivity to hyperparameters (Choshen et al., 2019) and code-level optimizations (Engstrom et al., 2020) present significant challenges in tuning for optimal performance and its tremendous success in Chat-GPT4 (OpenAI, 2023) has not been widely reproduced so far. Additionally, it often necessitates incorporating an extra reward model, a value network (known as a critic), and a reference model, potentially as large as the aligned LLM (Ouyang et al., 2022; Touvron et al., 2023). This imposes a significant demand on GPU memory resources. Thus, researchers have attempted to design alternative approaches for LLM alignment to resolve the aforementioned issues. Dong et al. (2023); Yuan et al. (2023); Touvron et al. (2023); Gulcehre et al. (2023) propose reward ranked finetuning (RAFT) (also known as the iterative finetuning, rejection sampling finetuning) by iteratively learning from the best-of-n policy (Nakano et al., 2021) to maximize the reward, which is a stable baseline with minimal hyper-parameter configuration and was applied to the alignment of LLaMA2 project. There is also a line of work focusing on deriving an algorithm from the KL-regularized formulation (Rafailov et al., 2023; Zhu et al., 2023b; Wang et al., 2023a; Liu et al., 2023a; Li et al., 2023a). Among them, Direct Preference Optimization (DPO) (Rafailov et al., 2023) has emerged as an attractive alternative approach to PPO with notable stability and competitive performance. The innovative idea of DPO is to train the LLMs directly as a reward model based on the offline preference dataset and bypassing the reward modeling. Similar to DPO, there are also other works aiming to optimize the LLMs directly from the preference data, including (Zhao et al., 2023; Azar et al., 2023), and has sparked considerable debate on whether reward modeling, as well as RL, is necessary for alignment. However, while these algorithms are partly inspired by mathematical principles and intuitions, a comprehensive theoretical analysis remains open.

Theoretical study of RLHF. The theoretical understanding of RLHF can be traced back to research on dueling bandits (e.g., Yue et al., 2012; Saha, 2021; Bengs et al., 2021), a simplified setting within the RLHF framework. Recently, many works have focused on the more challenging RLHF problem (also known as the preference-based RL). Xu et al. (2020); Novoseller et al. (2020); Pacchiano et al. (2021) delve into the study of tabular online RLHF, where the state space is finite and small. Moving beyond the tabular setting, Chen et al. (2022) provides the first results for online RLHF with general function approximation, capturing real-world problems with large state spaces. Wang et al. (2023c) presents a reduction-based framework, which can transform some sample-efficient algorithms for standard reward-based RL to efficient algorithms for online RLHF. Further advancements in algorithm designs are introduced by Zhan et al. (2023b); Wu & Sun (2023), encompassing the development of reward-free learning type algorithms and posterior sampling-based algorithms tailored for online RLHF. Initiating exploration into offline RLHF, Zhu et al. (2023a) presents a pessimistic algorithm that is provably efficient for offline RLHF. Additionally, Zhan et al. (2023a) and Li et al. (2023b) extend these investigations into the broader scope of general function approximation settings within offline RLHF. In comparison to these existing studies, our work introduces a new theoretical formulation and goal for RLHF, as well as novel problem settings, such as hybrid RLHF. The new mathematical formulation allows our framework to align more closely with recent advancements in LLMs, and we discuss the connections between our theoretical findings and practical algorithmic designs in Section 5.

Finally, concurrent to this work, Hoang Tran (2024) and Yuan et al. (2024) also consider variants of iterative DPO. We comment on the similarities and differences between our work and theirs as follows. Hoang Tran (2024) focuses on the batch online setting, which will be thoroughly developed in Theorem B.2 (as well as Algorithm 5) in this paper. We notice that Theorem B.2 essentially states that, for exploitation, we should choose the policies (LLMs) around $\pi_{r^{t}}$, i.e., the Gibbs distribution induced by $r^{t}$ (see Section 2 for formal definitions), and for exploration, we should increase the diversity by making two policies more different. In comparison, Hoang Tran (2024) sets the two policies as the best-of-4 policy and worst-of-4 policy induced by the a preference model PairRM-0.4B (Jiang et al., 2023) and $\pi_{r^{t}}$, which may be viewed as a reasonable practical approximation of the exploration-exploitation trade-off presented in Theorem B.2. Meanwhile, the resulting model from Hoang Tran (2024) achieves state-of-the-art (SOTA) performance in the AlpacaEval leaderboard even though the preference oracle is of only 0.4B parameters. This may also partially verify that the sample complexity of alignment depends on the complexity of reward/preference model space, which can be much smaller than that of the generator (also see Theorem B.2 for details). To summarize, the two works are complementary to each other, as Hoang Tran (2024) presents an exciting recipe to illustrate the effectiveness of batch-online iterative DPO, while our work focuses more on the theoretical side. Yuan et al. (2024) also proposes a variant of iterative DPO in the batch online setting but in a self-rewarding manner. The main difference between Yuan et al. (2024) and our work (as well as Hoang Tran (2024)) is that instead of optimizing against an external preference oracle, they adopt a clever idea by using the LLM itself as the reward model to provide preference signal, hence the name “self-rewarding”. We also remark that in addition to the practical algorithmic designs, our work also serves to establish the mathematical foundation of the RLHF, in offline, online, and hybrid settings.

Formulation of RLHF

In this section, we present the mathematical framework for the RLHF process, inspired by the standard LLM alignment workflow (Ouyang et al., 2022; Touvron et al., 2023).

Specifically, the LLM can take a prompt, denoted by $x\in\mathcal{X}$, and produce a response, denoted by $a=[w_{1},w_{2},\ldots]$, where $w_{i}$ is the $i$-th token generated by the model. Accordingly, we can take $\mathcal{X}$ as the state space of the contextual bandit and the $\mathcal{A}$ as the action space. Following Ouyang et al. (2022); Zhu et al. (2023a); Rafailov et al. (2023); Liu et al. (2023a), we assume that there exists a ground-truth reward function $r^{*}(x,a):\mathcal{X}\times\mathcal{A}\to$ and the preference satisfies the Bradley-Terry model (Bradley & Terry, 1952):

where $a^{1}\succ a^{2}$ means that $a^{1}$ is preferred to $a^{2}$, and $\sigma(z)=1/(1+\exp(-z))$ is the sigmoid function. We denote an LLM by a policy $\pi$ that maps $x$ to a distribution over $\mathcal{A}$.

In a typical LLM training pipeline, the tuning process begins with a pretrained LLM, which is subsequently fine-tuned using specialized and instructional data, yielding an initial LLM policy denoted as $\pi_{0}$. We will then align the LLM on RLHF data (prompt set), which we assume is taken from a distribution $x\sim d_{0}$. For preference learning, the way to gather information from the environment is to compare two different actions under the same state. Considering this, we assume that the agent can perform a pair of actions, aligning with precedents in existing literature (Novoseller et al., 2020; Pacchiano et al., 2021). In applications, we want the resulting LLM $\pi$ to be close to $\pi_{0}$, and our goal is to find a policy $\pi$ from some policy class $\Pi$ to maximize

where $\eta>0$ is the KL penalty coefficient. This formulation is widely studied in practice (Ziegler et al., 2019; Wu et al., 2021a; Ouyang et al., 2022; Rafailov et al., 2023; Liu et al., 2023a), and our paper aims to study its theoretical property.

Usually, we have a function class $\mathcal{F}$ for approximating the ground truth $r^{*}$. Following Pacchiano et al. (2021); Kong & Yang (2022); Zhu et al. (2023a), we make the following assumption for a clear presentation because it suffices to illustrate our ideas and the algorithmic design in this paper can also apply to the general case. The analysis also readily generalizes to general function class using standard complexity measures in RL theory literature (Russo & Van Roy, 2013; Gentile et al., 2022), which essentially state that there are some low-rank structures in reward model.

The most common way of reward modeling is Maximum Likelihood Estimation (MLE) (e.g., Ouyang et al., 2022; Bai et al., 2022a; Touvron et al., 2023).

Maximum Likelihood Estimation. A preference dataset $\mathcal{D}$ consists of numerous tuples, such as $(x,a^{1},a^{2},y)$, where $y$ is the preference signal. Specifically, $y=1$ means a preference for $a^{1}\succ a^{2}$, while $y=0$ indicates $a^{1}\prec a^{2}$. Given a dataset $\mathcal{D}=\{(x,a^{1},a^{2},y)\}$, we can write the log-likelihood function of the BT models as follows:

Accordingly, we take the policy class as \Pi:=\big{\{}\pi(\cdot|x)\propto\pi_{0}(\cdot|x)\cdot\exp\big{(}\frac{1}{\eta}\left\langle\theta,\phi(x,\cdot)\right\rangle\big{)}:\theta\in\Theta(B)\big{\}}. The goal is to design a sample-efficient algorithm, which finds a policy $\hat{\pi}\in\Pi$ so that the suboptimality $J(\pi)-J(\hat{\pi})<\epsilon$ with the number of samples polynomial in accuracy parameter $1/\epsilon$, feature dimension $d$, and other problem-dependent parameters, where $\pi$ is a comparator policy (e.g. $\pi^{*}$).

2 Preliminary

In this section, we present some useful technical tools and lemmas for subsequent analysis.

Value decomposition. We have the following lemma to decompose the value difference.

Given a comparator policy $\pi$, we can decompose the suboptimality of $\hat{\pi}$ as follows:

The equality can be verified directly by the definition of $J(\cdot)$ in Equation (2) and basic algebra. ∎

Policy improvement error. In standard RL setting, $\hat{\pi}$ is typically taken as a greedy policy of $\hat{r}$, leading to

In the KL-constrained case, since the policy cannot be greedy or deterministic, we need to additionally handle the policy improvement error. The following lemma provides such an estimation when our policy is obtained by calling the Oracle 2.2 with $\hat{r}$.

Suppose that $\pi,\hat{\pi}\in\Pi$ so that $\pi_{0},\pi,\hat{\pi}$ have the same support. If $\hat{\pi}$ is induced by calling Oracle 2.2 with $\hat{r}$, it holds that

We will provide the proof of the lemma in Appendix F. The analysis techniques are most similar to the policy gradient literature since they also consider the soft-max policies (Agarwal et al., 2021; Cai et al., 2020; Zanette et al., 2021b; Zhong & Zhang, 2023). The main difference is that in their iterative choices of policy, for choosing $\pi_{t}$, the reference policy they use is the policy of the last round, i.e., $\pi_{t-1}$, while we always use the SFT-model $\pi_{0}$ as our reference. We note that their algorithms essentially still use the non-KL-regularized reward as the target because though we prevent the policy from moving too far away in each individual step, the cumulative updates makes the reward estimations dominating in the final policy.

Covariance matrix. Given a preference dataset $\mathcal{D}$, a fixed $\lambda>0$, we denote $\Sigma_{\mathcal{D}}$ as the covariance matrix estimation:

Both the algorithmic design and analysis will be centered on the covariance matrix. For the readers that are not familiar with the eluder-type techniques (or elliptical potential lemma in this case), we provide a brief introduction to the high-level intuition in Appendix A.1.

Offline learning

In addition to adopting a pessimistic reward estimation, we may also use a modified target that is biased toward pessimism by penalizing the uncertainty as in Equation (4). Here we do not maintain a confidence set but use a modified target that is biased toward pessimism, similar to Xie et al. (2021a); Zhang (2022), which may be easier to approximate in practice (Liu et al., 2023b). Moreover, to handle the additional trade-off between the reward and the KL term, we also incorporate the KL divergence into the policy computation.

The full algorithm is presented in Algorithm 1 and is referred to as the offline Gibbs Sampling from Human Feedback (GSHF) because the output policy is the Gibbs distribution with some reward.

We also have the theoretical guarantee in Theorem 3.1.

We can combine the guarantee with dataset property, usually referred to as the coverage on the comparator policy $\pi$ (Jin et al., 2021b; Xie et al., 2021a), to obtain the concrete bound. See Proposition D.1 for an concrete example. The proof of the theorem is rather standard in offline learning based on the principle of pessimism but with a different analysis to handle the KL and the stochastic policy. We defer the proof of the theorem to Appendix C.

In comparison, the Option I achieves a sharper bound in the uncertainty bonus because the expectation is inside the norm and by Jensen’s inequality (Lemma G.1) we know that

Hybrid Learning with Batch Exploration

Beyond the offline learning, it is also common to query human feedback during the training process. For instance, Bai et al. (2022a); Touvron et al. (2023) typically iterate the RLHF process on a weekly cadence, where the fresh RLHF models are deployed to interact with crowdworkers and to collect new human preference data.

Non-symmetric algorithmic structure. The main technical challenge here is to decide the behavior policy pairs $(\pi_{t}^{1},\pi_{t}^{2})$. Our first idea is to adopt a non-symmetric structure in choosing $\pi_{t}^{1}$ and $\pi_{t}^{2}$. Specifically, we refer the $\pi_{t}^{1}$ as the main agent, which aims to learn a good policy so that the suboptimality gap $J(\pi^{*})-J(\pi_{t}^{1})$ is small. In contrast, the second agent, referred to as the enhancer, seeks to enhance the learning of the main agent by choosing appropriate $\pi_{t}^{2}$. The main advantage of such a non-symmetric structure is that we have a lot of freedoms to choose $\pi_{t}^{2}$ because we do not worry about the sub-optimality incurred by it. Using $\pi_{t}^{2}$ as an intermediate agent in Lemma 2.3, we have

The advantage of reward modeling. Theorem 4.2 and Theorem B.2 (for the online setting) reveal a key characteristic of reward modeling: the sample complexity is dependent on the complexity of the reward model rather than the generative models. For simple reward functions, such as sentiment or politeness evaluation, the required function class is substantially smaller compared to the generative model. This is corroborated by evidence showing that even compact models like BERT (Devlin et al., 2018) can yield accurate reward assessments. Besides, the reward function can also be recognized as a density ratio estimator, similar to the discriminator in Generative Adversarial Models (Goodfellow et al., 2014). The density ratio estimator can be applied to generative models without explicit likelihood (GAN, Energy-based models). Johnson & Zhang (2019) suggest that the distribution induced by the density ratio estimator is more stable than optimizing the generative model. Meanwhile, the density ratio estimator approximates the functional gradient in probability space, iterative reward modeling provides a provable path towards the target distribution. This may illustrate the advantage of the most popular RLHF framework used by Ouyang et al. (2022); Bai et al. (2022a); Touvron et al. (2023), in contrast to the idea of bypassing reward modeling (Rafailov et al., 2023; Zhao et al., 2023; Azar et al., 2023) and training based only on the offline dataset. In summary, reward learning typically exhibits lower sample dependency and results in a more stable induced distribution compared to that learned directly through generative modeling. Iterative reward modeling approximates the functional gradient in probability space. Furthermore, it can be effectively applied to universal generative models, even those without explicit likelihood (Goodfellow et al., 2014; Grathwohl et al., 2019; Du & Mordatch, 2019).

Practical Implementations of GSHF

In this section, we discuss how to practically implement the information-theoretical Algorithm 1 and Algorithm 2.

The main challenge to apply the theoretical algorithm lies in the Oracle 2.2, which is computationally intractable due to the exponential action space. To design an implementable algorithm, it is critical to approximate $\pi_{r}$ effectively.

In practice, the policy is represented by a deep neural network. In this case, one common choice (Ziegler et al., 2019; Wu et al., 2021a; Ouyang et al., 2022; Bai et al., 2022a) is to use the standard deep RL algorithms like PPO to optimize the regularized reward:

However, PPO is significantly less stable and sensitive to implementation as compared to SFT (Choshen et al., 2019; Engstrom et al., 2020). Recently, DPO (Rafailov et al., 2023) attracted significant attention due to its stability and easy implementation. Specifically, DPO chooses to train the LLM as a reward model, by optimizing the following loss:

where $a_{c}$, $a_{r}$ is the chosen/rejected response. It is shown that the optimal policy for the DPO loss in Equation (6) is identical to the one for the RLHF objective $\pi_{r}$, with $r$ as the MLE of Equation (3). To summarize, to move toward a practical approach from the theoretical algorithms, we may just replace the Oracle 2.2 with the practical RLHF algorithms (both deep RL methods or non-RL methods). In view of the simplicity and effectiveness of DPO, we will mainly investigate the performance of the proposed GSHF framework with DPO.

2 Multi-step Rejection Sampling for Offline Learning

To mitigate this issue and to make the algorithm more effective, we propose a multi-step approach to progressively achieve our ultimate target. Instead of using $\pi_{0}$ to approximate $\pi_{0}\exp(\frac{1}{\eta}r)$ directly, we divide the path into several steps by considering a sequence of distributions

One concern may be on the additional computations introduced by the multi-step approximations. However, in practice, the KL coefficient $\eta$ is also tuned as a hyper-parameter in an outer loop of the proposed framework (Huggingface, 2023) to achieve the best performance. The Algorithm 3 provides us with a sequence of models associated with different $\eta_{i}$, which exactly allows for further model selection via hyper-parameter tuning of $\eta$. In view of this, the Algorithm 3 does not introduce overhead in computation.

3 Algorithmic Simplicity and Data Coverage

We note that all the three settings: offline, online (Appendix B), and hybrid learning are complementary to each other and hold their own values. For instance, collecting new and online human feedback can be expensive for most of the developers and in this case, only offline learning is feasible. One appealing choice is to leverage AI feedback (Bai et al., 2022b), which is much cheaper than human feedback. However, for tasks with customized needs or requiring expertise, we may only query feedback from specific users or experts, whose preference is distinct from AI.

Experiments

In this section, we verify the effectiveness of the Algorithm 3 and Algorithm 4 by real-world RLHF experiments.

Model, and Task. We use the Open-LLaMA-3B-V2 (Geng & Liu, 2023) as the pretrained model and use the helpful subset of the Anthropic HH-RLHF dataset (Bai et al., 2022a) (see Table 4 for a sample example). We preprocess the dataset to get $103$K training set and $5$K test set, with details in Appendix H.1. We also sample a subset of the UltraFeedback (Cui et al., 2023), consisting of $5$K prompts, as another out-of-distribution (OOD) test set. Meanwhile, the UltraRM-13B (Cui et al., 2023) will be used as the ground truth reward model, also referred to as the gold reward, which is trained on a mixture of UltraFeedback, Anthropic HH-RLHF, and other open-source datasets based on LLaMA2-13B. For all the experiments, we fix the KL penalty in the learning target Equation (2) as $\eta=0.1$.

Stronger DPO Model with Gold RM for Model Selection. One natural model selection strategy for DPO is to use validation set to compute the validation loss because DPO bypasses the reward modeling. Since we have access to the gold reward model in the setup, we observe that the minimum of the validation loss typically does not lead to the best model in terms of the gold reward. Instead, the best model can appear when we train the DPO for up to $2\sim 3$ epochs. This is similar to the observation in Tunstall et al. (2023), where the authors found that overfitting the preference dataset within certain limit does not hurt the model performance (gold reward) and the strongest model was obtained with 3 epochs of DPO training. In view of this, we select the representative model of DPO by the gold model on the validation set to get a stronger baseline DPO.

2 Main Results

We present the main results in this subsection and defer implementation details to Appendix H. We report the gold rewards and the GPT4 evaluations compared to the DPO baseline in Table 1. As we can see, DPO, RSO, and GSHF significantly outperform the SFT baseline, and the GSHF algorithms further outperform the stronger baselines including both DPO and RSO in terms of gold reward, and GPT4 evaluations. In particular, the GSHF algorithms tend to be more robust in the face of OOD data, as they achieve a much smaller $\Delta$ compared to other RLHF algorithms.

In addition to the theoretical result provided in this paper, we may also intuitively justify the improvements achieved by the GSHF algorithm (as well as RSO) compared to DPO by noting that they use different data sources for the preference learning thus providing a better coverage of the state-action space. We would like to share some thoughts with more details between the coverage condition and the success of preference learning in Appendix E.

Reward-KL Trade-off. Since all the considered RLHF algorithms (except SFT) share the same KL-constraint reward optimization target in Equation (2), we first investigate the trade-off between the gold reward and the KL divergence achieved by the different RLHF algorithms and plot the curve in Figure 2. As we can see, both the Offline GSHF and the Hybrid GSHF significantly outperform the strong baselines DPO, and RSO by achieving a much higher reward, for a fixed KL level.

The Power of Exploration. We compare different iterations of Hybrid GSHF in Figure 3. For each iteration, we evaluate the models every 400 training steps and plot the representative models. Clearly, the previous iteration is strictly dominated by the subsequent one in terms of the frontier. This demonstrates the significant improvements achieved by further iterating DPO with online data. Notably, compared to offline DPO which uses more offline data than the iteration 1, leveraging online data proves to be far more efficient, as evidenced by the enhanced frontier of the reward-KL trade-off.

Performance Comparison Under Distribution Shift. We investigate the performance of the resulting models from different alignment algorithms under distribution shift. To this end, we sample a subset of the UltraFeedback (Cui et al., 2023), consisting of $5$K prompts, as our out-of-distribution (OOD) test set. The performance results of representative models are detailed in Table 1, and the trade-off between reward and KL divergence on this OOD test set is illustrated in Figure 4. It is observed that all models exhibit a decline in performance compared to the in-domain scenario. In comparison, the Hybrid GSHF and Offline GSHF are more stable in the face of the distribution shift because they achieve a smaller $\Delta$, which is the difference between in-domain and OOD rewards. Regarding the reward-KL trade-off, consistent with in-domain results, the GSHF algorithms outperform the baseline DPO and RSO models in producing a more efficient frontier. In particular, the Hybrid GSHF achieves the best performance, indicating the advantage of online exploration compared to the offline learning.

Performance Comparison Under Different Sampling Temperatures. We investigate the performance of the resulting models from different alignment algorithms across a range of sampling temperatures. We report the test gold reward with respect to the sampling temperature in Figure 5. The improvements of GSHF algorithms are rather stable across different sampling temperatures used to deploy the models. For all the models, a temperature of 0.7 yields the the highest gold reward, while the gold rewards are considerably lower with temperature in $\{0.2,0.5,1.0\}$. An exception is observed with the Offline RSO, which maintains robustness when the temperature is reduced from 1.0 to 0.7. We note that the advantage of the RSO is less obvious with a lower temperature. Conversely, both Offline GSHF and Hybrid GSHF models consistently surpass the baseline DPO and RSO models across various sampling temperatures. Notably, Hybrid GSHF shows more advantages over the Offline GSHF with a lower temperature, potentially indicating the benefits of online exploration.

Length Bias. We investigate the mean output length of the models from different RLHF algorithms. We observe that as the Hybrid GSHF iterates, the average output lengths increases: from 161 in the first iteration, to 243 in the second, and 263 in the third. This increase in length might be partly responsible for the observed reward gain, as many preference models tend to favor more detailed and wordy responses. In comparison, the average output lengths for DPO, RSO, and Offline GSHF are 241, 275, and 240, respectively. Though there is a trend towards longer responses in later iterations of the Hybrid GSHF model, we notice that the final output length of the Hybrid GSHF model does not significantly exceed that of DPO and RSO. In practice, however, the reward (signal) hacking is the fundamental issue of RLHF (Casper et al., 2023). Therefore, it may be beneficial to integrate additional strategies such as early stopping, replay, and a thorough validation process to ensure the selection of the most effective model during the training process.

Conclusion

In this paper, we formulate the real-world RLHF process as a reverse-KL regularized contextual bandit problem. Compared to existing theoretical RLHF frameworks, the proposed framework admits a stochastic optimal policy, that more accurately reflects the dynamics of foundation generative models and aligns closely with current alignment practices (Ouyang et al., 2022; Bai et al., 2022a; Rafailov et al., 2023). We design statistically efficient algorithms in offline, online, and hybrid settings, featuring the standard ideas of pessimism and optimism in the new framework, while also handling the distinct challenges of preference learning as well as the newly introduced KL constraint with distinct algorithmic designs.

The theoretical findings also sheds light on innovative pathways for practical algorithmic development, as we move toward implementations of the information-theoretical algorithms in Section 5. The practical implementations of the proposed algorithms outperform strong baselines like DPO and RSO in real-world alignment of LLMs.

References

Appendix A Notation Table and Backgrounds

To improve the readability of this paper, we provide a Table 2 for the notations used in this paper. We also provide an introduction to the eluder-type techniques and the rejection sampling for completeness.

Before we continue to prove the main results of this paper, we would like to briefly illustrate the high-level intuitions why the algorithmic design and analysis are centered on the covariance matrix. Given a preference dataset $\mathcal{D}$, and a fixed $\lambda>0$, we denote $\Sigma_{\mathcal{D}}$ as

Then, the in-sample error on the observed data in $\mathcal{D}$ is given by

where we additionally add a regularization term $\lambda\|\theta_{1}-\theta_{2}\|^{2}$. Meanwhile, if we test the hypothesis $(\theta_{1}-\theta_{2})$ on a newly observed data, the out-of-sample error would be given by $|\left\langle\theta_{1}-\theta_{2},\phi(x,a^{1})-\phi(x,a^{2})\right\rangle|.$ The ideal case would be that we can infer the out-of-sample error via the in-sample error, so we look at the ratio between them:

where we take a square root on the in-sample error to keep them being of the same order and use Cauchy-Schwarz inequality (Lemma G.2). Here, the $\|\phi(x,a^{1})-\phi(x,a^{2})\|_{\Sigma_{\mathcal{D}}^{-1}}$ is referred to as the elliptical potential in the literature of linear function approximation (Abbasi-Yadkori et al., 2011). The elliptical potential can be viewed as the uncertainty of $\phi(x,a^{1})-\phi(x,a^{2})$, given the historical samples in $\mathcal{D}$, and can be used to guide our exploration. The complexity of the reward model space is characterized by the following fact:

The ratio between the out-of-sample error and the in-sample error in the linear case can be readily generalized to the general function approximation using the variant of eluder dimension considered in Gentile et al. (2022); Zhang (2023); Ye et al. (2023); Agarwal et al. (2023), which essentially states that there is some low-rank structure in the reward model space so the generalization is limited (the elliptical potential cannot be large for too many times). Moreover, if we can effectively estimate the in-sample error from the preference data, by Lemma A.1, we can infer the out-of-sample error safely most of the time. Such an in-sample error estimation is provided in Lemma G.3. Essentially, the eluder-type complexity measures and techniques reduce the learning problem to an online supervised learning (in-sample error estimation and minimization) (Zhong et al., 2022).

A.2 Rejection Sampling

We briefly introduce the rejection sampling in this subsection. We first remark that in the literature, many papers use this terminology to refer best-of-n policy (Touvron et al., 2023), which can be different from the notion of rejection sampling here. Specifically, the best-of-n policy takes a base policy $\pi$ and a reward function $r$ as the input, and output a new policy $\widetilde{\pi}$: for each $x\in\mathcal{X}$, we sample $n$ independent policies from $\pi$ and output the one with the highest reward measured by $r$. In what follows, we introduce the rejection sampling.

Rejection sampling, a widely utilized method in Monte Carlo tasks, is designed to sample from a target distribution using samples from a proposal distribution and a uniform sampler (Neumann, 1951). This technique is applicable when the density ratio between the target distribution $q$ and the proposal distribution $p$ is bounded, satisfying $q(x)/p(x)\leq M$ for all $x\in\mathcal{X}$. In practical implementation, $n$ samples are drawn from the proposal distribution $p$. Each sample, denoted as $x\sim p$, is accepted with a probability $r=\frac{q(x)}{Mp(x)}$. This acceptance is determined by evaluating whether $u<r$, where $u$ is a number drawn from a uniform distribution $U$. The accepted samples $\widetilde{x}$ are then representative of the target distribution $q$.

The primary challenge in rejection sampling is its low acceptance rate, particularly problematic for high-dimensional data due to the curse of dimensionality, where the density ratio often scales with $\exp(d)$. This issue persists even in low-dimensional scenarios, as a large density ratio $M$ can drastically reduce acceptance rates. The method is most efficient when $p$ closely approximates $q$, leading to $M\approx 1$.

Appendix B (Batch) Online Learning with Enhancer

In this section, we develop the online framework of the KL-constraint contextual bandit, that is missing in the main paper.

The mathematical formulation of the online learning is almost the same as the hybrid case, except that we now start from scratch instead of the offline dataset. Consider the batch online setting of $T$ batches with fixed batch size $m$. At the beginning of each batch $t\in[T]$, An agent updates the policies $\pi_{t}^{1}$ and $\pi_{t}^{2}$. Then, $m$ prompts $\{x_{t,i}\}_{i=1}^{m}$ are sampled from $d_{0}$. Based on each prompt $x_{t,i}$, two responses $(a_{t,i}^{1},a_{t,i}^{2})$ are generated from two policies $(\pi_{t}^{1},\pi_{t}^{2})$, and a human preference signal $y_{t,i}\in\{0,1\}$ is yielded according to the ground-truth BT model.

We first consider the case of $m>1$, which leads to a more sparse update of the model. Our goal is also to design a sample-efficient algorithm, which finds a policy $\hat{\pi}$ so that the suboptimality $J(\pi^{*})-J(\hat{\pi})<\epsilon$ with the number of samples polynomial in the accuracy number $1/\epsilon$, feature dimension $d$, and other problem-dependent parameters. In practical applications, it is observed that the diversity of the outputs is critical, and the response pairs $(a^{1}_{t},a^{2}_{t})$ are recommended to be collected by different model variants with different temperature hyper-parameter (Touvron et al., 2023). To understand this choice, we recall the decomposition Lemma 2.3 and Lemma 2.4 to obtain for each batch $t\in[T]$

The main technical challenge is to relate the uncertainty of $\phi(x_{t},\pi_{t}^{1})-\phi(x_{t},\pi^{*})$ (analysis target) to the uncertainty of $\phi(x_{t},\pi_{t}^{1})-\phi(x_{t},\pi_{t}^{2})$ (the pair to collect data). Our algorithmic idea is built on optimism and non-symmetric structures. We present the complete algorithm in Algorithm 5. The main agent $\pi_{t}^{1}$ always takes the policy induced by $r^{t}$ from Oracle 2.2. On the other hand, the second agent $\pi_{t}^{2}$, referred to as the enhancer, seeks to maximize the uncertainty (similar to the practical choice of different model variants and temperature) for the fixed $\pi_{t}^{1}$, thus facilitating the learning of the main agent (similar idea was considered in the study of two-player zero-sum Markov game (Jin et al., 2021a; Huang et al., 2021; Xiong et al., 2022b)). In this case, the uncertainty compared to $\pi^{*}$ is upper bounded by that of $\pi_{t}^{2}$, which is referred to as the principle of optimism in the literature (Auer et al., 2002). Notably, in contrast to the case of Markov game (Jin et al., 2021a; Huang et al., 2021; Xiong et al., 2022b), the enhancer also converges to $\pi^{*}$ in terms of the metric of $J(\pi)$. We borrow the terminology of the main agent and enhancer to stress the non-symmetric algorithmic structure. Moreover, if we just regard the enhancer $\pi_{t}^{2}$ as an auxiliary policy and only care about the performance of $\pi_{t}^{1}$, there is no need to maintain the confidence set $\Pi_{t}$. Due to the realizability: $\pi^{*}\in\Pi$, we can construct $\pi_{t}^{2}$ as the solution of the following unconstrained problem:

where the uncertainty bonus will be specified later. Note that in Algorithm 5, we formulate that the agent first observes $m$ prompts and then establishes the enhancer. This is only for simplicity of analysis so that we can estimate the uncertainty and obtain the enhancer by maximizing the estimation. If we consider the standard online contextual bandit, we can first collect $m$ contexts, and estimate the uncertainty based on them. Then, for the next $m$ contexts, we interact with the environment in a strictly sequential manner using the policies determined by the first $m$ contexts. This will only roughly incur a constant factor $2$ in the final sample complexity.

To achieve optimism, we need to maintain a confidence set, that contains the $\pi^{*}$ for all iterations with high probability. The constructions of the confidence set are different compared to the dueling RL (Faury et al., 2020; Pacchiano et al., 2021) due to the reverse-KL regularized contextual bandit formulation, as well as the non-symmetric structure in our algorithm. We summarize the confidence set construction for the online setting in the following lemma.

For the linear model in Assumption 2.1, given the policy of the main agent $\pi_{t}^{1}$, we consider the following confidence set with \beta=O\big{(}\sqrt{\frac{d\log(T/\delta)}{\gamma^{2}m}}\big{)}:

Then, with probability at least $1-\delta$, we know that $\pi^{*}\in\Pi_{t}$ for all $t\in[T]$.

For any $\epsilon>0$, we set the batch size $m=d/(\gamma^{2}\epsilon^{2})$. Under Assumption 2.1 with the uncertainty estimator defined as

where the number of collected samples is at most mT=\widetilde{O}\Big{(}\frac{d^{2}}{\gamma^{2}\epsilon^{2}}\Big{)}.

Theorem B.2 reveals a key characteristic of reward modeling: the sample complexity is dependent on the complexity of the reward model rather than the generative models. For simple reward functions, such as sentiment or politeness evaluation, the required function class is substantially smaller compared to the generative model. We now present the proof of the theorem.

Recall the definition of the covariance matrix:

Then, by invoking Lemma G.3 for $\theta_{t}$ with $\Sigma_{\mathcal{D}}=m\Sigma_{t,m}$ and $\lambda^{\prime}=m\lambda$, we have with probability at least $1-\delta$, for any $t\in[T]$,

Now, by elliptical potential lemma (Lemma G.4), we have

Since each term on the left-hand side is positive, we know that there exists at least a $t_{0}\in[T]$, the value is smaller or equal than the average value:

We now consider the suboptimality at iteration $t_{0}$:

where the inequality uses the Cauchy-Schwarz inequality (Lemma G.2). Then, since the samples $\{x_{t,i}\}_{i=1}^{m}$ are i.i.d and for any $x\in\mathcal{X}$

we can use Chernoff bound (Theorem 2.16 of Zhang (2023)) to obtain that with probability at least $1-\delta/2$,

Similarly, we also get with probability at least $1-\delta/2$,

We take the two inequalities above back into Equation (B.1) to derive with that probability at least $1-3\delta$,

where the second inequality applies Lemma G.5 with $\lambda=\Omega(d\log(T/\delta)/m)$, and the last inequality uses Equation (B.1). By choosing $T$ satisfying that $T\geq d\log(T)$ and $\lambda=\Theta(d\log(T/\delta)/m\gamma^{2})$, we have

B.2 Sequential Online Setting

While we mainly care about finding a good model, with a slightly more involved analysis for the enhancer, we can also derive an upper bound for the average regret as in Pacchiano et al. (2021); Chen et al. (2022):

where we now discuss in the sequential case with $m=1$ in Algorithm 5. We consider two kinds of regrets: (1) cumulative suboptimality for the main policy $\pi_{t}^{1}$ compared to $\pi^{*}$:

where $\mathcal{D}^{1:t-1}=\cup_{s=1}^{t-1}\mathcal{D}^{s}$.

Under Assumption 2.1 with the uncertainty estimator defined in Equation (9), with $\lambda=\Omega(d\log(T/\delta)/(\gamma^{2}B^{2}))$ and \beta:=O\big{(}\sqrt{\frac{d\log(T/\delta)}{\gamma^{2}}}\big{)}, with probability at least $1-2\delta$, the regret of Algorithm 5 with $m=1$ satisfies

First, recalling the regret decomposition in Equation (B.1), we deduce that with probability at least $1-\delta$,

where the first inequality uses the Cauchy-Schwarz inequality, Lemma G.3 and reward $r\leq 1$ for any $r\in\mathcal{F}$, the second inequality uses $\pi^{*}\in\Pi_{t}$ according to Lemma B.1, and the last inequality uses the Cauchy-Schwarz inequality and Jensen’s inequality.

According to the concentration of the covariance matrix in Lemma G.5, since $\lambda=\Omega(d\log(T/\delta))$, we have with probability at least $1-\delta$, for any $t\in[T]$,

By taking the result above back into Equation (B.2), we get with probability at least $1-2\delta$,

We can deal with the Term $(\Delta_{t}^{2})$ by invoking Lemma B.1 with $\pi=\pi_{t}^{2}$ and using the definition of the confidence set:

Combining the above two inequalities and Equation (B.2), we have

Therefore, by combining the results above and Equation (15), we have

B.3 Construction of the Confidence Set

By the definition of the $\pi^{*}$ that $\pi^{*}$ is optimal at every context, for any $\pi_{t}^{1}\in\Pi$ and any $x_{t,i}\in\mathcal{X}$, we have

For Term (i), by Cauchy-Schwarz inequality and Lemma G.3 with $\Sigma_{\mathcal{D}}=m\Sigma_{t,m}$ and $\lambda^{\prime}=m\lambda$, we have

where \beta=O\big{(}\sqrt{\frac{d\log(T/\delta)}{\gamma^{2}m}}\big{)} and the additional $\log T$ factor is because of the union bound over the $T$ iterations. Meanwhile, by invoking Lemma 2.4 with $\pi=\pi^{*},~{}\hat{\pi}=\pi_{t}$, we obtain that

Taking respective upper bounds for Terms (i) and (ii) back into Equation (B.3) and summing over $i\in[m]$, we have

which implies that $\pi^{*}\in\Pi_{t}$. Therefore, we finish the proof of Lemma B.1. ∎

Appendix C Proof of the Offline Learning

For simplicity, we denote the LHS of Equation (19) as $(\star)$. We plugging this into the estimation of $J(\pi)-J(\hat{\pi})$:

where the first inequality is from the Equation (19) and the second inequality uses Cauchy-Schwarz inequality and Lemma G.3.

For Option II, we use the point-wise pessimism:

Then, we call Oracle 2.2 with $\hat{r}$ to get $\hat{\pi}$. By Lemma 2.3, we have

Since $\hat{r}$ is obtained from the Oracle 2.2 with $\hat{r}$, it follows from Lemma 2.4:

where we use Cauchy-Schwarz inequality in the last inequality.

Appendix D Proof of the Hybrid Setting

Under Assumption 2.1, assuming that there exists absolute constants $c^{\dagger}$ and $\alpha^{\ddagger}$ such that

where $\lambda_{j}$ denotes the $j$-th eigenvalue of $\Sigma^{\ddagger}$. It is not difficult to show that $\lambda_{j}\in[0,B^{2}]$, which further implies that

which concludes the proof of Proposition D.1. ∎

D.2 Sequential Hybrid Setting

Under Assumption 2.1, let $\lambda=d\log(T/\delta)/(\gamma^{2}B^{2})$ and \beta:=O\big{(}\sqrt{\frac{d\log(T/\delta)}{\gamma^{2}}}\big{)}. Under Assumption 4.1, with probability at least $1-2\delta$, the output policy of Algorithm 2 with $m=1$ satisfies

Define the following covariance matrices:

Similar to the proofs of the offline and online setting, we get the following decomposition: with probability at least $1-2\delta$,

For the term $P_{2}$, we can apply Lemmas G.4 and G.5 to obtain

By taking the upper bound of $P_{1}$ and $P_{2}$ back, we have

D.3 Proof of Theorem 4.2

We first restate the Theorem 4.2 for a slightly more general result.

by Cauchy-Schwarz inequality and Lemma G.3. It follows that

where we use $T\geq d\log(T)$ and $C>0$ is an absolute constant. Now we proceed to suppose that Assumption 4.1 holds. Then, we have

Plugging this estimation back and combining with the choices of parameters, we conclude the proof of Theorem D.3. ∎

Appendix E Discussion on the Practical Algorithmic Design

In this section, we investigate the connections between the proposed algorithms and the existing practical algorithms in the literature, including Direct Preference Optimization (DPO) (Rafailov et al., 2023) Rejection Sampling Optimization (RSO) (Liu et al., 2023a), and RewArd-ranked FineTuning (RAFT) (Dong et al., 2023).

where $a_{c}$ is the chosen response and $a_{r}$ is the rejected response. Given $x,a_{c},a_{r}$, fitting the model with the loss in Equation (21) yields a MLE for the preference probability (Lemma E.1) by training the LLM as a reward model. This process, however, necessitates considering the generation distributions of $a^{1}$ and $a^{2}$, which is missing in the original DPO paper. We now discuss the influence of the offline data distribution.

where $p^{\theta}$ is the preference model associated with $\pi_{\theta}$. Given $x,a^{1},a^{2}$, the following lemma demonstrates that $p^{\theta}=p^{*}$ uniquely minimizes the loss.

Given $x,a^{1},a^{2}$, we consider the preference learning for

Consider the population loss (when we have sufficiently many samples),

The solution satisfies $\pi_{\theta}({a}^{1}|x)/\pi_{\theta}({a}^{2}|x)=\pi^{*}(a^{1}|x)/\pi^{*}(a^{2}|x)$.

However, since the outpace $\mathcal{A}$ is exponentially large with respect to the sequence length, the ratio of using $\pi_{0}$ can be extremely large in the worst case, which may also lead to an inferior performance in practice, as shown in Liu et al. (2023a). On the other hand, (1) RSO uses rejection sampling to approximately sample data from $\pi_{r}$; (2) Offline GSHF improves RSO by adopting a more efficient multi-approach way to better approximate $\pi_{r}$ given the limited generation budget; (3) Hybrid GSHF uses both the offline dataset and the data from online exploration. These algorithms adopt different data sources for the preference learning thus exploring different parts of the state-action space.The improvements of the GSHF algorithms emphasize the importance of a more efficient data augmentation strategy and further exploration of the state-action space.

E.2 Iterative RLHF Training

The multi-step approximation in Algorithm 3 shares similar spirit with the iterative framework (i.e., the RAFT algorithm) proposed in Dong et al. (2023) and was also considered in Touvron et al. (2023) and Gulcehre et al. (2023). Our multi-step rejection sampling may be viewed as a generalization of that of RAFT, as we illustrate as follows.

RAFT starts from $\pi_{0}$ and aims to learn from the induced best-of-n policy (i.e., for each prompt $x$, we collect $n$ independent responses and output the one with highest reward). By standard concentration inequality, if the reward function is bounded by $M$, the upper bound of the best-of-$n$ policy satisfies

which increases at a rate of $\sqrt{\log n}$. Therefore, the marginal benefit of increasing $n$ diminishes quickly, which motivates the authors to adopt the iterative framework because the improved base policy will lead to an improved best-of-$n$ policy. In comparison, we decompose the target policies \pi_{0}\exp\big{(}\frac{1}{\eta_{N}}r\big{)} to several steps and we will use the improved policy associated with $\eta_{i}$ as the base policy to approximate that with $\eta_{i+1}$ with rejection sampling. The multi-step rejection sampling is far more efficient compared to using $\pi_{0}$ because the rejection rate is reduced, as we illustrate in Figure 6.

Another major difference is that Dong et al. (2023) only considers reward optimization without the KL constraint from the initial checkpoint. Therefore, they choose to train the model from the checkpoint obtained from the preceding iteration. On the other hand, we always start from the initial model at each iteration.

E.3 Heuristic Uncertainty Estimation and Implementation of Pessimism and Optimism

The uncertainty estimation for LLMs can be challenging due to the extremely large state-action space and a closed-form solution similar to the potential is unavailable in general.

Pessimistic MLE for Reward Modeling. The recent work (Coste et al., 2023) implements the principle of pessimism based on ensemble in two different ways, and demonstrate the effectiveness of them using real-world LLM alignment experiments. Specifically, to create an ensemble, the authors train $5$ independent reward models with different random seeds $\{r_{i}\}_{i=1}^{5}$. First, the authors consider worst-case optimization (Boyd & Vandenberghe, 2004), which gives a pessimistic reward estimation:

Second, the authors also consider a soft version of pessimism by penalizing the variance of estimation (Wu et al., 2021b):

where $\bar{r}(x,a)=\frac{1}{5}\sum_{i=1}^{5}r_{i}(x,a)$ and the $\lambda>0$ is a tuning parameter. It was observed that such a pessimistic RM can largely mitigate the issue of overfitting in RLHF. We refer interested readers to Coste et al. (2023) for details.

Optimistic Policy Selection for Enhancer. In comparison, selecting an appropriate optimistic policy for the enhancer to maximize the uncertainty with respect to the main agent $\pi_{t}^{1}=\pi_{r^{t}}$ is largely less explored in practical applications. The enhancer aims to maximize the uncertainty of the feature difference given in Equation (9). While there are works adopt an optimistic value estimation in practical DRL applications (Ciosek et al., 2019; Bai et al., 2020; Rashid et al., 2020), direct optimism in terms of the policy seems to be far more challenging. Meanwhile, we are in the face of distinct challenges from preference learning. These together call for new ideas for the practical implementations.

Essentially, the results in both hybrid learning and online learning presented in this paper emphasize the importance of sampling strategy for iterative RLHF. Although the optimistic enhancer is not readily available in practice, the theoretical insights behind such a choice of enhancer is that the enhancer should generate response so that the difference between it and that of the main agent is large, compared to the data collected so far, which should at least motivate the future algorithmic design in principle.

Since the advantages of pessimism in offline RLHF has been verified in a large amount of work (e.g., Christiano et al., 2017; Ziegler et al., 2019; Gao et al., 2023; Zhu et al., 2023a; Coste et al., 2023; Shin et al., 2023), we do not leverage pessimism in the experiments of this paper but focus on verify the effectiveness of the proposed multi-step rejection sampling. Moreover, as we cannot find a practical approximation for the optimistic enhancer, we hope that our theoretical insights can motivate future study in this direction to construct reliable and efficient uncertainty estimators for LLMs, especially for the implementation of an optimistic enhancer.

Appendix F Technical Lemma Proofs

Since $\hat{\pi}$ is induced by calling Oracle 2.2 with $\hat{r}$, we know that for any $x\in\mathcal{X}$,

where $Z(x)=\sum_{a\in\mathcal{A}}\pi_{0}(a|x)\exp(\frac{1}{\eta}\hat{r}(x,a))$ is the normalization constant. We can rewrite the reward function as

Plugging this reward reparameterization into the policy optimization error under $\hat{r}$, we have

Plugging the above equality into the LHS of the Lemma 2.4 completes the proof. ∎

The loss function can be reformulated as the KL divergence plus a constant term:

This implies that $p^{*}=p^{\theta}$ is the unique optimal solution for $p^{\theta}$. Moreover, if the condition $\pi_{\theta}({a}^{1}|x)/\pi_{\theta}({a}^{2}|x)=\pi^{*}(a^{1}|x)/\pi^{*}(a^{2}|x)$ is satisfied, the optimality of the solution is assured. ∎

Appendix G Technical Lemmas

See Proposition A.9 of Zhang (2023) for a proof. ∎

In particular, for a positive-definite matrix $\Sigma$, we can take $\left\langle u,\nu\right\rangle=\left\langle\Sigma^{1/2}u,\Sigma^{-1/2}\nu\right\rangle$ to get $\left\langle u,\nu\right\rangle\leq\|u\|_{\Sigma}\|\nu\|_{\Sigma^{-1}}$.

For a fixed $\lambda>0$, we denote $\Sigma_{\mathcal{D}}$ as

Assume that $\|\phi(x,a)\|\leq 1$ for all $(x,a)\in\mathcal{X}\times\mathcal{A}$ and $\|\theta\|\leq B$. Then, it follows that with probability at least $1-\delta$, we have

Further, if $\|x_{i}\|_{2}\leq L$ for all $i\in[T]$, then we have

Finally, if $\lambda_{\min}(\Lambda_{0})\geq\max(1,L^{2})$,

Given a loss functional with respect to $\pi(\cdot|x)$, written as

the minimizer of the loss functional is \pi^{*}(a|x)\propto\pi_{0}(a|x)\exp\Big{(}\frac{1}{\eta}r(x,a)\Big{)}, also known as Gibbs distribution.

Appendix H More Experiment Details

All the experiments are conducted using 8$\times$A40 (48G) with 600G RAM, and half-precision training (bf16). The implementations are based on open-source packages TRL (von Werra et al., 2020) and LMFlow (Diao et al., 2023), and the code will be publicly available on GitHub in the camera-ready version. The hyper-parameters used in the experiments are compactly provided in Table 7 and Table 8, with details described in the subsequent subsections.

Dataset preprocessing. We use the HH-RLHF dataset (Bai et al., 2022a) in our experiments, where each sample of the dartaset consists of a prompt $x$ (chat history between the Human and Assistant), and a chosen response $a_{c}$ and a rejected response $a_{r}$. We provide an example in Table 4 for readers’ reference. We delete the noisy samples (e.g., with the same chosen and rejected responses), and prompts longer than $400$ tokens, and eventually get $108$K prompts, which are divided into $103$K training set and $5$K test set. We also sample a subset of the UltraFeedback (Cui et al., 2023), consisting of $5$K prompts, as another out-of-distribution test set.

Rejection Sampling. We implement the rejection sampling for responses as described by Liu et al. (2023a). For each prompt, we initially generate a set of $K$ samples. Our objective is to extract preference pairs from these samples. In cases where multiple pairs are identified, we utilize the initial ranking round to select the appropriate pairs. Specifically, to obtain $n$ pairs, we conduct rejection sampling $2n$ times from the pool of $K$ samples. Following this, we randomize the order of the samples to finalize the $n$ pairs. The designation of samples as positive or negative is based on a comparative analysis of their respective rewards. It is important to note that in the context of rejection sampling, the coefficient corresponds to the $\eta$ parameter of the target distribution. Our implementation is grounded in the Python code outlined in Algorithm 1 (Liu et al., 2023a).

Multi-step approximation. We divide the path into three steps with $\eta\in\{0.1,0.3,0.5\}$ and use 25K prompts at each time. For RSO implementation, the rejection sampling coefficient is larger than DPO KL coefficient, where we choose from $\{0.5,1,2,3\}$ for better performance. Liu et al. (2023a) also suggest similar phenomenon in RSO.

Hybrid learning. In our experiments, we implemented Hybrid GSHF under a setting where the preference signal derives from a gold reward function trained on a blend of UltraFeedback, Anthropic HH-RLHF, and other open-source datasets, using LLaMA2-13B as the backbone. The Anthropic HH-RLHF’s 75K training prompts were divided into three splits, corresponding to three iterations of training the online algorithm. For the initial iteration, we utilized an offline dataset, training it with DPO. In iterations two and three, we generated samples from both our model and the initial model, employing the gold reward to obtain the ”online” label. Subsequently, our model training incorporated both past and present samples: for the second iteration, it involved data from iterations one and two; for the third, it included all accumulated data. Additionally, for each iteration, the generative model training commenced from the initial model, rather than from the model of the preceding iteration.

GPT4 Evaluation. We report the detailed GPT4 evaluation results in Table 3, where the model aligned with DPO is taken as the baseline. The test hyper-parameter is provided in Table 7. For GPT4 evaluation, we use the GPT-4-turbo model (gpt-4-1106-preview). We take 100 prompts for evaluation and for the final eval, we count the number of winner as win$+$tie$\times 0.5$.

Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user question displayed below. You should choose the assistant that follows the user’s instructions and answers the user’s question better. Your evaluation should consider factors such as the helpfulness, relevance, accuracy, depth, creativity, and level of detail of their responses. Begin your evaluation by comparing the two responses and provide a short explanation. Avoid any position biases and ensure that the order in which the responses were presented does not influence your decision. Do not allow the length of the responses to influence your evaluation. Do not favor certain names of the assistants. Be as objective as possible. After providing your explanation, output your final verdict by strictly following this format: [[A]] if assistant A is better, [[B]] if assistant B is better, and [[C]] for a tie.

Reward baseline. We mention in passing that we use the test reward of the initial model as the baseline when presenting the absolute values in Table 1 by convention (Gao et al., 2023; Dong et al., 2023).

H.2 Examples

We provide sample outputs of the models from different RLHF algorithms in Table 5 and Table 6.