Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

DPO Implementation › DPO Loss Computation

Michael BrenndoerferJanuary 2, 202614 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 284DPO Implementation › DPO Loss Computation

Reading section 3 of 11 · DPO Loss Computation · 14 minSection 3 of 11 · 14 min

Section 3 of 11 · 14 min read

DPO Implementation › DPO Loss ComputationOne section · Sections in this chapter3 / 11 ↓

Sections in this chapter11

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

With properly formatted data, we can now implement the DPO loss function. This is the heart of the algorithm: a single formula that converts pairwise preference comparisons into a differentiable training objective. Understanding every component of this formula mechanically and intuitively is what lets you debug when training goes wrong and tune when training goes right.

The DPO objective encourages the model to assign higher implicit rewards to chosen responses while constraining the model to stay close to the reference policy. DPO turns a reinforcement learning problem into a simple classification task: given a pair of responses, predict which one humans would prefer. This reframing is why DPO can be trained with standard supervised learning infrastructure. The reinforcement learning complexity has been absorbed into the loss function itself.

Think of the DPO loss as a fairness evaluation on a debate team. Two debaters (the chosen and rejected responses) give speeches in response to the same prompt. A judge (the DPO loss) scores each debater not on absolute merit, but relative to how much better the winning debater performs compared to a baseline version of themselves. If the winning debater improved a lot from baseline while the losing debater stayed the same, the judge rewards that generously. If both debaters performed at baseline, the judge sees no clear winner and pushes for more differentiation. This relative scoring is precisely what the log ratio achieves.

Recall from the DPO Derivation chapter that the loss is defined as:

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

where:

$\pi_\theta$ : the policy model being trained
$\pi_\text{ref}$ : the frozen reference model (usually the initial version of $\pi_\theta$ )
$\mathcal{D}$ : the dataset of preference triplets $(x, y_w, y_l)$
$\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}}$ : the expected value over preference triplets drawn from the dataset (computed as a batch average)
$x$ : the prompt or instruction
$y_w$ : the chosen (winning) response
$y_l$ : the rejected (losing) response
$\beta$ : a temperature parameter scaling the strength of the preference constraint
$\sigma$ : the logistic sigmoid function, $\sigma(z) = \frac{1}{1 + e^{-z}}$

To understand this formula intuitively, consider what each component accomplishes. The log ratio $\log \frac{\pi_\theta(y|x)}{\pi_\text{ref}(y|x)}$ measures how much more (or less) likely the policy model finds a response compared to the reference model. When this ratio is positive, the policy has increased its probability for that response relative to where it started. When negative, the policy has decreased its probability. The term $\beta \log \frac{\pi_\theta(y|x)}{\pi_\text{ref}(y|x)}$ represents the implicit reward the model assigns to a response, capturing the idea that preferred responses should see increased probability while rejected responses should see decreased probability.

The difference between these implicit rewards for the chosen and rejected responses measures how well the model has learned the preference. If this difference is large and positive, the model strongly prefers the chosen response. The sigmoid function $\sigma$ converts this difference into a probability, and by taking the negative log, we arrive at a cross-entropy loss that pushes the model to make this probability as high as possible.

The key insight is that DPO does not require the model to learn an explicit scalar reward for each response. Instead, the reward is implicitly encoded in how much the policy's probability for a response has moved relative to the reference. This encoding is exactly what emerges from the analytic solution to the RLHF optimization problem, which is why DPO and RLHF are solving the same underlying problem through mathematically equivalent means.

By minimizing this negative log-likelihood, DPO optimizes the policy so that the implicit reward for the chosen response $y_w$ is higher than for the rejected response $y_l$ . DPO steers the model toward human preferences while staying close to the reference policy using only supervised learning.

Out[11]:

Visualization

Line plot of the sigmoid function over reward margin with green shading for positive margins and red shading for negative margins. — Sigmoid transformation mapping reward margin to probability, with shaded regions showing correct and incorrect response rankings.

Line plot of the DPO loss over reward margin with annotations marking high-loss wrong-ranking and low-loss correct-ranking regions. — DPO loss as a function of reward margin, showing near-zero loss for correct rankings and exponentially increasing penalties for incorrect rankings.

Computing Sequence Log ProbabilitiesLink Copied

Before we can evaluate the DPO loss, we need a way to compute $\log \pi(y|x)$ , the log probability that a given language model assigns to a complete response sequence given a prompt. This computation is the atomic building block from which everything else is assembled.

The probability of an entire response sequence is not directly output by a language model. What the model outputs at each forward pass is a probability distribution over the vocabulary for the next token, conditioned on all previous tokens. To get the probability of a complete sequence we must chain these per-token probabilities together using the chain rule of probability.

To compute the log probability $\log \pi(y|x)$ for a response $y$ given prompt $x$ , we sum the log probabilities of each token conditioned on its history. Since autoregressive models generate text one token at a time, each token's probability depends on all previous tokens:

\log \pi(y|x) = \sum_{t=1}^{|y|} \log \pi(y_t | x, y_{<t})

where:

$y$ : the full response sequence
$x$ : the input prompt
$|y|$ : the total number of tokens in the response
$y_t$ : the token at position $t$
$y_{<t}$ : the sequence of tokens preceding $t$ (the history)
$\sum_{t=1}^{|y|}$ : the sum of log probabilities across all tokens in the sequence

This formula arises from the chain rule of probability. The probability of generating an entire sequence equals the product of generating each token given everything that came before it. In log space, the product of token probabilities becomes a sum, which is numerically stable and computationally convenient. Each term $\log \pi(y_t | x, y_{<t})$ represents the model's confidence in predicting token $y_t$ given the prompt $x$ and all previously generated tokens $y_{<t}$ .

This summation computes the total log probability of the sequence by aggregating the log probabilities of each token conditioned on its history. Longer sequences have lower log probabilities because they contain more terms. Comparing log ratios rather than raw log probabilities cancels out length effects when comparing the same response under different models. This is a subtle but important point: the DPO loss is a difference of log probabilities between two models, not a raw measure of absolute sequence likelihood.

The key insight is that because the same sequence $y$ appears in both the numerator (policy) and denominator (reference) of the log ratio, any length-based penalty cancels. A 200-token response has a very negative absolute log probability, but its log ratio can be close to zero if both the policy and reference assign similar probabilities to it. This is mathematically clean and computationally reassuring.

In[12]:

Code

import torch
import torch.nn.functional as F


def compute_log_probs(model, input_ids, attention_mask, response_mask):
    """
    Compute per-token log probabilities for response tokens.

    Args:
        model: Language model
        input_ids: Token IDs [batch_size, seq_len]
        attention_mask: Attention mask [batch_size, seq_len]
        response_mask: Mask showing response tokens [batch_size, seq_len]

    Returns:
        Per-sequence log probabilities [batch_size]
    """
    # Get model outputs (logits)
    with torch.no_grad() if not model.training else torch.enable_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.logits  # [batch_size, seq_len, vocab_size]

    # Shift for next-token prediction: logits[t] predicts token[t+1]
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()
    shift_mask = response_mask[:, 1:].contiguous()

    # Compute per-token log probabilities
    log_probs = F.log_softmax(shift_logits, dim=-1)

    # Gather log probs for actual tokens
    token_log_probs = log_probs.gather(
        dim=-1, index=shift_labels.unsqueeze(-1)
    ).squeeze(-1)  # [batch_size, seq_len-1]

    # Mask out non-response tokens and sum
    masked_log_probs = token_log_probs * shift_mask.float()
    sequence_log_probs = masked_log_probs.sum(dim=-1)  # [batch_size]

    return sequence_log_probs

This function handles the necessary shift operation for next-token prediction: since logits at position $t$ predict token at position $t+1$ , we need to align everything properly. This alignment is a common source of bugs in language model implementations, so it deserves careful attention. The model's output at each position represents a probability distribution over what the next token should be, not what the current token is. Therefore, to find the probability assigned to each actual token in the sequence, we must look at the logits from the preceding position.

After the shift, we apply F.log_softmax to convert raw logits into log probabilities, then use gather to extract the log probability of each actual token from the full vocabulary distribution. The mask is shifted by the same offset so it correctly identifies which gathered log probabilities correspond to response tokens. Finally, masked multiplication zeros out the prompt token contributions, and the sum aggregates the response log probability. Each of these operations has a specific purpose: removing any of them produces a silently incorrect implementation.

The Core DPO Loss FunctionLink Copied

Now we can implement the full DPO loss. This function brings together all the components we have discussed: computing log probabilities for both the policy and reference models on both chosen and rejected responses, calculating the log ratios that represent implicit rewards, and combining them through the sigmoid to produce a differentiable loss signal.

The function is structured to make the connection to the formula transparent. We compute four quantities: the policy's log probability of the chosen response, the policy's log probability of the rejected response, the reference model's log probability of the chosen response, and the reference model's log probability of the rejected response. From these four numbers we build two log ratios, one for each response. The DPO loss is then the negative log-sigmoid of the scaled difference of these ratios. Every term in the mathematical formula maps directly to a line of code.

In[13]:

Code

def dpo_loss(
    policy_model,
    reference_model,
    chosen_input_ids,
    chosen_attention_mask,
    chosen_response_mask,
    rejected_input_ids,
    rejected_attention_mask,
    rejected_response_mask,
    beta=0.1,
):
    """
    Compute DPO loss for a batch of preference pairs.

    Args:
        policy_model: The model being trained
        reference_model: Frozen reference model
        chosen_*: Tokenized chosen responses
        rejected_*: Tokenized rejected responses
        beta: Temperature parameter controlling deviation from reference

    Returns:
        loss: Scalar loss value
        metrics: Dictionary with logging information
    """
    # Compute log probs for policy model
    policy_chosen_logps = compute_log_probs(
        policy_model,
        chosen_input_ids,
        chosen_attention_mask,
        chosen_response_mask,
    )
    policy_rejected_logps = compute_log_probs(
        policy_model,
        rejected_input_ids,
        rejected_attention_mask,
        rejected_response_mask,
    )

    # Compute log probs for reference model (no gradients needed)
    with torch.no_grad():
        ref_chosen_logps = compute_log_probs(
            reference_model,
            chosen_input_ids,
            chosen_attention_mask,
            chosen_response_mask,
        )
        ref_rejected_logps = compute_log_probs(
            reference_model,
            rejected_input_ids,
            rejected_attention_mask,
            rejected_response_mask,
        )

    # Compute log ratios
    chosen_log_ratio = policy_chosen_logps - ref_chosen_logps
    rejected_log_ratio = policy_rejected_logps - ref_rejected_logps

    # DPO loss: -log(sigmoid(beta * (chosen_ratio - rejected_ratio)))
    logits = beta * (chosen_log_ratio - rejected_log_ratio)
    loss = -F.logsigmoid(logits).mean()

    # Compute metrics for monitoring
    with torch.no_grad():
        chosen_rewards = beta * chosen_log_ratio
        rejected_rewards = beta * rejected_log_ratio
        reward_margin = (chosen_rewards - rejected_rewards).mean()
        accuracy = (chosen_rewards > rejected_rewards).float().mean()

    metrics = {
        "loss": loss.item(),
        "reward_margin": reward_margin.item(),
        "accuracy": accuracy.item(),
        "chosen_reward": chosen_rewards.mean().item(),
        "rejected_reward": rejected_rewards.mean().item(),
    }

    return loss, metrics

These metrics help monitor training effectively and each tells a different story about what the model is learning:

Reward margin: The average difference between implicit rewards for chosen and rejected responses. This should increase during training as the model learns to more strongly prefer the chosen responses. A growing reward margin indicates the model is successfully learning the preference signal.
Accuracy: The fraction of examples where the policy assigns higher implicit reward to the chosen response. Well-trained models should approach high accuracy on the training set. However, reaching 100% accuracy too quickly may indicate overfitting.
Individual rewards: Tracking both chosen and rejected rewards helps diagnose issues like reward hacking. Ideally, the chosen reward should increase modestly while the rejected reward decreases or stays stable. If both rewards increase dramatically, the model may be drifting too far from the reference policy.

A word on the torch.no_grad() context for the reference model computations: this is not optional. The reference model is frozen and we never backpropagate through it. Wrapping its forward pass in torch.no_grad() saves both memory and computation by preventing PyTorch from building a computational graph for those operations. Forgetting this wrapper is a common mistake that does not cause errors but significantly increases memory usage and training time.

Numerical Stability ConsiderationsLink Copied

The DPO loss involves log probabilities that can become very negative for long sequences. A sequence of 100 tokens where each token has an average log probability of -3 would yield a sequence log probability of -300. While this does not cause issues when computing ratios, the subtraction of two very negative numbers to form the ratio must be handled carefully in floating-point arithmetic.

A few practices help maintain numerical stability throughout the computation:

In[14]:

Code

def dpo_loss_stable(
    policy_chosen_logps,
    policy_rejected_logps,
    ref_chosen_logps,
    ref_rejected_logps,
    beta=0.1,
    label_smoothing=0.0,
):
    """
    Numerically stable DPO loss computation.

    Supports optional label smoothing for regularization.
    """
    # Compute log ratios (these are differences, not absolute values)
    chosen_log_ratio = policy_chosen_logps - ref_chosen_logps
    rejected_log_ratio = policy_rejected_logps - ref_rejected_logps

    logits = beta * (chosen_log_ratio - rejected_log_ratio)

    if label_smoothing > 0:
        # Soft labels: slightly prefer chosen but allow some uncertainty
        # This can help with noisy preference labels
        smooth_loss = (
            -F.logsigmoid(logits) * (1 - label_smoothing)
            - F.logsigmoid(-logits) * label_smoothing
        )
        loss = smooth_loss.mean()
    else:
        loss = -F.logsigmoid(logits).mean()

    return loss

Working with log ratios rather than raw probabilities ensures numerical stability. When we compute $\log \pi_\theta(y|x) - \log \pi_\text{ref}(y|x)$ , the very negative values from long sequences largely cancel out, leaving a ratio that reflects how differently the two models view the same response. This cancellation is what keeps the numbers in a manageable range.

Using F.logsigmoid(x) rather than torch.log(torch.sigmoid(x)) is another important stability choice. The naive implementation computes sigmoid(x) first, which saturates to exactly 0.0 for very negative values, and then takes the log of 0, creating -inf. The logsigmoid function avoids this by using a numerically stable formula equivalent to $-\log(1 + e^{-x})$ for positive $x$ and $x - \log(1 + e^x)$ for negative $x$ .

Label smoothing can be helpful when your preference labels are noisy, as discussed in the Human Preference Data chapter. Setting a small smoothing value (0.01-0.1) prevents the model from becoming overconfident about potentially mislabeled examples. The smoothing works by mixing in a small probability of the "wrong" label, which acts as a form of regularization that improves generalization when the training data contains annotation errors. In practice, most DPO training runs use no label smoothing, reserving it for datasets with known annotation noise.