DPO Derivation › Code Implementation
Language AI HandbookAll chapters DPO Derivation › Code ImplementationOne section · Sections in this chapter8 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement the DPO loss function and verify our derivation with code.
import torch
import torch.nn.functional as F
def dpo_loss(
policy_chosen_logps: torch.Tensor,
policy_rejected_logps: torch.Tensor,
ref_chosen_logps: torch.Tensor,
ref_rejected_logps: torch.Tensor,
beta: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Compute the DPO loss for a batch of preference pairs.
Args:
policy_chosen_logps: Log probs of chosen responses under policy [batch_size]
policy_rejected_logps: Log probs of rejected responses under policy [batch_size]
ref_chosen_logps: Log probs of chosen responses under reference [batch_size]
ref_rejected_logps: Log probs of rejected responses under reference [batch_size]
beta: Temperature parameter controlling deviation from reference
Returns:
loss: Scalar DPO loss
chosen_rewards: Implicit rewards for chosen responses [batch_size]
rejected_rewards: Implicit rewards for rejected responses [batch_size]
"""
# Compute implicit rewards (log-ratios scaled by beta)
chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)
# DPO loss: negative log-sigmoid of reward difference
logits = chosen_rewards - rejected_rewards
loss = -F.logsigmoid(logits).mean()
return loss, chosen_rewards, rejected_rewardsNow let's verify with our worked example:
# Values from our worked example
policy_chosen_logps = torch.tensor([-42.1])
policy_rejected_logps = torch.tensor([-39.5])
ref_chosen_logps = torch.tensor([-45.2])
ref_rejected_logps = torch.tensor([-38.7])
loss, chosen_rewards, rejected_rewards = dpo_loss(
policy_chosen_logps,
policy_rejected_logps,
ref_chosen_logps,
ref_rejected_logps,
beta=0.1,
)Chosen implicit reward: 0.310 Rejected implicit reward: -0.080 Reward difference: 0.390 Preference probability: 0.596 DPO loss: 0.517
The implicit rewards align with our derivation: the chosen response has a positive reward (0.310), while the rejected response has a negative reward (-0.080). The positive difference (0.390) results in a preference probability of 0.596. This indicates the model correctly prefers the chosen response, though the probability is only moderately high.
Let's also visualize how the loss and preference probability change as the model learns:
The loss approaches zero as the reward difference increases, meaning the model strongly prefers chosen responses. Conversely, when the model prefers rejected responses (negative reward difference), the loss grows large.
Let's also implement a function to compute the gradients and see how the implicit weighting works:
def dpo_loss_with_grad_weights(
policy_chosen_logps: torch.Tensor,
policy_rejected_logps: torch.Tensor,
ref_chosen_logps: torch.Tensor,
ref_rejected_logps: torch.Tensor,
beta: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Compute DPO loss and the implicit gradient weights.
The gradient weight is sigma(-reward_diff), which is the probability
the model assigns to the WRONG preference.
"""
chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)
logits = chosen_rewards - rejected_rewards
loss = -F.logsigmoid(logits).mean()
# Gradient weight: probability of wrong preference
grad_weight = torch.sigmoid(-logits)
return loss, grad_weight# Test with different scenarios
scenarios = [
("Model strongly correct", -20.0, -83.5, -45.2, -38.7),
("Model weakly correct", -42.1, -50.6, -45.2, -38.7),
("Model uncertain", -45.0, -38.5, -45.2, -38.7),
("Model weakly wrong", -47.0, -25.5, -45.2, -38.7),
("Model strongly wrong", -75.5, -19.0, -45.2, -38.7),
]
results = []
for name, pc, pr, rc, rr in scenarios:
loss, grad_weight = dpo_loss_with_grad_weights(
torch.tensor([pc]),
torch.tensor([pr]),
torch.tensor([rc]),
torch.tensor([rr]),
beta=0.1,
)
results.append((name, loss.item(), grad_weight.item()))Scenario Loss Grad Weight -------------------------------------------------- Model strongly correct 0.001 0.001 Model weakly correct 0.201 0.182 Model uncertain 0.693 0.500 Model weakly wrong 1.701 0.818 Model strongly wrong 5.007 0.993
Notice how the gradient weight automatically adjusts based on how wrong the model is. In the "Model strongly correct" case, the gradient weight is near zero (0.000), meaning there's little to learn. Conversely, when the model is "Strongly wrong", the gradient weight approaches one (0.993), pushing hard to correct the mistake.
