Kahneman-Tversky Optimization

Setup

For a prompt $x$, completion $y$, trainable policy $\pi_\theta$, and reference policy $\pi_{\mathrm{ref}}$, define the KTO implied reward:

$$\begin{aligned} r_\theta(x,y) &= \log\frac{\pi_\theta(y\mid x)} {\pi_{\mathrm{ref}}(y\mid x)},\\ z_0 &= \mathrm{KL}(\pi_\theta(\cdot\mid x)\|\pi_{\mathrm{ref}}(\cdot\mid x)),\\ \delta &= r_\theta(x,y)-z_0,\\ v_D &= \lambda_D\,\sigma(\beta\delta),\\ v_U &= \lambda_U\,\sigma(-\beta\delta),\\ \mathcal L_y &= \lambda_y-v_y. \end{aligned}$$

With $z_0$ treated as fixed for the update and $s=\sigma(\beta\delta)$, the label-dependent derivative through $r_\theta$ is

$$\begin{aligned} \frac{\partial \mathcal L_D}{\partial r_\theta} &= -\lambda_D\beta s(1-s),\\ \frac{\partial \mathcal L_U}{\partial r_\theta} &= \lambda_U\beta s(1-s). \end{aligned}$$

For an autoregressive language model, both log-probabilities are sequence log-probabilities: sums over the completion tokens.

Each training example has a binary label $\ell\in\{D,U\}$, where $D$ means desirable and $U$ means undesirable.

The KL reference point

KTO does not compare $y$ to one rejected completion. It compares $y$ to the policy's average reference-relative reward:

$$z_0 = \mathrm{KL}(\pi_\theta(\cdot\mid x)\|\pi_{\mathrm{ref}}(\cdot\mid x)).$$

Equivalently, this is the current policy's expected implied reward at prompt $x$:

$$z_0 = \mathbb E_{y'\sim\pi_\theta(\cdot\mid x)} [r_\theta(x,y')].$$

The important scalar is

$$\delta = r_\theta(x,y)-z_0.$$

A desirable output should have positive $\delta$. An undesirable output should have negative $\delta$.

In practical KTO training, $z_0$ is usually estimated from mismatched outputs in the same microbatch, clamped to be nonnegative, and treated as a stop-gradient quantity. It controls loss saturation; it is not itself the thing KTO tries to optimize through.

For a microbatch of size $m\ge 2$, define $\rho_i$ as the policy/reference log-ratio of a shifted output $y_{j(i)}$ under prompt $x_i$:

$$\rho_i = \log \frac{\pi_\theta(y_{j(i)}\mid x_i)} {\pi_{\mathrm{ref}}(y_{j(i)}\mid x_i)}.$$

A toy version of the paper's reference estimate is then

$$\hat z_0=\max(0,\operatorname{mean}_i \rho_i).$$

The point is not to reuse the labeled pair $(x_i,y_i)$, because those completions were deliberately selected as good or bad and can have unrepresentative rewards.

Label-dependent value and loss

KTO uses a logistic value function. For a desirable example,

$$v_D = \lambda_D\,\sigma(\beta\delta).$$

For an undesirable example,

$$v_U = \lambda_U\,\sigma(-\beta\delta).$$

The default KTO loss subtracts this value from the label's class weight:

$$\mathcal L_y = \lambda_y - v_y.$$

Over the dataset, this is:

$$\mathcal L_{\mathrm{KTO}} = \mathbb E_{(x,y,\ell)\sim D} [\lambda_\ell-v_\ell(x,y)].$$

So the two label cases are:

$$\mathcal L_D = \lambda_D(1-\sigma(\beta\delta)),$$

and

$$\mathcal L_U = \lambda_U(1-\sigma(-\beta\delta)).$$

Gradient descent pushes desirable examples upward in reference-relative reward and undesirable examples downward. If $z_0$ is treated as fixed for the update, let $s=\sigma(\beta\delta)$. Then

$$\frac{\partial \mathcal L_D}{\partial r_\theta} = -\lambda_D\beta s(1-s),$$

while

$$\frac{\partial \mathcal L_U}{\partial r_\theta} = \lambda_U\beta s(1-s).$$

Both gradients have their largest magnitude near $\delta=0$, then saturate.

Here $\beta$ is KTO's value-function saturation parameter. It has a practical anchoring effect similar to DPO's $\beta$, but in KTO it is introduced directly to control how quickly utility saturates.

The weights $\lambda_D$ and $\lambda_U$ are class/utility weights, not a universal moral setting. They are often chosen to account for the ratio of desirable to undesirable examples. Increasing $\lambda_U$ emphasizes undesirable examples; increasing $\lambda_D$ emphasizes desirable examples.

How this differs from DPO

DPO uses a pairwise margin:

$$r_\theta(x,y_w)-r_\theta(x,y_\ell).$$

KTO uses a pointwise margin:

$$r_\theta(x,y)-z_0.$$

That is the bridge. DPO asks whether the winner beats the loser by enough. KTO asks whether a single labeled output sits on the correct side of a KL-derived baseline.

Limits of the mechanism

KTO is natural when feedback is already binary or when desirable and undesirable examples are imbalanced. If preference data is clean, consistent, and pairwise, DPO may be the better fit. If labels are noisy or intransitive, KTO's saturation can be useful because extreme examples receive small updates.

When KTO is trained from preference pairs, a common conversion is to treat the preferred response as desirable and the rejected response as undesirable. That is a simplifying assumption, not a law of the method. Naturally binary feedback is the cleaner setting for KTO.

The same property can fail. A hard desirable example with very negative implied reward may be ignored instead of learned. Lower $\beta$ and more epochs can reduce this underfitting risk, but they do not remove it.

A natural next step is Process Reward Models: instead of labeling a whole completion as desirable or undesirable, ask what changes when feedback attaches to intermediate reasoning steps.