The Three Stages of RLHF

The Motivation: From Prediction to Instruction-Following

As we've discussed, a base LLM trained solely on next-token prediction over a vast internet corpus is a powerful knowledge repository, but it's not inherently an instruction-following agent. Its objective function is misaligned with a user's goal of getting a direct, helpful answer. Getting a base model to perform a specific task often requires clever "prompt engineering" to trick it into the desired format.

The video clearly illustrates the problem: the base model completes the text in a statistically likely way (like continuing a forum post), while the aligned model follows the implicit instruction to answer the question. RLHF is the machinery that bridges this gap.

The Three Stages of RLHF

The RLHF pipeline is a sophisticated, three-act process that progressively shapes a model's behavior. It combines supervised learning with reinforcement learning to instill human preferences into the model.

The three stages are:

1. Supervised Fine-Tuning (SFT)
2. Reward Model (RM) Training
3. Reinforcement Learning (RL) Optimization

This entire workflow can be visualized as a sequence of transformations, starting with a pre-trained model and culminating in an aligned one.

Let's now briefly walk through each of these stages.

Now, let's unpack each of those steps with a more technical lens.

Stage 1: Supervised Fine-Tuning (SFT)

The journey begins not with reinforcement learning, but with a more familiar technique: supervised fine-tuning.

• Objective: To adapt the general-purpose pre-trained LLM to the format of instruction-following. The goal is to teach it what a "good answer" looks like and to establish a strong starting point for the subsequent RL stage.
• Process: A small, high-quality dataset of prompt-response pairs is created. Human labelers are given various prompts and asked to write out the ideal, high-quality response. The pre-trained LLM is then fine-tuned on this dataset.
• Mechanism: This is standard supervised learning. The training objective is the same causal language modeling objective used in pre-training—minimizing the cross-entropy loss to predict the next token. Crucially, the loss is typically calculated only on the tokens of the response, not the prompt (a technique known as loss masking).

The output of this stage is an SFT model. This model is already significantly better at following instructions than the base model, but its knowledge of human preferences is limited by the size and diversity of the demonstration dataset.

Stage 2: Reward Model (RM) Training

This is where human feedback is explicitly modeled. The goal is to create a scalable proxy for human judgment, so we don't need a human to evaluate every single output during the final training phase.

• Objective: To train a model that takes a prompt and a response and outputs a single scalar score representing "goodness" or human preference.
• Process:
  1. A prompt is selected from a dataset.
  2. The SFT model from Stage 1 generates multiple different responses (e.g., A, B, C, D) for that prompt.
  3. A human labeler ranks these responses from best to worst (e.g., D > B > A > C).
  4. This ranking is broken down into a dataset of pairwise comparisons: (Prompt, Winner Response, Loser Response). For the ranking D > B > A > C, this yields pairs like (D, B), (D, A), (D, C), (B, A), etc.
  5. A separate model, the Reward Model (RM), is trained on this preference data. The RM is typically initialized from the SFT model, but its final layer is replaced with a linear head that outputs a single scalar value.
• Mechanism: The RM is trained to assign a higher score to the "winner" response than the "loser" response for any given pair. This is framed as a classification problem, often using a pairwise ranking loss based on the Bradley-Terry model. The loss function aims to maximize the margin between the scores of preferred and dispreferred responses.

The output of this stage is a trained Reward Model. This model encapsulates human preferences and can now be used as an automated, differentiable reward function for the final stage.

Stage 3: RL Optimization with PPO

In the final stage, we use the Reward Model to fine-tune the SFT model. This is where reinforcement learning enters the picture.

• Objective: To adjust the weights of the SFT model (now called the "policy") to generate responses that maximize the score given by the Reward Model.
• Process: The training is an active loop:
  1. A prompt is sampled from the dataset.
  2. The current policy (the LLM we are training) generates a response.
  3. The Reward Model evaluates this (prompt, response) pair and returns a scalar reward.
  4. This reward signal is used to update the policy's weights through an RL algorithm, most commonly Proximal Policy Optimization (PPO).
• Mechanism: PPO is a policy gradient algorithm that encourages updates that increase the expected reward. However, there is a critical constraint. The objective function includes a penalty term to prevent the policy from diverging too far from the original SFT model. This penalty is the Kullback-Leibler (KL) divergence between the current policy's output distribution and the SFT model's output distribution.

The final objective function is a combination of the reward and this KL penalty:

Where:

• $r_{\varphi }(x,y)$ is the score from the Reward Model.
• ${\pi }_{\text{RL}}$ is the policy being trained.
• ${\pi }_{\text{SFT}}$ is the frozen SFT model from Stage 1, acting as a reference.
• $\beta$ is a coefficient controlling the strength of the KL penalty.

This KL penalty is crucial. It acts as a regularizer, addressing the "reward hacking" failure mode we discussed previously. Without it, the policy might learn to generate gibberish that happens to trick the RM into giving a high score. The penalty pulls the model back towards the sane, coherent distribution of the SFT model, balancing alignment with fluency.

The output of this final stage is the fully aligned, RLHF-tuned model, ready for evaluation and deployment.

Conclusion

In this lesson, we have mapped out the high-level architecture of the Reinforcement Learning from Human Feedback pipeline. We've seen how it systematically combines different training paradigms to instill complex human preferences into a large language model.

The key takeaways are:

• RLHF is a three-stage process designed to align LLMs with human intent, making them more helpful and harmless than their base pre-trained versions.
• Stage 1 (SFT): Supervised fine-tuning on a small set of high-quality demonstrations teaches the model the format of following instructions, providing a strong foundation.
• Stage 2 (RM Training): A Reward Model is trained on human-ranked responses to create a differentiable and scalable proxy for human preferences.
• Stage 3 (RL Optimization): The SFT model is optimized using a reinforcement learning algorithm (PPO) to maximize the score from the Reward Model, while a KL-divergence penalty prevents it from diverging too far from the coherent SFT policy, mitigating reward hacking.

You now have a bird's-eye view of the entire process. In our next lesson, we will zoom in on Stage 2, taking a deeper look at how the reward model is trained, the structure of preference data, and the mathematical underpinnings of the Bradley-Terry model used in the loss function.