Preference Alignment: RLHF & DPO Mathematics | AI & Machine Learning Engineering

Pre-trained foundation models are pure text completion engines optimized to predict the next token across internet corpora. Consequently, they mirror internet distributions, generating factually dubious hallucinations, toxic outputs, or evasive responses.

Transforming raw foundation models into helpful, safe assistants requires Preference Alignment. While Classical RLHF requires complex reinforcement learning loops with separate reward models, Direct Preference Optimization (DPO) analytically bypasses reinforcement learning altogether.

DPO Paradigm ShiftLLM Alignment

Zero PPO Required

Analytically substitutes the optimal policy into the Bradley-Terry model to align LLMs using simple binary cross-entropy.

1. Classical RLHF & The Bradley-Terry Preference Model#

In human preference learning, human annotators are presented with a prompt $x$ and two candidate model responses: a winning response $y_w$ and a losing response $y_l$ ( $y_w \succ y_l$ ).

Under the Bradley-Terry (1952) preference model, the probability that response $y_w$ is preferred over $y_l$ given latent scalar reward model $r(x, y)$ is:

P(y_w \succ y_l | x) = \sigma\Big( r(x, y_w) - r(x, y_l) \Big) = \frac{1}{1 + \exp\Big( -(r(x, y_w) - r(x, y_l)) \Big)}

Architecture Diagram

Rendering diagram...

The reward model parameters $\psi$ are trained by minimizing negative log-likelihood:

\mathcal{L}_R(\psi) = -\mathbb{E}_{(x, y_w, y_l)}\left[ \ln \sigma\Big( r_\psi(x, y_w) - r_\psi(x, y_l) \Big) \right]

The PPO Policy Optimization Objective#

Once the reward model $r_\psi$ is trained, policy $\pi_\theta$ is optimized using Proximal Policy Optimization (PPO) against an anchor reference model $\pi_{\text{ref}}$ :

\max_{\theta} \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_\theta}\left[ r_\psi(x, y) - \beta D_{KL}\Big( \pi_\theta(y | x) \,\|\, \pi_{\text{ref}}(y | x) \Big) \right]

Where $\beta$ penalizes drift away from the original model to prevent reward hacking (exploiting weaknesses in reward model $r_\psi$ ).

2. Mathematical Derivation of Direct Preference Optimization (DPO)#

Rafailov et al. (2023) observed that the KL-regularized RL objective possesses a closed-form analytical optimum, enabling direct optimization without training a separate reward model or executing PPO loops.

Step 1: The Optimal Policy Solution

Consider the constrained optimization problem over all candidate policies $\pi$ :

\max_{\pi} \mathbb{E}_{x \sim \mathcal{D}} \left[ \mathbb{E}_{y \sim \pi(y|x)}[r(x, y)] - \beta D_{KL}\Big( \pi(y | x) \,\|\, \pi_{\text{ref}}(y | x) \Big) \right]

Expanding the expectation and KL divergence for a single prompt $x$ :

\sum_{y} \pi(y | x) r(x, y) - \beta \sum_{y} \pi(y | x) \ln\left( \frac{\pi(y | x)}{\pi_{\text{ref}}(y | x)} \right)

= -\beta \sum_y \pi(y | x) \left[ \ln\left( \frac{\pi(y | x)}{\pi_{\text{ref}}(y | x)} \right) - \frac{1}{\beta} r(x, y) \right]

Using calculus of variations subject to constraint $\sum_y \pi(y | x) = 1$ , the optimal policy satisfies:

\pi^*(y | x) = \frac{1}{Z(x)} \pi_{\text{ref}}(y | x) \exp\left( \frac{1}{\beta} r(x, y) \right)

Where partition function $Z(x) = \sum_y \pi_{\text{ref}}(y | x) \exp\left( \frac{1}{\beta} r(x, y) \right)$ normalizes the distribution.

Step 2: Invert the Relation to Isolate the Reward Function

Take the natural logarithm of both sides:

\ln \pi^*(y | x) = \ln \pi_{\text{ref}}(y | x) + \frac{1}{\beta} r(x, y) - \ln Z(x)

Multiply by $\beta$ and isolate reward $r(x, y)$ :

r(x, y) = \beta \ln\left( \frac{\pi^*(y | x)}{\pi_{\text{ref}}(y | x)} \right) + \beta \ln Z(x)

This remarkable equation proves that the ground-truth reward is uniquely determined by the log-ratio of the optimal policy to the reference policy, plus a prompt-dependent normalization constant.

Step 3: Substitute into the Bradley-Terry Model

Substitute this analytical reward formulation into the Bradley-Terry preference probability:

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

Evaluate the reward difference $r(x, y_w) - r(x, y_l)$ :

r(x, y_w) - r(x, y_l) = \left[ \beta \ln\left( \frac{\pi(y_w | x)}{\pi_{\text{ref}}(y_w | x)} \right) + \beta \ln Z(x) \right] - \left[ \beta \ln\left( \frac{\pi(y_l | x)}{\pi_{\text{ref}}(y_l | x)} \right) + \beta \ln Z(x) \right]

Notice that the intractable partition function term $\beta \ln Z(x)$ cancels out completely:

r(x, y_w) - r(x, y_l) = \beta \ln\left( \frac{\pi_\theta(y_w | x)}{\pi_{\text{ref}}(y_w | x)} \right) - \beta \ln\left( \frac{\pi_\theta(y_l | x)}{\pi_{\text{ref}}(y_l | x)} \right)

Step 4: Formulate the Final DPO Loss

Applying Maximum Likelihood Estimation over preference dataset $\mathcal{D} = \{(x, y_w, y_l)\}$ :

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

3. Analysis of the DPO Gradient Update#

To understand how DPO updates network parameters $\theta$ , we compute the gradient with respect to $\theta$ :

\nabla_\theta \mathcal{L}_{\text{DPO}}(\theta) = -\beta \cdot \sigma\left( \beta \ln \frac{\pi_\theta(y_l | x)}{\pi_{\text{ref}}(y_l | x)} - \beta \ln \frac{\pi_\theta(y_w | x)}{\pi_{\text{ref}}(y_w | x)} \right) \cdot \left[ \nabla_\theta \ln \pi_\theta(y_w | x) - \nabla_\theta \ln \pi_\theta(y_l | x) \right]

The gradient has two intuitive factors:

Directional Push/Pull: $\nabla_\theta \ln \pi_\theta(y_w | x) - \nabla_\theta \ln \pi_\theta(y_l | x)$ increases the likelihood of preferred response $y_w$ while decreasing the likelihood of rejected response $y_l$ .
Implicit Weighting: The leading scalar weight scales with how incorrectly the current model scores the pair. If the model already strongly favors $y_w$ , the gradient evaluates close to zero. If the model incorrectly prefers $y_l$ , the gradient scales to its maximum.

4. From-Scratch Implementation: Zero Dependencies#

The following implementations compute the Bradley-Terry Likelihood and the Direct Preference Optimization (DPO) Loss from scratch in pure Python and pure Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import math
from typing import List, Tuple

class DPOPrimitives:
    @staticmethod
    def _sigmoid(z: float) -> float:
        if z < -700.0: return 0.0
        if z > 700.0: return 1.0
        return 1.0 / (1.0 + math.exp(-z))

    @classmethod
    def dpo_loss(
        cls,
        policy_chosen_logps: float,      # ln pi_theta(y_w | x)
        policy_rejected_logps: float,    # ln pi_theta(y_l | x)
        reference_chosen_logps: float,   # ln pi_ref(y_w | x)
        reference_rejected_logps: float, # ln pi_ref(y_l | x)
        beta: float = 0.1
    ) -> Tuple[float, float, float]:
        # Compute log-ratios for winning and losing responses
        # pi_logratios = ln(pi_theta / pi_ref)
        chosen_logratio = policy_chosen_logps - reference_chosen_logps
        rejected_logratio = policy_rejected_logps - reference_rejected_logps

        # Logit discrepancy: beta * (chosen_logratio - rejected_logratio)
        logits = beta * (chosen_logratio - rejected_logratio)

        # Loss = -ln(sigmoid(logits))
        # Numerically stable: -ln(1 / (1 + exp(-logits))) = ln(1 + exp(-logits))
        if logits > 0:
            loss = math.log(1.0 + math.exp(-logits))
        else:
            loss = -logits + math.log(1.0 + math.exp(logits))

        # Implicit rewards
        chosen_reward = beta * chosen_logratio
        rejected_reward = beta * rejected_logratio

        return loss, chosen_reward, rejected_reward

if __name__ == "__main__":
    # Case 1: Model currently prefers the winning response (low loss)
    loss1, r_w1, r_l1 = DPOPrimitives.dpo_loss(
        policy_chosen_logps=-2.5,
        policy_rejected_logps=-4.0,
        reference_chosen_logps=-3.0,
        reference_rejected_logps=-3.0,
        beta=0.1
    )
    print("Case 1 (Policy correctly favors winning response):")
    print(f"  DPO Loss: {loss1:.6f}")
    print(f"  Implicit Reward (Winning):  {r_w1:.4f}")
    print(f"  Implicit Reward (Losing):   {r_l1:.4f}")

    # Case 2: Model incorrectly prefers the losing response (high loss)
    loss2, r_w2, r_l2 = DPOPrimitives.dpo_loss(
        policy_chosen_logps=-4.5,
        policy_rejected_logps=-1.5,
        reference_chosen_logps=-3.0,
        reference_rejected_logps=-3.0,
        beta=0.1
    )
    print("\nCase 2 (Policy incorrectly favors losing response):")
    print(f"  DPO Loss: {loss2:.6f}")
    print(f"  Implicit Reward (Winning):  {r_w2:.4f}")
    print(f"  Implicit Reward (Losing):   {r_l2:.4f}")

5. Active Knowledge Checks#

Concept check

In the mathematical derivation of Direct Preference Optimization (DPO), what cancellation allows the loss to evaluate without computing the partition function Z(x)?

Concept check

Why does DPO eliminate the need for an explicit actor-critic reinforcement learning loop (PPO) during LLM alignment?