Alternatives: Direct Preference Optimization (DPO)

While Reinforcement Learning from Human Feedback (RLHF) using Proximal Policy Optimization (PPO) has proven effective for aligning language models with human preferences, the process can be complex and sometimes unstable. It typically involves multiple stages: supervised fine-tuning (SFT), training a separate reward model (RM) on human preference data, and then fine-tuning the SFT model using reinforcement learning guided by the RM. This multi-stage pipeline introduces several hyperparameters and potential points of failure, particularly during the RL phase which can be sensitive to implementation details and prone to issues like reward hacking.

Direct Preference Optimization (DPO) offers a more streamlined approach to preference alignment, bypassing the need for explicit reward model training and the complexities of reinforcement learning altogether. DPO directly optimizes the language model on the preference data using a simple classification objective.

The core idea behind DPO stems from a mathematical relationship between the optimal policy sought by RLHF and the underlying (implicit) reward function. Recall that standard reward modeling often uses the Bradley-Terry model to link pairwise preferences $(y_w, y_l)$ for a prompt $x$ (where $y_w$ is preferred over $y_l$) to a latent reward function $r$:

$P(y_w \succ y_l | x) = \sigma(r(x, y_w) - r(x, y_l))$

Here, $\sigma$ is the logistic function. The goal of RLHF is to find a policy $\pi_{RL}$ that maximizes the expected reward $E[r(x, y)]$ while staying close to a reference policy $\pi_{ref}$ (usually the SFT model), controlled by a KL divergence penalty:

$\max_{\pi_{RL}} E_{(x,y) \sim \pi_{RL}}[r(x, y)] - \beta D_{KL}(\pi_{RL}(y|x) || \pi_{ref}(y|x))$

DPO uses the analytical solution to this constrained optimization problem. It can be shown that the optimal policy $\pi_{RL}$ has the form:

$\pi_{RL}(y|x) = \frac{1}{Z(x)} \pi_{ref}(y|x) \exp\left(\frac{1}{\beta} r(x, y)\right)$

where $Z(x)$ is a partition function ensuring the probabilities sum to one. This equation connects the optimal policy, the reference policy, and the reward function. By substituting this relationship back into the Bradley-Terry preference model, the reward function $r(x, y)$ can be eliminated, expressing the preference probability directly in terms of policies:

$P(y_w \succ y_l | x) = \sigma\left( \beta \log \frac{\pi_{RL}(y_w|x)}{\pi_{ref}(y_w|x)} - \beta \log \frac{\pi_{RL}(y_l|x)}{\pi_{ref}(y_l|x)} \right)$

DPO trains the language model $\pi_\theta$ directly to satisfy this preference model, using $\pi_\theta$ in place of the unknown optimal policy $\pi_{RL}$. The objective is to maximize the log-likelihood of the observed human preferences under this model. This leads to the DPO loss function:

$\mathcal{L}_{DPO}(\pi_\theta; \pi_{ref}) = - E_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \log \sigma \left( \beta \log \frac{\pi_\theta(y_w|x)}{\pi_{ref}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{ref}(y_l|x)} \right) \right]$

Here, $\mathcal{D}$ is the dataset of preference triplets $(x, y_w, y_l)$, $\pi_\theta$ is the language model being optimized, $\pi_{ref}$ is the frozen reference SFT model, and $\beta$ is a hyperparameter that implicitly controls how much the optimized policy $\pi_\theta$ deviates from the reference policy $\pi_{ref}$. A higher $\beta$ places more weight on the preference data, potentially leading to larger deviations.

Implementation and Usage

Implementing DPO is significantly simpler than RLHF. It requires:

1. A preference dataset $\mathcal{D} = \{(x_i, y_{w,i}, y_{l,i})\}$.
2. A base SFT model to serve as the reference policy $\pi_{ref}$. This model's weights are kept frozen during DPO training.
3. A copy of the SFT model, $\pi_\theta$, whose weights will be updated.

The training loop involves standard supervised learning:

• Sample a batch of preference triplets $(x, y_w, y_l)$ from $\mathcal{D}$.
• Compute the log-probabilities of the chosen ($y_w$) and rejected ($y_l$) responses under both the current model $\pi_\theta$ and the reference model $\pi_{ref}$.
• Calculate the DPO loss $\mathcal{L}_{DPO}$ using these log-probabilities.
• Perform a backward pass and update the weights of $\pi_\theta$ using an optimizer like AdamW.

Here's a PyTorch snippet for calculating the core part of the DPO loss within a training step, assuming you have obtained the log-probabilities:

import torch
import torch.nn.functional as F

# Assume log_probs_policy and log_probs_ref contain log probabilities
# Shape: (batch_size,) for both chosen (w) and rejected (l) responses
# log_probs_policy_w, log_probs_policy_l: Log probs from the model being
#                                         trained (pi_theta)
# log_probs_ref_w, log_probs_ref_l: Log probs from the frozen reference
#                                   model (pi_ref)
# beta: Hyperparameter (e.g., 0.1)

def dpo_loss(log_probs_policy_w, log_probs_policy_l,
             log_probs_ref_w, log_probs_ref_l, beta):
    """Calculates the DPO loss for a batch of preferences."""

    # Calculate log ratios
    log_ratio_w = log_probs_policy_w - log_probs_ref_w
    log_ratio_l = log_probs_policy_l - log_probs_ref_l

    # Difference in log ratios scaled by beta
    diff_scaled = beta * (log_ratio_w - log_ratio_l)

    # Calculate loss using logistic sigmoid
    loss = -F.logsigmoid(diff_scaled)

    # Average loss over the batch
    return loss.mean()

# Example usage within a training loop
# batch = get_preference_batch()
# # (prompts, chosen_responses, rejected_responses)
# outputs_policy = policy_model(prompts,
#                               chosen_responses,
#                               rejected_responses)
# with torch.no_grad():
#     outputs_ref = ref_model(prompts,
#                             chosen_responses,
#                             rejected_responses)
#
# # Extract log probabilities (details depend on model implementation)
# log_probs_policy_w = get_log_probs(outputs_policy, chosen_responses)
# log_probs_policy_l = get_log_probs(outputs_policy, rejected_responses)
# log_probs_ref_w = get_log_probs(outputs_ref, chosen_responses)
# log_probs_ref_l = get_log_probs(outputs_ref, rejected_responses)
#
# loss = dpo_loss(log_probs_policy_w, log_probs_policy_l,
#                 log_probs_ref_w, log_probs_ref_l, beta=0.1)
# loss.backward()
# optimizer.step()

Advantages and Disadvantages

Advantages:

• Simplicity: DPO requires only a single stage of training after SFT, eliminating the need for reward model fitting and complex RL algorithms.
• Stability: It avoids the optimization instabilities often encountered with PPO in RLHF, such as sensitivity to KL divergence coefficient and reward scaling.
• Efficiency: It is generally less computationally expensive than PPO-based RLHF as it doesn't require sampling generations from the policy during training or training a separate reward model.

Disadvantages:

• Performance: While often achieving comparable results to RLHF, DPO might sometimes underperform compared to a meticulously tuned PPO pipeline, especially when the underlying reward signal is complex or noisy.
• Hyperparameter Tuning: The hyperparameter $\beta$ still needs careful tuning, similar to the KL coefficient in PPO.
• Data Requirements: Like RLHF, DPO relies on a high-quality preference dataset.

In summary, DPO presents a simpler alternative to the standard RLHF pipeline. Its stability and ease of implementation make it an attractive option for aligning language models with human preferences, especially when computational resources or RL expertise are limited. It represents a significant simplification in the process of making LLMs more helpful, honest, and harmless.