What it actually does
A pretrained model can produce a plausible continuation but has no notion of which of two plausible continuations anyone would rather receive, and the reason is structural: next-token likelihood rewards imitating the training distribution, which contains excellent and terrible text in unknown proportion. Preferences are the missing supervision, and they are cheap in a specific form — people are unreliable at scoring a response out of ten and fairly reliable at saying which of two they prefer. So collect comparisons, fit something that explains them, and push the model toward what wins. Two routes from there. Classic RLHF fits a reward model to the comparisons and then optimizes the policy against it with reinforcement learning, held near the original by a KL penalty. DPO observes that the optimal policy under that penalty has a closed form, inverts it to express the reward in terms of the policy itself, and lands on a logistic loss over preference pairs with no reward model and no sampling. The idea that makes both work is also the one that makes them fail: because only differences are ever observed, the optimization has a direction it can run in forever, and the KL term is the only thing holding it. Release that and you get a model that scores superbly and is worse.
- Human preferences are consistent enough to be modelled by Bradley-Terry — transitive, and driven by the response rather than by presentation order
- The annotators' preferences are the ones worth optimizing, which is a value judgement embedded in the dataset rather than a property of the method
- The reward model or implicit reward remains valid over the region the policy reaches, which is exactly what the KL term is there to enforce and exactly what over-optimization violates
- The preference data is on-policy for the model being tuned; the further the generating policy is from it, the less of the data means anything, and DPO cannot sample to correct it
- The reference policy is already competent, since neither method teaches capability — both only reweight behaviour the model can already produce
Margin objective
A logistic loss on a margin: the difference between the implicit reward of the preferred response and that of the dispreferred one, with nothing in the objective that cares about either value on its own. That is the structural fact the whole entry turns on. Classic RLHF writes the same problem in two steps — fit a reward model to the comparisons, then maximize expected reward minus beta times the KL divergence from the reference — and the optimum of that second step has a closed form, a reference policy exponentially tilted by the reward. Inverting that identity expresses the reward as a log-ratio of policies, substituting it back into the Bradley-Terry likelihood yields exactly the formula above, and the reinforcement-learning problem has become a classification problem. Three consequences worth carrying. Because only differences are observed, a reward model is identified only up to a shift per prompt, so its absolute values are meaningless and comparing them across prompts is a mistake people make constantly. Because the objective only raises a margin, it can be satisfied by lowering the preferred response's likelihood as long as the dispreferred one falls faster — the widely reported pathology where both log-probabilities decline while the loss improves, and probability mass drains into sequences nobody labelled. And because the KL term is the only anchor, beta is not a regularization detail but the entire specification of how far a proxy may be trusted.
- the preferred and dispreferred responses to the same prompt — always a pair, never an absolute score
- the policy being trained and the frozen reference it started from
- the KL strength; small values permit a large departure from the reference, which is where the failures live
- the implicit reward — DPO's central identity, a reward expressed entirely through the policy
- the logistic function, from the Bradley-Terry model of pairwise choice
How it is minimized
PPO against a learned reward model with a KL penalty, or DPO's gradient descent on the pairwise logistic loss with no sampling at all
- the implicit reward, beta times the log-ratio of policy to reference on that response
- a per-pair weight near one when the model currently has the pair backwards and near zero once it is confidently right
- the ordinary likelihood gradient for the preferred response — raised
- the same for the dispreferred response — lowered, and nothing bounds how far
The gradient is worth reading slowly, because it explains both why DPO is stable and why it drifts. It is the ordinary likelihood gradient on the winner minus the ordinary likelihood gradient on the loser, scaled by how badly the current implicit reward has the pair ordered — an automatic curriculum that concentrates on pairs still being got wrong and ignores ones already settled. That is the good half. The bad half is that only the difference is controlled, so the update is indifferent between raising the winner and sinking the loser, and in practice it sinks the loser; the reported symptom is a healthy-looking loss with both log-probabilities falling, which means probability mass is moving to sequences that appeared in no label. Against that, PPO with a reward model keeps a property DPO gives up: it samples from the current policy, so it can reward responses the preference set never contained, which matters whenever the dataset was collected from a materially different model. The cost is operational — four models resident, a value head, rollouts, and a clip range and KL coefficient that both need tuning — and it is a real cost rather than a fashion, which is why DPO displaced it for most work. The honest summary is that DPO is a better default and a worse ceiling: reliable, cheap, one meaningful hyperparameter, and structurally unable to discover anything outside its dataset.
beta
The KL strength, and the only thing bounding how far the policy may chase a proxy. Too small and reward hacking is guaranteed; too large and nothing moves (0.01 to 0.5, commonly 0.1)
reference policy
What the KL is measured against. Almost always the supervised-fine-tuned checkpoint, and a mismatch between it and the policy that generated the preference data is DPO's main silent failure (the SFT checkpoint)
learning rate
Low, because the reference anchor is the only thing preventing collapse and a large step escapes it in a handful of updates (1e-7 to 5e-6)
preference batch size
Pairs per step. Larger batches damp the noise in individual annotations, which is substantial at typical agreement rates (32 to 512 pairs)
PPO clip range
RLHF only. Bounds the policy ratio per update; loose values make the run unstable rather than fast (0.1 to 0.3)
reward-model ensemble size
RLHF only. Disagreement across an ensemble is the cheapest available detector of the policy leaving the reward model's support (1 to 5)
Four failures, and the first is not a bug in any implementation. Reward hacking is Goodhart's law on schedule: optimize a learned proxy hard enough and true quality rises, peaks, and then declines while the proxy keeps improving, so the measured reward is a rising line over a model that is getting worse. The peak sits at a finite KL distance from the reference, which is why beta is the specification of the whole procedure and why KL from reference is the single most informative thing to plot. Second, the annotator ceiling: pairwise human agreement on realistic comparisons typically runs between seventy and eighty percent, so a reward model cannot be meaningfully more accurate than that, and every downstream number inherits the ceiling — a reward model reporting ninety-five percent validation accuracy has learned an artefact, most often length. Third, length and style bias, the most reproducible artefact in this area: annotators prefer longer, more confident, better-formatted answers, the reward model learns that faithfully, and the policy discovers it immediately, so verbosity is the first thing any preference-optimized model acquires. Fourth, DPO's likelihood displacement: because only the margin is constrained, both chosen and rejected log-probabilities commonly fall together, which satisfies the objective while moving mass to unlabelled sequences — monitor the two log-probabilities separately, since the loss curve cannot show it.
Classic RLHF holds four models at once — policy, frozen reference, reward model and value head — and is dominated by rollout generation rather than by gradient steps, since each update needs fresh samples from the current policy. DPO holds two, policy and reference, needs two forward passes per response, and does no sampling at all; the reference passes depend on nothing being trained, so they can be computed once offline and cached, which removes the second model from the training loop entirely. Reward-model training itself is ordinary supervised learning, two forward passes per pair, and is the cheapest stage by a wide margin.
How, where, and why it is used
Time-Series Forecasting
Not applicableForecasts are scored against values that arrive on their own, so substituting a learned proxy for human preference replaces a ground truth with an approximation of one — and every failure mode here comes from optimizing a proxy that a real target was available for.
Anomaly Detection
Not applicablePreference optimization changes what a model prefers to emit and produces no measure of how unusual an input is; the reward is an ordering over responses to the same prompt, identified only up to a shift, so it cannot be read as a density or a score.
Optimization
Primary fitThe entire technique is a constrained optimization and a case study in optimizing a proxy. Maximize expected reward under a KL trust region around a reference policy, where the reward is itself learned from noisy pairwise labels — and the constraint, not the objective, is what determines whether the result is any good.
- KL-constrained policy optimization, where the constraint is the specification rather than a regularization detail
- Goodhart's law with a measurable turning point: true quality against proxy reward as a function of distance from the reference
- An algebraic reparameterization that converts a reinforcement-learning problem into a classification loss with no sampling
- Learning from comparisons rather than scores, which identifies the objective only up to a per-prompt shift
The most instructive optimization story in this reference, because the proxy is visibly a proxy and the divergence between it and the truth can be plotted. Several lessons transfer well beyond alignment. A learned objective is only trustworthy within the region where its training data lived, so unconstrained optimization against it is guaranteed to leave that region — the KL term is not conservatism, it is the statement of where the proxy is valid. The DPO derivation is a clean demonstration that a change of variables can eliminate an entire optimization stage, since the closed-form optimum of the constrained problem can be inverted and substituted back rather than reached by iteration. And the objective depending only on differences shows up everywhere once noticed: it means absolute reward values carry no information, which invalidates any comparison of reward across prompts and a great deal of reported tooling that does exactly that.
- Plot true quality against KL from the reference rather than against training step; the turning point is the result and step count hides it
- Normalize or centre reward per prompt before any comparison, since the Bradley-Terry model fixes it only up to a per-prompt shift
- Hold out preference pairs from a different annotator cohort, because agreement rate is the ceiling on everything downstream
- Track chosen and rejected log-probabilities separately; a margin improving while both fall is the failure the loss curve cannot show
Win rate against the reference policy judged by something other than the reward being optimized — human review or an independent judge — reported against KL distance rather than against step. The critical measurement is where that curve turns: proxy reward rises monotonically past that point while true quality falls, and a run reported at its final step rather than its peak is reporting the wrong model.
- Reporting the proxy reward as the result, which rises exactly while quality declines
- Treating beta as a regularization detail rather than as the statement of where the proxy can be trusted
- Comparing absolute reward values across prompts, which the preference model never identified
- Evaluating with a judge related to the reward model, which shares its blind spots and confirms them
Natural Language
Primary fitCollect pairwise comparisons on model responses to real prompts, fit a reward model or optimize the policy directly, and hold the result near the supervised-fine-tuned checkpoint with a KL term. This is the stage that turns a model which can produce a plausible answer into one that produces the answer people wanted.
- Instruction-following and helpfulness tuning after supervised fine-tuning
- Harmlessness and refusal behaviour, where the preference data encodes a policy nobody can write down as rules
- Tone, format and house-style conformance that supervised data cannot specify by example alone
- Task-specific preference tuning — summarization faithfulness, code review style, support register
The standard final stage of post-training, and the reason is that the target is genuinely unspecifiable any other way: nobody can write the loss function for "a helpful answer", but almost anyone can pick the better of two. The domain-specific caveats are sharp. Verbosity bias is the most reproducible finding in the area — annotators prefer longer and more confident answers, the reward model learns it perfectly, and the policy exploits it within a few hundred steps, so length-controlled evaluation is mandatory rather than fastidious. Preference data ages badly: it reflects the model that generated it, so a dataset collected from last quarter's checkpoint is partly off-policy for this quarter's, and DPO is more sensitive to that than PPO because it cannot sample its way out. And alignment is a distribution over annotators, not a fact — who labelled the comparisons is part of the model's specification and belongs in its documentation as much as the training corpus does.
- Report length-controlled win rates, because a raw win rate is measuring verbosity as much as quality
- Collect comparisons from the policy being tuned rather than from an older checkpoint, or accept a distribution gap DPO cannot close
- Record annotator cohort and agreement rate with the dataset; they are the ceiling on the reward model and on everything after it
- Deduplicate near-identical pairs, which inflate apparent agreement and concentrate the gradient on a narrow slice of behaviour
Pairwise win rate against the reference, length-controlled, judged independently of the reward model — an automated judge sharing the reward model's lineage confirms its biases rather than testing them. Alongside it, a capability benchmark the preference data never touched, since preference tuning routinely costs a few points of reasoning or factual accuracy and nothing in the objective notices.
- Uncontrolled win rates, which reward length and formatting more than substance
- Preference data generated by a different checkpoint than the one being tuned
- A judge model related to the reward model, which makes the evaluation agree with the failure
- No capability benchmark, so the alignment tax goes unmeasured and is discovered by users
Computer Vision
AdaptedApply the same pairwise machinery to image generation: collect comparisons between two generations from the same prompt, and either fit an aesthetic or alignment reward model or optimize the diffusion model directly with a DPO-style loss over the denoising likelihood.
- Aesthetic and prompt-adherence tuning of text-to-image models
- Safety and content-policy behaviour expressed through comparisons rather than rules
- Brand or house-style conformance for a specific deployment
- Reducing characteristic artefacts that are obvious to a viewer and invisible to a pixel metric
A real and increasingly standard fit, with one structural complication that makes it harder than the language case. The likelihood of a diffusion sample is not directly available — it is a bound over a sampling trajectory — so the log-ratio at the centre of the loss becomes an approximation over sampled timesteps rather than an exact quantity, which adds variance the language version does not have. The motivation is strong anyway, because image quality is precisely the sort of target that is easy to compare and impossible to specify: no pixel metric agrees with a viewer about which of two images is better. The failure modes rhyme with the language ones and are more visible. Reward hacking produces a recognizable house style — oversaturated, over-smoothed, centrally composed — that scores well and collapses diversity, and because mode collapse is obvious by eye here, this is the setting where the cost of over-optimization is easiest to see.
- Sample timesteps for the likelihood ratio rather than assuming a single step represents the trajectory
- Hold a fixed prompt suite, including prompts unrelated to the preference data, to watch diversity rather than only quality
- Separate aesthetic preference from prompt adherence in the data; a single comparison conflates them and the model optimizes the easier one
- Keep beta higher than the language default, since the noisier ratio makes a loose constraint fail faster
Human comparison against the reference on a fixed prompt suite, plus an explicit diversity measure across seeds — reward-hacked image models score well per image and produce the same image repeatedly, and no per-image metric can see that.
- Optimizing until every output shares one house style, which scores well and is a collapse
- Conflating aesthetic quality with prompt adherence in a single comparison
- Ignoring the extra variance from approximating the diffusion likelihood ratio
- Evaluating only on prompts resembling the preference data, which cannot show lost diversity
Recommendation & Ranking
AdaptedThe Bradley-Terry model underneath all of this is the same one classical learning-to-rank has used for decades, so pairwise implicit feedback — this item was clicked over that one — fits the reward-modelling half directly. The policy half applies where the system generates something rather than selecting it.
- Pairwise learning-to-rank from implicit feedback, which is reward modelling by another name
- Generated surfaces — explanations, summaries, message copy — tuned by comparison rather than by click-through alone
- Conversational recommendation, where the response is generated and preference data is the natural supervision
- Offline preference data used to pre-train a ranker before online exploration
Worth including because it shows the lineage honestly: the pairwise logistic objective here is the same one RankNet used, and the interesting part is what the modern framing adds rather than any novelty in the loss. What it adds is the KL anchor and a language of proxy over-optimization that ranking has always needed and rarely names — a ranker optimized hard against a click proxy degrades in exactly the way a policy optimized against a reward model does. The boundary is worth being clear about. Ranking systems usually select from a catalogue rather than generate, and for selection the policy-optimization half has nothing to do; the reward-modelling half is simply supervised ranking, already well served by established methods. And implicit feedback is confounded by exposure in a way human comparisons are not: the loser of a click pair may have been positioned worse rather than been worse, which the Bradley-Terry model cannot distinguish and will happily fit.
- Correct for position and exposure before treating a click pair as a preference, since the model cannot separate them
- Keep the pair within a single impression, or the comparison is confounded by context rather than reflecting quality
- Anchor a generated surface to its reference policy with a KL term, exactly as in the language case
- Treat the click proxy as a proxy and watch for the same turnover, because it behaves like one
Online tests decide, as always in this domain. Offline pairwise accuracy on held-out impressions is worth having as a regression guard, with the same caveat that it measures agreement with logged behaviour — including the incumbent ranker's exposure decisions — rather than quality.
- Treating a click pair as a preference without correcting for position
- Applying the policy-optimization half to a system that selects rather than generates
- Optimizing a click proxy without watching for the turnover that a learned proxy always has
- Comparing items across impressions, where the Bradley-Terry shift is not identified
What it costs to run
Classic RLHF is the expensive path and most of the cost is not gradients: four models resident, and generation of fresh rollouts at every step, which dominates wall-clock and makes the run sensitive to sampling throughput rather than to training throughput. DPO is a different budget entirely — two forward passes per response, no sampling, and the reference passes precomputable offline, which in practice brings a preference run within reach of the same hardware that did the supervised fine-tune. The dominant cost in both cases is neither: it is collecting the comparisons. Illustrative, not a measured benchmark.
Nothing at all. Both methods produce an ordinary policy with the same architecture and the same latency as the model that went in, which is worth stating because it is unusual — the entire cost of this stage is paid at training time. The one change a deployment notices is behavioural: response length typically increases, which raises token cost per request even though per-token latency is unchanged.
Set by preference-data collection, which is slow and human-bound, so the realistic cadence is weeks to months rather than continuous. The trap is that preference data ages against the model: a dataset collected from an earlier checkpoint is partly off-policy for the current one, and DPO has no mechanism to correct for that because it never samples. Budget fresh comparisons from the current policy alongside any retrain rather than reusing the existing set.
- KL divergence from the reference policy, which is the axis every failure in this technique moves along and the one plot worth having
- Chosen and rejected log-probabilities tracked separately, since a margin can improve while both fall and the loss cannot show it
- Response-length distribution, because verbosity is the first thing a preference-optimized model acquires and the easiest to mistake for quality
- Capability benchmarks the preference data never touched, to measure the alignment tax rather than assume it is zero
- Reward-model ensemble disagreement in RLHF, which rises as the policy leaves the region the reward model was fitted on
- Refusal and over-refusal rates together, since harmlessness tuning moves both and only one is usually watched
- A reward model is identified only up to a per-prompt shift, so absolute reward values mean nothing and comparing them across prompts is a mistake tooling encourages
- Reward hacking is not an implementation bug. Proxy reward rises monotonically while true quality peaks and declines, so a run reported at its final step reports the wrong checkpoint
- The reference policy is part of the artefact. Change it and beta means something different, and a KL budget tuned against one reference does not transfer
- Verbosity is acquired within a few hundred steps and reads as quality in every uncontrolled evaluation, including most automated judges
- A judge model sharing lineage with the reward model agrees with its biases, which makes the evaluation confirm the failure instead of detecting it
- DPO can lower the likelihood of chosen and rejected responses together, moving probability mass to sequences nobody labelled — visible only if both are logged
- Preference data encodes the annotator cohort. Who labelled it is part of the model's specification and should be documented as carefully as the training corpus
- The alignment tax is real: a few points of reasoning or factual accuracy commonly go, and nothing in the objective notices or reports it
Pros and cons, in context
Optimizes a target that cannot be written down as a loss, using comparisons people can actually make reliably
The reason this stage exists. Absolute quality ratings are noisy and poorly calibrated between annotators; pairwise choices are far more stable, and Bradley-Terry is the standard way to turn the second into a model
DPO removes the reward model, the value head and the sampling loop by an exact algebraic substitution
Not an approximation but an identity under the KL-constrained objective. It turns four resident models into two and a reinforcement-learning loop into a supervised one, which is why it displaced PPO for most work
The KL term makes "how far may the model move" an explicit, tunable quantity
Rare and valuable: most techniques leave the trust in a proxy implicit. Here it is one number, and plotting quality against KL distance turns over-optimization into something you can see rather than suspect
Inference cost is completely unchanged — the output is an ordinary policy
The whole expense is at training time, so a preference-tuned model drops into existing serving with no architectural change. Token cost may rise through longer responses, which is behavioural rather than structural
The DPO gradient weights each pair by how wrong the model currently is on it
An automatic curriculum that costs nothing: settled pairs contribute almost no gradient while contested ones dominate, which is part of why the method is stable at low learning rates
Reward hacking is guaranteed rather than possible: optimize the proxy far enough and true quality turns over while the measured reward keeps rising
The peak sits at a finite KL distance, so the method has no natural stopping point and the stopping criterion has to come from an independent evaluation. Reporting the final checkpoint reports the wrong model
The reward model cannot exceed the annotators' agreement rate, typically seventy to eighty percent on realistic comparisons
A low ceiling that everything downstream inherits, and a reward model reporting much higher validation accuracy has usually learned length or formatting rather than quality
Verbosity and formatting bias are learned faithfully and exploited immediately
The most reproducible artefact in this area. It makes uncontrolled win rates unusable as evidence, and length-controlled evaluation is the minimum bar rather than a refinement
DPO can drive the likelihood of preferred and dispreferred responses down together
The objective constrains only their difference, so this satisfies it perfectly while moving mass to unlabelled sequences. Invisible in the loss curve and visible only if both log-probabilities are logged
DPO cannot explore, so it is bounded by the responses its dataset already contains
The cost of removing the sampling loop. Where the preference data came from a materially different checkpoint, the gap cannot be closed by training longer, and PPO's on-policy sampling is the reason it still has a role
Preferences are a distribution over annotators, not a fact, and the dataset is a value judgement
Two cohorts produce two different models from the same procedure. This belongs in the model's documentation, and it usually is not there
The alignment tax is commonly a few points of reasoning or factual accuracy
Nothing in the objective protects a capability the preference data never mentioned, so it must be measured on held-out benchmarks or it will be discovered in production
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""RLHF and DPO, transcribed from the objective.
No ML library, and a deliberately tiny world: a handful of prompts with
a handful of candidate responses each, so every quantity in the theory
can be computed exactly rather than estimated. That is what makes the
three claims below checkable instead of quoted.
RLHF: max_pi E[r(x, y)] - beta * KL(pi || pi_ref)
optimum: pi*(y|x) proportional to pi_ref(y|x) * exp(r(x, y)/beta)
DPO: -log sigmoid( beta*log(pi/pi_ref)[y_w] - beta*log(pi/pi_ref)[y_l] )
Claim 1: the closed form above IS the optimum. tilted_policy() computes
it directly; dpo_train() reaches it by gradient descent on a
classification loss that never mentions reinforcement learning.
Claim 2: the reward is identified only up to a per-prompt shift.
bradley_terry_log_likelihood() is invariant to adding a constant to
every response of a prompt -- so absolute reward values carry no
information, and any tool comparing them across prompts is reading
noise.
Claim 3: reward hacking is not a bug. overoptimization_curve() fits a
reward model to limited comparisons, sweeps beta, and reports true
reward against KL from the reference. It rises, peaks, and falls, while
the proxy rises throughout.
"""
import math
import random
Distribution = list[float]
def softmax(logits: list[float]) -> Distribution:
largest = max(logits)
exponentials = [math.exp(value - largest) for value in logits]
total = sum(exponentials)
return [value / total for value in exponentials]
def kl_divergence(policy: Distribution, reference: Distribution) -> float:
"""KL(policy || reference), the term that anchors the whole method.
Asymmetric on purpose and in the direction that matters: it is
infinite where the policy puts mass that the reference does not, so
it forbids the policy from inventing behaviour rather than merely
discouraging it.
"""
total = 0.0
for probability, base in zip(policy, reference):
if probability <= 0.0:
continue
if base <= 0.0:
return float("inf")
total += probability * math.log(probability / base)
return total
def tilted_policy(reference: Distribution, rewards: list[float], beta: float) -> Distribution:
"""The closed-form optimum of the KL-constrained problem.
pi*(y|x) proportional to pi_ref(y|x) * exp(r(x, y) / beta)
Worth staring at, because DPO is nothing more than this identity
read backwards. Solve for r and you get beta * log(pi*/pi_ref) plus
a term constant in y -- which is exactly the implicit reward, and
exactly why the constant cancels in a difference.
Note what beta does here. Large beta flattens the exponent and the
optimum stays near the reference; small beta concentrates all mass
on the highest-reward response, reference be damned. Beta is not a
regularization detail, it is the statement of how far the proxy is
trusted.
"""
tilted = []
for probability, reward in zip(reference, rewards):
tilted.append(probability * math.exp(reward / beta))
total = sum(tilted)
return [value / total for value in tilted]
class PreferencePair:
def __init__(self, prompt: int, winner: int, loser: int) -> None:
self.prompt = prompt
self.winner = winner
self.loser = loser
def bradley_terry_log_likelihood(rewards: list[list[float]], pairs: list[PreferencePair]) -> float:
"""The likelihood a reward model is fitted to.
P(y_w beats y_l) = sigmoid(r(y_w) - r(y_l))
Only the DIFFERENCE appears. Add any constant to every response of
one prompt and this value does not move, which is the formal
statement that reward is identified up to a per-prompt shift --
and the reason a dashboard comparing mean reward across prompts is
displaying an arbitrary offset.
"""
total = 0.0
for pair in pairs:
margin = rewards[pair.prompt][pair.winner] - rewards[pair.prompt][pair.loser]
# log sigmoid(margin), written stably.
total += -math.log1p(math.exp(-margin)) if margin > 0 else margin - math.log1p(math.exp(margin))
return total / len(pairs)
def fit_reward_model(
n_prompts: int, n_responses: int, pairs: list[PreferencePair], steps: int, lr: float
) -> list[list[float]]:
"""Ordinary supervised learning on pairwise labels.
Tabular, so there is no function approximation to hide behind: the
reward model is one number per (prompt, response) and the only
signal it ever sees is which of two was preferred. Everything it
knows about a response it never saw is nothing, which is precisely
why optimizing against it out of distribution is unsafe.
"""
rewards = [[0.0 for _ in range(n_responses)] for _ in range(n_prompts)]
for _ in range(steps):
gradients = [[0.0 for _ in range(n_responses)] for _ in range(n_prompts)]
for pair in pairs:
margin = rewards[pair.prompt][pair.winner] - rewards[pair.prompt][pair.loser]
# d/dmargin of log sigmoid(margin) is sigmoid(-margin):
# large when the pair is currently ordered wrongly, near
# zero once it is confidently right. An automatic curriculum
# that costs nothing.
weight = 1.0 / (1.0 + math.exp(margin))
gradients[pair.prompt][pair.winner] += weight
gradients[pair.prompt][pair.loser] -= weight
for prompt in range(n_prompts):
for response in range(n_responses):
rewards[prompt][response] += lr * gradients[prompt][response] / len(pairs)
return rewards
def dpo_loss_and_gradient(
logits: list[list[float]],
reference: list[Distribution],
pairs: list[PreferencePair],
beta: float,
):
"""The DPO objective, computed rather than described.
The implicit reward is beta * log(pi_theta / pi_ref), and the loss
is a logistic loss on the difference between the winner's and the
loser's. No reward model appears anywhere; no response is ever
sampled.
"""
policies = [softmax(row) for row in logits]
gradients = [[0.0 for _ in row] for row in logits]
total_loss = 0.0
for pair in pairs:
policy = policies[pair.prompt]
base = reference[pair.prompt]
implicit_winner = beta * math.log(policy[pair.winner] / base[pair.winner])
implicit_loser = beta * math.log(policy[pair.loser] / base[pair.loser])
margin = implicit_winner - implicit_loser
total_loss += math.log1p(math.exp(-margin)) if margin > 0 else -margin + math.log1p(math.exp(margin))
# sigma(r_loser - r_winner): how wrong the implicit reward
# currently is on this pair. The whole gradient is scaled by it.
weight = 1.0 / (1.0 + math.exp(margin))
# d/dlogit_j of log pi(y) is [j == y] - pi(j), the ordinary
# softmax gradient. So the update is the likelihood gradient of
# the winner MINUS that of the loser -- and nothing in it cares
# which of the two moves, only that they separate. That is the
# source of the likelihood-displacement pathology: sinking the
# loser satisfies the objective exactly as well as raising the
# winner, and empirically it is what happens.
for response in range(len(policy)):
winner_term = (1.0 if response == pair.winner else 0.0) - policy[response]
loser_term = (1.0 if response == pair.loser else 0.0) - policy[response]
gradients[pair.prompt][response] -= beta * weight * (winner_term - loser_term)
count = len(pairs)
for row in gradients:
for index in range(len(row)):
row[index] /= count
return total_loss / count, gradients
def dpo_train(
reference: list[Distribution], pairs: list[PreferencePair], beta: float, steps: int, lr: float
):
"""Gradient descent on a classification loss, arriving at an
RL optimum.
The policy is initialized AT the reference, which is what the theory
assumes and what production does. Every log-ratio starts at zero, so
the implicit reward starts at zero everywhere and the first update
is driven entirely by the labels.
"""
logits = [[math.log(max(probability, 1e-12)) for probability in row] for row in reference]
history = []
for _ in range(steps):
loss, gradients = dpo_loss_and_gradient(logits, reference, pairs, beta)
for prompt, row in enumerate(logits):
for response in range(len(row)):
row[response] -= lr * gradients[prompt][response]
history.append(loss)
return [softmax(row) for row in logits], history
def log_probability_report(policy: list[Distribution], pairs: list[PreferencePair]):
"""Chosen and rejected log-probabilities, separately.
The loss curve cannot distinguish "the winner went up" from "the
loser went down twice as far", and the second is the reported
failure. Logging both is the only way to see it, and it is the
cheapest diagnostic in this entire entry.
"""
chosen = 0.0
rejected = 0.0
for pair in pairs:
chosen += math.log(max(policy[pair.prompt][pair.winner], 1e-12))
rejected += math.log(max(policy[pair.prompt][pair.loser], 1e-12))
count = len(pairs)
return {
"mean_chosen_logprob": chosen / count,
"mean_rejected_logprob": rejected / count,
"margin": (chosen - rejected) / count,
}
def sample_preferences(
true_rewards: list[list[float]], n_pairs: int, seed: int
) -> list[PreferencePair]:
"""Noisy annotators, sampled from Bradley-Terry.
Labels come from sigmoid(r_a - r_b) rather than from argmax,
because that is what the model assumes and what people actually do:
disagree at a rate set by how close the two options are. Agreement
on genuinely close pairs is near chance, and no reward model can
exceed that.
"""
rng = random.Random(seed)
pairs = []
for _ in range(n_pairs):
prompt = rng.randrange(len(true_rewards))
n_responses = len(true_rewards[prompt])
first, second = rng.sample(range(n_responses), 2)
gap = true_rewards[prompt][first] - true_rewards[prompt][second]
probability_first_wins = 1.0 / (1.0 + math.exp(-gap))
if rng.random() < probability_first_wins:
pairs.append(PreferencePair(prompt, first, second))
else:
pairs.append(PreferencePair(prompt, second, first))
return pairs
def overoptimization_curve(
n_prompts: int = 6, n_responses: int = 8, n_pairs: int = 160, seed: int = 0
):
"""Goodhart's law, plotted.
A true reward exists here and is never shown to the fitting
procedure. A reward model is fitted to a limited, noisy sample of
comparisons drawn from it, then the policy is tilted against the
LEARNED reward at a sweep of beta values, and both rewards are
measured under the result.
What comes back: proxy reward rises monotonically as beta falls,
true reward rises, peaks, and declines. The peak is at a finite KL
distance, so there is no "train until converged" here -- the
stopping point is a measurement, and it must come from an
evaluation the optimizer never saw.
"""
rng = random.Random(seed)
true_rewards = [
[rng.gauss(0.0, 1.0) for _ in range(n_responses)] for _ in range(n_prompts)
]
reference = [softmax([rng.gauss(0.0, 0.5) for _ in range(n_responses)]) for _ in range(n_prompts)]
pairs = sample_preferences(true_rewards, n_pairs, seed + 1)
learned = fit_reward_model(n_prompts, n_responses, pairs, steps=600, lr=2.0)
curve = []
for beta in [2.0, 1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125]:
proxy_total = 0.0
true_total = 0.0
kl_total = 0.0
for prompt in range(n_prompts):
policy = tilted_policy(reference[prompt], learned[prompt], beta)
kl_total += kl_divergence(policy, reference[prompt])
for response in range(n_responses):
proxy_total += policy[response] * learned[prompt][response]
true_total += policy[response] * true_rewards[prompt][response]
curve.append(
{
"beta": beta,
"kl_from_reference": kl_total / n_prompts,
"proxy_reward": proxy_total / n_prompts,
"true_reward": true_total / n_prompts,
}
)
return curve
def shift_invariance_check(rewards: list[list[float]], pairs: list[PreferencePair], shift: float):
"""The non-identifiability, demonstrated rather than claimed.
Add a constant to every response of every prompt and the Bradley-
Terry likelihood is unchanged to floating-point noise. The reward
model has no opinion about absolute value and never did.
"""
before = bradley_terry_log_likelihood(rewards, pairs)
shifted = [[value + shift for value in row] for row in rewards]
after = bradley_terry_log_likelihood(shifted, pairs)
return {"before": before, "after": after, "difference": after - before}
Reward-model fitting is O(steps · pairs); DPO training is O(steps · pairs · responses) because the softmax gradient touches every response of the pair's prompt; the closed-form tilt is a single O(responses) pass per prompt. Illustrative, not a measured benchmark: the shape worth noticing is that the closed form and the gradient descent reach the same policy, so the entire reinforcement-learning apparatus in the tabular case is replaceable by one exponential — which is the identity DPO exploits at scale, where the closed form is unreachable.
How this language actually wants it written.
"""Preference optimization with the silent failures made loud.
The literal version accepted all of these without complaint:
* a beta tuned against one reference policy, applied to another,
where it specifies a different trust region entirely;
* an annotator cohort agreeing at chance, whose labels carry no
signal and whose reward model will still report a number;
* training to convergence, which by construction returns a
checkpoint past the point where true quality turned over.
Each becomes an exception, a budget or a guard here. What is added
beyond defence is diagnostic: log_probability_report() and
length_controlled_win_rate() are the two measurements that separate
this working from this appearing to work.
"""
from __future__ import annotations
import contextlib
import math
from collections.abc import Iterator, Sequence
from dataclasses import dataclass, field
from typing import NamedTuple, Protocol
import numpy as np
import numpy.typing as npt
FloatArray = npt.NDArray[np.float64]
# Below this pairwise agreement the cohort is not distinguishable from
# coin flipping, and a reward model fitted to it is fitting noise.
MIN_ANNOTATOR_AGREEMENT = 0.55
class PreferenceError(Exception):
"""Base for every recoverable failure in preference optimization."""
class DegeneratePair(PreferenceError):
"""A comparison that carries no information."""
class UninformativeAnnotations(PreferenceError):
"""Agreement at or near chance; there is no signal to fit."""
class ReferenceMismatch(PreferenceError):
"""The policy is being anchored to a reference it did not start from.
Beta is defined relative to a specific reference, so this is not a
bookkeeping complaint: the same beta against a different reference
is a different trust region and a different objective.
"""
class PreferencePair(NamedTuple):
"""A comparison, with the cohort that produced it.
The cohort travels with the pair because the cohort's agreement
rate is the ceiling on the reward model and on everything after it.
Leaving it in a README is how that ceiling gets forgotten.
"""
prompt_id: str
chosen: str
rejected: str
cohort: str
chosen_tokens: int
rejected_tokens: int
@dataclass(frozen=True)
class ReferencePolicy:
"""The frozen anchor, fingerprinted.
Pinned by identity rather than by value: a reference is a specific
checkpoint, and two checkpoints that score identically still define
different KL geometries.
"""
checkpoint_id: str
log_probs: dict[tuple[str, str], float]
def log_prob(self, prompt_id: str, response: str) -> float:
key = (prompt_id, response)
if key not in self.log_probs:
raise ReferenceMismatch(
f"reference {self.checkpoint_id} has no log-probability for "
f"{prompt_id!r}/{response[:32]!r}; the preference data was generated "
"by a different policy than the one anchoring this run"
)
return self.log_probs[key]
class Policy(Protocol):
"""What the optimizer needs from a policy, and nothing more."""
def log_prob(self, prompt_id: str, response: str) -> float: ...
def parameters(self) -> FloatArray: ...
def set_parameters(self, values: FloatArray) -> None: ...
@dataclass(frozen=True)
class DpoConfig:
"""Validated once, at construction."""
beta: float
learning_rate: float
reference_checkpoint_id: str
def __post_init__(self) -> None:
if not 0.0 < self.beta <= 1.0:
raise PreferenceError(
f"beta must lie in (0, 1]; got {self.beta}. Beta is the statement of "
"how far the proxy may be trusted, not a regularization detail"
)
if self.learning_rate <= 0.0:
raise PreferenceError(f"learning rate must be positive, got {self.learning_rate}")
def validate_pairs(pairs: Sequence[PreferencePair]) -> None:
"""Guard clauses, so the interesting path is the training loop."""
if not pairs:
raise PreferenceError("no preference pairs")
for pair in pairs:
if pair.chosen == pair.rejected:
raise DegeneratePair(
f"{pair.prompt_id}: chosen and rejected are identical; the margin is "
"zero by construction and the gradient is exactly nothing"
)
def cohort_agreement(
pairs: Sequence[PreferencePair], second_pass: Sequence[PreferencePair]
) -> float:
"""Agreement between two independent labellings of the same pairs.
This number is the ceiling on the reward model, and a reward model
reporting validation accuracy well above it has learned an
artefact — length, most often — rather than quality.
"""
if len(pairs) != len(second_pass):
raise PreferenceError("labelling passes cover different pair counts")
if not pairs:
raise UninformativeAnnotations("no pairs to measure agreement over")
agreed = sum(
1 for first, second in zip(pairs, second_pass) if first.chosen == second.chosen
)
rate = agreed / len(pairs)
if rate < MIN_ANNOTATOR_AGREEMENT:
raise UninformativeAnnotations(
f"cohort agreement {rate:.3f} is below {MIN_ANNOTATOR_AGREEMENT}; "
"these labels are close to coin flips and the reward model will fit noise"
)
return rate
@dataclass
class KlBudget:
"""The stopping criterion, as an object rather than a step count.
True quality peaks at a finite KL distance from the reference and
declines after it, so "train to convergence" returns the wrong
checkpoint by construction. This stops on measured distance and
keeps the checkpoint from the best independent evaluation, not the
last one.
"""
max_kl: float
best_score: float = -math.inf
best_parameters: FloatArray | None = None
best_kl: float = 0.0
_stopped: bool = field(default=False, init=False)
def observe(self, kl: float, score: float, parameters: FloatArray) -> bool:
if kl < 0.0:
raise ValueError(f"KL cannot be negative, got {kl}")
if score > self.best_score:
self.best_score = score
self.best_parameters = parameters.copy()
self.best_kl = kl
if kl > self.max_kl:
self._stopped = True
return not self._stopped
@contextlib.contextmanager
def preference_training(policy: Policy, budget: KlBudget) -> Iterator[KlBudget]:
"""Restore the best checkpoint on the way out, always.
A context manager rather than a restore call at the end, because
the run will eventually raise — an evaluation times out, a batch is
malformed — and a restore that depends on that not happening is a
restore that will be skipped exactly when the run is salvageable.
"""
try:
yield budget
finally:
if budget.best_parameters is not None:
policy.set_parameters(budget.best_parameters)
def implicit_reward(
policy: Policy, reference: ReferencePolicy, pair_side: tuple[str, str], beta: float
) -> float:
"""beta * log(pi_theta / pi_ref) on one response.
DPO's central identity. Absolute values are meaningless — the
Bradley-Terry model identifies reward only up to a per-prompt shift
— so this is only ever used inside a difference.
"""
prompt_id, response = pair_side
return beta * (policy.log_prob(prompt_id, response) - reference.log_prob(prompt_id, response))
def dpo_loss(
policy: Policy,
reference: ReferencePolicy,
pairs: Sequence[PreferencePair],
config: DpoConfig,
) -> float:
"""The pairwise logistic loss, computed stably.
logaddexp rather than log(1 + exp(-margin)): the naive form
overflows for the confidently-wrong pairs, which are precisely the
ones carrying the whole gradient early in a run.
"""
validate_pairs(pairs)
if reference.checkpoint_id != config.reference_checkpoint_id:
raise ReferenceMismatch(
f"config anchors to {config.reference_checkpoint_id}, reference is "
f"{reference.checkpoint_id}; beta specifies a different trust region "
"against a different anchor"
)
margins = np.array(
[
implicit_reward(policy, reference, (pair.prompt_id, pair.chosen), config.beta)
- implicit_reward(policy, reference, (pair.prompt_id, pair.rejected), config.beta)
for pair in pairs
]
)
return float(np.mean(np.logaddexp(0.0, -margins)))
class LogProbReport(NamedTuple):
mean_chosen: float
mean_rejected: float
margin: float
both_falling: bool
def log_probability_report(
policy: Policy, pairs: Sequence[PreferencePair], previous: LogProbReport | None = None
) -> LogProbReport:
"""Chosen and rejected log-probabilities, separately.
The objective constrains only their difference, so it is satisfied
equally well by raising the winner and by sinking the loser faster
— and in practice it sinks the loser, moving probability mass to
sequences nobody labelled. The loss curve cannot show this. These
two numbers can, and nothing else will.
"""
chosen = np.array([policy.log_prob(p.prompt_id, p.chosen) for p in pairs])
rejected = np.array([policy.log_prob(p.prompt_id, p.rejected) for p in pairs])
mean_chosen = float(chosen.mean())
mean_rejected = float(rejected.mean())
both_falling = (
previous is not None
and mean_chosen < previous.mean_chosen
and mean_rejected < previous.mean_rejected
)
return LogProbReport(
mean_chosen=mean_chosen,
mean_rejected=mean_rejected,
margin=mean_chosen - mean_rejected,
both_falling=both_falling,
)
def length_controlled_win_rate(
wins: Sequence[bool], policy_tokens: Sequence[int], reference_tokens: Sequence[int]
) -> dict[str, float]:
"""A win rate with the length effect regressed out.
Annotators prefer longer answers, reward models learn that
faithfully, and policies exploit it within a few hundred steps. A
raw win rate is therefore measuring verbosity as much as quality,
which makes it evidence of very little on its own.
The residual below is not a causal correction and is not offered as
one; it is the minimum honest presentation — raw and adjusted side
by side, with the length gap stated so a reader can judge.
"""
if not (len(wins) == len(policy_tokens) == len(reference_tokens)):
raise ValueError("win, policy-length and reference-length counts differ")
if not wins:
raise ValueError("no comparisons")
outcomes = np.asarray(wins, dtype=np.float64)
length_gap = np.asarray(policy_tokens, dtype=np.float64) - np.asarray(
reference_tokens, dtype=np.float64
)
design = np.column_stack([np.ones_like(length_gap), length_gap])
coefficients, *_ = np.linalg.lstsq(design, outcomes, rcond=None)
return {
"raw_win_rate": float(outcomes.mean()),
"length_adjusted_win_rate": float(coefficients[0]),
"mean_length_gap_tokens": float(length_gap.mean()),
"slope_per_token": float(coefficients[1]),
}
What changed: The reference policy stops being an array passed alongside the data and becomes a pinned, fingerprinted object, because beta is meaningless without it — a KL budget tuned against one reference does not transfer to another, and the literal version would happily train against whatever was passed. Every recoverable failure names the values that caused it: a preference pair whose two responses are identical, an annotator cohort whose agreement rate is at chance and therefore carries no signal, a reference probability of zero where the policy is being asked to place mass, and a beta outside the range where the objective is numerically meaningful. Preference pairs become a NamedTuple carrying the annotator cohort, so the ceiling that cohort imposes travels with the data instead of living in a README. The stopping criterion becomes an explicit object rather than a step count, because the peak of true quality sits at a finite KL distance and a run that trains to convergence returns the wrong checkpoint by construction — the KLBudget here stops on a measured distance and records the checkpoint at the peak. Training runs inside a context manager that guarantees the reference model is released and the best checkpoint is restored, on the exception path as well as the ordinary one. And the two diagnostics that catch this failing are functions rather than advice: separate chosen and rejected log-probabilities, which is the only view that reveals likelihood displacement, and a length-controlled win rate, because an uncontrolled one measures verbosity.
Same arithmetic as the literal version — one implicit-reward difference per pair, a logistic loss over the batch — with the per-pair Python loop replaced by array operations for the reductions. Illustrative, not a measured benchmark: the substantive change is that a mismatched reference, an annotator cohort at chance, a degenerate pair and a run trained past its KL budget are now failures the caller must handle, and the best checkpoint is restored on the exception path rather than lost with the run.
Faster, and the justification for each change.
"""Preference optimization, shaped around one observation.
The reference policy is frozen. Its log-probabilities depend on nothing
being trained, so they can be computed once for the whole dataset and
cached -- which removes an entire model from the training loop, halves
resident memory, and halves the forward work per step. That is the
largest win available here and it is structural rather than numerical.
Everything else is array shaping: chosen and rejected scored in one
padded batch, the per-token gather vectorized, and the sums fused into
the loss so no per-token intermediate lands.
What none of this changes: the peak of true quality still sits at a
finite KL distance, so a faster run reaches the wrong checkpoint sooner
unless the budget from the previous stage travels with it.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import numpy.typing as npt
FloatArray = npt.NDArray[np.float32]
IdArray = npt.NDArray[np.int32]
@dataclass(frozen=True)
class BatchShape:
max_pairs: int
max_tokens: int
class ReferenceCache:
"""Reference sequence log-probabilities, computed once.
Keyed by the reference checkpoint id, so changing the anchor
invalidates the cache instead of silently reusing numbers from a
different geometry -- which would leave beta specifying a trust
region around a policy no longer in the run.
The values are float64 on purpose. They are summed over sequences
that can run to thousands of tokens, and this is the one place where
float32 accumulation loses enough precision to move a margin.
"""
def __init__(self, reference_checkpoint_id: str) -> None:
self.reference_checkpoint_id = reference_checkpoint_id
self._chosen: npt.NDArray[np.float64] | None = None
self._rejected: npt.NDArray[np.float64] | None = None
def populate(
self, chosen_logprobs: npt.NDArray[np.float64], rejected_logprobs: npt.NDArray[np.float64]
) -> None:
if chosen_logprobs.shape != rejected_logprobs.shape:
raise ValueError("chosen and rejected caches must cover the same pairs")
self._chosen = np.ascontiguousarray(chosen_logprobs)
self._rejected = np.ascontiguousarray(rejected_logprobs)
def lookup(self, indices: IdArray) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
if self._chosen is None or self._rejected is None:
raise RuntimeError(
"reference cache empty; run the offline pass before training or the "
"reference model has to stay resident for the whole run"
)
return self._chosen[indices], self._rejected[indices]
class PreferenceScratch:
"""Every buffer a step needs, allocated once at peak shape.
A training loop that allocates per step spends a measurable share of
a cheap stage inside the allocator, and both shapes are known from
the batch configuration before the first step runs.
"""
def __init__(self, shape: BatchShape) -> None:
self.shape = shape
rows = 2 * shape.max_pairs # chosen and rejected in one batch
self.token_ids = np.zeros((rows, shape.max_tokens), dtype=np.int32)
self.mask = np.zeros((rows, shape.max_tokens), dtype=np.float32)
self.gathered = np.zeros((rows, shape.max_tokens), dtype=np.float32)
self.sequence_logprob = np.zeros(rows, dtype=np.float64)
def pack(self, sequences: list[IdArray]) -> int:
"""Length-bucketed padding into the preallocated block.
Padding to the longest sequence in the batch wastes work in
proportion to the length spread, which is why callers bucket by
length before packing. Unbucketed batches are where this layout
quietly gives back most of what it gained.
"""
rows = len(sequences)
if rows > self.token_ids.shape[0]:
raise ValueError(f"batch of {rows} exceeds scratch rows {self.token_ids.shape[0]}")
longest = max(len(sequence) for sequence in sequences)
if longest > self.shape.max_tokens:
raise ValueError(f"sequence of {longest} tokens exceeds {self.shape.max_tokens}")
self.token_ids[:rows, :longest] = 0
self.mask[:rows, :longest] = 0.0
for row, sequence in enumerate(sequences):
length = len(sequence)
self.token_ids[row, :length] = sequence
self.mask[row, :length] = 1.0
return longest
def sequence_log_probs(
token_log_probs: FloatArray, token_ids: IdArray, mask: FloatArray, out: npt.NDArray[np.float64]
) -> npt.NDArray[np.float64]:
"""Per-sequence log-probability, without a loop over sequences.
token_log_probs is (rows, time, vocab) log-softmax output. One
take_along_axis pulls the realized token's log-probability at every
position, the mask zeroes the padding, and the sum runs along time.
The previous stage did this one sequence at a time in Python, which
dominated its step cost entirely.
"""
gathered = np.take_along_axis(token_log_probs, token_ids[:, :, None], axis=2)[..., 0]
np.multiply(gathered, mask, out=gathered)
# float64 accumulation: the sum runs over thousands of tokens and
# float32 error here is large enough to move a margin.
np.sum(gathered, axis=1, dtype=np.float64, out=out)
return out
def fused_dpo_loss_and_margin(
policy_chosen: npt.NDArray[np.float64],
policy_rejected: npt.NDArray[np.float64],
reference_chosen: npt.NDArray[np.float64],
reference_rejected: npt.NDArray[np.float64],
beta: float,
) -> tuple[float, npt.NDArray[np.float64], dict[str, float]]:
"""Implicit rewards, margin and loss in one pass.
The intermediate implicit rewards are never materialized separately:
the difference of differences collapses algebraically, which is both
fewer temporaries and fewer chances for a sign to go astray.
margin = beta * [(pi_c - ref_c) - (pi_r - ref_r)]
logaddexp rather than log(1 + exp(-margin)) because the naive form
overflows on exactly the confidently-wrong pairs that carry the
gradient early in a run.
"""
margin = beta * ((policy_chosen - reference_chosen) - (policy_rejected - reference_rejected))
loss = float(np.mean(np.logaddexp(0.0, -margin)))
# The diagnostics that matter, computed from arrays already in hand
# rather than from a second pass: the objective constrains only the
# difference, so both log-probabilities falling together satisfies
# it while draining mass to unlabelled sequences.
diagnostics = {
"mean_chosen_logprob": float(policy_chosen.mean()),
"mean_rejected_logprob": float(policy_rejected.mean()),
"mean_margin": float(margin.mean()),
"accuracy": float((margin > 0.0).mean()),
}
return loss, margin, diagnostics
def sequence_kl_estimate(
policy_logprob: npt.NDArray[np.float64], reference_logprob: npt.NDArray[np.float64]
) -> float:
"""A low-variance KL estimate from the log-ratios already computed.
k3 = exp(-d) - 1 + d, d = log pi_ref - log pi_theta
Always non-negative, unbiased, and far less noisy than the naive
mean of -d, which is negative roughly half the time on small
batches and makes the one plot that matters unreadable.
It costs nothing here because both terms are already in hand -- and
KL from the reference is the axis every failure in this technique
moves along, so it should be on every step.
"""
log_ratio = reference_logprob - policy_logprob
return float(np.mean(np.expm1(log_ratio) - log_ratio))
def bucket_by_length(lengths: npt.NDArray[np.int32], batch_size: int) -> list[IdArray]:
"""Group pairs of similar length before packing.
Padding cost is the spread within a batch, not the mean length, so
sorting first turns a mostly-padding block into a mostly-tokens one.
The cost is that batch composition is no longer random, which
correlates gradient noise across steps in a way a shuffled loader
avoids -- shuffle the buckets, not the rows inside them.
"""
order = np.argsort(lengths, kind="stable").astype(np.int32)
return [order[start : start + batch_size] for start in range(0, len(order), batch_size)]
def overoptimization_sweep(
proxy_reward: FloatArray, true_reward: FloatArray, kl: FloatArray
) -> dict[str, float]:
"""Where the curve turned, from logged checkpoints.
Cheap, and the single most useful summary of a preference run: the
proxy is monotone in KL while the truth is not, so the argmax of the
true reward names the checkpoint to ship and the final step names
the one that would have been shipped by default.
"""
if not (proxy_reward.shape == true_reward.shape == kl.shape):
raise ValueError("reward and KL series must cover the same checkpoints")
if proxy_reward.size == 0:
raise ValueError("no checkpoints logged")
peak = int(np.argmax(true_reward))
return {
"peak_checkpoint": peak,
"peak_kl": float(kl[peak]),
"peak_true_reward": float(true_reward[peak]),
"final_true_reward": float(true_reward[-1]),
"final_proxy_reward": float(proxy_reward[-1]),
"true_reward_lost_by_training_on": float(true_reward[peak] - true_reward[-1]),
}
What changed: The largest win in this technique is structural rather than numerical, and it comes from noticing that the reference policy is frozen: its log-probabilities depend on nothing being trained, so they can be computed once over the whole dataset and cached, which removes an entire model from the training loop and halves both the resident memory and the forward work per step. Everything else follows from shaping the remaining work as arrays. Chosen and rejected responses are packed into a single padded batch and scored in one forward pass rather than two, since they differ only in which rows they occupy; the per-token log-probability gather becomes one take_along_axis over a padded (batch, time) block with a mask, eliminating the Python loop over sequences that dominated the previous stage; and the sequence sums, the implicit-reward difference and the logistic loss are fused into one pass over the token block so the per-token intermediates are never materialized. Token ids are int32 and every float buffer is float32 and C-contiguous, which is what keeps the gather and the reductions on a single kernel rather than silently upcasting mid-expression. Scratch is allocated once at the maximum batch-by-sequence shape and sliced, so a training loop performs no per-step allocation — and because the reference cache is keyed by reference checkpoint, a changed anchor invalidates it rather than silently reusing the wrong numbers.
Eliminate Python-level loops over samples
The per-sequence log-probability gather becomes one take_along_axis over a padded token block with a validity mask
Trade-off: Padding to the longest sequence in the batch wastes work proportional to the length spread, so batches must be length-bucketed or the saving erodes
Batch work to amortize interpreter overhead
Chosen and rejected responses are scored in one forward pass rather than two, since they differ only in which rows of the batch they occupy
Trade-off: Peak activation memory doubles per step, which is usually what limits batch size in preference training in the first place
Fuse adjacent operations so intermediates are never materialized
Sequence sums, the implicit-reward difference and the logistic loss collapse into one pass, so the per-token arrays never land
Trade-off: The fused form hides the per-token quantities, and per-token log-probabilities are exactly what you want when diagnosing a run that has gone wrong
Pre-allocate output arrays and use in-place operations
The token block, mask and gather scratch are sized once at the maximum batch-by-sequence shape, so a training step allocates nothing
Trade-off: Buffers stay at peak shape for the process lifetime and are not reentrant, so a worker cannot share them across threads
Ensure contiguous memory layout and a single dtype
int32 ids and float32 log-probabilities in C-contiguous blocks keep the gather and the reductions on one kernel instead of upcasting mid-expression
Trade-off: float32 accumulation over long sequences loses precision in the sequence sum, so the running total is the one place that must stay float64
Per step: one forward pass over 2·pairs padded sequences instead of two passes of pairs each, one take_along_axis over a (rows, time) block, and O(rows) reductions — with the reference model absent from the loop entirely once its log-probabilities are cached. Illustrative, not a measured benchmark: the useful shape is that the cache removes a whole model rather than a constant factor, and that padding waste scales with the length spread inside a batch, which is why bucketing matters more here than any change to the reductions.