Model Alignment: RLHF & Direct Preference Optimization
Raw pre-trained foundation models are pure next-token probability predictors. They lack an intrinsic understanding of user intent, safety, truthfulness, or societal norms.Alignment is the mathematical process of steering model outputs to be Helpful, Harmless, and Honest (the 3 Hs).
1. The Classical 3-Stage RLHF Pipeline (InstructGPT / PPO)
Pioneered by Christiano et al. (2017) and popularized by OpenAI's InstructGPT (Ouyang et al., 2022), the traditional Reinforcement Learning from Human Feedback (RLHF) framework functions across three sequential stages:
Supervised Fine-Tuning (SFT)
A raw base model is fine-tuned on curated instruction-response demonstrations to produce a capable initial conversational policy π_SFT.
Reward Modeling (RM)
Human annotators review model outputs and rank them: y_w ≻ y_l (preferred winning response vs losing response). A scalar reward network r_θ(x, y) is trained via the Bradley-Terry preference loss:ℒ_RM = -𝔼 [ log σ(r(x, y_w) - r(x, y_l)) ]
PPO Policy Optimization
The policy π_θ is updated via Proximal Policy Optimization (PPO) to maximize predicted reward while adding a KL-divergence penalty to prevent policy drift away from π_SFT.
max_π 𝔼_(x~D, y~π) [ r_θ(x, y) - β · D_KL( π_θ(y | x) ‖ π_ref(y | x) ) ]
The Problem with PPO: Training requires holding 4 separate models concurrently in GPU memory: the Actor (being trained), the Critic (value network), the Reference model (frozen baseline for KL), and the Reward model. This causes extreme memory bloat and hyperparameter instability.
2. Direct Preference Optimization (DPO): The Modern Standard
Introduced by Rafailov et al. (Stanford, 2023), DPO revolutionized model alignment by deriving an exact mathematical equivalence that eliminates the reward model and PPO loop entirely.
The Mathematical Core of DPO
Rafailov et al. proved that under the Bradley-Terry preference model, the optimal reward function can be expressed in closed form directly using the policy probabilities:r(x, y) = β · log( π_θ(y | x) / π_ref(y | x) ) + β · log Z(x)
ℒ_DPO(θ; π_ref) = -𝔼_(x, y_w, y_l) [ log σ( β · log( π_θ(y_w|x) / π_ref(y_w|x) ) - β · log( π_θ(y_l|x) / π_ref(y_l|x) ) ) ]
Why DPO Won: DPO fits preferences using standard supervised binary cross-entropy loss over log-likelihood ratios. It requires zero sampling during training, is mathematically exact, completely stable, and runs on a fraction of the GPU memory of PPO. Models like Llama 3, Mistral, and Zephyr align using DPO.
3. Constitutional AI & RL from AI Feedback (RLAIF)
Human labeling is slow, expensive, and subject to human bias and psychological fatigue. Anthropic introduced Constitutional AI (RLAIF) (Bai et al., 2022), replacing human annotators with AI feedback guided by an explicit constitution:
1. Critique and Revision (SL-CAI)
The model is prompted to evaluate its own initial response against explicit constitutional principles (e.g. "Please choose the response that is most harmless and least likely to promote violence") and generate a self-revised response.
2. AI Preference Ranking (RL-CAI)
A frontier evaluator model (e.g. Claude 3.5 Sonnet) compares response pairs and generates automated preference rankings, creating synthetic preference datasets at 1/100th the cost of human contractors.
4. Production Code: Training DPO with Hugging Face TRL
The following Python snippet demonstrates aligning a model on preference pairs using Hugging Face DPOTrainer:
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import DPOTrainer
from peft import LoraConfig
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
# Load model and reference model (using LoRA so reference model uses base weights)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="auto"
)
peft_config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
# Example Pairwise Preference Data: prompt, chosen (winning), rejected (losing)
training_args = TrainingArguments(
output_dir="./llama3-dpo-aligned",
per_device_train_batch_size=2,
learning_rate=5e-5,
lr_scheduler_type="cosine",
max_steps=100,
bf16=True,
logging_steps=10
)
dpo_trainer = DPOTrainer(
model=model,
ref_model=None, # None means it uses the frozen base model with LoRA disabled as ref
args=training_args,
beta=0.1, # Temperature parameter beta in DPO loss
train_dataset=dataset["train"],
tokenizer=tokenizer,
peft_config=peft_config,
max_prompt_length=512,
max_length=1024
)
dpo_trainer.train()
dpo_trainer.save_model("./final_aligned_model")
print("Model successfully aligned via Direct Preference Optimization.")Module Key Takeaways
- Pre-training produces an unaligned text predictor; alignment guides model behavior toward human preferences.
- Traditional RLHF uses PPO with a separate reward model and KL penalty, which is resource-intensive and unstable.
- DPO directly optimizes log-likelihood ratios of preferred vs rejected responses, eliminating the reward model and PPO entirely.
- Constitutional AI (RLAIF) uses structured principles and automated AI feedback to scale alignment efficiently.
Explore Related Tools
Hand-picked utilities and calculators related to this tool.
