Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

DPO Variants › Kahneman-Tversky Optimization (KTO)

Michael BrenndoerferJanuary 3, 202613 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 285DPO Variants › Kahneman-Tversky Optimization (KTO)

Reading section 3 of 12 · Kahneman-Tversky Optimization (KTO) · 13 minSection 3 of 12 · 13 min

Section 3 of 12 · 13 min read

DPO Variants › Kahneman-Tversky Optimization (KTO)One section · Sections in this chapter3 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

While DPO and IPO improve upon RLHF's computational complexity, they share a basic data requirement: paired preferences. Each training example must contain a prompt together with both a chosen and rejected response. This pairing constraint creates practical challenges that are easy to underestimate if you have only worked with carefully curated benchmark datasets.

Think of the difference between a controlled taste test and a product review. In a taste test, each participant explicitly compares two options side by side and says which they prefer. Product reviews, by contrast, arrive as standalone ratings: a customer gives three stars or five stars to a single item, with no explicit comparison to an alternative. Most of the feedback you collect from deployed language models looks much more like product reviews than taste tests, and forcing it into paired-comparison format requires either discarding most of your data or constructing artificial comparisons that may not reflect real user preferences.

KTO, introduced by Ethayarajh et al. (2024), eliminates the paired data requirement by designing an objective that works directly with unpaired binary feedback. The result is a method that can absorb the full breadth of available user signal without the preprocessing overhead of converting it into an unnatural format.

The Unpaired Feedback ProblemLink Copied

Real-world human feedback often comes in unpaired form, and this mismatch between how feedback is collected and how preference optimization algorithms expect data creates significant friction in practical alignment pipelines. Consider the forms that user feedback naturally takes:

Binary ratings: users click thumbs up or thumbs down on individual responses, with no stated alternative.
Flagging systems: users report problematic outputs without giving a corrected version.
Implicit signals: engagement metrics, session lengths, and follow-up question rates indicate whether a response was helpful, but never identify a counterfactual.

Converting this abundant unpaired feedback into DPO's paired format requires either discarding data or artificially constructing pairs. Discarding wastes useful signal. Constructing pairs introduces artifacts: if you pair a thumbs-up response to a randomly selected thumbs-down response, the pair may not reflect a real comparison because the two responses came from different conversations with different prompts and different contexts.

Consider a scenario where you have 10,000 thumbs-up ratings and 5,000 thumbs-down ratings, but these ratings come from different conversations. To use DPO, you would need to either match these into pairs somehow, losing most of your data, or generate new responses to create artificial comparisons. KTO avoids this entirely by treating each labeled response as an independent training example.

Inspiration from Prospect TheoryLink Copied

KTO draws inspiration from Kahneman and Tversky's prospect theory, which describes how humans make decisions under uncertainty. This connection to behavioral economics is more than a naming convention: it gives principled guidance for how to weight different types of feedback. Two key insights from behavioral economics directly inform KTO's design.

Reference dependence is the first insight. People evaluate outcomes relative to a reference point, not in absolute terms. A gain of $100 feels different depending on whether you expected $0 or $200. For alignment, this suggests that the quality of a model response should be measured relative to some baseline expectation, not in absolute terms. KTO implements this by measuring implicit rewards relative to the expected KL divergence across the training distribution, creating a dynamic baseline that adapts as training progresses.

Loss aversion is the second insight. Losses loom larger than equivalent gains. Losing $100 feels psychologically worse than gaining $100 feels good, and behavioral experiments consistently estimate that losses are weighted roughly 2x more heavily than gains in human value assessments. For alignment, this suggests that suppressing bad outputs may be more important than promoting good ones. You might forgive a bland response, but you remember a harmful or embarrassingly incorrect one. KTO incorporates this asymmetry by letting different loss weights for desirable and undesirable examples.

KTO incorporates both principles into its loss function, treating "desirable" and "undesirable" responses asymmetrically in a way that mirrors how humans experience quality differences.

Out[5]:

Visualization

Line plot of the prospect theory value function compared to a symmetric reference, with shaded red and green regions showing the loss and gain domains respectively. — Prospect theory's value function compared to a symmetric baseline. Losses (negative x-axis) produce steeper decreases in perceived value than equivalent gains produce increases, which shows the empirical phenomenon of loss aversion. KTO applies this asymmetry to alignment training.

The KTO Loss FunctionLink Copied

The construction of KTO's loss function proceeds in several stages, each motivated by the behavioral economics principles described above. We begin by defining the implicit reward, which is the raw signal that KTO turns into a learning objective.

For a response $y$ to prompt $x$ with binary label $z \in \{\text{desirable}, \text{undesirable}\}$ , KTO defines:

r_\theta(x, y) = \log \frac{\pi_\theta(y|x)}{\pi_{\text{ref}}(y|x)}

where:

$r_\theta(x, y)$ : the implicit reward assigned to response $y$ given prompt $x$
$\pi_\theta(y|x)$ : the probability of the response under the current policy
$\pi_{\text{ref}}(y|x)$ : the probability of the response under the reference model

This is the same implicit reward used in DPO. It measures how much more likely the current policy makes this response compared to the reference. A positive implicit reward means the policy has learned to favor this response; a negative implicit reward means the policy has learned to disfavor it. The important property is that this quantity is well-defined for a single response, without needing a comparison partner.

KTO then defines a reference point that implements the reference dependence principle from prospect theory:

z_0 = \mathbb{E}_{x' \sim \mathcal{D}} \left[ \text{KL}(\pi_\theta(\cdot|x') \| \pi_{\text{ref}}(\cdot|x')) \right]

where:

$z_0$ : the reference point (baseline) for evaluation
$\pi_\theta, \pi_{\text{ref}}$ : the policy and reference models
$\mathcal{D}$ : the training dataset distribution
$\mathbb{E}_{x' \sim \mathcal{D}}$ : the expectation over prompts $x'$ sampled from $\mathcal{D}$
$\text{KL}$ : the Kullback-Leibler divergence measuring the drift of the policy from the reference

The reference point $z_0$ is the expected KL divergence between policy and reference across the training distribution. In practice, this is estimated from a running average during training. The reference point defines what counts as "above average" versus "below average" performance. Rather than using an arbitrary fixed threshold, KTO adapts the reference point to the current state of training, creating a dynamic baseline that evolves as the model improves.

With the implicit reward and reference point defined, KTO constructs a value function that differs based on whether the response is desirable:

v(x, y) = \begin{cases} \sigma(\beta \cdot r_\theta(x, y) - z_0) & \text{if } z = \text{desirable} \\ \sigma(z_0 - \beta \cdot r_\theta(x, y)) & \text{if } z = \text{undesirable} \end{cases}

where:

$v(x, y)$ : the value assigned to the response, bounded between 0 and 1
$\sigma$ : the sigmoid function
$z$ : the label showing if the response is desirable or undesirable
$\beta \cdot r_\theta(x, y) - z_0$ : the shifted implicit reward used for desirable examples

This asymmetric definition is the mathematical implementation of reference dependence. For desirable responses, we ask whether the implicit reward exceeds the reference point. For undesirable responses, we ask whether the implicit reward falls below the reference point. The sigmoid function squashes these comparisons into the range (0, 1), creating a smooth value measure that is easy to differentiate.

The loss weights desirable and undesirable examples differently, implementing the loss aversion principle:

\mathcal{L}_{\text{KTO}} = \mathbb{E}_{(x, y, z)} \left[ \lambda_z \cdot (1 - v(x, y)) \right]

where:

$\mathcal{L}_{\text{KTO}}$ : the KTO loss function
$\mathbb{E}_{(x, y, z)}$ : the expectation over the dataset of labeled examples
$\lambda_z$ : the weighting factor specific to the label type
$v(x, y)$ : the value computed by the value function

The weighting factors are $\lambda_{\text{desirable}} = \lambda_d$ and $\lambda_{\text{undesirable}} = \lambda_u$ , typically with $\lambda_u > \lambda_d$ to implement loss aversion. By setting $\lambda_u > \lambda_d$ , we tell the model that failing to suppress a bad response is worse than failing to promote a good response. This asymmetry shows the empirical finding that people are more bothered by failures than impressed by equivalent successes.

Understanding the Value FunctionLink Copied

The asymmetric value function encodes prospect theory's core insights, and understanding its behavior illuminates why KTO works. The function turns implicit rewards into values differently depending on the label. This creates distinct learning signals for positive and negative feedback.

For desirable responses, the sigmoid argument is $\beta \cdot r_\theta - z_0$ . The model is rewarded with a low loss when the implicit reward exceeds the reference point. Intuitively, the model should push desirable responses to have higher implicit rewards than average. When $\beta \cdot r_\theta > z_0$ , the sigmoid output is greater than 0.5, meaning the value is high and the loss $(1 - v)$ is low.

For undesirable responses, the sigmoid argument is $z_0 - \beta \cdot r_\theta$ . The model is rewarded when the implicit reward falls below the reference point. When $\beta \cdot r_\theta < z_0$ , the sigmoid output is greater than 0.5, meaning the value is high and the loss is low. The model learns to associate undesirable responses with below-average implicit rewards.

The reference point $z_0$ is the dividing line between "gains" and "losses" in the prospect theory sense. This relative framing means KTO does not need paired comparisons; it learns to push good responses up and bad responses down relative to an adaptive baseline. The elegance of this approach is that it converts an inherently comparative problem into a classification problem, which is exactly the form that unpaired feedback gives.

Out[6]:

Visualization

Line plot showing sigmoid-shaped value functions for desirable and undesirable responses, with a vertical dashed line marking the reference point that separates the two response classes. — KTO value functions for desirable and undesirable responses. The reference point at the vertical dashed line divides the implicit reward space: desirable responses gain value above it, while undesirable responses gain value below it.

Line plot showing KTO loss curves for desirable and undesirable responses as functions of implicit reward, illustrating the opposing slopes around the reference point. — KTO loss functions for desirable and undesirable responses. Each curve is minimized in the region where the model correctly orders the implicit reward relative to the reference point.

Practical Implementation DetailsLink Copied

KTO requires tracking the reference point $z_0$ during training. This running average must be estimated from the current batch of data and updated incrementally as training proceeds.

In[7]:

Code

import torch


def kto_loss(
    policy_logps: torch.Tensor,  # log π_θ(y|x) for batch
    reference_logps: torch.Tensor,  # log π_ref(y|x) for batch
    is_desirable: torch.Tensor,  # binary mask: 1 for desirable, 0 for undesirable
    kl_reference: float,  # z_0: running average KL
    beta: float = 0.1,
    lambda_d: float = 1.0,  # weight for desirable
    lambda_u: float = 1.0,  # weight for undesirable (often > lambda_d)
) -> torch.Tensor:
    """
    Compute KTO loss for a batch of (potentially unpaired) examples.
    """
    # Implicit reward: log ratio
    implicit_reward = policy_logps - reference_logps

    # Value function differs by desirability
    desirable_mask = is_desirable.bool()

    values = torch.zeros_like(policy_logps)
    values[desirable_mask] = torch.sigmoid(
        beta * implicit_reward[desirable_mask] - kl_reference
    )
    values[~desirable_mask] = torch.sigmoid(
        kl_reference - beta * implicit_reward[~desirable_mask]
    )

    # Weighted loss
    weights = torch.where(desirable_mask, lambda_d, lambda_u)
    loss = weights * (1 - values)

    return loss.mean()

The running average KL is typically computed as:

In[8]:

Code

def update_kl_reference(
    kl_reference: float, batch_kl: float, momentum: float = 0.99
):
    """Update running average of KL divergence."""
    return momentum * kl_reference + (1 - momentum) * batch_kl

The momentum parameter controls how quickly the reference point adapts to changes in the policy. A high momentum (close to 1) means the reference point changes slowly. This gives a stable baseline. A low momentum means it tracks the current policy more aggressively, which can be destabilizing if the policy is changing rapidly at the start of training.

KTO's AdvantagesLink Copied

KTO offers several practical benefits for real-world alignment pipelines:

Data efficiency: Uses all available binary feedback without discarding unpaired examples.
Natural data format: Matches how most user feedback is collected in deployed products.
Principled asymmetry: The loss-aversion weighting shows empirical findings about human judgment rather than an arbitrary design choice.
Stable training: The reference point gives a grounding mechanism similar to IPO's target, preventing unbounded optimization.