KTO (Kahneman-Tversky Optimization), explained

Kahneman-Tversky Optimization (KTO) aligns a language model from unpaired binary feedback: each training example is a prompt, one response, and a label saying whether that response is desirable or undesirable. There is no need for two responses to the same prompt. The loss borrows its shape from prospect theory, the model of human decision-making by Daniel Kahneman and Amos Tversky: each response's implied reward is compared against a reference point, gains and losses are scored by a saturating value function, and losses can be weighted differently from gains. KTO sits where DPO sits in the post-training stack, after or sometimes instead of SFT, and is the method of choice when your feedback is thumbs up or thumbs down rather than side-by-side comparisons.

Why KTO exists#

Pairwise preference data is expensive. Someone has to see two responses to the same prompt and pick one. Production systems, on the other hand, generate a steady stream of single-response signals: a thumbs-down button, a user regenerating an answer, a support ticket marked resolved, a unit test passing or failing. These are abundant and cheap, but they do not fit DPO's input format.

Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky and Douwe Kiela posted KTO: Model Alignment as Prospect Theoretic Optimization in February 2024 (ICML 2024). The paper makes two arguments. First, successful alignment losses such as DPO and PPO-Clip already encode human-like biases; the authors formalize these as human-aware losses (HALOs) and prove DPO and PPO-Clip belong to the class. Second, if those biases are why they work, you can design a loss around them directly, and the result needs only a binary signal. They report that KTO matches or exceeds DPO from 1B to 30B parameters (Pythia 1.4B to 12B and Llama 7B to 30B), still matches DPO after discarding up to 90% of the desirable examples, and, for sufficiently strong pretrained Llama models (13B and 30B), can skip SFT without losing quality.

The prospect theory behind it#

Prospect theory (Kahneman and Tversky, 1979; Tversky and Kahneman, 1992) says people evaluate outcomes relative to a reference point rather than in absolute terms, are risk-averse for gains and risk-seeking for losses, and feel losses more than equal gains. Tversky and Kahneman's value function for an outcome $z$ relative to reference point $z_0$ is

with median estimates $\alpha =0.88$ (diminishing sensitivity) and $\lambda =2.25$ (loss aversion) across individuals. KTO keeps the three ideas (a reference point, concavity in gains and convexity in losses, and asymmetric weights) but swaps the power function for a logistic sigmoid, which is bounded and easier to optimize.

The KTO objective#

For a prompt $x$ and response $y$, define the implied reward exactly as in DPO, the log-ratio between the policy and a frozen reference:

The reference point is the KL divergence between policy and reference on the prompt distribution,

the average implied reward of what the policy would generate. The value of an example is

and the loss is

where ${\lambda }_y$ is ${\lambda }_D$ or ${\lambda }_U$ depending on the label (it only makes the loss non-negative), $\beta$ controls how sharply value saturates, and ${\lambda }_D,{\lambda }_U$ weight desirable and undesirable examples.

Why it works

A desirable example earns value by having its implied reward exceed the reference point; an undesirable example earns value by falling below it. Because $z_0$ is the policy's own average log-ratio, the model cannot win by raising every response's probability bluntly: doing so raises $z_0$ along with the rewards and the margin stays flat. The paper puts it this way: the model is forced to learn exactly what makes an output desirable, so that its reward can rise while the KL term stays flat. The reference point is not backpropagated through; it only sets where the sigmoid saturates.

The gradient and saturation

For a desirable example with $z=r_{\theta }-z_0$,

and symmetrically for undesirable examples with the sign flipped. The factor $\sigma (\beta z)\sigma (-\beta z)$ peaks at $z=0$ with value $1/4$ and vanishes in both directions. That is the paper's Proposition 4.1: as an example's implied reward goes to $\pm \infty$, its gradient contribution goes to zero. Examples the model finds too easy are ignored, and so are examples it finds impossibly hard, such as a "desirable" response the model considers extremely unlikely. The authors argue this is partly a blessing: real feedback is noisy, and a desirable example with a hugely negative implied reward may simply be mislabeled. The cost is potential underfitting on hard but genuine examples, which they suggest countering with a lower $\beta$ and more epochs.

KTO next to DPO

DPO's per-pair loss is $-\log \sigma (\beta [r_{\theta }(x,y_w)-r_{\theta }(x,y_l)])$: the reference point for the chosen response is the rejected response's reward, so the two must come from the same prompt. KTO replaces "the other response" with "the policy's average response", which is what frees it from pairs. Three other differences matter in practice. DPO's loss is concave in the margin everywhere (risk-seeking in the paper's terms), while KTO is concave for gains and convex for losses. KTO can weight the two classes independently. And their behavior on contradictory data differs: the paper's Theorem 4.3 shows that with contradictory preferences between two outputs, in proportion $p$ to $1-p$, the optimal DPO policy can assign more probability to the minority-preferred output when the reference model is poorly aligned, whereas the KTO update with ${\lambda }_D={\lambda }_U$ increases the relative likelihood of the majority-preferred one.

The paper's ablations support the loss shape: replacing KTO's symmetric $1-\sigma (\cdot )$ with DPO's $-\log \sigma (\cdot )$ cost 9.4 points on BBH and 11.0 on GSM8K, and making the value function risk-neutral (identity) collapsed BBH performance entirely.

Estimating the reference point#

Computing $z_0$ properly would mean sampling from ${\pi }_{\theta }$ during training, which is slow. KTO instead uses a biased but cheap estimate from the microbatch. Take $m$ examples $(x_1,y_1),\dots ,(x_m,y_m)$, shift the responses by one to create mismatched pairs $(x_1,y_2),(x_2,y_3),\dots ,(x_m,y_1)$, and compute

Mismatched responses are used because the labeled responses were chosen to be canonically good or bad and have unrepresentatively large rewards. Clamping at zero adds upward bias but reduces variance. This costs one extra forward pass on each model per microbatch and is why KTO needs a microbatch size of at least 2. The paper notes that when KTO follows SFT on the same data used as desirable feedback, the estimate quickly approaches zero, and you can skip it and set ${\hat{z}}_0=0$.

Worked example#

A microbatch of four examples, $\beta =0.1$, ${\lambda }_D={\lambda }_U=1$. The mismatched log-ratios are $0.8,-0.3,1.1,0.2$; their mean is 0.45, so ${\hat{z}}_0=0.45$.

Examples 2 and 4 are on the wrong side of the reference point and carry the most loss. With $\beta =0.1$ all four are still near the middle of the sigmoid, so gradient factors are close to the 0.025 maximum. An example would need $|r-z_0|$ of roughly 30 nats before its gradient factor fell below a third of that.

Now the class weights. Suppose you have 400 desirable and 3,600 undesirable examples. The paper recommends choosing ${\lambda }_D,{\lambda }_U$ so that

With ${\lambda }_U=1$: ${\lambda }_D\times 400/3600\in [1,1.5]$, so ${\lambda }_D\in [9,13.5]$. The paper's own example: at a 1:10 ratio, ${\lambda }_U=1$ and ${\lambda }_D\in [10,15]$.

A toy run on imbalanced feedback#

We trained KTO on a log-linear policy over 200 responses with 16 shared features (true reward a linear function of the features, standardized). We sampled 4,000 responses from the reference, scored each with the true reward plus Gaussian noise (standard deviation 0.5), and labeled the top 10% desirable: 400 desirable, 3,600 undesirable. Because the toy has one prompt, the reference point can be computed exactly as the policy's KL (detached), rather than estimated from mismatched pairs. Adam, learning rate 0.02, batch 64, $\beta =0.1$, 3,000 steps, three seeds averaged. The reference policy's expected true reward is −0.02 and the best response scores 2.45.

Unweighted KTO on 1:9 data does what the paper warns about: the implied rewards of desirable and undesirable examples both collapse, by 46 and 54 nats. The model learns mostly "push everything in the data down" and expected reward stalls at 0.61. Setting ${\lambda }_D=10$, inside the recommended band, nearly triples the gain (1.77). Overshooting to ${\lambda }_D=30$ keeps desirable rewards positive but pushes undesirable ones down less and ends slightly lower. The recommended band was the best of the three here, though the sensible range is wide.

Two more probes from the same setup. Training ${\lambda }_D=10$ three times longer (9,000 steps) left KL at 2.37 and expected reward at 1.73; raising the learning rate fivefold gave 2.23 and 1.67. KTO plateaus, which is the saturation of Figure 1 at work. Lowering $\beta$ to 0.03 moved it further (KL 3.24, expected reward 1.99), consistent with the paper's advice to lower $\beta$ if KTO underfits. For comparison, DPO on 400 pairs built from the same examples (each desirable response paired with a random undesirable one) reached the best response, 2.45, at a KL of 5.61. In this toy the true reward is exactly what the labels reveal, so drift is rewarded; on a real model that much drift is not free. Read the comparison as "KTO moves less and stops on its own", not as a quality ranking.

The algorithm#

Input: pi_sft (or a strong base model), examples D = {(x, y, label)}, beta, lambda_D, lambda_U
pi_ref = frozen copy; pi_theta = trainable copy
for each microbatch of m >= 2 examples:
    r_i  = logp(pi_theta, x_i, y_i) - logp(pi_ref, x_i, y_i)
    kl_i = logp(pi_theta, x_i, y_{i+1}) - logp(pi_ref, x_i, y_{i+1})     # mismatched
    z0   = max(0, mean(kl_i)).detach()        # (sync the mean across devices)
    v_i  = lambda_D * sigmoid(beta * (r_i - z0))   if desirable
           lambda_U * sigmoid(beta * (z0 - r_i))   if undesirable
    loss = mean(lambda_{y_i} - v_i)
    step optimizer

Data

Each row is (prompt, completion, boolean label). Natural sources: thumbs feedback from a deployed product, task success signals, rule or test outcomes, a judge model's accept/reject call, or SFT data mixed with known-bad samples. Preference datasets convert trivially by treating every chosen response as desirable and every rejected one as undesirable, which is how the paper ran most experiments; it also tested a "one-y-per-x" variant that keeps only one response per prompt, removing any trace of pairing, and KTO still beat DPO on its UltraFeedback setup. Imbalance is fine as long as you set the weights. For score-valued feedback, the paper suggests thresholding and weighting high-magnitude examples more.

Hyperparameters#

The paper's recommended settings on UltraFeedback all use 5e-6. Examples: Llama-3 8B base with SFT then KTO used $\beta =0.05$; KTO directly on the base model used 0.10; Llama-3 8B Instruct used 0.25; Qwen2.5 3B Instruct used 0.50. The learning-rate guidance has a diagnostic attached: if the rate is too aggressive, the implied rewards of desirable and undesirable examples both decline. Ideally only the undesirable ones fall while desirable ones stay flat or rise.

Compute and memory#

Memory matches DPO: a trainable policy (about 16 bytes per parameter with Adam in mixed precision) plus a frozen reference (2 bytes per parameter in bf16), or no extra reference with LoRA. Compute per example is lower than DPO per pair because each example has one response instead of two, but the reference-point estimate adds one forward pass of each model on the mismatched completions. Per labeled example: one forward and backward on the policy, one gradient-free policy forward on the mismatched completion, and two reference forwards. Counting a forward pass as a third of forward plus backward, that is about 2x an SFT step on one response, against about 2.7x for a DPO pair. Two KTO examples built from one pair's responses therefore cost somewhat more than the DPO pair, unless you can set ${\hat{z}}_0=0$ (KTO after SFT on the same data) and precompute reference log-probs, which brings each example down to roughly 1.3x SFT. The real saving is upstream: binary labels are far cheaper to collect than comparisons.

The paper also tested a reference-free variant that assumes a uniform reference. It saves memory and, with ${\lambda }_D=1.75$, beat DPO on some tasks and lost on others, while trailing standard KTO and being more sensitive to the loss-aversion weights.

Failure modes#

Ignoring class balance. As in the toy run, unweighted KTO on skewed data pushes everything down. Check the per-class mean implied reward: desirable should be flat or rising.

Learning rate set like DPO. At DPO-style rates around 5e-7, KTO's smaller reference-adjusted rewards produce small gradients, and runs can look stuck. The paper's default of 5e-6 is ten times higher.

Underfitting from saturation. If both class rewards plateau early and downstream metrics barely move, lower $\beta$ and train longer.

Broken reference point. With a microbatch of 1 per device, or a mean that is not synchronized across devices, ${\hat{z}}_0$ is noisy or wrong. Log it. It should start near 0 and grow slowly as the policy moves.

Label semantics. "Undesirable" must mean undesirable for that prompt. Mixing in generic bad text that is off-topic for the prompt gives the model an easy, useless signal.

What a healthy run looks like

At step zero the policy equals the reference, every implied reward is 0, ${\hat{z}}_0$ is 0, and each example's loss is ${\lambda }_y(1-\sigma (0))={\lambda }_y/2$. That is a useful first check: with unit weights the initial loss must be exactly 0.5. From there, the undesirable class should carry the visible progress. Its mean implied reward falls steadily and its loss drops below the desirable class's. The desirable class's mean implied reward should stay roughly flat or rise slowly; the paper notes it may stay flat when the reference has already been through SFT or earlier alignment, because the model already places high probability on those responses. ${\hat{z}}_0$ grows slowly from zero and should remain small relative to the class rewards; if it tracks the desirable reward upward one-for-one, the model is raising probability bluntly and the margin is not improving. Finally, because KTO saturates, the loss flattens rather than trending to zero, and a run that has flattened will not improve by training longer at the same $\beta$.

Metrics to watch: loss split by class, mean implied reward per class, ${\hat{z}}_0$, KL on generated samples, output length, and a task metric.

When to use KTO#

In POFE, KTO is the default suggestion when a project's feedback arrives as production thumbs data, with the class weights computed automatically from the label counts.

Minimal implementation#

import torch

def kto_loss(logp, ref_logp, logp_mis, ref_logp_mis, desirable,
             beta=0.1, lam_d=1.0, lam_u=1.0):
    """logp, ref_logp: summed log-probs of each labeled completion, shape [B].
    logp_mis, ref_logp_mis: same for mismatched completions (x_i, y_{i+1}).
    desirable: bool tensor [B]."""
    r = logp - ref_logp
    with torch.no_grad():
        z0 = (logp_mis - ref_logp_mis).mean()
        if torch.distributed.is_initialized():
            torch.distributed.all_reduce(z0); z0 /= torch.distributed.get_world_size()
        z0 = z0.clamp(min=0)
    v = torch.where(desirable,
                    lam_d * torch.sigmoid(beta * (r - z0)),
                    lam_u * torch.sigmoid(beta * (z0 - r)))
    lam = torch.where(desirable, torch.full_like(r, lam_d), torch.full_like(r, lam_u))
    return (lam - v).mean(), {"z0": z0.item(),
                              "r_desirable": r[desirable].mean().item(),
                              "r_undesirable": r[~desirable].mean().item()}

In TRL, KTOTrainer takes a dataset with prompt, completion and boolean label columns, and exposes beta, desirable_weight and undesirable_weight on KTOConfig:

from trl import KTOConfig, KTOTrainer

args = KTOConfig(output_dir="kto-out", beta=0.1, learning_rate=5e-6,
                 desirable_weight=10.0, undesirable_weight=1.0,
                 per_device_train_batch_size=8, num_train_epochs=1, bf16=True)
trainer = KTOTrainer(model=model, args=args, train_dataset=binary_feedback,
                     processing_class=tokenizer)
trainer.train()

Further reading#

The KTO paper's Sections 3 and 4 contain the HALO definition, the loss and the hyperparameter guidance used throughout this article.

• Ethayarajh et al. (2024), KTO: Model Alignment as Prospect Theoretic Optimization. HALOs, the loss, hyperparameter guidance and theory.
• Kahneman and Tversky (1979), Prospect Theory: An Analysis of Decision under Risk, and Tversky and Kahneman (1992), Advances in Prospect Theory: Cumulative Representation of Uncertainty.
• Rafailov et al. (2023), Direct Preference Optimization. The paired baseline KTO is compared against.
• Tunstall et al. (2023), Zephyr. The SFT model and UltraFeedback setup used in several KTO experiments.
• Meng et al. (2024), SimPO. Includes tuned KTO baselines against other preference losses.

Frequently asked questions#

References#

1. Ethayarajh, K., Xu, W., Muennighoff, N., Jurafsky, D., Kiela, D. (2024). KTO: Model Alignment as Prospect Theoretic Optimization. arxiv.org/abs/2402.01306
2. Kahneman, D., Tversky, A. (1979). Prospect Theory: An Analysis of Decision under Risk. doi.org/10.2307/1914185
3. Tversky, A., Kahneman, D. (1992). Advances in Prospect Theory: Cumulative Representation of Uncertainty. doi.org/10.1007/BF00122574
4. Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., Finn, C. (2023). Direct Preference Optimization: Your Language Model is Secretly a Reward Model. arxiv.org/abs/2305.18290
5. Tunstall, L. et al. (2023). Zephyr: Direct Distillation of LM Alignment. arxiv.org/abs/2310.16944
6. Meng, Y., Xia, M., Chen, D. (2024). SimPO: Simple Preference Optimization with a Reference-Free Reward. arxiv.org/abs/2405.14734
7. von Werra, L. et al. (Hugging Face) (2020). TRL: Transformer Reinforcement Learning (KTOTrainer). github.com/huggingface/trl

Related reading