Aligning AI: An Introduction to RLHF and DPO
After an LLM completes its pre-training phase on large corpora of text, it becomes excellent at predicting the next token. However, a pre-trained base model is not necessarily a helpful assistant; it may complete prompts by repeating them, generating low-quality content, or outputting toxic or unsafe language.
To transform base models into helpful, honest, and harmless conversational agents, researchers perform alignment. Two primary paradigms govern this process: Reinforcement Learning from Human Feedback (RLHF) and Direct Preference Optimization (DPO). In this article, we will break down the mechanics and differences between these two methodologies.
The Baseline: Supervised Fine-Tuning (SFT)
Before applying preference alignment, models undergo Supervised Fine-Tuning (SFT). During SFT, the model is trained on curated datasets of instruction-response pairs (e.g., "Write a python function..." followed by high-quality code). SFT teaches the model the format of instructions and answers, but it fails to address nuanced human preferences or prevent subtle, undesirable behaviors.
Reinforcement Learning from Human Feedback (RLHF)
RLHF has been the standard alignment framework for state-of-the-art models like GPT-4 and Claude. It is a multi-stage process that leverages reinforcement learning to optimize the model’s outputs.
Aligning AI: RLHF and DPO alignment pipeline flowchart
1. Training the Reward Model
First, human annotators are presented with a prompt x and multiple model-generated responses (e.g., y_1, y_2). Annotators rank these responses from best to worst, creating preference pairs where y_w is the winning (preferred) response and y_l is the losing (dispreferred) response.
A Reward Model R_\phi is trained on this data. It outputs a scalar score representing how much humans would favor a given response. The loss function for training the reward model is:
\mathcal{L}_{\text{RM}}(\phi) = - \mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \log \sigma \left( R_\phi(x, y_w) - R_\phi(x, y_l) \right) \right]
Where \sigma is the sigmoid function. This maximizes the score difference between the winning and losing completions.
2. Reinforcement Learning via PPO
Once the reward model is trained, the SFT model is updated using Proximal Policy Optimization (PPO). The agent is the language model policy \pi_\theta, and the state space is the sequence of tokens. The policy generates a completion, the reward model scores it, and PPO updates the policy weights to maximize the expected reward.
To prevent the model from drifting too far from the original SFT distribution and generating gibberish (a phenomenon called reward hacking), a Kullback-Leibler (KL) divergence penalty is added to the objective:
\text{Objective}(\theta) = \mathbb{E}_{(x, y) \sim \pi_\theta} \left[ R_\phi(x, y) \right] - \beta \, \mathbb{D}_{\text{KL}} \left( \pi_\theta(y | x) \,\|\, \pi_{\text{ref}}(y | x) \right)
Where \pi_{\text{ref}} is the frozen SFT model, and \beta controls the strength of the KL penalty.
Challenges of RLHF
RLHF is notoriously unstable. It requires running four large models in memory simultaneously during training: the active policy \pi_\theta, the reference model \pi_{\text{ref}}, the reward model R_\phi, and a value model used by PPO. This makes it resource-heavy and sensitive to hyperparameters.
Direct Preference Optimization (DPO)
Developed by researchers at Stanford, Direct Preference Optimization (DPO) simplifies alignment by eliminating the need for a separate reward model and reinforcement learning step.
DPO mathematically proves that the RLHF objective can be optimized directly using a simple binary cross-entropy loss over preference pairs. The key insight is that the implicit reward function can be expressed directly in terms of the policy \pi_\theta and reference policy \pi_{\text{ref}}:
R^*(x, y) = \beta \log \frac{\pi_\theta(y | x)}{\pi_{\text{ref}}(y | x)}
Substituting this back into the preference loss yields the DPO loss function:
\mathcal{L}_{\text{DPO}}(\theta) = - \mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \log \sigma \left( \beta \log \frac{\pi_\theta(y_w | x)}{\pi_{\text{ref}}(y_w | x)} - \beta \log \frac{\pi_\theta(y_l | x)}{\pi_{\text{ref}}(y_l | x)} \right) \right]
How DPO Works Intuitively
The loss function forces the active policy \pi_\theta to increase the likelihood of the preferred response y_w relative to the reference policy \pi_{\text{ref}}, while simultaneously decreasing the likelihood of the dispreferred response y_l. This achieves the same mathematical optimum as RLHF but with single-stage supervised training.
Implementing DPO with the TRL Library
Below is a Python implementation utilizing the Hugging Face trl (Transformer Reinforcement Learning) library to train a model using DPO.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOTrainer, DPOConfig
from datasets import load_dataset
# 1. Load Tokenizer and Models
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="auto"
)
# DPO requires a reference model. If left None, DPOTrainer will clone
# the training model and freeze it to act as the reference.
ref_model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
# 2. Load and Format Dataset
# DPO datasets need columns: 'prompt', 'chosen', and 'rejected'
dataset = load_dataset("Intel/orca_dpo_pairs", split="train[:1000]")
# 3. Configure DPO Trainer Parameters
training_args = DPOConfig(
output_dir="./dpo_results",
beta=0.1, # Temperature hyperparameter for DPO (implicit reward scale)
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
learning_rate=5e-6,
logging_steps=10,
max_length=512,
max_prompt_length=256,
remove_unused_columns=False,
fp16=False,
bf16=True
)
# 4. Initialize and Run Trainer
dpo_trainer = DPOTrainer(
model=model,
ref_model=ref_model,
args=training_args,
train_dataset=dataset,
tokenizer=tokenizer,
)
dpo_trainer.train()
Comparison: RLHF vs. DPO
| Feature | RLHF | DPO | | :--- | :--- | :--- | | Number of Stages | 3 (SFT, Reward Model, PPO) | 2 (SFT, DPO) | | Active Models in VRAM| 4 (Policy, Reference, Reward, Value) | 2 (Policy, Reference) | | Stability | Low (sensitive to PPO hyperparams) | High (standard supervised loss) | | Compute Overhead | Very High | Moderate | | Optimality | Iterative approximations | Analytically exact |
By bypassing the instabilities and memory demands of PPO, DPO has become the preferred alignment strategy for open-source AI teams, facilitating fast and cost-effective alignment on downstream consumer hardware.