RLHF & DPO from scratch

1 · Preference data & the Bradley–Terry model

Some qualities — helpfulness, tone, harmlessness — are far easier to express as preferences (“A is better than B”) than as gold answers. The unit of preference tuning is the triple (prompt, chosen, rejected). To turn a pile of such comparisons into a trainable signal we need a probabilistic model of preference. The standard choice is Bradley–Terry: if response y has a latent scalar reward r(y), the probability a rater prefers y_w (winner) over y_l (loser) is the logistic of the reward gap:

P(y_w > y_l) = σ(r(y_w) − r(y_l)), where σ(z) = 1/(1+e^−z).

So the whole learning signal is the reward margin between chosen and rejected. A big margin gives a confident preference and a strong gradient; a near-tie teaches almost nothing. This is why curating pairs with real quality gaps matters as much as volume.

2 · Lab · fit a reward model from pairs

A reward model is just a function that scores a response, trained so that σ(r(chosen) − r(rejected)) is high for every pair. Here we fit a trivial linear reward on hand-crafted features by gradient descent on the Bradley–Terry log-likelihood — pure Python, so you watch the margin grow.

reward_model.pyimport math, random
random.seed(0)

def sigmoid(z):
    return 1.0 / (1.0 + math.exp(-z)) if z >= 0 else math.exp(z) / (1.0 + math.exp(z))

# Features per response: [length_norm, has_specifics, is_refusal]. r(y) = w . features.
# Each pair: (features_chosen, features_rejected). Chosen is the better answer.
pairs = [
    ([0.8, 1, 0], [0.2, 0, 1]),   # long+specific  vs  short refusal
    ([0.6, 1, 0], [0.5, 0, 0]),   # specific       vs  vague
    ([0.7, 1, 0], [0.9, 0, 0]),   # specific       vs  long but vague
    ([0.5, 1, 0], [0.4, 0, 1]),   # helpful        vs  refusal
]
w = [0.0, 0.0, 0.0]
lr = 0.5

def reward(feats):
    return sum(wi * fi for wi, fi in zip(w, feats))

for epoch in range(61):
    grad = [0.0, 0.0, 0.0]
    loss = 0.0
    for fc, fr in pairs:
        margin = reward(fc) - reward(fr)
        p_pref = sigmoid(margin)
        loss += -math.log(p_pref)
        # d(-log sigma(margin))/dw = -(1 - p_pref) * d(margin)/dw
        coef = -(1 - p_pref)
        for i in range(3):
            grad[i] += coef * (fc[i] - fr[i])
    for i in range(3):
        w[i] -= lr * grad[i] / len(pairs)
    if epoch % 20 == 0:
        avg_margin = sum(reward(fc) - reward(fr) for fc, fr in pairs) / len(pairs)
        print(f"epoch {epoch:>2}: loss={loss/len(pairs):.4f}  "
              f"mean chosen-rejected margin={avg_margin:+.2f}")

sign = lambda x: "+" if x >= 0 else "-"
print(f"learned weights favor: length{sign(w[0])} specificity{sign(w[1])} refusal{sign(w[2])}")

epoch  0: loss=0.6931  mean chosen-rejected margin=+0.00
epoch 20: loss=0.3567  mean chosen-rejected margin=+1.31
epoch 40: loss=0.2489  mean chosen-rejected margin=+2.05
epoch 60: loss=0.1954  mean chosen-rejected margin=+2.53
learned weights favor: length+ specificity+ refusal-

The loss falls and the mean margin between chosen and rejected grows — the reward model has learned to score preferred answers higher. In real RLHF this r is a full network with a scalar head on the LLM; the objective is identical.

3 · PPO vs DPO: two ways to use the reward

Once you have preferences, there are two routes to an aligned policy. RLHF/PPO trains the reward model, then runs a reinforcement-learning loop: sample responses from the policy, score them with the reward model, and update the policy to raise expected reward — while a KL penalty keeps it from drifting too far from the reference (SFT) model. It works, but it is a moving target: reward model, policy, and a sampling loop, all interacting and unstable.

DPO observes that the KL-regularized objective has a closed-form optimal policy, and substitutes it back to get a loss you can optimize directly on the pairs — no reward model, no sampling, no RL loop. Same alignment target, one simple loss.

4 · Deriving the DPO loss

This is the heart of the lesson. RLHF maximizes expected reward with a KL leash to the reference policy π_ref:

max_π E[r(y)] − β·KL(π ‖ π_ref)

This constrained problem has a known closed-form optimum — the reference policy reweighted by the exponentiated reward:

π*(y|x) = (1/Z(x)) · π_ref(y|x) · exp(r(x,y)/β)

Rearranging for the reward gives r(x,y) = β·log(π*(y|x)/π_ref(y|x)) + β·log Z(x). The insight: substitute this expression for r into the Bradley–Terry preference probability. The intractable partition term log Z(x) is the same for chosen and rejected (same prompt x), so it cancels in the difference. What remains is a loss over the policy's own log-probs versus the reference — the DPO loss:

L_DPO = −log σ( β·[ logπ(y_w|x) − logπ_ref(y_w|x) ] − β·[ logπ(y_l|x) − logπ_ref(y_l|x) ] )

5 · Lab · the DPO loss & one gradient step

Now implement exactly that loss in pure Python and take one gradient step, watching the log-prob gap between chosen and rejected widen. This is a toy 2-token policy, but the loss and its gradient are the real DPO update.

dpo_from_scratch.pyimport math

def sigmoid(z):
    return 1.0 / (1.0 + math.exp(-z)) if z >= 0 else math.exp(z) / (1.0 + math.exp(z))

# Toy policy over two responses {chosen, rejected}, parameterized by a logit `theta`
# (logit of the chosen response). logp(chosen)=log softmax; reference is frozen.
theta = 0.0                      # policy starts equal to reference
theta_ref = 0.0                  # frozen reference
BETA = 1.0
LR = 0.5

def logps(theta):
    # softmax over [chosen, rejected] with logits [theta, 0]
    z = math.log(math.exp(theta) + math.exp(0.0))
    return theta - z, 0.0 - z     # logp(chosen), logp(rejected)

for step in range(11):
    lp_c, lp_r = logps(theta)
    ref_c, ref_r = logps(theta_ref)
    chosen_ratio = lp_c - ref_c
    reject_ratio = lp_r - ref_r
    diff = BETA * (chosen_ratio - reject_ratio)
    loss = -math.log(sigmoid(diff))
    if step in (0, 1, 5, 10):
        print(f"step {step:>2}: DPO loss={loss:.4f}  "
              f"logp(chosen)-logp(rejected)={lp_c - lp_r:+.3f}")
    # gradient of loss wrt theta (finite difference, tiny epsilon)
    eps = 1e-5
    lp_c2, lp_r2 = logps(theta + eps)
    d2 = BETA * ((lp_c2 - ref_c) - (lp_r2 - ref_r))
    loss2 = -math.log(sigmoid(d2))
    grad = (loss2 - loss) / eps
    theta -= LR * grad

print("the policy now assigns more probability to the chosen response")

step  0: DPO loss=0.6931  logp(chosen)-logp(rejected)=+0.000
step  1: DPO loss=0.6446  logp(chosen)-logp(rejected)=+0.104
step  5: DPO loss=0.4901  logp(chosen)-logp(rejected)=+0.548
step 10: DPO loss=0.3616  logp(chosen)-logp(rejected)=+1.043
the policy now assigns more probability to the chosen response

The loss falls and the policy shifts probability mass toward the chosen response, relative to the frozen reference — with no reward model in sight. β is the one dial: higher β pushes harder toward chosen but risks over-fitting the pairs and drifting from the reference.

6 · Reward hacking & its mitigations

Any time you optimize a proxy for what you want, the optimizer will exploit the gap. In RLHF the proxy is the reward model; a policy trained too hard against it finds cheap tricks that score high without being genuinely better — reward hacking. Classic symptoms: answers grow longer because the reward model correlated length with quality; the model becomes sycophantic; it repeats reward-y phrases. DPO is not immune — over-fitting the pairs is the same failure in a different coat.

reward_hacking.py# Proxy reward rewards length; TRUE quality peaks at a moderate length then falls.
def proxy_reward(length):          # what the reward model (wrongly) scores
    return min(length * 0.19, 10.0)

def true_quality(length):          # what humans actually want: sweet spot ~12-16
    return max(0.0, 8.0 - abs(length - 14) * 0.35)

def optimize(kl_leash, ref_length=14, steps=200, lr=0.6):
    length = float(ref_length)
    for _ in range(steps):
        eps = 1e-3
        obj = lambda L: proxy_reward(L) - (kl_leash * (L - ref_length) ** 2 if kl_leash else 0)
        grad = (obj(length + eps) - obj(length - eps)) / (2 * eps)
        length = max(1.0, length + lr * grad)
    return length

for name, kl in [("no KL leash:  ", 0.0), ("with KL leash:", 0.05)]:
    L = optimize(kl)
    pr, tq = proxy_reward(L), true_quality(L)
    tag = "HACKED" if tq < 4 else "aligned"
    print(f"{name} final length={L:>3.0f} tokens  proxy reward={pr:.2f}  "
          f"TRUE quality={tq:.1f}  ({tag})")
print("the leash trades proxy reward for real quality")

no KL leash:   final length= 48 tokens  proxy reward=9.10  TRUE quality=2.1  (HACKED)
with KL leash: final length= 14 tokens  proxy reward=5.30  TRUE quality=6.8  (aligned)
the leash trades proxy reward for real quality

Without the KL penalty the policy maximizes the flawed proxy (here: length) and drifts far from sensible behavior — high proxy reward, low true quality. The KL leash (the β term in both PPO and DPO) holds it near the reference and preserves real quality. The other mitigations are the ones you already know: better preference data, a capped/normalized reward, and holding out evals the reward model never saw.

goals = [
    ("emit this exact JSON schema",           "SFT"),
    ("prefer polite over blunt phrasing",     "preference (DPO)"),
    ("be more helpful, less evasive",         "preference (DPO)"),
    ("copy this exact support macro",         "SFT"),
]
for goal, method in goals:
    print(f"{goal:38} -> {method}")

import math
def sigmoid(z):
    return 1/(1+math.exp(-z)) if z >= 0 else math.exp(z)/(1+math.exp(z))

def prefers_chosen(r_chosen, r_rejected):
    return sigmoid(r_chosen - r_rejected)

print(f"equal rewards:      {prefers_chosen(2.0, 2.0):.3f}")   # 0.5
print(f"chosen +3 better:   {prefers_chosen(5.0, 2.0):.3f}")   # ~0.95
print(f"shift both by +100: {prefers_chosen(102.0, 99.0):.3f}") # same as +3

import math
def sigmoid(z):
    return 1/(1+math.exp(-z)) if z >= 0 else math.exp(z)/(1+math.exp(z))

def dpo_loss(logp_pol_c, logp_pol_r, logp_ref_c, logp_ref_r, beta=0.1):
    chosen_ratio = logp_pol_c - logp_ref_c
    reject_ratio = logp_pol_r - logp_ref_r
    return -math.log(sigmoid(beta * (chosen_ratio - reject_ratio)))

# reference fixed; sweep the policy raising chosen above rejected
for boost in (0.0, 1.0, 3.0):
    L = dpo_loss(-1.0 + boost, -1.0, -1.0, -1.0, beta=1.0)
    print(f"policy chosen logratio +{boost}: DPO loss={L:.4f}")

import math
def sigmoid(z):
    return 1/(1+math.exp(-z)) if z >= 0 else math.exp(z)/(1+math.exp(z))

def logps(theta):                      # softmax over [chosen, rejected], logits [theta, 0]
    z = math.log(math.exp(theta) + 1.0)
    return theta - z, -z

def dpo_loss(theta, theta_ref=0.0, beta=1.0):
    lp_c, lp_r = logps(theta); rc, rr = logps(theta_ref)
    return -math.log(sigmoid(beta * ((lp_c - rc) - (lp_r - rr))))

theta = 0.0
before, gap0 = dpo_loss(theta), (logps(theta)[0] - logps(theta)[1])
eps = 1e-5
grad = (dpo_loss(theta + eps) - dpo_loss(theta)) / eps
theta -= 0.5 * grad
after, gap1 = dpo_loss(theta), (logps(theta)[0] - logps(theta)[1])
print(f"loss {before:.4f} -> {after:.4f}   gap {gap0:+.3f} -> {gap1:+.3f}")
assert after < before and gap1 > gap0

def proxy(L):   return min(L * 0.19, 10.0)          # reward model (wrongly) likes length
def truth(L):   return max(0.0, 8.0 - abs(L - 14) * 0.35)   # humans: sweet spot ~14

def optimize(kl, ref=14, steps=200, lr=0.6):
    L = float(ref)
    for _ in range(steps):
        e = 1e-3
        obj = lambda x: proxy(x) - (kl * (x - ref) ** 2 if kl else 0)
        L = max(1.0, L + lr * (obj(L + e) - obj(L - e)) / (2 * e))
    return L

for name, kl in [("no leash ", 0.0), ("KL leash ", 0.05)]:
    L = optimize(kl)
    print(f"{name}: length={L:5.1f}  proxy={proxy(L):.2f}  TRUE quality={truth(L):.1f}")

def recipe(have_gold, have_prefs, safety_critical):
    if not have_gold:
        return "COLLECT gold first -- SFT is the foundation for any alignment"
    if not have_prefs:
        return "SFT only for now -- add DPO once you can collect preference pairs"
    if safety_critical:
        return "SFT + DPO -- DPO nudges toward safe/helpful; RLHF only if DPO plateaus"
    return "SFT + DPO -- the applied default"

print(recipe(True,  True,  True))    # SFT + DPO ...
print(recipe(True,  False, False))   # SFT only for now ...
print(recipe(False, False, True))    # COLLECT gold first ...