The RLHF Pipeline: Components and Workflow

Reinforcement Learning from Human Feedback (RLHF) provides a structured approach to steer Large Language Models (LLMs) towards desired behaviors, such as helpfulness, honesty, and harmlessness, by incorporating human judgments directly into the training loop. While the concept is straightforward, the implementation involves a multi-stage pipeline that combines elements of supervised learning, preference modeling, and reinforcement learning. Understanding this workflow is essential to applying RLHF effectively.

The typical RLHF pipeline, illustrated below, consists of three primary stages, although variations exist:

1. Model Preparation (Optional but Common): Starting with a base pre-trained LLM and optionally fine-tuning it using supervised methods.
2. Reward Model Training: Creating a model that learns to score LLM outputs based on human preferences.
3. RL Policy Fine-Tuning: Using the reward model to further fine-tune the LLM via reinforcement learning to maximize preferred outputs.

High-level overview of the RLHF pipeline stages and their interactions. SFT is often used but optional. The SFT model frequently serves as the reference policy ($\pi_{ref}$) during RL fine-tuning.

Let's examine each stage in more detail.

Stage 1: Model Preparation (Base Model and Optional SFT)

The process typically begins with a capable pre-trained LLM. This model has learned general language patterns from extensive amounts of text data but may not be specifically adapted to follow instructions or exhibit desired behaviors like helpfulness.

Often, an intermediate step called Supervised Fine-Tuning (SFT) is performed. In this stage, the pre-trained model is fine-tuned on a smaller, high-quality dataset of input prompts and desired outputs (demonstrations). This dataset is curated to exemplify the target behavior (e.g., helpful instruction following, specific conversational styles).

The SFT stage aims to adapt the model to the distribution of prompts it will likely encounter and teach it the basic format and style of desired responses. The resulting model from this stage, let's call it $\pi_{SFT}$, often serves as the starting point for the subsequent RLHF stages. Specifically, it becomes the initial policy for RL fine-tuning and the reference model ($\pi_{ref}$) used for regularization during RL training. While beneficial, SFT is not strictly mandatory for RLHF; one could potentially start reward modeling and RL fine-tuning directly from the base pre-trained model.

Stage 2: Reward Model Training

The core idea of RLHF is to learn human preferences. This is achieved by training a separate model, the reward model ($r_\theta$), to predict which LLM output a human would prefer.

1. Data Collection:

   • A set of prompts ($x$) is selected (often drawn from the distribution used in SFT or expected in deployment).
   • The current LLM (usually the SFT model $\pi_{SFT}$) generates two or more responses ($y_1, y_2, ...$) for each prompt $x$.
   • Human annotators are presented with the prompt $x$ and a pair of responses (e.g., $y_i, y_j$) and asked to choose which response they prefer (or indicate a tie, or that both are unacceptable). Common practice focuses on pairwise comparisons.
   • This process yields a dataset of preference tuples: $(x, y_w, y_l)$, where $y_w$ is the preferred ("winning") response and $y_l$ is the less preferred ("losing") response for prompt $x$.

2. Model Architecture: The reward model typically shares the same architecture as the LLM being trained (or uses the base pre-trained architecture). The final layer, however, is modified to output a single scalar value representing the predicted reward (preference score) instead of token probabilities.

3. Training Objective: The reward model $r_\theta(x, y)$ is trained to assign a higher score to the preferred response $y_w$ than the less preferred response $y_l$. A common approach uses a loss function derived from the Bradley-Terry model, which models the probability that $y_w$ is preferred over $y_l$: $P(y_w \succ y_l | x) = \sigma(r_\theta(x, y_w) - r_\theta(x, y_l))$ where $\sigma$ is the sigmoid function. The model is trained by minimizing the negative log-likelihood of the human preference labels over the dataset: $\mathcal{L}(\theta) = - \mathbb{E}_{(x, y_w, y_l) \sim D} [\log \sigma(r_\theta(x, y_w) - r_\theta(x, y_l))]$ This loss encourages the reward model $r_\theta$ to output a larger difference in scores between the winning and losing responses.

The quality and calibration of this reward model are significant factors influencing the success of the entire RLHF process. We'll discuss challenges in reward modeling later in this chapter.

Stage 3: RL Policy Fine-Tuning

With a trained reward model $r_\theta$ that supposedly captures human preferences, the final stage uses reinforcement learning to fine-tune the LLM policy $\pi_\phi$ (initially $\pi_{SFT}$) to maximize the rewards assigned by $r_\theta$.

1. The RL Setup:

   • Policy: The LLM itself, $\pi_\phi(y|x)$, parameterized by $\phi$.
   • Action Space: The vocabulary of tokens the LLM can generate.
   • Observation Space: The sequence of tokens generated so far given the input prompt $x$.
   • Reward Function: Provided by the learned reward model $r_\theta(x, y)$ applied to the final generated sequence $y$ for a given prompt $x$.

2. Optimization Process: An RL algorithm, most commonly Proximal Policy Optimization (PPO), is used to update the policy parameters $\phi$. The process generally involves:

   • Sampling prompts $x$ from a dataset (often the SFT prompt distribution).
   • Generating a response $y$ for each prompt $x$ using the current policy $\pi_\phi(y|x)$.
   • Calculating the reward for the prompt-response pair $(x, y)$ using the reward model: $r = r_\theta(x, y)$.
   • Updating the policy parameters $\phi$ using the PPO algorithm to increase the likelihood of generating high-reward responses.

3. PPO Objective with KL Penalty: A critical aspect of RLHF is preventing the policy $\pi_\phi$ from deviating too drastically from the original SFT model $\pi_{SFT}$ (or the base pre-trained model if SFT was skipped). Over-optimizing solely for the learned reward $r_\theta$ can lead to generating nonsensical or repetitive text that happens to score highly (a phenomenon known as "reward hacking") or catastrophic forgetting of general language capabilities. To mitigate this, a penalty term based on the Kullback-Leibler (KL) divergence between the current policy $\pi_\phi$ and the reference policy $\pi_{ref}$ (usually $\pi_{SFT}$) is added to the reward. The objective optimized by PPO in RLHF is typically: $\text{Objective}(\phi) = \mathbb{E}_{(x,y) \sim D_{\pi_\phi}} [r_\theta(x, y) - \beta \text{KL}(\pi_\phi(y|x) || \pi_{ref}(y|x))]$ Here, $D_{\pi_\phi}$ is the distribution of prompt-response pairs obtained by sampling prompts $x$ and generating responses $y \sim \pi_\phi(y|x)$. The coefficient $\beta$ controls the strength of the KL penalty. This term encourages the policy $\pi_\phi$ to maximize the reward predicted by $r_\theta$ while staying relatively close to the behavior of the reference model $\pi_{ref}$, preserving language fluency and mitigating reward hacking.

The output of this stage is the final aligned LLM, $\pi_{\phi^*}$, which has been tuned to generate responses that are preferred according to the learned reward model, balanced by the need to maintain coherence and stay close to the initial capabilities learned during pre-training and SFT.

This three-stage process (SFT -> Reward Modeling -> RL Fine-Tuning) forms the backbone of many successful RLHF implementations. Subsequent sections will explore the details, challenges, and variations associated with each stage, particularly reward modeling and PPO optimization.