The RLHF Pipeline Overview

While Supervised Fine-Tuning (SFT) provides a strong foundation for instruction following, aligning a model more closely with complex human preferences often requires Reinforcement Learning from Human Feedback (RLHF). The RLHF process involves training a model based on which outputs humans judge to be better, allowing for optimization towards qualities like helpfulness and harmlessness that are difficult to capture solely through supervised examples. The standard RLHF pipeline typically consists of three main stages:

Initial Model Preparation (Pre-training & SFT): Start with a capable pre-trained language model. This model is then fine-tuned using supervised learning on a dataset of high-quality prompt-response pairs (as discussed in Chapter 25). This SFT step adapts the model to follow instructions and generate responses in the desired style and format. This SFT model serves as the starting point for the subsequent RLHF stages.
Reward Model (RM) Training: The core idea is to learn a model that can predict human preferences.
- Data Collection: Select a set of diverse prompts. For each prompt, generate several responses using the SFT model (or multiple model variants). Human labelers are then presented with pairs (or sometimes more) of these responses and asked to choose which one they prefer based on criteria like helpfulness, accuracy, or safety.
- RM Architecture: The reward model ( $RM$ ) is typically another language model (often initialized from the SFT model or a smaller pre-trained model) whose final layer is replaced with a linear layer outputting a single scalar value representing the predicted preference score. It takes a prompt and a response as input and outputs this score.
- Training Objective: The RM is trained on the collected pairwise preference data. Given a prompt $p$ and two responses $y_w$ (winner) and $y_l$ (loser) where $y_w$ was preferred over $y_l$ by a human, the RM is trained to assign a higher score to $y_w$ than to $y_l$ . A common loss function is the pairwise ranking loss: $\mathcal{L}_{RM} = - \mathbb{E}_{(p, y_w, y_l) \sim D} [\log(\sigma(r_{\theta}(p, y_w) - r_{\theta}(p, y_l)))]$ Here, $r_{\theta}(p, y)$ is the scalar score output by the reward model with parameters $\theta$ for prompt $p$ and response $y$ , $\sigma$ is the sigmoid function, and $D$ is the dataset of human preferences. This objective maximizes the probability that the preferred response $y_w$ gets a higher score.
Reinforcement Learning (RL) Fine-Tuning: The trained RM now acts as a proxy for human preferences, providing reward signals to fine-tune the SFT model further using an RL algorithm, typically Proximal Policy Optimization (PPO).
- RL Setup: The SFT model acts as the initial policy ( $\pi_{SFT}$ ) in the RL framework. The action space consists of the possible tokens the model can generate, and the state is the sequence of tokens generated so far.
- Optimization Loop: The process iteratively samples prompts $p$ from the dataset. The current RL policy ( $\pi_{RL}$ ) generates a response $y$ . The RM assigns a reward $r = r_{\theta}(p, y)$ to the generated response. The PPO algorithm then uses this reward signal to update the weights of the policy $\pi_{RL}$ .
- KL Divergence Penalty: A significant challenge in RL fine-tuning is that the policy might learn to generate outputs that maximize the RM score but deviate significantly from the original SFT model's distribution, potentially leading to nonsensical or repetitive text ("reward hacking"). To mitigate this, PPO in RLHF incorporates a Kullback-Leibler (KL) divergence penalty term into the reward function or objective. The objective aims to maximize the RM score while staying close to the initial SFT policy: $\text{Objective} = \mathbb{E}_{(p,y) \sim \pi_{RL}} [r_{\theta}(p, y) - \beta \cdot \text{KL}(\pi_{RL}(\cdot|p) || \pi_{SFT}(\cdot|p))]$ The term $\text{KL}(\pi_{RL}(\cdot|p) || \pi_{SFT}(\cdot|p))$ measures the divergence between the token distributions predicted by the RL policy and the original SFT policy for prompt $p$ . The hyperparameter $\beta$ controls the strength of this penalty, preventing the RL policy from straying too far from the learned distribution of the SFT model, thus maintaining coherence and language quality.

The following diagram illustrates this multi-stage process:

The three stages of the standard RLHF pipeline: Supervised Fine-Tuning (SFT), Reward Model (RM) training based on human preferences, and Reinforcement Learning (PPO) fine-tuning guided by the RM and a KL penalty against the original SFT model.

This pipeline allows the model to learn from comparative human feedback, refining its ability to generate responses that align better with desired attributes than what can be achieved through SFT alone. The following sections will detail the practical aspects of collecting preference data, training the reward model, and implementing the PPO fine-tuning step.

Was this section helpful?

References

Deep Reinforcement Learning from Human Preferences, Paul F. Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei, 2017 Advances in Neural Information Processing Systems 30, Vol. 30 DOI: 10.48550/arXiv.1706.03741 - Introduces the method of learning a reward function from human comparisons and then using this learned reward function to train deep reinforcement learning agents. This work laid a significant foundation for modern RLHF.
Proximal Policy Optimization Algorithms, John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, Oleg Klimov, 2017 arXiv DOI: 10.48550/arXiv.1707.06347 - Describes the Proximal Policy Optimization (PPO) algorithm, an effective and widely used policy gradient method for reinforcement learning, which is a core algorithm in the RL fine-tuning stage of RLHF.
Training language models to follow instructions with human feedback, Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Christiano, Jan Leike, Ryan Lowe, 2022 Advances in Neural Information Processing Systems 35, Vol. 35 - This paper introduces the InstructGPT model and details the complete three-stage RLHF pipeline-including Supervised Fine-Tuning (SFT), Reward Model (RM) training using human preferences, and Reinforcement Learning fine-tuning with PPO and a KL penalty-for aligning large language models.