How RLHF Alignment Works

Introduction

Reinforcement Learning from Human Feedback (RLHF) is the training paradigm that transforms a language model capable of text completion into one that follows instructions, refuses harmful requests, and generates responses humans prefer. First applied at scale by Ziegler et al. (2019) and popularized by InstructGPT (Ouyang et al., 2022), RLHF bridges the gap between a model optimized for next-token prediction on internet text and one optimized for human utility and safety. The core intuition is that it is far easier for humans to evaluate a response than to specify exactly what a good response looks like — comparing two outputs and picking the better one is a tractable labeling task, while writing down a complete rubric for “a good AI response” is nearly impossible.

RLHF consists of three sequential phases: supervised fine-tuning (SFT) on human demonstrations, reward model training on human preference comparisons, and reinforcement learning optimization of the policy against the reward model. Each phase has distinct failure modes and hyperparameter sensitivities, and the entire pipeline is notoriously difficult to stabilize — PPO (Proximal Policy Optimization) applied to LLMs is one of the most complex training setups in modern ML. Newer alternatives like Direct Preference Optimization (DPO) and Constitutional AI (CAI) have emerged to address some of these complexities.

Step-by-Step: The Three Phases of RLHF

flowchart TD
    A["1. Supervised Fine-Tuning (SFT)"]
    B["2. Reward Model Training"]
    C["3. PPO Optimization with KL Divergence Penalty"]
    D["4. Constitutional AI"]
    A --> B
    B --> C
    C --> D

Phase 1: Supervised Fine-Tuning (SFT)

The process begins with the pre-trained base language model π_base — trained on internet-scale text to predict next tokens but with no instruction-following behavior. SFT adapts this model for instruction following by fine-tuning on a dataset of (prompt, high-quality response) pairs written or curated by human demonstrators.

Mathematically, SFT is standard supervised learning: maximize the log-likelihood of the demonstration responses given the prompts:

L_SFT(θ) = E_{(x,y) ~ D_demo} [ Σ_t log π_θ(y_t | x, y_{<t}) ]

Where x is the prompt, y is the demonstration response, and π_θ is the model being trained. The SFT dataset typically contains 10K-100K (prompt, response) pairs covering a diverse range of tasks: summarization, question answering, creative writing, code generation, and so on. The resulting model π_SFT can follow instructions but hasn’t yet been optimized for preference — it generates plausible demonstrations, not necessarily the best possible outputs.

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")

# SFT dataset format: each entry has "prompt" and "completion" fields
# Only the completion tokens contribute to the loss (not the prompt)
response_template = "[/INST]"
collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)

trainer = SFTTrainer(
    model=model,
    train_dataset=sft_dataset,
    dataset_text_field="text",
    data_collator=collator,
)
trainer.train()
sft_model = trainer.model

Phase 2: Reward Model Training (Bradley-Terry Preference Model)

With π_SFT in hand, human raters generate preference comparison data: for each prompt x, the SFT model generates multiple candidate responses y_w (the “winner”) and y_l (the “loser”), and human annotators choose which they prefer. This produces a dataset of triplets: D_preference = {(x, y_w, y_l)}.

The Bradley-Terry model is the standard probabilistic framework for learning from pairwise comparisons. It models the probability that response y_w is preferred over y_l as a function of scalar reward scores:

P(y_w ≻ y_l | x) = σ(r_φ(x, y_w) - r_φ(x, y_l))
               = exp(r_φ(x, y_w)) / (exp(r_φ(x, y_w)) + exp(r_φ(x, y_l)))

Where σ is the sigmoid function, and r_φ is the reward model (a language model with an added scalar head on the final token). Training minimizes the binary cross-entropy loss:

L_RM(φ) = -E_{(x,y_w,y_l) ~ D_pref} [
    log σ(r_φ(x, y_w) - r_φ(x, y_l))
]

This trains the reward model to assign higher scalar scores to preferred responses. The reward model is initialized from π_SFT (same architecture) and has a linear projection added on top of the final hidden state to output a single scalar reward value.

from trl import RewardTrainer, RewardConfig

reward_config = RewardConfig(
    output_dir="./reward-model",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    learning_rate=1e-5,
    max_length=512,
)

reward_trainer = RewardTrainer(
    model=reward_model,
    tokenizer=tokenizer,
    args=reward_config,
    train_dataset=preference_dataset,
    # Dataset must have "chosen" and "rejected" columns
)
reward_trainer.train()

Phase 3: PPO Optimization with KL Divergence Penalty

The final phase uses Proximal Policy Optimization (PPO) to optimize the SFT policy π_θ to maximize rewards from the reward model, while staying close to the SFT policy via a KL divergence penalty. The full objective is:

Objective(θ) = E_{x ~ D, y ~ π_θ(·|x)} [
    r_φ(x, y)                                      -- Reward maximization
    - β · KL(π_θ(·|x) || π_SFT(·|x))             -- KL penalty
]

where:
  KL(π_θ || π_SFT) = Σ_t π_θ(y_t|x,y_{<t}) · log(π_θ(y_t|x,y_{<t}) / π_SFT(y_t|x,y_{<t}))

The KL penalty (controlled by coefficient β) is critical: without it, the policy would “reward hack” the reward model — finding degenerate outputs that score high on the learned reward function but are clearly bad (e.g., generating an extremely long response full of filler text that the reward model spuriously assigns high scores to, since it wasn’t trained on such edge cases). The KL penalty keeps the policy within a “trust region” of the original SFT model.

The PPO update involves four models loaded simultaneously:

1. Active policy π_θ: Being trained, generates responses
2. Reference policy π_SFT: Frozen, computes KL denominator
3. Reward model r_φ: Frozen, scores responses
4. Value function V_ψ: Estimates expected future rewards (critic)

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

ppo_config = PPOConfig(
    model_name="./sft-model",
    learning_rate=1.41e-5,
    batch_size=256,
    mini_batch_size=16,
    gradient_accumulation_steps=1,
    ppo_epochs=4,
    kl_penalty="kl",
    init_kl_coef=0.2,          # Initial β (KL coefficient)
    target_kl=6.0,             # Adaptive β target
    adap_kl_ctrl=True,        # Adaptively adjust β during training
)

model = AutoModelForCausalLMWithValueHead.from_pretrained("./sft-model")
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained("./sft-model")
# ref_model weights are frozen

ppo_trainer = PPOTrainer(
    config=ppo_config,
    model=model,
    ref_model=ref_model,
    tokenizer=tokenizer,
    dataset=prompt_dataset,
)

# Training loop
for batch in ppo_trainer.dataloader:
    # 1. Generate responses with active policy
    query_tensors = batch["input_ids"]
    response_tensors = ppo_trainer.generate(query_tensors, max_new_tokens=256)
    
    # 2. Score responses with reward model
    rewards = [reward_model(q, r) for q, r in zip(query_tensors, response_tensors)]
    
    # 3. PPO update with KL-penalized reward
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
    ppo_trainer.log_stats(stats, batch, rewards)

Phase 4: Constitutional AI — An Alternative to Human Preference Data

Constitutional AI (Anthropic, 2022) provides an alternative to human preference labeling in Phase 2. Instead of human raters choosing between outputs, a set of constitutional principles (a written constitution, e.g., “Do not generate content that could help someone build weapons”) is used to train the reward model using AI feedback (RLAIF — Reinforcement Learning from AI Feedback):

1. Critique step: The SFT model generates a response, then critiques it against each constitutional principle: “Does this response violate the principle ‘do not assist with harmful activities’? Explain why.”
2. Revision step: The model revises the response to address the critique
3. Preference generation: The original and revised responses form a preference pair (revised is preferred if the critique identified a violation)
4. Reward model training: Proceed as in Phase 2, but using AI-generated preference pairs instead of human labels

This dramatically scales preference data collection — you can generate millions of preference pairs without human labelers.

# Pseudo-code for Constitutional AI feedback generation
constitution = [
    "Do not provide instructions for creating weapons.",
    "Be helpful, harmless, and honest.",
    "Acknowledge uncertainty rather than stating false information.",
]

def generate_constitutional_feedback(prompt, response, model):
    feedback_pairs = []
    for principle in constitution:
        critique_prompt = f"""
        Human: {prompt}
        Assistant: {response}
        
        Critique the above response against this principle: "{principle}"
        Does it violate the principle? If so, how?
        """
        critique = model.generate(critique_prompt)
        
        revision_prompt = f"""
        Original response: {response}
        Critique: {critique}
        
        Revise the response to address the critique:
        """
        revised = model.generate(revision_prompt)
        
        if "violates" in critique.lower():
            feedback_pairs.append({"chosen": revised, "rejected": response})
    
    return feedback_pairs

Key Takeaways

• SFT establishes the baseline instruction-following capability: By training on high-quality (prompt, response) demonstrations via standard cross-entropy loss, the SFT model learns the format and general approach of helpful responses before preference optimization begins.
• The Bradley-Terry model transforms pairwise comparisons into a differentiable scalar reward: By modeling P(y_w ≻ y_l) = σ(r(y_w) - r(y_l)), human preference judgments become a log-likelihood training signal for the reward model.
• The KL divergence penalty prevents reward hacking: Without the β · KL(π_θ || π_SFT) penalty term, PPO will find degenerate high-reward strategies that fool the reward model but produce clearly bad outputs; β controls the exploration-exploitation trade-off.
• PPO requires 4 models in memory simultaneously: The active policy, reference policy, reward model, and value function must all fit in GPU memory, making RLHF one of the most memory-intensive ML training setups — typically requiring 8× A100 80GB GPUs per training run.
• Constitutional AI replaces human preference labels with AI-generated critiques: By using a language model to apply a written constitution of principles, RLAIF can generate preference datasets at scale without human annotators, making it practical for iterative alignment improvements.