RLHF & AI Alignment

1. The Alignment Problem

The fundamental challenge in deploying large language models (LLMs) is ensuring they produce outputs that are helpful, harmless, and honest (the HHH criteria). RLHF provides a framework for aligning model behavior with human preferences without requiring explicit supervision for every possible output.

1.1 Why RLHF?

Pre-training with next-token prediction (self-supervised learning) optimizes for:

This objective captures statistical patterns in the training data but does not directly encode human preferences about which outputs are better when multiple completions are plausible. Fine-tuning on demonstrations (SFT) helps but is limited by the quality and diversity of human-written examples.

RLHF bridges this gap by learning a reward model from human preference comparisons and then optimizing the language model policy against this reward using reinforcement learning.

2. The RLHF Pipeline

The standard RLHF pipeline consists of three phases:

Phase 1: Supervised Fine-Tuning (SFT)

Given a pre-trained model , fine-tune on a curated dataset of high-quality demonstrations :

Phase 2: Reward Model Training

Collect human preference data: for each prompt , sample two completions and ask a human annotator to indicate which is preferred. This yields a dataset:

where is the preferred (winning) completion and is the dispreferred (losing) completion.

Phase 3: RL Optimization (PPO)

Optimize the policy using Proximal Policy Optimization (PPO) to maximize the learned reward while staying close to the SFT policy via a KL penalty.

3. Reward Modeling

3.1 The Bradley-Terry Model

The standard approach models human preferences using the Bradley-Terry pairwise comparison model. Given a reward function , the probability that completion is preferred over for prompt is:

where is the logistic sigmoid function:

3.2 Reward Model Loss

The reward model is trained by maximizing the likelihood of the observed preferences:

This is equivalent to minimizing the pairwise ranking loss. The reward model is typically initialized from the SFT model with a scalar head appended to the final transformer layer.

3.3 Implementation

import torch
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification

class RewardModel(torch.nn.Module):
    def __init__(self, model_name, max_length=512):
        super().__init__()
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=1
        )
        self.max_length = max_length

    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.logits.squeeze(-1)

    def compute_loss(self, chosen_ids, chosen_mask, rejected_ids, rejected_mask):
        chosen_rewards = self.forward(chosen_ids, chosen_mask)
        rejected_rewards = self.forward(rejected_ids, rejected_mask)
        loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
        return loss, chosen_rewards.mean(), rejected_rewards.mean()

3.4 Reward Hacking

A critical failure mode is reward hacking (or reward over-optimization), where the policy learns to exploit artifacts in the reward model rather than genuinely improving quality. The代理 policy can achieve high reward scores by generating outputs that the reward model rates highly but humans would find poor.

Quantitatively, if is the true human reward and is the learned reward, the Goodhart's Law effect manifests as:

where the gap increases with the KL divergence from the reference policy.

4. Proximal Policy Optimization (PPO) for RLHF

4.1 RLHF Objective

The RLHF optimization objective is:

where controls the KL penalty strength, preventing the policy from deviating too far from the SFT model. The KL term can be expanded as:

4.2 PPO Algorithm

PPO uses a clipped surrogate objective to ensure stable policy updates:

where is the estimated advantage function and is the clipping parameter (typically 0.2).

In the RLHF context, each token generation step is treated as an action in the RL formulation:

• State : the prompt plus previously generated tokens
• Action : the next token to generate
• Reward: at the final token, 0 otherwise

4.3 PPO Training Loop

import torch
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

config = PPOConfig(
    learning_rate=1.4e-5,
    batch_size=64,
    mini_batch_size=16,
    ppo_epochs=4,
    kl_penalty="kl",
    init_kl_coef=0.2,
    target_kl=6.0,
)

ppo_trainer = PPOTrainer(
    config=config,
    model=AutoModelForCausalLMWithValueHead.from_pretrained("sft-model"),
    ref_model=AutoModelForCausalLMWithValueHead.from_pretrained("sft-model"),
    tokenizer=tokenizer,
)

for batch in dataloader:
    query_tensors = batch["input_ids"]

    # Generate responses
    response_tensors = ppo_trainer.generate(query_tensors, max_new_tokens=256)

    # Compute rewards
    reward_scores = reward_model(
        batch["input_ids"],
        torch.cat(response_tensors, dim=0)
    )

    # Run PPO step
    stats = ppo_trainer.step(
        query_tensors, response_tensors, reward_scores
    )

4.4 The KL Penalty Implementation

The per-token KL penalty is computed as:

The total reward at each step combines the task reward and the KL penalty:

where indicates the final token position.

5. Direct Preference Optimization (DPO)

5.1 Motivation

PPO-based RLHF is complex to implement and train, requiring:

1. Training a separate reward model
2. Running a costly RL loop with multiple models in memory
3. Careful hyperparameter tuning for stable training

DPO (Rafailov et al., 2023) provides a simpler alternative by deriving a closed-form mapping between optimal policies and reward functions under the RLHF objective.

5.2 Theoretical Foundation

The optimal policy under the RLHF objective satisfies:

where is the partition function.

Inverting this relationship, we can express the reward in terms of the optimal policy:

5.3 DPO Loss Derivation

Substituting the reward expression into the Bradley-Terry preference model:

Note that the partition function cancels out in the difference! The DPO loss is:

This is a remarkable result: we can optimize the policy directly on preference data without training a separate reward model.

5.4 DPO Gradient Analysis

The gradient of the DPO loss with respect to is:

where is the implicit reward.

Key insight: the gradient upweights examples where the current policy assigns higher probability to the dispreferred response — precisely the cases where the model needs the most correction.

5.5 DPO Implementation

import torch
import torch.nn.functional as F

def dpo_loss(policy_logps_w, policy_logps_l,
             ref_logps_w, ref_logps_l, beta=0.1):
    """
    Compute the DPO loss.

    Args:
        policy_logps_w: log π_θ(y_w|x) for preferred completions
        policy_logps_l: log π_θ(y_l|x) for dispreferred completions
        ref_logps_w: log π_ref(y_w|x) for preferred completions
        ref_logps_l: log π_ref(y_l|x) for dispreferred completions
        beta: temperature parameter
    """
    logits = beta * (
        (policy_logps_w - ref_logps_w) -
        (policy_logps_l - ref_logps_l)
    )
    loss = -F.logsigmoid(logits).mean()
    return loss

# Training loop
for batch in dataloader:
    # Compute log probabilities under policy and reference
    policy_logps_w = compute_logprobs(policy_model, batch["chosen"])
    policy_logps_l = compute_logprobs(policy_model, batch["rejected"])
    ref_logps_w = compute_logprobs(ref_model, batch["chosen"])
    ref_logps_l = compute_logprobs(ref_model, batch["rejected"])

    loss = dpo_loss(
        policy_logps_w, policy_logps_l,
        ref_logps_w, ref_logps_l, beta=0.1
    )
    loss.backward()
    optimizer.step()

5.6 DPO vs RLHF Comparison

6. Constitutional AI (CAI)

6.1 Overview

Constitutional AI (Bai et al., 2022) reduces reliance on human annotators by using an AI system to both generate and evaluate training data according to a set of explicit constitutional principles.

6.2 Two-Phase Process

Phase 1: Supervised Learning from AI Feedback (SL-AF)

1. Generate initial responses using the base model
2. Ask the model to critique its own response against constitutional principles
3. Ask the model to revise its response based on the critique
4. Use the revised responses as SFT training data

Phase 2: RL from AI Feedback (RL-AF)

1. Generate pairs of responses
2. Ask an AI annotator (based on the same model) to select the preferred response based on constitutional principles
3. Train a reward model on these AI-generated preferences
4. Optimize the policy using RL

6.3 Constitutional Principles

Example principles might include:

1. "Choose the response that is least likely to be considered harmful or offensive"
2. "Choose the response that is most helpful while remaining truthful"
3. "Choose the response that answers the question directly and concisely"
4. "Choose the response that demonstrates the most nuanced understanding"

6.4 RLAIF (RL from AI Feedback)

RLAIF generalizes Constitutional AI by using AI feedback more broadly:

where the AI annotator selects preferences based on learned or specified criteria. The key insight is that for many alignment criteria, AI systems can provide consistent and scalable feedback that approximates human judgment.

Recent work has shown that RLAIF can achieve performance comparable to RLHF with human annotators, while being significantly more scalable and cost-effective.

7. Advanced Topics

7.1 Iterative DPO and Online DPO

Standard DPO uses a fixed offline dataset, which can lead to distribution shift. Iterative DPO addresses this by:

1. Training the policy with DPO on the current dataset
2. Generating new completions with the updated policy
3. Obtaining new preference labels (human or AI)
4. Adding to the training dataset and repeating

Online DPO performs this process at the batch level, generating new examples on-the-fly during training.

7.2 IPO and KTO

Identity Preference Optimization (IPO) addresses a weakness in DPO by not assuming the Bradley-Terry model:

Kahneman-Tversky Optimization (KTO) works with individual thumbs-up/thumbs-down labels rather than pairwise comparisons:

where is based on the KL-divergence from the reference and is a weighting function that accounts for loss aversion.

7.3 Multi-Objective Alignment

Real-world alignment requires balancing multiple objectives (helpfulness, harmlessness, honesty). This can be formulated as a multi-objective optimization problem:

subject to KL constraints. Methods like MOO-DPO extend DPO to the multi-objective setting using techniques from multi-objective optimization.

8. SVG Diagrams

RLHF Pipeline

9. Key Metrics and Evaluation

9.1 Win Rate

The primary metric: comparing the aligned model's outputs to a baseline (typically SFT) via human or AI preference:

9.2 KL Divergence

Monitored during training to ensure the policy doesn't drift too far:

9.3 Reward Model Accuracy

The reward model's accuracy on held-out preference data:

10. Open Problems

1. Scalable oversight: How do we align models that are smarter than their human evaluators?
2. Reward hacking mitigation: Better techniques for preventing reward over-optimization
3. Multi-turn alignment: Extending RLHF to conversational settings with multiple turns
4. Cross-cultural alignment: Whose values should we align to?
5. Formal verification: Can we prove alignment properties?

References

1. Ouyang et al. (2022). "Training language models to follow instructions with human feedback." NeurIPS.
2. Bai et al. (2022). "Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback." arXiv.
3. Rafailov et al. (2023). "Direct Preference Optimization: Your Language Model is Secretly a Reward Model." NeurIPS.
4. Christiano et al. (2017). "Deep Reinforcement Learning from Human Preferences." NeurIPS.
5. Schulman et al. (2017). "Proximal Policy Optimization Algorithms." arXiv.
6. Bai et al. (2022). "Constitutional AI: Harmlessness from AI Feedback." arXiv.