LLM Alignment & RLHF

Why Alignment Matters

A pretrained LLM (e.g., trained to predict next tokens on the internet) is not automatically:

• Helpful: it predicts plausible text, not useful responses
• Harmless: it may reproduce harmful, biased, or dangerous content from training data
• Honest: it may confabulate facts confidently

Alignment is the set of techniques that make a model's behavior match human preferences and values.

The Post-Training Pipeline

Pretraining          SFT                  RLHF / DPO
(predict tokens) → (learn to follow) → (optimize for human preference)
     Base LLM     →  Instruct LLM    →       Aligned LLM
  (GPT-3 base)      (text-davinci-003)    (ChatGPT / Claude / Llama-3-Instruct)

Stage 1 — Supervised Fine-Tuning (SFT): fine-tune on high-quality (prompt, ideal response) pairs written by human contractors. Teaches format, instruction-following style.

Stage 2 — Reward Model: train a model to predict which of two responses humans prefer.

Stage 3 — RL from Human Feedback: optimize the SFT model against the reward model using PPO.

Stage 1: Supervised Fine-Tuning (SFT)

Starting from a pretrained base model, fine-tune on a curated dataset of:

$\mathcal{D}_{SFT} = \{(x_i, y_i^*)\}$

where $x_i$ is a prompt and $y_i^*$ is a high-quality human-written response.

Training objective: standard cross-entropy next-token prediction on the responses:

$\mathcal{L}_{SFT} = -\sum_t \log p_\theta(y_t^* \mid x, y_{<t}^*)$

What SFT Teaches

• Response format (assistant-style, markdown, structured output)
• Instruction following (do X, don't do Y)
• Refusal of clearly harmful requests
• Basic helpful behavior

Limitation of SFT Alone

SFT can only teach what's explicitly in the labeled data. It cannot:

• Generalize well to novel instruction types
• Optimize for subtle quality distinctions
• Learn from ranking/preference signals (which response is better)

This motivates the reward model + RL stage.

Stage 2: Reward Model Training

The reward model $r_\phi(x, y)$ predicts how much a human would prefer response $y$ to prompt $x$.

Data Collection

For each prompt $x$, generate $k$ responses from the SFT model. Human raters rank them (or compare pairs):

$\mathcal{D}_{RM} = \{(x, y_w, y_l)\}$

where $y_w$ is the preferred (winner) and $y_l$ is the less-preferred (loser) response.

Bradley-Terry Loss

$\mathcal{L}_{RM} = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}}\left[\log \sigma(r_\phi(x, y_w) - r_\phi(x, y_l))\right]$

This maximizes the probability that the preferred response gets a higher score.

Architecture: typically the SFT model with the final token's hidden state projected to a scalar reward.

Goodhart's Law Warning

"When a measure becomes a target, it ceases to be a good measure." — Goodhart's Law

The reward model is an imperfect proxy. Optimizing too hard against it causes reward hacking: the LLM finds outputs that score high on the RM but are not actually good (verbose, sycophantic, confident-but-wrong).

KL penalty (used in PPO) limits how far the policy drifts from the SFT model to mitigate this.

Stage 3: PPO-Based RLHF

The RLHF objective maximizes reward while staying close to the SFT policy via a KL penalty:

$\mathcal{L}_{RLHF}(\theta) = \mathbb{E}_{x \sim \mathcal{D},\ y \sim \pi_\theta(\cdot|x)}\left[r_\phi(x, y) - \beta\, \text{KL}(\pi_\theta(\cdot|x) \,\|\, \pi_{SFT}(\cdot|x))\right]$

• $r_\phi(x, y)$: scalar reward from the reward model
• $\beta$: KL penalty coefficient (typically 0.1–0.5) — prevents reward hacking
• $\pi_\theta$: the policy (LLM) being optimized
• $\pi_{SFT}$: frozen SFT model as reference

Why PPO for LLMs?

PPO (Proximal Policy Optimization) clips the policy update ratio to keep updates conservative:

$\mathcal{L}^{PPO} = \mathbb{E}_t\left[\min\left(r_t(\theta)\hat{A}_t,\ \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_t\right)\right]$

For LLMs: the "action" is generating a token, the "episode" is generating the full response, the "reward" is given at the end (or per-token via the RM).

PPO-RLHF Components in Practice

Component	Role
Policy model	The LLM being trained ($\pi_\theta$)
Reward model	Scores full responses
Value model	Estimates expected future reward (PPO critic)
Reference model	Frozen SFT model for KL computation

Requires 4 large models in memory simultaneously — very GPU-intensive. InstructGPT used this for GPT-3.5; Claude 2 and Llama 2-Chat also used PPO-RLHF.

DPO — Direct Preference Optimization

Rafailov et al. (2023) — a simpler alternative that eliminates the reward model and PPO entirely.

Key Insight

The optimal policy under the RLHF objective has a closed-form relationship to the reward model:

$r(x, y) = \beta \log \frac{\pi^*(y|x)}{\pi_{SFT}(y|x)} + \beta \log Z(x)$

Substituting this into the Bradley-Terry preference loss and canceling $Z(x)$ (which is prompt-dependent and cancels in pairwise comparison):

$\mathcal{L}_{DPO}(\theta) = -\mathbb{E}_{(x, y_w, y_l)}\left[\log \sigma\!\left(\beta \log \frac{\pi_\theta(y_w|x)}{\pi_{SFT}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{SFT}(y_l|x)}\right)\right]$

In plain English: increase the log-probability of the preferred response relative to the SFT baseline, and decrease it for the rejected response — directly, without a reward model.

DPO vs PPO Comparison

	PPO-RLHF	DPO
Reward model	Required	Not needed
Complexity	4 models in memory	2 models (policy + reference)
Training stability	Tricky (RL)	Stable (supervised-style)
Sample efficiency	Needs online sampling	Offline from fixed dataset
Quality	Slightly better on complex tasks	Comparable or better on many tasks
Adoption	InstructGPT, Llama 2	Llama 3, Mistral, Qwen, most OSS models

DPO has become the dominant method for open-source alignment due to its simplicity and effectiveness.

DPO Variants

• IPO (Identity Preference Optimization): removes the $\sigma$ link function, more robust
• ORPO (Odds Ratio Preference Optimization): combines SFT + DPO in a single stage
• SimPO: reference-free DPO variant using average log-prob normalization

import torch
import torch.nn.functional as F
from torch.optim import AdamW

def compute_log_probs(model, input_ids, attention_mask, labels):
    """Compute sum of log-probs for non-padding tokens."""
    with torch.no_grad() if not model.training else torch.enable_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    logits = outputs.logits[:, :-1, :]     # (B, T-1, vocab)
    target = labels[:, 1:]                 # shift right
    log_probs = F.log_softmax(logits, dim=-1)
    token_log_probs = log_probs.gather(
        dim=-1, index=target.unsqueeze(-1)
    ).squeeze(-1)
    mask = (target != -100).float()
    return (token_log_probs * mask).sum(dim=-1)  # (B,)


def dpo_loss(
    policy_model,
    reference_model,
    chosen_ids, chosen_mask, chosen_labels,
    rejected_ids, rejected_mask, rejected_labels,
    beta: float = 0.1,
):
    """
    DPO loss:  -log σ(β * (log π_θ(y_w|x)/π_ref(y_w|x) - log π_θ(y_l|x)/π_ref(y_l|x)))
    """
    # Policy log-probs
    policy_chosen_lp   = compute_log_probs(policy_model, chosen_ids, chosen_mask, chosen_labels)
    policy_rejected_lp = compute_log_probs(policy_model, rejected_ids, rejected_mask, rejected_labels)

    # Reference log-probs (frozen)
    reference_model.eval()
    with torch.no_grad():
        ref_chosen_lp   = compute_log_probs(reference_model, chosen_ids, chosen_mask, chosen_labels)
        ref_rejected_lp = compute_log_probs(reference_model, rejected_ids, rejected_mask, rejected_labels)

    # Log-ratios
    chosen_ratio   = policy_chosen_lp   - ref_chosen_lp    # log π/π_ref for winner
    rejected_ratio = policy_rejected_lp - ref_rejected_lp  # log π/π_ref for loser

    # DPO loss
    logits = beta * (chosen_ratio - rejected_ratio)
    loss = -F.logsigmoid(logits).mean()

    # Diagnostics
    reward_chosen   = (beta * chosen_ratio).mean().item()
    reward_rejected = (beta * rejected_ratio).mean().item()
    reward_margin   = reward_chosen - reward_rejected
    accuracy = (logits > 0).float().mean().item()

    return loss, {
        'reward_chosen': reward_chosen,
        'reward_rejected': reward_rejected,
        'reward_margin': reward_margin,
        'dpo_accuracy': accuracy,
    }


# ---- Training sketch ----
# from transformers import AutoModelForCausalLM
# policy_model    = AutoModelForCausalLM.from_pretrained("sft_checkpoint")
# reference_model = AutoModelForCausalLM.from_pretrained("sft_checkpoint")
# for param in reference_model.parameters():
#     param.requires_grad_(False)
#
# optimizer = AdamW(policy_model.parameters(), lr=5e-7)
# for batch in dataloader:
#     loss, metrics = dpo_loss(policy_model, reference_model, **batch, beta=0.1)
#     loss.backward(); optimizer.step(); optimizer.zero_grad()
#     print(f"Loss: {loss:.4f} | Accuracy: {metrics['dpo_accuracy']:.3f} | Margin: {metrics['reward_margin']:.3f}")

print("DPO implementation ready.")
print("Key metric to monitor: dpo_accuracy (should trend toward 1.0)")
print("Key risk: if reward_margin grows too large → overfitting to preference data")

Constitutional AI (CAI)

Anthropic (2022) — an alternative alignment approach that reduces dependence on human-labeled preference data.

The Two Phases

Phase 1 — SL-CAI (Supervised):

1. Generate harmful/unhelpful responses to red-team prompts
2. Ask the model to critique its own response against a set of principles (the "constitution")
3. Ask the model to revise its response based on the critique
4. Fine-tune on the revised responses

Phase 2 — RL-CAI:

1. Generate pairs of responses to harmful prompts
2. Use the model itself as a preference rater (asks which response better follows the constitution)
3. Train a reward model from these AI-labeled preferences
4. Apply PPO or DPO against this reward model

The Constitution

A list of principles the model should follow, e.g.:

• "Choose the response that is least likely to be harmful"
• "Choose the response that is most honest and does not involve deception"
• "Prefer responses that support human autonomy and democratic values"

Why CAI Matters

• Scalable: reduces human labeling bottleneck — AI generates preference labels
• Transparent: the principles are explicit and auditable
• Foundation for RLAIF (RL from AI Feedback): use a capable AI (e.g., Claude Opus) to rate responses instead of humans
• Powers Claude's alignment approach alongside human feedback

Alignment Failure Modes

Problem	Description
Sycophancy	Model agrees with users even when wrong
Reward hacking	Finds high-reward outputs that aren't actually good
Jailbreaking	Users bypass safety training with adversarial prompts
Over-refusal	Model refuses benign requests out of excessive caution
Value lock-in	Current human preferences baked in; may not generalize

Summary

• SFT teaches format and instruction-following from human demonstrations — necessary but not sufficient for alignment
• Reward model: trained on human preference pairs using Bradley-Terry loss; scores full responses
• PPO-RLHF: optimize the LLM policy against the reward model with a KL penalty to prevent reward hacking — used in InstructGPT, GPT-3.5, Llama 2
• DPO: reparameterizes RLHF to train directly on preference pairs without a reward model or RL loop — simpler, now dominant in open-source models
• Constitutional AI / RLAIF: use the model itself (with a constitution or a stronger AI) to generate preference labels, scaling beyond human labeling
• Alignment is an active research area — sycophancy, reward hacking, and jailbreaking remain unsolved challenges