Policy Filtration in RLHF to Fine-Tune LLM
for Code Generation

Reinforcement Learning from Human Feedback (RLHF) is a key technique to align large language models (LLMs) with human values and preferences (Christiano et al., 2017; Ziegler et al., 2019; Ouyang et al., 2022). RLHF has been proven to be an essential process for LLMs to produce more helpful, harmless, and honest responses (Bai et al., 2022). Despite various non-RL algorithms such as DPO (Rafailov et al., 2024) are proposed, state-of-the-art applications such as ChatGPT/GPT-4 (OpenAI, 2023), Claude (Anthropic, 2023), and Gemini (Team et al., 2023) adopt the RL algorithm (e.g., PPO) for policy optimization. The key challenge of RLHF is the inaccuracy of the intermediate reward model. While there are researchers investigate how to learn reliable reward models (see e.g., Wang et al., 2024), we focus on how to learn better policy under the guidance of such inaccurate reward models.

We observe that, though the reward model gives inaccurate rewards in general, it can be more reliable in specific regions (e.g., when it gives high rewards) than the others. The observation is based on the simple experiment: We use a policy model fine-tuned for code generation to generate a set of responses for prompts in the HumanEval dataset. Later, we score these responses using a reward model trained with the common recipe (see Ouyang et al., 2022, and also Section 2) and compare them with the actual scores. We find that, across different sets of samples, the reward model is more reliable when it gives high or low rewards than when it gives moderate rewards. (This property also holds on other datasets and tasks. See Appendix A for more experiment results and further discussion.) Considering that RLHF updates the policy solely based on the reward signal, this observation motivates us to filter out the samples with possibly unreliable rewards aiming to improve RLHF by increasing the signal-to-noise ratio on training samples.

Based on this motivation, we propose a simple modification to the standard PPO-based RLHF algorithm (Ouyang et al., 2022), Policy Filtration for PPO (PF-PPO), that learns a filtered version of the policy using PPO. Specifically, we generate $N$ samples for each prompt, score these samples using the reward model, and use a filtered subset of these samples for subsequent policy training. We design filtering strategies to improve the reliability of the reward model on the filtered samples by maximizing the coefficient of determination ($R^{2}$) between the rewards and actual scores on these filtered samples. We show that the reward model can evaluate more accurately on these filtered samples, thus providing better training signal and improving the performance of the policy. Our method is also connected with reject sampling that filters out responses with low rewards during inference to yield a better response. Reject sampling is a simple but surprisingly strong inference-time strategy, whereas we adopt similar filtration in an RL algorithm.

Empirically, we show that PF-PPO can greatly improve the performance of LLMs on code generation tasks, which is challenging since complex logic behind these tasks makes the reward model inaccurate in general. We conduct extensive ablation studies to validate the design of our algorithm. Moreover, we illustrate the effectiveness of our algorithm by fine-tuning LLMs that achieves new sota on HumanEval and LeetCode Contest benchmarks across 7-billion-parameter LLMs. To evaluate whether PF-PPO can be effective on more challenging coding tasks, we create the LeetCode Contest benchmark that includes competition-level coding tasks for human experts. We find that the policy filtration technique can result in even more significant improvement on this challenging benchmark.

Limitation of reward model. The outcome of RLHF highly relies on the quality of the reward model. Unfortunately, the reward model can hardly provide accurate scores due to 1) the mis-specified reward modeling to represent human preferences (Lambert et al., 2023; Pitis, 2023); 2) the presence of incorrect and ambiguous preferences in the dataset (Ouyang et al., 2022; Bai et al., 2022), and 3) the poor generalization ability of the reward model (McKinney et al., 2023). The inaccuracy of reward model is attributed as one major cause of reward hacking and hallucination in LLMs (Kalai & Vempala, 2024). While there are previous papers try to improve the accuracy of the reward model itself (Wang et al., 2024; Coste et al., 2023; Zhang et al., 2024), the objective of our paper is to design a better RLHF algorithm in the face of inaccurate reward models. Moreover, Bai et al. (2022) also mentioned that using the output of the reward model directly in the RLHF process may not be a good choice. A possible solution is to penalize the outputs with low rewards more to improve the worst-case responses but they did not further implement this.

Reject sampling. Reject sampling (or best-of-N sampling) is a popular and effective inference-time strategy to enhance the response of an LLM by generating $N$ responses and select the best one according to a reward model (Nakano et al., 2021; Cobbe et al., 2021). This trick can yield good responses while keeping a tight KL constraint to the original policy. Inspired by its effectiveness in inference, researchers also try to involve this trick in policy optimization. For example, RAFT (Dong et al., 2023), BOND (Sessa et al., 2024) and vBoN (Amini et al., 2024) learn a policy that distills the best-of-$N$ policy using supervised fine-tuning losses. In a boarder sense, the rank information of the $N$ samples can also be leveraged. For example, RRHF (Yuan et al., 2023) and PRO (Song et al., 2024) train the policy using the combination of a ranking loss and a SFT loss (w.r.t. the best response) based on $N$ responses for each prompt. However, these algorithms do not adopt an elaborate RL algorithm, while state-of-the-art language models adopts RL algorithms in alignment, benefiting from the generalization power of the reward model especially in reasoning tasks (Ivison et al., 2024). Unlike these algorithms, we adopt the idea of reject sampling in the sampling phase of an RL algorithm instead of using supervised learning losses.

RLHF algorithms in the face of inaccurate reward models. One key challenge in RLHF is the inaccuracy of reward model, which can lead to reward over-optimization (Gao et al., 2023; Skalse et al., 2022; Chaudhari et al., 2024). Optimization with a policy constraint (e.g., a KL divergence between the target policy and the reference policy) is a remedy frequently used in not only RL-based algorithms (Ouyang et al., 2022; Wu et al., 2023; Zhu et al., 2023) but also direct policy optimization algorithms (Rafailov et al., 2024; Zhao et al., 2023; Liu et al., 2023). Going beyond policy constraint, Moskovitz et al. (2023) only maximize rewards up to a threshold to avoid excessive deviation from a pre-trained policy. In this paper, we not only rely on the policy constraint to optimize in the face of inaccurate rewards but also try to avoid using samples with unreliable rewards.

Notations. We use $[a,b]$ to denote the set $\{a,a+1,\cdots,b\}$ and use $[b]$ as the shorthand for $[1,b]$. We use $\oplus$ to denote the concatenation on tokens, and use $x_{a:b}$ as the shorthand for the concatenation $(x_{a}\oplus x_{a+1}\oplus\cdots\oplus x_{b})$. We use $c_{i}$ and $y_{i}$ to indicate the $i$-th token in the context $c$ (including task instruction, prompt, inputs, etc.) and the response $y$ respectively.

MDP formulation. We adopt a Markov decision process (MDP) formulation for RLHF. Specifically, language generation is formulated as an MDP $M=(\mathcal{S},\mathcal{A},P,R)$ with states $s\in\mathcal{S}$, actions $a\in\mathcal{A}$, transition probabilities $P\in\Delta(\mathcal{S})^{\mathcal{S}\times\mathcal{A}}$, and the next-state-based reward function $R:\mathcal{S}\to[0,1]$. Given a context $c$ with $T_{c}$ tokens, on each step $t\in[T_{c}+1,T]$¹¹1We fix the index of the terminal state to be the maximum length $T$. To adapt responses of different lengths, we left pad the context $c$., the language model $\pi_{\theta}(a_{t}|s_{t})$ selects a token $a_{t}=y_{t-T_{c}}$ based on the state $s_{t}:=(c_{1:T_{c}}\oplus y_{1:t-T_{c}-1})$. Then, the language model enters the next state $s_{t+1}:=(c_{1:T_{c}}\oplus y_{1:t-T_{c}})$ until the language model completes the response $y_{1:T-T_{c}}$. For simplicity, we will also use contextual-bandit-style notations, e.g., we denote the language generation process as $y\sim\pi_{\theta}(\cdot|c)$.

RLHF. Reinforcement learning with human feedback (RLHF) is an important process to address objective mismatch between the next-token-prediction objective in pre-training and our expectation of LLMs to follow the instructions and assist humans to complete various tasks. We briefly review the pipeline of RLHF.

• •

  Supervised fine-tuning. In the supervised fine-tuning (SFT) phase, a pre-trained LLM is fine-tuned with a high-quality supervised dataset collected for specific downstream tasks. Typically, the LLM is fine-tuned with a maximum likelihood loss, and we denote the output of this phase as $\pi^{\text{SFT}}$. While subsequent RLHF procedure is necessary for training high-quality LLMs, this phase alone can also yield an LLM that reasonably follows human instructions (see e.g., Longpre et al., 2023).

• •

  Reward model learning. In the reward model learning phase, we learn a reward model $r_{\phi}(y|c)\in[-1,1]$ parameterized by $\phi$ that scores the response $y$ to the context $c$ based on collected preference data $\mathcal{D}_{\text{HF}}:=\{(c,y^{w},y^{l})\}$ specifying that $y^{w}$ is a preferred response to $c$ than $y^{l}$. The reward model is initialized by $\pi^{\text{SFT}}$ with an additional output layer. A preference model links the reward model with the preference data, and Bradley-Terry model (Bradley & Terry, 1952) is a common choice:

  $$\mathbb{P}(y^{w}\succ y^{l}|c)=\sigma(R_{\phi}(y^{w}|c)-R_{\phi}(y^{l}|c)),$$

  (1)

  where $\sigma$ is the sigmoid function. The learning objective of reward model is to maximize the log-probability on preference data:

  $$\max_{\phi}\mathbb{E}_{(c,y_{w},y_{l})\sim\mathcal{D}_{\text{HF}}}\left[\log\mathbb{P}(y_{w}\succ y_{l}|c)\right].$$

  (2)

• •

  RL fine-tuning. In this stage, we fine-tune the language model $\pi_{\theta}$ to maximize the rewards given by the reward model with a policy constraint. The optimization problem is formulated as

  $$\max_{\theta}\mathbb{E}_{c}\mathbb{E}_{y\sim\pi_{\theta}(\cdot|c)}\left[r_{\phi}(y|c)-\beta D_{\text{KL}}(\pi_{\theta}(\cdot|c)||\pi^{\text{SFT}}(\cdot|c))\right].$$

  (3)

  The second term prevents the learned policy deviating too much from the SFT model, and this is a popular technique to alleviate reward over-optimization (Jaques et al., 2019; Stiennon et al., 2020).

PPO. Proximal policy optimization (PPO) (Schulman et al., 2017) is an RL algorithm that uses a clipped version of the policy gradient for more conservative and stable learning. It becomes a standard algorithm for RL fine-tuning in RLHF that optimizes the modified (cumulative) reward

where the reward model gives sparse rewards and the policy constraint yields dense rewards. PPO is an on-policy algorithm where the policy gradient is estimated based on the samples collected by the current policy $\pi_{\theta}$.

Our method is motivated by the observation that the reward model is more reliable for the responses assigned with high/low rewards (cf. Figure 1). Consequently, we conjecture that, if we wrap the policy with proper filtration during policy optimization of RLHF, the reward model can avoid yielding unreliable rewards and thus give better signal to guide policy learning.

Policy filtration. Given an unfiltered policy model $\pi_{\theta}(y|c)$ that generates responses $y$ to the context $c$, we denote the corresponding filtered policy as $\mu_{\theta}(y|c)$. We consider a family of policy filtration, from which we can sample responses to the context $c$ as follows: We first sample $N$ responses from $\pi_{\theta}(\cdot|c)$ and rank them by the reward model $R_{\phi}$, obtaining $y_{1},\cdots,y_{N}$ with $R_{\phi}(y_{1}|c)\geq\cdots\geq R_{\phi}(y_{N}|c)$. Then, given a weight vector $\mathbf{w}=(w_{1},\cdots,w_{N})$ satisfying $\sum_{i\in[N]}w_{i}=1$, we sample a one-hot vector $\mathbf{z}=(z_{1},\cdots,z_{N})$ from the categorical distribution parameterized by $\mathbf{w}$ such that $\mathbb{P}[z_{i}=1]=w_{i}$. At last, the filtered policy $\mu_{\theta}(\cdot|c)$ yields the response selected by $\mathbf{z}$ following $y=\sum_{i\in[N]}z_{i}y_{i}$.

We can define several filtered policies under this family. Specifically, we obtain the best-of-$N$ (BoN), best-random (BR), and best-worst (BW) filtered policy by setting the weight vector to $\mathbf{w}^{\text{BoN}}=(1,0,\cdots,0)$, $\mathbf{w}^{\text{BR}}=\left(\dfrac{1}{2},\dfrac{1}{2(N-1)},\cdots,\dfrac{1}{2(N-1)}\right)$, and $\mathbf{w}^{\text{BW}}=\left(\dfrac{1}{2},0,\cdots,0,\dfrac{1}{2}\right)$ respectively.

Training objective. Since our target is to learn a good filtered policy $\mu_{\theta}$, we consider the follow objective:

In practice, use the samples collected by the filtered policy $\pi_{\theta}$ as if they were collected by $\mu_{\theta}$ in the original PPO algorithm. This leads to Policy Filtration Proximal Policy Optimization (PF-PPO) listed in Algorithm 2, which is an algorithm that only modifies the sampling process of PPO.

Weight choice. By defining different weight vectors $\mathbf{w}$, we can obtain different policy filtering strategies for PF-PPO. Our objective is to choose a weight vector $\mathbf{w}$ such that the accuracy of the reward model on the responses generated by the filtered policies can be maximized. To measure this accuracy, we calculate the coefficient of determination (aka R-squared or $R^{2}$) (Draper, 1998) between the rewards and the actual scores of the responses generated by the policy. $R^{2}$ measures how well the actual scores can be predicted by the rewards with a linear model. Specifically, given a set of responses $\{(c_{i},y_{i})\}$ sampled from the filtered policy $y_{i}\sim\mu_{\theta}(\cdot|c_{i})$, we can collect the corresponding reward $R_{i}:=R_{\phi}(y_{i}|c_{i})$ and the actual score $s_{i}$. Then, we fit a linear model $f$ to predict the actual score based on the reward and denote the predicted score as $\hat{s}_{i}=f(R_{i})$. The R-squared is calculated as $1-\frac{\sum_{i}(s_{i}-\hat{s}_{i})^{2}}{\sum_{i}(s_{i}-\bar{s})^{2}}$ where $\bar{s}$ is the average of actual scores. Since PF-PPO optimizes the policy based on the rewards on these responses, how well these rewards indicate the actual performance is closely related to the final performance of our algorithm. We find $R^{2}$ well correlates with the final performance and can imply the level of reward over-optimization of the subsequent RLHF algorithm, therefore serving as a useful metrics to determine the weight vector used in PF-PPO.

To select a weight vector, we first checkpoint three policies $\pi_{\theta}$ collected from different stages of a standard RLHF process and collect responses using filtered policies $\mu_{\theta}$ in combination with different policy filtering strategies. Then, we group the responses with similar rewards, record the average actual score and reward for each group, and calculate the $R^{2}$ by treating each group as a sample point. We exam how different policy filtering strategies can improve the reliability of the rewards on the responses generated by the corresponding filtered policies.

We present the results in Table 1. We observe that best-random (BR) and best-worst (BW) can improve the reliability of the given reward model on sampled responses compared with unfiltered policy. The BoN strategy does not improve the $R^{2}$, which indicates that learning a BoN filtered policy may not result in good performance in RL, although learning for a best-of-$N$ policy using supervised learning presents good performance (Sessa et al., 2024).

In this paper, we propose a new reinforcement learning with human feedback (RLHF) method, Policy Filtration for Proximal Policy Optimization (PF-PPO), aimed at mitigating the adverse effects of reward noise. When training the reward model using the Bradley-Terry approach, the reward signal is generally more reliable in the high or low reward regions but less reliable in the moderate reward regions. Motivated by this observation, we adopt a rank-based method to selectively use sample from these reliable regions more in PPO to improve the quality of the signal provided by the reward model. We conduct comprehensive experiments on code generation tasks, demonstrating that PF-PPO outperforms existing baselines. Additionally, we analyze PF-PPO, standard PPO, and PPO with multiple responses in details and show that filtering samples with unreliable rewards can improve the performance of the outcome policy.