KTO: Unpaired Preference Learning
KTO aligns language models using only binary good/bad labels per response, avoiding the paired (chosen, rejected) format that makes preference data expensive and brittle to collect.
Collecting preference data is a bottleneck in practice. DPO and RLHF both require pairs - a chosen response and a rejected response for the same prompt. That pairing requirement is harder to satisfy than it sounds: annotators must evaluate two responses simultaneously, and even small timing or framing differences between the two can introduce noise. In many real deployments, you already have logs of model outputs with thumbs-up/thumbs-down ratings, but those ratings were never collected as paired comparisons. They are inherently unpaired.
Kahneman-Tversky Optimisation (KTO), introduced by Ethayarajh et al. at ICML 2024, is built precisely for this regime. It trains directly on binary desirability signals - each example is simply (prompt, response, label) where label is either desirable or undesirable - and matches or exceeds DPO's performance on models from 1B to 30B parameters.
The prospect theory framing
The name is not decorative. Kahneman and Tversky's prospect theory models how humans actually perceive utility under uncertainty, rather than how a rational agent would. Two key properties matter here:
- Reference-point sensitivity: people evaluate outcomes relative to a reference point, not in absolute terms.
- Loss aversion: losses loom larger than equivalent gains.
KTO argues that existing alignment objectives implicitly encode some of these biases, and that making them explicit - through a Human-Aware Loss function (HALO) - leads to better behaviour. The framework unifies several existing objectives: you can derive DPO, IPO, and related methods as special cases of the HALO family, each with different implicit utility functions.
What makes KTO distinct is that it picks the utility function that actually matches the prospect theory literature, rather than one that merely happens to train well.
The loss function
Each training example has a prompt x, a response y, and a binary label. For a desirable example the loss is:
L_KTO = λ_D * σ( β * (r_θ(x,y) - z₀) )
For an undesirable example the loss is:
L_KTO = λ_U * σ( β * (z₀ - r_θ(x,y)) )
Where:
| Symbol | Meaning |
|---|---|
r_θ(x,y) |
log π_θ(y|x) - log π_ref(y|x), the implicit reward |
z₀ |
KL reference point: E[log π_θ(y'|x) - log π_ref(y'|x)] over sampled y' |
β |
KL penalty coefficient; controls how far policy can stray from reference |
λ_D, λ_U |
loss weights for desirable and undesirable examples respectively |
σ |
sigmoid function |
The z₀ term is what makes this a reference-point model in the prospect theory sense. Rather than measuring reward in absolute terms, KTO measures reward relative to the expected reward under the current policy on the same prompt. A response that looks good in isolation but is merely average for that prompt does not receive a strong positive signal.
In practice z₀ is estimated per mini-batch by averaging the implicit reward over responses sampled from the current policy for the same prompts. The TRL implementation does this automatically via the KL batch inside each training step.
A worked example. Suppose after a gradient step:
- r_θ(x, y_good) = 0.8, z₀ = 0.3 for prompt x.
- For the desirable label: σ(β * (0.8 - 0.3)) = σ(β * 0.5) > 0.5, so the loss contribution is low and the model is rewarded.
- For a bad response where r_θ(x, y_bad) = 0.1: σ(β * (0.3 - 0.1)) = σ(β * 0.2) > 0.5, contributing a stronger push away from that response.
Why no pairing is needed
DPO's loss is defined over pairs (y_w, y_l) for the same prompt. The gradient signal is essentially: "make y_w more likely than y_l." This requires knowing which of two responses is preferred.
KTO only needs to know: "is this response good or bad, independent of any comparison?" Each example contributes its own gradient. Chosen and rejected data can come from entirely different sources, different annotators, or different time periods. The reference point z₀ provides the implicit comparison, not a paired rejected response.
This matters because: - Annotation is cheaper: a binary rating takes seconds; a pairwise ranking requires reading two responses. - Existing feedback logs can be used directly (helpfulness ratings, upvotes, thumbs-down flags). - The dataset does not go stale in the same way: a new good response does not need a contemporaneous bad counterpart.
Practical setup with TRL
The KTO trainer in TRL (now under trl.experimental.kto as of v1.0) expects a dataset with columns prompt, completion, and label (boolean). A minimal training loop:
from datasets import load_dataset
from trl.experimental.kto import KTOConfig, KTOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
dataset = load_dataset("trl-lib/kto-mix-14k", split="train")
config = KTOConfig(
output_dir="my-kto-model",
beta=0.1,
desirable_weight=1.0,
undesirable_weight=1.0,
learning_rate=1e-6,
)
trainer = KTOTrainer(
model=model,
args=config,
processing_class=tokenizer,
train_dataset=dataset,
)
trainer.train()
Key hyperparameters to calibrate:
- beta: KL penalty strength. Default 0.1. Lower beta lets the policy stray further from the reference model but risks reward hacking; higher beta keeps behaviour close to SFT.
- desirable_weight / undesirable_weight: if your dataset has 70% positive labels, upweight the negative class so the effective ratio of weighted positives to weighted negatives stays between 1:1 and 4:3.
- learning rate: should not exceed 1e-6 for beta=0.1. The KL estimate is noisy, and a high learning rate amplifies that noise into unstable gradients.
- per-step batch size: TRL recommends at least 4 per step. The z₀ KL estimate is computed over the mini-batch, so very small batches produce high-variance reference points.
When it falls down
Imbalanced labels with naive weights. If the dataset is 90% desirable and you leave both weights at 1.0, the model effectively trains almost entirely on positive signal. It will learn to be less bad at rejections without learning what "bad" looks like with enough resolution. The fix is adjusting undesirable_weight upward, but getting this ratio wrong degrades performance noticeably.
Small batch sizes corrupt z₀. The KL reference point z₀ is estimated from the mini-batch. With batch size 2 or 4, the estimate has high variance. The model can receive misleading gradient directions: a genuinely good response may appear below the noisy reference point and get penalised. TRL explicitly warns against this.
Very small models. The paper reports that KTO's gains over DPO are most pronounced at 13B parameters and above. At 1B-7B, the methods perform comparably. If you are working with a 1B model, KTO's simpler data format still helps, but do not expect a clear performance advantage over DPO.
Only-desirable or only-undesirable data. Technically the trainer accepts single-polarity datasets, but the quality degrades. With only desirable data, KTO effectively becomes a weighted SFT with a KL penalty; with only undesirable data it pushes the model away from bad outputs without grounding what good looks like. You need both polarities for the reference point to be meaningful.
Noisy labels at scale. Because each example stands on its own (no contrastive pressure from a paired counterpart), label noise has a direct effect on gradient direction. Mislabelled desirable examples will push the policy toward bad outputs. With DPO, a mislabelled pair merely inverts the preference direction; with KTO, a mislabelled example is an unchecked positive or negative signal with no implicit correction from a paired example.
Majority-preference encoding. The utility function optimises toward the majority consensus in the training data. Minority-preference groups whose desirable responses were labelled undesirable by a majority of annotators are systematically underserved. This is a structural limitation shared with most RLHF methods, but KTO makes it more explicit because there is no pairwise comparison to partially average it out.
Further reading
- Ethayarajh, K., Xu, W., Muennighoff, N., Jurafsky, D., Kiela, D. (2024). "KTO: Model Alignment as Prospect Theoretic Optimisation." ICML 2024. https://arxiv.org/abs/2402.01306
- Rafailov, R., Sharma, A., Mitchell, E., et al. (2023). "Direct Preference Optimization: Your Language Model is Secretly a Reward Model." https://arxiv.org/abs/2305.18290
- TRL KTO Trainer documentation. https://huggingface.co/docs/trl/kto_trainer
- ContextualAI/HALOs - the reference implementation released with the paper. https://github.com/ContextualAI/HALOs
7 flashcards for this concept
Click a card to reveal the answer.