KTO Math

From paired preferences to a single bit

Preference tuning grew up on pairwise data. RLHF trains a reward model on pairs (a winner y_w and a loser y_l for the same prompt x); DPO collapses that into one loss but still consumes the same pair. The pair is the expensive part — getting two responses to one prompt and asking which is better is slow, and much real-world feedback is not shaped that way: a user upvotes one reply, flags another, abandons a third, each an isolated verdict on a single output.

KTO throws the pairing requirement away. Its unit of data is one (x, y) plus one bit: y is desirable or undesirable given x. You can have a thousand desirable examples and fifty undesirable ones, from entirely different prompts, and KTO still trains. This is the data-flexibility win — the format matches how feedback is actually logged (accepts, rejects, ratings binarized around a threshold) rather than forcing it into curated A-vs-B pairs. The question is how to build a sensible gradient from an unpaired signal, and the answer comes from behavioural economics.

The Kahneman-Tversky value function

Prospect theory’s central claim is that humans judge outcomes not on absolute wealth but on gains and losses relative to a reference point. Its value function v(z) passes through that point at z = 0, is concave for gains (z > 0 — the second won dollar thrills less than the first) and convex for losses (z < 0), and is steeper on the loss side: a loss hurts more than an equal gain pleases — loss aversion.

KTO’s insight is that a language model has its own notion of gain and loss: how much more (or less) probability it puts on an output than a frozen reference model. Turn that into the value-function input z, pass it through an S-shaped curve, and you get a per-example objective that makes desirable outputs feel like gains and undesirable ones like losses — without ever comparing two outputs. The logistic sigmoid σ(z) = 1 / (1 + e^(-z)) supplies the S-curve: concave for z > 0, convex for z < 0, exactly the gains/losses geometry prospect theory describes.

The implied reward

KTO, like DPO, never trains a separate reward model. It reads the reward straight off the policy through the KL-regularized-RL solution: the optimal policy’s log-ratio against the reference is the reward up to scale. Define the implied reward of output y for prompt x as

r_θ(x, y) = β · log( π_θ(y|x) / π_ref(y|x) )

Here π_θ is the model being trained, π_ref is the frozen starting checkpoint (usually the SFT model), and β > 0 controls how hard the KL leash pulls the policy back toward the reference. When the policy raises an output’s probability above the reference, r_θ > 0; when it lowers it, r_θ < 0. This scalar is all KTO needs from the network, computed from ordinary log-probabilities in one forward pass through each model. The next step decides what counts as ‘above’ — the reference point.

The reference point: the z_0 baseline

A reward is meaningless in isolation; prospect theory says value is measured against a reference. KTO’s reference point is z_0, an estimate of the typical KL divergence between the policy and the reference over the current batch:

z_0 = KL( π_θ(y’|x) || π_ref(y’|x) )
     ≈ max( 0,  (1/m) Σ_i  r_θ(x_i, y’_i) )

In practice it is estimated cheaply: shuffle outputs against prompts within the batch so each y’ is mismatched, average the implied rewards of those pairs, and clamp at zero. This z_0 is detached from the gradient — it shifts the origin of the value function but is not itself optimized, which keeps training stable. Read it as a floating baseline for how far the policy has drifted from the reference on average. An output whose implied reward sits above z_0 is a gain; one below it is a loss. That single comparison, r_θ − z_0, is the argument the value function shapes.

The KTO value function and loss

Now assemble the value. For a desirable output we want its reward pushed above the baseline (a gain); for an undesirable one, below (a loss). The sigmoid turns the signed distance from z_0 into a bounded value in (0, 1):

v_KTO(x, y) = σ( r_θ(x, y) − z_0 )   if y is desirable
v_KTO(x, y) = σ( z_0 − r_θ(x, y) )   if y is undesirable

In both cases v is near 1 when the model does the right thing (desirable reward high, undesirable reward low) and near 0 otherwise. The loss is the shortfall from a perfect value, weighted per class:

L_KTO = E_(x,y ~ D) [ w(y) · ( 1 − v_KTO(x, y) ) ]

where  w(y) = λ_D  if y desirable
       w(y) = λ_U  if y undesirable

Minimizing 1 − v drives v → 1: for desirable examples the gradient raises log π_θ(y|x), for undesirable ones it lowers it — each example moved on its own, no pair in sight.

Gains, losses, and loss aversion via λ_D, λ_U

The two weights λ_D (desirable) and λ_U (undesirable) are where KTO encodes prospect theory’s asymmetry — and where it earns its robustness to imbalance. The sigmoid curve itself is symmetric; loss aversion does not come from the curve, it comes from making λ_U weigh undesirable examples more heavily than λ_D weighs desirable ones, so a ‘loss’ moves the loss function more than an equal ‘gain.’

These same knobs absorb class imbalance. With n_D desirable and n_U undesirable examples, if n_U dwarfs n_D the raw loss is swamped by undesirable examples and the model just learns to suppress everything. KTO counters this by choosing weights so the two classes contribute comparably — the authors recommend keeping

(λ_D · n_D) / (λ_U · n_U)  ∈  [ 1 , 4/3 ]

So with ten times as many undesirable as desirable examples, you raise λ_D (or lower λ_U) to rebalance. DPO has no equivalent lever: its pairs are, by construction, one winner and one loser each, so it cannot express ‘I have mostly negative feedback.’

A worked numeric example

Take β = 0.1 and a batch-estimated reference point z_0 = 0.1. Consider one desirable and one undesirable example.

Desirable y+:   log π_θ = -18 , log π_ref = -22
  log-ratio = 4  ->  r_θ = β · 4 = 0.4
  v = σ( r_θ - z_0 ) = σ(0.3) = 0.574
  loss (λ_D = 1)   = 1 - 0.574 = 0.426

Undesirable y-:  log π_θ = -20 , log π_ref = -21
  log-ratio = 1  ->  r_θ = 0.1
  v = σ( z_0 - r_θ ) = σ(0) = 0.5
  loss (λ_U = 4/3) = 1.333 · 0.5 = 0.667

The desirable example sits above the baseline (r_θ = 0.4 > z_0 = 0.1), so its value 0.574 exceeds a half and its loss is modest; the gradient still nudges log π_θ(y+|x) higher to push v toward 1. The undesirable example is exactly at the baseline (r_θ = z_0): the model still gives it slightly more probability than the reference, which is wrong, so v = 0.5 and the heavier λ_U makes its loss the larger. Minimizing it lowers log π_θ(y-|x), dragging r_θ below z_0 so σ(z_0 − r_θ) → 1.

KTO vs DPO, precisely

Both methods share the implied-reward machinery (r_θ = β log(π_θ / π_ref)) and avoid a separate reward model. The divergence is entirely in the data and the objective’s reference.

DPO’s loss −log σ( r_θ(x,y_w) − r_θ(x,y_l) ) contrasts two rewards for the same prompt. KTO contrasts one reward against a population baseline z_0. That is the whole conceptual shift: DPO asks ‘is this better than that?’, KTO asks ‘is this better than average?’ — and the second question can be answered from an unpaired thumbs-up.

Practical notes and pitfalls

KTO’s data flexibility is its headline advantage: you can train on logged production feedback (accepts, rejects, ratings binarized around a threshold) with no pairing step, and on skewed label distributions that would break a pairwise pipeline. On a CPU-bound small-model workflow this matters twice over — each step still costs the same two forward passes as DPO (π_θ and the frozen π_ref), but you skip generating and annotating matched pairs and start from whatever feedback you already have.

The pitfalls are specific. First, get the λ_D/λ_U ratio right for your class balance — leaving them at 1:1 on skewed data is the most common way KTO underperforms, collapsing toward blanket suppression or approval. Second, β is the same KL-leash trade-off as in DPO: too small and the model drifts; too large and it barely learns. Third, the z_0 estimate is only as stable as the batch is large, so the shuffled-pair estimate wants a reasonable batch size. Get those three right and KTO delivers DPO-competitive alignment from cheaper, more abundant data.