DPO Meets PPO: Reinforced Token Optimization for RLHF

Introduction

Reinforcement Learning from Human Feedback (RLHF) has emerged as a key technique for aligning foundation models with human values and preferences (Christiano et al., 2017; Ziegler et al., 2019). It has been pivotal in enabling Large Language Models (LLMs) to produce more helpful, harmless, and honest responses (Bai et al., 2022), as demonstrated in significant applications such as ChatGPT (OpenAI, 2023), Claude (Anthropic, 2023), and Gemini (Team et al., 2023). The classical RLHF pipeline (Ziegler et al., 2019; Ouyang et al., 2022) consists of two steps: (i) Reward training from human feedback, where the learner learns the reward function based on preference data, typically through Maximum Likelihood Estimation (MLE). (ii) Reward-based RL training, where the learner employs the seminal deep RL algorithm Proximal Policy Optimization (PPO; Schulman et al., 2017) to optimize the reward learned in the previous step.

Despite the success of this framework in the aforementioned powerful closed-source LLMs, the training of PPO is known to be unstable and sample-inefficient (Choshen et al., 2019) compared to supervised learning. PPO frequently fails to maintain a consistent average response length or experiences sudden drops in reward value. Moreover, the superior performance of PPO also relies on the code-level optimization and an appropriate configuration of the hyper-parameters (Engstrom et al., 2020), while the training stability issue further prohibits us from achieving the best performance of the PPO. So far, the success of PPO has not been widely reproduced, especially in the open-source community with rather limited resources. While researchers have made efforts to propose alternative approaches to the PPO algorithm, with notable examples like rejection sampling fine-tuning (Dong et al., 2023; Gulcehre et al., 2023), direct preference learning algorithms (Rafailov et al., 2023; Zhao et al., 2023; Azar et al., 2023), there is little evidence that these newly proposed approaches alone can make the state-of-the-art LLMs. Therefore, improving the performance of the PPO algorithm in the context of RLHF is still an important research direction that is largely under-explored.

After examining the open-source implementation of PPO, we identify that one potential reason for the sub-optimal performance of PPO is the mismatch between the formulation of RLHF and the nature of PPO. Specifically, in the existing framework (Ouyang et al., 2022; Bai et al., 2022), RLHF is formulated as a bandit, where the entire response sentence is considered to be an action, and the reward is sentence-level, evaluating only the overall quality of the response. However, PPO is designed for multi-step RL problems modeled as Markov decision processes (MDPs), requiring a token-wise reward assignment to each step. In typical implementations of PPO (e.g., the TRL package from huggingfacehttps://github.com/huggingface/trl), besides the regularization reward function assigned to each token to ensure the fine-tuned LLM stays close to the supervised fine-tuning (SFT) model, the learned sentence-level reward is only distributed to the last token, while other tokens receive zero learned reward. See (2.3) for the formal mathematical description. Clearly, there is a separation in terms of the assignment strategies of the regularization reward and the learned reward. Meanwhile, while it is generally believed that a fine-grained characterization with token-wise feedback can provide more information, in practice, it is also challenging to collect effective token-wise feedback for human conversations and use it in the MLE process. Consequently, the construction of token-wise reward signals also remains largely under-explored in the literature of RLHF.

In this work, we aim to address the aforementioned issues by developing an RLHF framework with a fine-grained token-wise reward characterization, establishing the mathematical foundation, and advancing practical algorithmic designs. The key contributions of this work are summarized as follows.

We propose a framework that models RLHF as an MDP, offering a more precise token-wise characterization of the LLM’s generation process. Furthermore, we provide theoretical insights into why the token-wise MDP formulation is superior to the previous sentence-level bandit formulation of RLHF.

Under the MDP formulation of RLHF, we introduce Reinforced Token Optimization (RTO), which extracts token-wise reward signals from offline preference data and subsequently performs RL training with respect to the learned token-wise rewards. Using MLE as the token-wise reward learning oracle, we prove that RTO can learn a near-optimal policy in a sample-efficient manner.

Moving toward the practical implementation of RTO, we adopt a novel token-wise reward extraction approach from direct preference optimization (DPO; Rafailov et al., 2023). By assigning this DPO-based token-wise reward function to each token and then optimizing with PPO, RTO outperforms existing baselines such as PPO and DPO in the task of dialogue.

In summary, under the MDP formulation of RLHF, we develop a new principled RLHF algorithm, RTO, that leverages token-wise reward signals derived from offline preference data using DPO, and subsequently performs PPO training to optimize the token-wise rewards. The pipeline of RTO is visualized in Figure 1.

2 Related Works

We review the works that are mostly related to our project in this subsection. Due to the space constraint, we refer interested readers to the survey (Casper et al., 2023) for a more comprehensive overview of RLHF.

The classic RLHF framework is established in Christiano et al. (2017); Ziegler et al. (2019) and further developed in Ouyang et al. (2022); Bai et al. (2022), where the latter can be viewed as the results of the preliminary versions of Chat-GPT and Claude. PPO (Schulman et al., 2017) is the default choice for all these projects and its effectiveness has been showcased in the resulting revolutionary foundation language models. However, as we mentioned in the introduction, tuning the PPO algorithm to its best performance requires extensive efforts and resources are often unavailable to the open-source community. Motivated by this, researchers have made efforts to develop alternative approaches to the PPO algorithm. As a direct extension of the best-of-n inference (Nakano et al., 2021), rejection sampling fine-tuning is proposed by Dong et al. (2023); Gulcehre et al. (2023); Wang et al. (2024), which prompts the LLM to generate $n$ responses per prompt and uses a learned reward function to rank the responses and fine-tune the model on those with high rewards. Besides, inspired by the reward-conditioned training in RL literature (Chen et al., 2021), Hu et al. (2023); Yang et al. (2024a) develop conditional SFT to avoid the reward learning. Another line of work aims to skip the reward modeling step and may be referred to as the direct preference learning approach (Zhao et al., 2023; Rafailov et al., 2023; Azar et al., 2023; Tang et al., 2024). Among them, the direct preference optimization (DPO) algorithm is the most popular one, mostly due to its innovative idea: your language model is secretly a reward model. In particular, according to the reward benchmark (Lambert et al., 2024), the DPO-aligned algorithm often admits a competing ranking accuracy as a reward function. We will formally discuss the principle of DPO in Appendix C.1, which also partly motivates our methods. After these, there are also many tasks that consider the variants of this direct preference learning approach by increasing the training steps (Xiong et al., 2023; Hoang Tran, 2024) and consider the more general preference signal sources (Ye et al., 2024; Rosset et al., 2024). Although all these recently proposed algorithms achieve promising results, there is little evidence that these algorithms alone without PPO can make state-of-the-art LLMs. Therefore, understanding PPO and improving its performance in the context of foundation model alignment is still an important research direction.

Theoretical study of RLHF.

The theoretical study of RLHF may date back to the dueling bandit and dueling RL (e.g., Yue et al., 2012; Saha, 2021; Faury et al., 2020; Bengs et al., 2021; Pacchiano et al., 2021; Chen et al., 2022; Zhu et al., 2023; Wang et al., 2023; Zhan et al., 2023a, b), where the reward maximization problem is considered in the face of preference signals, instead of the absolute reward signals. However, the reward maximization framework admits a greedy and deterministic optimal policy, which deviates from the principle of generative AI. Meanwhile, instead of the original reward function, the most widely used learning target is a Kullback-Leibler (KL)-regularized one. In recognition of the above issues, Xiong et al. (2023) first formally formulates the RLHF as the reverse-KL constrained contextual bandit in offline, online, and hybrid settings, and proposes sample-efficient algorithms in different settings accordingly. Beyond the reward-based framework under the Bradley-Terry model, Azar et al. (2023); Ye et al. (2024) consider the RLHF under a general preference oracle, and motivate the algorithmic design in a KL-regularized minimax game between two LLMs. In particular, Azar et al. (2023) proposes the first sample-efficient planning algorithm, and Ye et al. (2024) designs the sample-efficient learning algorithms in offline and online settings. Notably, as these studies of the KL-regularized framework align with the practical applications closely, the theoretical insights naturally motivate practically powerful algorithms like GSHF (Xiong et al., 2023), Nash-MD (Azar et al., 2023), and DNO (Rosset et al., 2024). However, we remark that Xiong et al. (2023); Azar et al. (2023); Ye et al. (2024) are still confined to the bandit setting, thus differing from the MDP formulation presented in this paper.

Improving PPO in the context of RLHF.

Although some works (e.g., Uesato et al., 2022; Lightman et al., 2023; Yang et al., 2024b) use token-wise or step-wise information to enhance the performance of LLMs, such as their reasoning ability, we will not discuss them in detail here. Instead, we will focus on comparing our work with others that aim to improve the PPO in RLHF. In particular, Li et al. (2023a) and Ahmadian et al. (2024) state that the PPO is not the best fit for RLHF because of the sentence-level reward and deterministic transition, and argue that the reinforce-style (Williams, 1992) algorithms perform better. Wu et al. (2024) proposes to construct several separate reward functions for different goals and use the linear combination of them to guide the PPO training, but the separate models are still confined to the sentence level. Similarly, Jang et al. (2023) extends the PPO to the multi-objective optimization scenario, but still uses the sentence-level modeling. Chan et al. (2024) shares similar insights that aim to improve PPO via a dense reward. They still follow the two-staged RLHF framework to model the reward function via MLE of the Bradley-Terry model and assume that the learned reward is based on the transformer (Vaswani et al., 2017). Then, they propose to use the attention value to redistribute the final scalar reward on a token level. In comparison, while sharing similar insights about using a token-wise reward, our techniques to obtain the dense signal and mathematical motivation are fundamentally different.

Concurrent work.

During the preparation of this work, there is a concurrent and independent work (Rafailov et al., 2024) that also provides a token-wise MDP formulation for RLHF. Their work shares the same insight as ours, namely that “DPO implicitly optimizes the token-wise reward”. Based on this insight, they improve the efficiency of search-based algorithms. In contrast, we propose a new algorithm RTO that leverages the token-wise reward functions to enhance the performance of PPO. In addition, our work provides a theoretical foundation for the unique advantages of token-wise MDP and its sample-efficient learning.

3 Notation

Preliminaries

In this section, we introduce the standard RLHF paradigm. Let $x\in\mathcal{X}$ denote the prompt sampled from a distribution $\rho\in\Delta(\mathcal{X})$, and $y=(y_{1},y_{2},\dots,y_{h},\dots)$ be the corresponding response, which is a sequence of tokens generated by LLMs, where $y_{i}$ represents the $i$-th token. In practice, it is widely assumed (Christiano et al., 2017; Ziegler et al., 2019; Bai et al., 2022; Ouyang et al., 2022; Touvron et al., 2023) that the preference signal is generated according to the Bradley-Terry (BT) model (Bradley and Terry, 1952):

where $\sigma(z)=1/(1+\exp(-z))$ is the sigmoid function, and $r$ is a ground-truth reward function defined at the sentence level. In other words, the reward function $r$ only evaluates the overall performance of the entire response. The classical RLHF pipeline (Ziegler et al., 2019; Ouyang et al., 2022) typically consists of two steps: reward training from human feedback and reward-based RL training. In the first step, the learner is given a dataset $\mathcal{D}=\{(x,y^{w},y^{l})\}$, where $y^{w}$ denotes the preferred response over the $y^{l}$. The reward function is learned through Maximal Likelihood Estimation (MLE) on this dataset $\mathcal{D}$:

where $\pi$ is the current policy to be improved. However, it is well known that sparse rewards can make learning more difficult compared to dense rewards (Andrychowicz et al., 2017). One natural solution is to design dense token-wise rewards used for PPO training, but this is beyond the scope of the current bandit formulation for RLHF and motivates us to provide a framework with more fine-grained token-wise characterization that enables the use of token-wise rewards.

Formulation for RLHF: From Bandit to MDP

In this section, we introduce our MDP formulation for RLHF. Section 3.1 describes how to characterize RLHF using token-wise MDPs in the context of LLMs. Section 3.2, we provide the learning objective under this framework. Lastly, Section 3.3 demonstrates the advantages of the token-wise MDP formulation compared to the sentence-wise bandit formulation.

We model the RLHF problem as a Markov decision process (MDP), which is denoted as a tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},r,\rho,H)$. Here $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $\mathcal{P}:\mathcal{S}\times\mathcal{A}\rightarrow\Delta(\mathcal{S})$ is the transition kernel, $r$ denotes the reward function, $\rho$ signifies the initial state distribution and $H$ is the maximal number of interaction steps. A (Markov) policy in MDPs $\pi:\mathcal{S}\rightarrow\Delta(\mathcal{A})$ is a mapping from state to a distribution over actions. The interaction between the environment $\mathcal{M}$ and the agent can be described as follows. Initially, the starting state $s_{1}$ is sampled from the initial distribution $\rho$. At the $h$-th step, the agent observes the state $s_{h}$ and selects an action $a_{h}$ based on its policy. The environment then transits to the next state $s_{h+1}$, which is sampled from the distribution $\mathcal{P}(\cdot{\,|\,}s_{h},a_{h})$. This interaction continues until a certain ending condition is satisfied, which will be triggered within $H$ steps.

In our MDP formulation for RLHF, we also model the preference signal using BT model (Bradley and Terry, 1952), but replace the sentence-level reward function in (2.1) with token-wise reward functions. In specific, for any trajectory pair $\tau^{1}=\{(s_{h}^{1},a_{h}^{1})\}_{h=1}^{H}$ and $\tau^{2}=\{(s_{h}^{2},a_{h}^{2})\}_{h=1}^{H}$In fact, these two trajectories can have different lengths, say $\tau^{1}=\{(s_{h}^{1},a_{h}^{1})\}_{h=1}^{H_{1}}$ and $\tau^{2}=\{(s_{h}^{2},a_{h}^{2})\}_{h=1}^{H_{2}}$ with $1\leq H_{1},H_{2}\leq H$. These trajectories can be extended to length $H$ by assuming that the state ending with EoS is absorbing and yields zero reward. This modification is to simplify the mathematical formulation and does not affect the problem modeling in (3.1). For the sake of clarity, the following theoretical discussion may focus on length-$H$ trajectories., the preference is specified by

Compared to literature that formulates the RLHF problem as a contextual dueling bandit, a subtle difference is that the policy in the contextual dueling bandit maps a prompt to a distribution over sentences, which does not capture the autoregressive nature of LLMs. In contrast, our MDP formulation precisely captures this nature. We defer the discussion of these two types of policies in Section C.2. More importantly, the main difference is that the reward function in the MDP formulation is defined on a token level, which contrasts significantly with the sentence-level reward in the contextual dueling bandit. We discuss the advantages of token-level rewards in Section 3.3.

2 Learning Objective

Different from classical RL literature, where the sole goal is to maximize the reward function, the objective of RLHF is to maximize the reward function while ensuring that the learned policy does not deviate too much from the reference model (e.g., SFT model) too much. Inspired by this and the formulation of entropy-regularized MDPs (Williams and Peng, 1991; Ziebart, 2010), for any policy $\pi$, we define its corresponding regularized value-function by

Our learning objective is to find a near-optimal policy $\widehat{\pi}$, and its optimality gap is measured by the following suboptimality gap:

3 Advantages of Token-Wise MDP over Sentence-Wise Bandit

Intuitively, the distinction between token-based and trajectory-based rewards reflects the difference between sparse and dense reward settings. In the sparse reward scenario, exploration proves to be more challenging. To illustrate this, we focus on the deterministic MDP with an action set size of $A=|\mathcal{A}|$. We employ an autoregressive policy $\pi^{*}$ to represent the policy of a powerful LLM, such as GPT-4. Fixing a prompt $x$, given responses $(y^{1}=y_{1:H}^{1},y^{2}=y_{1:H}^{2})$, the evaluation provided by $\pi^{*}$ is

By comparing this with the BT models of bandit in (2.1) and of our MDP formulation in (3.1), we observe that the sentence-wise reward $r_{s}$ and token-wise as $r_{t}$ can be specified by

Intuitively, the responses that powerful LLMs tend to choose have higher rewards. In addition, it is straightforward to show that $r_{s}(x,y)=\sum_{h=1}^{H}r_{t}((x,y_{1:h-1}),y_{h})$. We also make the following natural assumption.

There exists a response $y=y_{1:H}$ satisfying $\pi^{*}(y{\,|\,}x)\geq A^{-\xi}$.

By the pigeon-hole principle, there must be a response $y$ such that $\pi^{*}(y{\,|\,}x)\geq A^{-H}$, implying that $\xi\leq H$. In practice, $\xi$ is usually much smaller than $H$ because the language model tends to choose the optimal response rather than making a random guess. Now, we define the interaction protocol and the sample complexity. The learner can determine a response $y=y_{1:H}$ and receive either $r_{s}(x,y)$ or $\{r_{t}((x,y_{1:h-1}),y_{h})\}_{h=1}^{H}$, depending on whether the sentence-level reward or the token-wise reward is used. The sample complexity is defined as the number of responses and corresponding reward signals that need to be gathered to find the optimal response $y^{*}=y_{1:H}^{*}$ with length $H$.

Suppose Assumption 3.1 holds. In the setting where only the sentence-wise reward $r_{s}$ in (3.6) is accessible, finding the optimal response $y^{*}$ requires a sample complexity of $A^{H}$. However, if token-reward signals $r_{t}$ in (3.6) are available, there exists an algorithm that can find the optimal policy with sample complexity $A^{\min\{\xi+1,H\}}$.

If only the sentence-level reward $r_{s}$ is available, the learner must try every possible response and determine the optimal one by ranking the collected sentence-level reward signals, resulting in a sample complexity of $A^{H}$. Instead, we consider a binary tree with depth $H+1$, where each node is indexed by some token sequence $y_{1:h}$ and has $A$ children $\{(y_{1:h},y_{h+1})\}_{y_{h+1}\in\mathcal{A}}$. All $A^{H}$ leaf nodes denote a unique prompt-response pair $(x,y_{1:H})$. We define a set of nodes as

Our key observation is that $|\mathcal{N}|\leq A^{\xi+1}$. We also maintain a node set $\mathcal{N}^{\prime}$. Initially, we set $\mathcal{N}^{\prime}=\emptyset$. If $\mathcal{N}^{\prime}$ is updated, we delete all paths containing some node in $\mathcal{N}^{\prime}$. Each query of a new path (response with length $H$) will identify an additional node in $\mathcal{N}$. Then we add it to $\mathcal{N}^{\prime}$ and delete all paths containing some node in $\mathcal{N}^{\prime}$. This operation ends after at most $A^{\xi+1}$ iterations. Finally, ranking all gathered rewards identifies the optimal $y^{*}=y_{1:H}^{*}$. Together with the fact that there exists as most $A^{H}$ nodes, we finish the proof of Proposition 3.2. To facilitate understanding, we visualize a simplified learning process in Figure 2. ∎

Since $\xi\ll H$ typically holds in practice, the gap between $A^{H}$ and $A^{\min\{\xi+1,H\}}$ is deemed large. Hence, Proposition 3.2 reveals the significant separation of sample complexity between two types of reward signals, providing theoretical insights into the superiority of the token-wise MDP formulation over the sentence-wise bandit formulation.

Reinforced Token Optimization

Motivated by Section 3, we tackle RLHF by treating it as an MDP problem. Under this MDP framework, we aim to develop an algorithmic framework that fully utilizes the token-level information. To this end, we develop the Reinforced Token Optimization (RTO) algorithm. At a high level, RTO consists of two main steps: (i) token-wise reward learning, where RTO learns a token-wise reward based on the preference data; and (ii) optimizing token-wise reward through RL training methods such as PPO. In Section 4.1, we provide a theoretically grounded version of RTO with guaranteed sample complexity. To align more closely with practice, we present a practical implementation of RTO in Section 4.2.

We focus on the offline setting and assume the access to an offline dataset $\mathcal{D}=\{(\tau^{w},\tau^{l})\}$ that contains several trajectory pairs, where $\tau^{w}=\{(s_{h}^{w},a_{h}^{w})\}_{h=1}^{H}$ is preferred over $\tau^{l}=\{(s_{h}^{l},a_{h}^{l})\}_{h=1}^{H}$. Each pair of trajectories shares the same initial state/prompt (i.e., $s_{1}^{w}=s_{1}^{l}$), but differs in the subsequent tokens. We also assume that the reward function is linear, and our following results are ready to be extended to general function approximation (Chen et al., 2022; Wang et al., 2023; Zhan et al., 2023a).

Following the standard reward learning pipeline (Ouyang et al., 2022), we learn the reward function via maximum likelihood estimation (MLE). Specifically, if we parametrize the reward function by $\theta$, then the MLE is given by

Suppose Assumption 4.1 holds. For $\beta>0$, $\lambda>0$, $\delta\in(0,1)$, if we choose $\varrho=\widetilde{\mathcal{O}}(\sqrt{d})$ (see (A.2)), then the output policy $\widehat{\pi}$ of Algorithm 1 satisfies

The first term in Theorem 4.2 measures how well the offline dataset covers the trajectory generated by the policy $\pi_{\beta}^{*}$. Typically, this term decreases at a rate of $|\mathcal{D}|^{-1/2}$ under the mild partial coverage assumption (Jin et al., 2021; Uehara and Sun, 2021; Xiong et al., 2022; Zhu et al., 2023; Zhan et al., 2023a), where $|\mathcal{D}|$ is the size of the offline dataset. The second KL term is always negative, and it arises from the goal of learning a regularized value. We also remark that our algorithm relies on the known transition kernel to compute the exact optimal policy with respect to $\widehat{r}$. While this is natural in the context of large language models, we provide insights on how to extend our findings to stochastic regularized MDPs and the variant of our RTO algorithm in Appendix B.

There have also been previous works (Pacchiano et al., 2021; Chen et al., 2022; Wang et al., 2023; Li et al., 2023b; Zhan et al., 2023a) studying RLHF under the MDP framework, also known as dueling RL and preference-based RL. However, these works do not consider the KL constraint, which is an essential component of RLHF. Furthermore, they do not explicitly emphasize the superiority of the MDP framework over the contextual dueling bandit problem in the context of LLMs, and their proposed algorithms lack practical implementation. In contrast, we will provide a practical implementation of our algorithm, demonstrating the practicality of our approach.

2 Practical Implementation

In this subsection, we shift our focus to developing a practical version of RTO. The key challenge in implementing RTO in Algorithm 1 lies in learning the token-wise reward to be optimized from the offline data. In the most popular frameworks outlined in Instruct-GPT (Ouyang et al., 2022), Claude (Bai et al., 2022), and LLaMA2 (Touvron et al., 2023) projects replace the last layer of the LLM with a linear layer for a scalar output and maximize the log-likelihood as in (2.2). However, this approach gives only a sentence-level reward. To bridge the gap in the literature, we present our practical version of RTO in Algorithm 2, which features a novel calculation of token-wise reward. Our key observation is that, given a trajectory $\tau=\{(s_{h},a_{h})\}_{h=1}^{H}$, we have

Building upon this result and combining it with the definition of the BT model in (3.1), for any trajectory pair $\{\tau^{j}=\{(s_{h}^{j},a_{h}^{j})\}_{h=1}^{H}\}_{j=1}^{2}$ satisfying $s_{1}^{1}=s_{1}^{2}$, we have

where $x$ is the prompt, $y_{1:h-1}$ is the tokens generated so far, and $y_{h}$ is the token chosen at the current step. In contrast to the previous PPO implementation with sparse reward in (2.3), we will assign the token-wise reward function defined in (4.5) to each step. Formally, for any $h$, we define

Experiments

In this section, we conduct real-world alignment experiments to verify the effectiveness of RTO. We provide the experimental setups and experimental results in Sections 5.1 and 5.2, respectively.

We study the performance of our model on the single-turn dialogue generation task (Bai et al., 2022). Given a text sequence ($x$) representing dialogue history between the user and the assistant, the goal of the task is to generate a helpful response ($y$) as the answer. For this purpose, we utilize the helpful subset of the Anthropic Helpful and Harmless (HH-RLHF) dialogue datasethttps://huggingface.co/datasets/Anthropic/hh-rlhf (Bai et al., 2022). Each sample of the HH-RLHF dataset is accompanied by a history and two alternative responses, with preferences annotated by humans. We provide an example of the HH-RLHF dataset in Appendix D.1.

Model and Baselines.

where $(x,y=y_{1:H})$ is the prompt-response pair, $\pi$ is the current policy, and $(\beta_{1},\beta_{2})$ are tuning hyperparameters. In other words, we use the DPO to extract a sentence-level reward from the preference data, assign it to the last token, and fine-tune the model using the PPO algorithm. We refer to this baseline as DPPO. For our proposed RTO, we use the DPO model to derive a token-wise reward model (LABEL:eq:prac:5), and train the policy to align human preference using PPO, as detailed in Algorithm 2. The training configurations of all aforementioned models are given in Appendix D.2.

Evaluation.

We use two metrics to evaluate the alignment performance of different methods: oracle reward evaluation and GPT-4 evaluation. For oracle reward evaluation, we employ an open-sourced reward model https://huggingface.co/weqweasdas/RM-Mistral-7B as oracle, which is trained from Mistral-7B (Jiang et al., 2023) and achieves one of the highest accuracy on the Anthropic Helpful and Harmless dialogue task (Bai et al., 2022). For each pair of models, given the same prompt, we generate the responses using both models and calculate the rewards of the responses by the oracle reward model. We then compare the rewards of the two models and report the win rates between them. The GPT-4 evaluation, on the other hand, leverages the capabilities of GPT-4 itself and has been demonstrated to correlate with human evaluations (Rafailov et al., 2023) well. Given the two responses for the same prompt using two models, we ask GPT-4 about which one is better and calculate the win rates, following Rafailov et al. (2023). The prompt for GPT-4 evaluation is provided in Table 6 in Appendix D.3. For each evaluation by the oracle reward model, 400 dialogue histories from the test dataset are sampled, and for each evaluation by GPT-4, 100 dialogue histories are sampled.

2 Experimental Results

The experiment results of our proposed method and baselines are detailed in Table 1. This table meticulously presents the win rates between different models, assessed through both the oracle reward and the GPT-4.

From these results, we can see that the model trained by RTO achieves win rates over $50\%$ against all other baselines, especially compared to DPO, evaluated by both the oracle reward model and GPT-4. This highlights the effectiveness of the RTO algorithm in alignment tasks. Furthermore, the model trained by RTO gets a win rate of $61.1\%$ evaluated by the oracle reward model and a win rate of $56\%$ evaluated by GPT-4 over the DPPO algorithm. This implies that the token-wise reward mechanism significantly improves the performance of the RL algorithm in training models. To further investigate the benefits of the token-wise reward mechanism in the optimization process, we compare the estimated reward during the training period in Figure 3. In this figure, the x-axis represents the training iterations (1 epoch roughly corresponds to 160 PPO training iterations). The y-axis represents the reward given by the implicit reward model derived from the DPO model (the reward model used in training) per batch. As we can see, in one epoch, the reward of the model trained by RTO can achieve about $0.1$, while the reward of the model trained by DPPO is roughly $-0.25$. The results demonstrate that the token-wise reward mechanism significantly enhances the training process, leading to a remarkably higher reward. All these empirical findings demonstrate the token-wise reward mechanism’s advantage in improving model performance.

Conclusion

In this work, we suggest that the suboptimal performance of open-source implementations of PPO may be attributed to their reliance on sentence-level rewards, which neglect valuable token-wise information. To tackle this problem caused by the limitations of the previous bandit framework for RLHF, we propose an MDP formulation for RLHF that better characterizes token-wise information, along with theoretical insights demonstrating its superiority. Building upon this formulation, we introduce a novel algorithm called Reinforced Token Optimization (RTO), which leverages token-wise rewards to improve the policy. RTO is shown to be both provably sample-efficient and practical. Our practical implementation involves a novel token-wise reward learning approach via DPO, followed by optimization using PPO. This innovative combination of DPO and PPO allows RTO to effectively utilize token-level information and significantly improve the performance of baselines. Furthermore, our research opens up several intriguing future research directions, such as designing alternative methods for learning token-wise rewards beyond DPO and exploring other effective algorithms for optimizing token-level rewards besides PPO.

References

Appendix A Proof of Theorem 4.2

Recall that the visitation measure of policy $\pi$ is

Under this notation, we can rewrite the value function in (3.2) as

For simplicity, we will use the shorthand $d^{*}=d^{\pi_{\beta}^{*}}$.

Our proof relies on the following standard MLE analysis.

It holds with probability $1-\delta$ that

where $C$ is an absolute constant and $\Upsilon=1/(2+\exp(-2HLB)+\exp(2HLB))$.

See e.g., Faury et al. (2020); Pacchiano et al. (2021); Zhu et al. (2023) for a detailed proof. ∎

Back to the proof of Theorem 4.2, we first decompose the suboptimality gap defined in (3.5) as

Then we analyze these three terms respectively.

Recall that the pessimistic reward $\widehat{r}$ defined in (4.2) takes the form

where the first inequality is obtained by Cauchy-Schwarz inequality, and the last inequality follows from Lemma A.1.

Term (ii).

Similar to the derivation of (A), we have

where the first inequality uses Cauchy-Schwarz inequality, and the last inequality is implied by Lemma A.1.

Term (iii).

To handle this term, we introduce the following performance difference lemma for MDP with KL constraint.

For any reward function $r$ and policy pair $(\pi,\pi^{\prime})$, it holds that

When $\beta=0$, the regularized MDP becomes the standard MDP, and Lemma A.2 reduces to the standard performance difference lemma (Kakade and Langford, 2002). Applying Lemma A.2 to Term (iii) in (A), we have

where the second equality follows from the fact that $\widehat{\pi}$ is the optimal policy with respect to $V_{\beta}^{\pi}(s;\widehat{r})$ and the expression of optimal policy $\widehat{\pi}(a{\,|\,}s)=\exp\{(Q_{\beta}^{\widehat{\pi}}(s,a;\widehat{r})-V_{\beta}^{\widehat{\pi}}(s;\widehat{r}))/\beta\}$ in (3.4), and the last equality is obtained by the definition of KL divergence.

Finishing the Proof.

Plugging (A), (A), and (A) into (A), we obtain that

which finishes the proof of Theorem 4.2. ∎

If we do not have access to the exact optimal policy $\widehat{\pi}$ with respect to $\widehat{r}$, we can use the policy optimization algorithms to find a near-optimal optimal policy $\widetilde{\pi}$. In such case, Term (iii) in (A) becomes $V_{\beta}^{\pi_{\beta}^{*}}(\rho;\widehat{r})-V_{\beta}^{\widetilde{\pi}}(\rho;\widehat{r})=V_{\beta}^{\pi_{\beta}^{*}}(\rho;\widehat{r})-V_{\beta}^{\widehat{\pi}}(\rho;\widehat{r})+V_{\beta}^{\widehat{\pi}}(\rho;\widehat{r})-V_{\beta}^{\widetilde{\pi}}(\rho;\widehat{r})$, and we need to handle the additional error term $V_{\beta}^{\widehat{\pi}}(\rho;\widehat{r})-V_{\beta}^{\widetilde{\pi}}(\rho;\widehat{r})$. This type of error analysis has been established for NPG (Agarwal et al., 2021; Cen et al., 2022) and PPO (Cai et al., 2020; Wu et al., 2022; Zhong and Zhang, 2024).

A.1 Proof of Lemma A.2

Without loss of generality, we assume that the initial state is a fixed state $s_{1}\in\mathcal{S}$. For simplicity, we also omit the dependency of $r$ in the regularized Q-function and value function. First, we have

Plugging this into Term ($\star$) of (A.7), we have

where we use $d_{h}^{\pi}(s)$ to denote the visitation measure at the $h-$th step. Meanwhile, we rewrite ($\star\star$) in (A.7) as

Plugging (A.1) and (A.9) into (A.7), we have

Appendix B Variants of Reinforced Token Optimization

Different from Algorithm 1 where the learner constructs a pessimistic reward estimation and then outputs its corresponding optimal policy. Indeed, we can also perform pessimistic planning with respect to the value function to find the near-optimal policy:

Suppose Assumption 4.1 holds. For $\beta>0$, $\lambda>0$, $\delta\in(0,1)$, if we choose $\varrho=\widetilde{\mathcal{O}}(\sqrt{d})$ (see (A.2)), then the output policy $\widehat{\pi}$ of (B.1) satisfies

By Lemma A.1, we know that $\theta^{*}\in\Theta$ with probability $1-\delta$. This implies that

Plugging this into the definition of the suboptimality gap in (3.5), we have

Now we introduce the notation of $\widehat{\theta}$:

Under this notation, we further obtain that

where the second inequality uses Cauchy-Schwarz inequality, and the last inequality is obtained by Lemma A.1. Therefore, we conclude the proof of Theorem B.1. ∎

In (B.1), we assume that the transition kernel is known so that we can compute the state distribution $d^{\pi}$ induced by the policy $\pi$. Although this is natural in LLMs, we briefly sketch the extension to the unknown transition setting. Following Zhan et al. (2023a), which is inspired by previous works on standard reward-based RL theory (Uehara and Sun, 2021; Liu et al., 2022; Zhong et al., 2022; Liu et al., 2023; Huang et al., 2024), we can also construct a confidence set for the transition kernel

where $P(\tau)$ is the probability of observing the trajectory $\tau$ under the transition $P$ and $\zeta$ is a tuning parameter. With a proper choice of $\zeta$, one can also show that $\mathcal{P}\in\Theta_{\mathcal{P}}$ with high probability. Then we can perform the following pessimistic planning

where $d_{P}^{\pi}$ denotes the state distribution induced by policy $\pi$ under the environment $P$. Combining the analysis of Theorem B.1 and previous work on offline RL (Uehara and Sun, 2021; Zhan et al., 2023a), we can also establish a similar result to Theorem B.1, but with an additional estimation error for the transition kernel part. As this part is standard and not the focus of our work, we omit it for simplicity.

Appendix C Additional Discussions

Direct Preference Optimization (DPO) is a representative algorithm of the direct preference learning algorithm (Rafailov et al., 2023; Zhao et al., 2023; Azar et al., 2023; Tang et al., 2024). From a high level, these type of algorithms aim to skip the reward modeling and learn directly from the preference data, hence the name direct preference learning. In this section, we introduce the mathematical principle of DPO for completeness.

We first recall that in the original two-staged learning paradigm, we aim to optimize the following KL-regularized target:

One notable feature of this KL-constrained optimization problem is that it admits a closed-form solution, as summarized in the following lemma.

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

Therefore, for any fixed reward function $r$, it leads to a closed-form policy:

Interestingly, while the DPO is derived from the sentence-level reward function and BT model, the implicit reward naturally gives a token-wise characterization of the prompt-response pair and can be leveraged as a dense reward signal for the PPO training.

C.2 Autoregressive Policy

For the policy defined in a contextual dueling bandit setting, it maps from a prompt to a complete sentence. For ease of presentation, we call this type of policy the predetermined policy since it determines the entire sentence regardless of the generation process. In contrast, the Markov policy defined in the MDP formulation generates responses autoregressively: it considers not only the prompt but also the tokens generated so far. By definition, the Markov policy is at least as good as the policy that determines the whole sentence based solely on the prompt. In deterministic MDPs, the optimal action sequence is predetermined given the initial state, which demonstrates the equivalence of these two types of policies. However, for stochastic MDPs, the Markov policy is strictly more expressive than the predetermined policy. The transition can be stochastic for various reasons. For example, if the LLM uses an external search engine, the next state $s_{h+1}$ depends not only on the current tokens $(x,y_{1:h})$ but also on the text generated by the external search engine $\pi^{\prime}(\cdot{\,|\,}x,y_{1:h})$, making it stochastic. Moreover, RLHF may have applications in other scenarios, such as robotics (Christiano et al., 2017), where the transition kernel is stochastic. To clarify, we distinguish these two types of policies in the following proposition.

There exists an MDP such that the value of any predetermined policy is at least $0.5$ less than that of optimal Markov/autoregressive policy.

We construct an MDP $\mathcal{M}$ with state space $\mathcal{S}=\{s_{0},s_{1},s_{2}\}$, action space $\mathcal{A}=\{a_{1},a_{2}\}$, horizon $H=2$, fixed initial state $s_{0}$. The reward $r$ and transition kernel $\mathcal{P}$ are given by

It is straightforward to see that the optimal autoregressive policy achieves a value of $1$. In contrast, any predetermined policy only achieves a value of $0.5$. This completes the proof. ∎

Appendix D Additional Experimental Details

Prompt: Human: What does ugly Christmas sweater mean? Assistant: Chosen response: The ugly Christmas sweater is a popular meme that’s become pretty big over the past few years. The concept is simple: you buy a cheap ugly Christmas sweater, you wear it on Christmas day, and you post a photo of yourself in the sweater on social media with some funny caption or message. Rejected response: It means something that a person might wear as a Christmas sweater. And by “Christmas sweater”, I assume you mean a sweater that people might wear around Christmas time, and not a regular sweater worn in the winter, and not a wool sweater.

D.2 Training Configurations

We provide the training configuration of SFT, DPO, PPO, DPPO, and RTO below. In the table of the training configuration of the standard PPO algorithm, we also present the configuration of training the reward model used in the PPO algorithm in this table.

D.3 Evaluation Details

We employ the implicit reward model, derived from the DPO model, as the Oracle reward model. This evaluation is conducted for both dialogue generation and summarization tasks. To assess the performance across models, we use top-p sampling and set $p=0.99$ and temperature $\tau=0.9$ to generate completions for 400 prompts from the test set. These samples are then compared based on their rewards to calculate the win rate of one model over another.

Evaluation via GPT-4:

Following the previous work (Rafailov et al., 2023), for evaluations utilizing GPT-4, completions are sampled by top-p sampling method with temperature of $\tau=0.9$ and $p=0.99$ for 100 prompts. To mitigate any positional bias inherent in GPT-4’s responses, we ensure that the order of completions within each pair is randomized. The version of the GPT-4 we used is GPT-4-0613, and the specific prompt utilized for GPT-4 evaluation is detailed as follows.