Direct Preference Optimization: A Complete Guide

August 14, 2024

1 4 minutes read

import torch
import torch.nn.functional as F
class DPOTrainer:
    def __init__(self, model, ref_model, beta=0.1, lr=1e-5):
        self.model = model
        self.ref_model = ref_model
        self.beta = beta
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
    
    def compute_loss(self, pi_logps, ref_logps, yw_idxs, yl_idxs):
        """
        pi_logps: policy logprobs, shape (B,)
        ref_logps: reference model logprobs, shape (B,)
        yw_idxs: preferred completion indices in [0, B-1], shape (T,)
        yl_idxs: dispreferred completion indices in [0, B-1], shape (T,)
        beta: temperature controlling strength of KL penalty
        Each pair of (yw_idxs[i], yl_idxs[i]) represents the indices of a single preference pair.
        """
        # Extract log probabilities for the preferred and dispreferred completions
        pi_yw_logps, pi_yl_logps = pi_logps[yw_idxs], pi_logps[yl_idxs]
        ref_yw_logps, ref_yl_logps = ref_logps[yw_idxs], ref_logps[yl_idxs]
        # Calculate log-ratios
        pi_logratios = pi_yw_logps - pi_yl_logps
        ref_logratios = ref_yw_logps - ref_yl_logps
        # Compute DPO loss
        losses = -F.logsigmoid(self.beta * (pi_logratios - ref_logratios))
        rewards = self.beta * (pi_logps - ref_logps).detach()
        return losses.mean(), rewards
    def train_step(self, batch):
        x, yw_idxs, yl_idxs = batch
        self.optimizer.zero_grad()
        # Compute log probabilities for the model and the reference model
        pi_logps = self.model(x).log_softmax(-1)
        ref_logps = self.ref_model(x).log_softmax(-1)
        # Compute the loss
        loss, _ = self.compute_loss(pi_logps, ref_logps, yw_idxs, yl_idxs)
        loss.backward()
        self.optimizer.step()
        return loss.item()
# Usage
model = YourLanguageModel()  # Initialize your model
ref_model = YourLanguageModel()  # Load pre-trained reference model
trainer = DPOTrainer(model, ref_model)
for batch in dataloader:
    loss = trainer.train_step(batch)
    print(f"Loss: {loss}")

Challenges and future directions

Although DPO offers significant advantages over traditional RLHF approaches, there are still challenges and areas for further research:

a) Scalability to larger models:

As language models continue to grow in size, efficiently applying DPO to models with hundreds of billions of parameters remains an open challenge. Researchers are investigating techniques such as:

Efficient fine-tuning methods (e.g. LoRA, prefix tuning)
Distributed training optimizations
Gradient control and mixed precision training

Example of using LoRA with DPO:

from peft import LoraConfig, get_peft_model
class DPOTrainerWithLoRA(DPOTrainer):
    def __init__(self, model, ref_model, beta=0.1, lr=1e-5, lora_rank=8):
        lora_config = LoraConfig(
            r=lora_rank,
            lora_alpha=32,
            target_modules=["q_proj", "v_proj"],
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM"
        )
        self.model = get_peft_model(model, lora_config)
        self.ref_model = ref_model
        self.beta = beta
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
# Usage
base_model = YourLargeLanguageModel()
dpo_trainer = DPOTrainerWithLoRA(base_model, ref_model)

b) Adaptation to multiple tasks and few shots:

Developing DPO techniques that can efficiently adapt to new tasks or domains with limited preference data is an active area of research. Approaches being explored include:

Meta-learning frameworks for rapid adaptation
Prompt-based fine-tuning for DPO
Transfer learning from general preference models to specific domains

c) Dealing with ambiguous or conflicting preferences:

Real-world preference data often contains ambiguities or conflicts. Improving DPO’s robustness to such data is critical. Possible solutions include:

Probabilistic preference modeling
Active learning to resolve ambiguities
Aggregation of preferences for multiple agents

Example of probabilistic preference modeling:

class ProbabilisticDPOTrainer(DPOTrainer):
    def compute_loss(self, pi_logps, ref_logps, yw_idxs, yl_idxs, preference_prob):
        # Compute log ratios
        pi_yw_logps, pi_yl_logps = pi_logps[yw_idxs], pi_logps[yl_idxs]
        ref_yw_logps, ref_yl_logps = ref_logps[yw_idxs], ref_logps[yl_idxs]
        
        log_ratio_diff = pi_yw_logps.sum(-1) - pi_yl_logps.sum(-1)
        loss = -(preference_prob * F.logsigmoid(self.beta * log_ratio_diff) +
                 (1 - preference_prob) * F.logsigmoid(-self.beta * log_ratio_diff))
        return loss.mean()
# Usage
trainer = ProbabilisticDPOTrainer(model, ref_model)
loss = trainer.compute_loss(pi_logps, ref_logps, yw_idxs, yl_idxs, preference_prob=0.8)  # 80% confidence in preference

d) Combining DPO with other coordination techniques:

Integrating DPO with other coordination approaches could lead to more robust and capable systems:

Constitutional AI principles for explicit constraint satisfaction
Debate and recursive reward modeling for complex preference elicitation
Inverse reinforcement learning for inferring underlying reward functions

Example of combining DPO with Constitutional AI:

class ConstitutionalDPOTrainer(DPOTrainer):
    def __init__(self, model, ref_model, beta=0.1, lr=1e-5, constraints=None):
        super().__init__(model, ref_model, beta, lr)
        self.constraints = constraints or []
    def compute_loss(self, pi_logps, ref_logps, yw_idxs, yl_idxs):
        base_loss = super().compute_loss(pi_logps, ref_logps, yw_idxs, yl_idxs)
        
        constraint_loss = 0
        for constraint in self.constraints:
            constraint_loss += constraint(self.model, pi_logps, ref_logps, yw_idxs, yl_idxs)
        
        return base_loss + constraint_loss
# Usage
def safety_constraint(model, pi_logps, ref_logps, yw_idxs, yl_idxs):
    # Implement safety checking logic
    unsafe_score = compute_unsafe_score(model, pi_logps, ref_logps)
    return torch.relu(unsafe_score - 0.5)  # Penalize if unsafe score > 0.5
constraints = [safety_constraint]
trainer = ConstitutionalDPOTrainer(model, ref_model, constraints=constraints)

Practical considerations and best practices

When implementing DPO for real-world applications, keep the following tips in mind:

A) Data quality: The quality of your preference data is crucial. Make sure your dataset:

Covers a wide range of inputs and desired behavior
Has consistent and reliable preference annotations
Balances different types of preferences (e.g. factuality, safety, style)

B) Hyperparameter tuning: Although DPO has fewer hyperparameters than RLHF, tuning is still important:

β (beta): controls the trade-off between preference satisfaction and deviation from the reference model. Start with values around 0.1-0.5.
Learning rate: Use a lower learning rate than standard fine-tuning, usually in the range of 1st-6 to 1st-5.
Batch Size: Larger batch sizes (32-128) often work well for learning preferences.

C) Iterative refinement: DPO can be applied iteratively:

Train an initial model using DPO
Generate new responses using the trained model
Collect new preference data on these responses
Retrain using the expanded dataset

Performance for instant preference optimization

This view examines the performance of LLMs such as GPT-4 compared to human judgments on various training techniques, including Direct Preference Optimization (DPO), Supervised Fine-Tuning (SFT), and Proximal Policy Optimization (PPO). The table shows that the results of GPT-4 are increasingly in line with human preferences, especially in summary tasks. The level of agreement between GPT-4 and human reviewers demonstrates the model’s ability to generate content that resonates with human reviewers almost as closely as human-generated content.

Case studies and applications

To illustrate the effectiveness of DPO, let’s look at some real-world applications and some of their variations:

Iterative DPO: Developed by Snorkel (2023), this variant combines rejection sampling with DPO, allowing a more refined training data selection process. By repeating multiple rounds of preference sampling, the model is better able to generalize and avoid overfitting to noisy or biased preferences.
IPO (Iterative preference optimization): Introduced by Azar et al. (2023), IPO adds a regularization term to prevent overfitting, which is a common problem in preference-based optimization. This extension allows models to maintain a balance between retaining preferences and retaining generalization capabilities.
KTO (Optimization of knowledge transfer): A more recent variant of Ethayarajh et al. (2023), KTO does not use binary preferences at all. Instead, it focuses on transferring knowledge from a reference model to the policy model, optimizing for a smoother and more consistent alignment with human values.
Multimodal DPO for cross-domain learning by Xu et al. (2024): An approach that applies DPO to different modalities (text, image and audio), demonstrating its versatility in tailoring models to human preferences for different data types. This research highlights the potential of DPO in creating more comprehensive AI systems that can perform complex, multimodal tasks.

Conclusion

Direct preference optimization represents a significant advance in tailoring language models to human preferences. Its simplicity, efficiency and effectiveness make it a powerful tool for researchers and practitioners alike.

By harnessing the power of Direct Preference Optimization and keeping these principles in mind, you can create language models that not only exhibit impressive capabilities, but also closely align with human values and intentions.

Source link