Implementing Direct Preference Optimization (DPO) on Llama-3 Models Without the Computation Costs of RLHF

Puntos clave

1. Understand the math. DPO derives the optimal policy from preference data without an explicit reward model.
2. Prepare preference data. You need pairs of chosen and rejected responses for the same prompt.
3. Set the beta parameter. Beta controls how far the policy moves from the reference; start at 0.1.
4. Use a reference model. The reference is the SFT model frozen; the policy starts from it and diverges.
5. Train with a small learning rate. DPO is sensitive to learning rate; use 1e-6 to 5e-7.
6. Monitor KL divergence. If the policy diverges too far from the reference, reduce beta or learning rate.
7. Evaluate on held-out preferences. Check that the policy prefers chosen over rejected on unseen data.
8. Iterate on data quality. DPO amplifies data quality issues; clean, diverse preferences matter more than quantity.

Architecture

The DPO loss compares the log probability of chosen and rejected responses under the policy and reference models. The policy is updated to increase the margin between chosen and rejected, weighted by how much the reference already prefers the chosen response.

Production configuration

import torch
import torch.nn.functional as F

def dpo_loss(policy, ref, chosen, rejected, beta=0.1):
    pi_chosen = policy.log_prob(chosen)
    pi_rejected = policy.log_prob(rejected)
    ref_chosen = ref.log_prob(chosen)
    ref_rejected = ref.log_prob(rejected)
    logits = beta * ((pi_chosen - pi_rejected) - (ref_chosen - ref_rejected))
    loss = -F.logsigmoid(logits).mean()
    return loss

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained('meta-llama/Llama-3-8B')
ref_model = AutoModelForCausalLM.from_pretrained('meta-llama/Llama-3-8B')
ref_model.eval()

for batch in dataloader:
    loss = dpo_loss(model, ref_model, batch['chosen'], batch['rejected'])
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

pip install trl
python -c "
from trl import DPOTrainer
trainer = DPOTrainer(
  model=model,
  ref_model=ref_model,
  args=training_args,
  train_dataset=dataset,
  tokenizer=tokenizer,
  beta=0.1,
)
trainer.train()
"

Edge cases from production

• Low-quality preference data produces a worse model than SFT; curate carefully.
• DPO can overfit to the preference distribution; evaluate on diverse prompts.
• Large beta values cause training instability; start small and increase gradually.
• Reference model must match the policy's starting point exactly; use the same checkpoint.

Preguntas frecuentes