KTO: Model Alignment as Prospect Theoretic Optimization

In plain words

Most ways of teaching a language model what people prefer need comparisons: a person looks at two answers to the same question and picks the better one. Those comparisons are slow and costly to collect. This paper asks whether a cruder signal is enough: for each answer, only a yes or no on whether it is good or bad.

The method, KTO, borrows from prospect theory, a model of how people judge gains and losses. People weigh losses more heavily than equal gains, and they judge outcomes against a reference point rather than in absolute terms. KTO scores each answer against a reference point estimated from other answers in the same batch. It pushes good answers up and bad answers down, and two weights let you balance the two kinds when one is much rarer.

The appeal is that thumbs-up and thumbs-down data is cheaper to get than rankings. The authors report that, on the same data, this matches or beats DPO, a popular comparison-based method, across model sizes from 1 billion to 30 billion parameters.

TL;DR

• KTO aligns language models from binary labels: each output is marked desirable or undesirable, so paired preferences are not needed §4.1.
• The loss maximises a prospect-theory utility of each output's log-ratio against a reference point, estimated from mismatched outputs in the same microbatch §4.1.
• Author-reported: KTO matches or exceeds DPO at 1B to 30B parameters; on Zephyr-β-SFT, GSM8K rises from 40.0 (DPO) to 53.5 §4.3 tab. 2.
• Needs a reference model (or a reference-free variant), a microbatch of at least 2, and the loss-aversion weights tuned for class balance §4.2.
• An official, Apache-2.0 repository exists, but its default branch is a revised version; the original release is on a legacy branch. See Paper vs code.

Problem

• Current strong alignment methods, RLHF and DPO, need pairwise preferences, which are scarce and expensive to collect §1.
• RLHF samples from the model and optimises with PPO, which is slow and often unstable, especially in distributed training §2.
• DPO maximises the likelihood of observed preferences. The authors argue that the utility functions implied by existing methods differ from the human value function measured in prospect-theory experiments §1 §3.2.
• Binary feedback is more abundant, cheaper and faster to collect, but the authors expected it to be weaker per example; their experiments test this §1 §4.3.
• Assumptions: a frozen reference model π_ref (usually an SFT model) is available to compute log-probabilities, or a reference-free approximation is accepted §4.1 §4.3. Turning preference pairs into binary labels assumes the chosen output is always desirable and the rejected output always undesirable; the authors call this naive §4.1 Data.

Method

The authors first show that DPO and PPO-Clip belong to a family they call human-aware losses (HALOs): losses built from an implied reward and a utility function applied to it §3.2. KTO is then derived as a HALO that uses a prospect-theory value function instead of a log-likelihood objective §4.1.

The starting point is the Kahneman-Tversky value function for gains and losses, with median parameters α = 0.88 and λ = 2.25 across individuals eq. 4 §3.1. KTO changes three things. It replaces the power-law curve with a logistic function σ, which avoids numerical instability. It adds a risk-aversion parameter β. It splits the loss-aversion coefficient into λ_D for desirable and λ_U for undesirable outputs §4.1.

Steps:

1. Implied reward. For output y given input x, r(x, y) is the log-ratio of the policy to the reference model over the whole output, with no length normalisation (the normalising factor is 1) §4.1 §3.2 trainers.py.
2. Reference point. The ideal z0 is the KL divergence between policy and reference. The paper estimates it per microbatch: each input x_i is paired with the output of the next example in the batch (cyclically), the log-ratios of these mismatched pairs are averaged, and the mean is clamped at zero. The estimate is not backpropagated through §4.1, "KL Estimate".
3. Value. A desirable output gets λ_D σ(β(r − z0)). An undesirable output gets λ_U σ(β(z0 − r)) §4.1.
4. Loss. The loss for each example is λ_y − v(x, y), averaged over the batch eq. 8 §4.1. Since λ_y is a constant, the gradient equals that of λ_y (1 − σ(·)).
5. Class balance. The authors suggest choosing λ_D and λ_U so that λ_D n_D / (λ_U n_U) lies between 1 and 3/2, where n_D and n_U count desirable and undesirable examples eq. 9 §4.2.
6. Data. Desirable outputs come from desirable labels. For preference data, the chosen output is treated as desirable and the rejected output as undesirable §4.1 Data. Keeping one output per input removes pairing entirely §4.3.
7. Reference-free variant. Assume π_ref is uniform over outputs. Then r − z0 reduces to log π_θ(y|x) minus the entropy H of π_θ(·|x) §4.3. This avoids loading π_ref but trails standard KTO §4.3 tab. 2.
8. SFT-free variant. If SFT was run on the same desirable data and that SFT model is π_ref, z0 can be set to 0 to skip the extra forward pass §4.1. The authors say the estimate then tends to be near zero anyway.

Pseudocode, written from the description above. The paper has no algorithm box, so there is no alg. N to cite; the citations point to the equations and sections the code implements.

# KTO loss for one microbatch  [§4.1] [eq. 8]
# Input:  m >= 2 (input, output, desirable?) examples
# Output: scalar loss

def kto_loss(batch, policy, reference, beta, lambda_d, lambda_u):
    # Sequence log-ratio of each output, summed over tokens  [§4.1]
    logratios = (
        log_prob(policy, batch.inputs, batch.outputs)
        - log_prob(reference, batch.inputs, batch.outputs)
    )  # (m,)
    # Mismatched outputs: input i gets the output of example i+1, cyclic
    # [§4.1, "KL Estimate"]
    mismatched = roll(batch.outputs, shift=-1)  # (m,)
    mismatched_logratios = (
        log_prob(policy, batch.inputs, mismatched)
        - log_prob(reference, batch.inputs, mismatched)
    )  # (m,)
    # Shared reference point: clamped mean of the mismatched log-ratios.
    # No gradient flows through it  [§4.1]
    kl = max(0, mean(mismatched_logratios)).detach()  # scalar
    # Desirable outputs: lambda_D * (1 - sigmoid(beta * (r - z0)))  [§4.1]
    desirable_loss = lambda_d * (1 - sigmoid(beta * (logratios - kl)))
    # Undesirable outputs: lambda_U * (1 - sigmoid(beta * (z0 - r)))  [§4.1]
    undesirable_loss = lambda_u * (1 - sigmoid(beta * (kl - logratios)))
    loss = where(batch.desirable, desirable_loss, undesirable_loss)
    return mean(loss)

Key equations:

L_\text{KTO}(\pi_\theta, \pi_\text{ref}) = \mathbb{E}_{x,y \sim D} [ \lambda_y - v(x, y) ] eq. 8

\frac{\lambda_D n_D}{\lambda_U n_U} \in \left[ 1, \frac{3}{2} \right] eq. 9

Value function and reference-point estimate, as stated in §4.1 (unnumbered in the paper):

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

\hat{z}_0 = \max\left(0, \frac{1}{m} \sum_{i=1}^{m} \log \frac{\pi_\theta(y_{i+1}|x_i)}{\pi_\text{ref}(y_{i+1}|x_i)}\right) §4.1, "KL Estimate"; index i+1 is cyclic

Background equations used by the method: the Bradley-Terry preference model eq. 1, the RLHF objective eq. 2, the DPO loss eq. 3 and the closed-form optimal policy with its implied reward eq. 7 §2 §3.2.

Hyperparameters

Table 1 lists the settings the authors recommend for KTO on UltraFeedback. It is the main table of settings; the others report results. Its caption says all sweeps used AdamW, an effective batch size of 32, gradient norm clipping at 10, and λ_D = λ_U = 1 tab. 1.

Source: tab. 1. Caption notes: smaller or larger models may benefit from a smaller or larger learning rate; the best λ_D, λ_U depend on the desirable:undesirable ratio.

Results

Author-reported. Setups:

• Main comparison (§3.3, §4.3): Pythia 1.4B, 2.8B, 6.9B, 12B and Llama 7B, 13B, 30B, trained on Anthropic-HH, OpenAssistant and SHP. Winrates against SFT targets, judged by GPT-4-0613 on helpfulness, harmlessness and conciseness §3.3 fig. 2 fig. 3.
• Zephyr-β-SFT on UltraFeedback, 1 epoch. The base model was finetuned on UltraChat, not UltraFeedback tab. 2 tab. 5.
• Mistral-7B on OpenAssistant, winrate vs SFT target tab. 3.
• Hardware and run time: not stated.

Headline numbers:

Full tables, as reported:

Table 2 (Zephyr-β-SFT on UltraFeedback, 1 epoch; EM unless noted):

Source: tab. 2. Bold marks the best value per column in the paper's table.

Table 3 (Mistral-7B on OpenAssistant; winrate vs SFT target with 90% binomial CI):

Source: tab. 3. The one-y-per-x setting cuts training data by 72% tab. 3 §4.3.

Table 4 (appendix; the caption does not name the base model):

Source: tab. 4. Caption: KTO trained only on desirable examples (the SFT data) gives results similar to SFT when β is small, with λ_D fixed at 1.

Table 5 (Zephyr-β-SFT on UltraFeedback, 1 epoch, all methods see the same data; AlpacaEval 2 with GPT-4 Turbo CoT judge at temperature 0.7):

Source: tab. 5. Bold marks are as printed in the paper.

Figure results (described, not reproduced):

• fig. 2: HALO-based DPO and offline PPO outperform non-HALO SLiC and CSFT in GPT-4-judged winrate against SFT targets. The gap is significant only at 13B and above; only the HALO-aligned Llama-13B and Llama-30B reach or exceed a 50% winrate.
• fig. 3: SFT+KTO is competitive with SFT+DPO from 1B to 30B. KTO alone beats DPO alone on Llama-7B, 13B and 30B; the gap is significant at 7B and 30B (p < 0.01 after correction). No significant difference on Pythia.
• fig. 4: Without SFT, DPO-aligned models tend to produce long outputs that hallucinate further conversation turns. KTO keeps the average response length roughly unchanged. Sample generations are in tab. 6.
• fig. 5: On Llama-7B, KTO still matches or beats DPO after up to 90% of the desirable examples are discarded, with the undesirable examples untouched.
• fig. 6 (appendix): AlpacaEval 2 winrates for Mistral-7B on OpenAssistant; KTO beats SFT and DPO at all sampling temperatures.

Human and judge agreement app. D: 256 OpenAssistant prompts; 214 comparisons each for DPO and KTO after dropping domain-specific questions. Individual human judgments agreed with GPT-4 on 68.7% for KTO and 65.9% for DPO.

Paper vs code

The official repository is ContextualAI/HALOs (license Apache-2.0). The paper's footnote links it. Code was read at code_commit 4831988 (default branch). The README says the code was "significantly revised" after the paper and that the original code is on a legacy branch README.md. The legacy branch was read via WebFetch, not cloned, so the line numbers below refer to the default branch only.

Differences and agreements:

• Gradient clipping. Paper sweeps clip the gradient norm at 10 tab. 1. The code default is 1.0 base_model.yaml, applied in trainers.py. Which value produced the paper's numbers is not stated.
• Learning rate and optimizer. The paper recommends 5e-6 with AdamW §4.2 tab. 1. The default config matches: lr 5e-6 config.yaml and AdamW config.yaml. Agreement.
• Weight decay and warmup. The paper says nothing about either. The code defaults are weight decay 0.01 config.yaml and warmup over 10% of training config.yaml.
• Reference point and the z_0 = 0 option. The paper says that after SFT on the same desirable data, z_0 may be set to zero to skip the extra forward pass §4.1. The code always builds the mismatched inputs dataloader.py and always runs the forward pass under no_grad trainers.py. I found no switch that sets z_0 to zero.
• Scope of the KL average. The paper uses one reference point per microbatch §4.1. The code sums the mismatched log-ratios and their counts across all processes with accelerator.reduce, then divides trainers.py. With more than one process, the reference point is therefore shared across processes, not per microbatch. The clamp at zero matches the paper trainers.py.
• Gradient through z_0. The paper does not backpropagate through z_0 §4.1. The code detaches it trainers.py trainers.py. Agreement.
• The loss. The desirable term is w_D (1 − σ(β(r − KL))) trainers.py and the undesirable term is w_U (1 − σ(β(KL − r))) trainers.py. These equal eq. 8 with the constant λ_y dropped. Default weights are 1.0 kto.yaml kto.yaml. The default β is 0.1 kto.yaml. Agreement with the paper's defaults.
• Sequence reward. Token log-ratios are summed, not averaged, by default trainers.py. This matches the paper's normalising factor of 1 §4.1.
• Microbatch. The code raises an error when the microbatch times the number of processes is 1 or less dataloader.py. This matches the paper's requirement of at least 2 §4.2.
• Preference data. The code treats chosen outputs as desirable and rejected outputs as undesirable when the data are pairs, and it can subsample one output per input dataloader.py. Agreement with the paper's data handling §4.1 Data.
• Humanline option. The code has an extra humanline option that clamps log-ratios config.yaml trainers.py. It is off by default config.yaml and is not described in the paper.
• Legacy branch. The original code, read via WebFetch, uses the same loss form: a mean of the mismatched policy minus reference log-probabilities, detached, clamped at zero, and averaged across processes. Its top-level config, as summarised by WebFetch, sets lr 5e-7 with RMSprop, which fits the paper's note that comparison runs used DPO defaults §4.2. I did not verify the legacy KTO loss config.

No discrepancies found in the core loss, the reward, the KL estimate or the λ weights.

Pitfalls

• Reference point bias. The KL estimate is biased. Mismatched outputs are not samples from the policy, and clamping at zero adds an upward bias while reducing variance §4.1. The authors argue that human reference points are biased in a similar way §4.1.
• Microbatch composition sets the reference point. The estimate uses the other examples in the same microbatch, so shuffling and microbatch size change it. The paper does not state the microbatch size used in its runs.
• Learning rate is the most sensitive knob. The authors say aligned performance is more sensitive to lr than to any other hyperparameter, and that the good KTO lr is usually 2x to 10x the good DPO lr §4.2. Diagnose with the implied rewards: if both desirable and undesirable rewards fall, lr is too high; ideally only the undesirable rewards should fall §4.2.
• β does not transfer from DPO. Optimal DPO settings "do not necessarily transfer" to KTO, even for β §4.2. Small β (0.01 to 0.10) suits a reference already tuned on the same data; large β (0.10 to 0.50) suits a strong reference §4.2.
• Class imbalance. The loss weights must follow eq. 9. If desirable examples are discarded, λ_D rises to compensate; with 90% dropped, the example uses λ_D = 13.33 §4.3. Choose λ by the ratio rule, not by the number of examples alone.
• λ ratio is an empirical rule. The interval was found empirically. The authors say that in toxicity-type tasks, where the downside matters more, the reverse may work better §4.2.
• Extra cost per step. KTO needs a second policy forward pass and a reference forward pass on the mismatched outputs §4.1 trainers.py. A reference model must be held in memory unless the reference-free variant is used, and that variant trails standard KTO on some tasks §4.3 tab. 2. The code can cache reference log-probabilities, but the default is off config.yaml.
• Desirable-only data behaves like SFT when β is small. Training on desirable examples alone with small β gives results similar to SFT tab. 4.
• Preference pairs as binary labels. Splitting a pair into a desirable and an undesirable example is called naive by the authors; they leave better decompositions to future work §4.1 Data.
• Scored feedback. The paper gives no HALO for scores. The suggested route is a weighting function plus a threshold for desirability §4.1 Data.
• Theory is conditional. The claims about intransitive preferences hold for a loss-neutral value function (λ_D = λ_U), and the worst-case argument depends on the reference model and p as stated in the theorem §4.4.
• Small models. Up to 7B parameters, alignment gave almost no gain over SFT for the methods compared in §3.3 §3.3.
• Judge dependence. Most winrates are GPT-4-judged. Individual human judgments agreed with GPT-4 on 68.7% of comparisons for KTO and 65.9% for DPO app. D.
• Missing details. Hardware, run time, warmup, weight decay, Adam betas, maximum sequence length and the SLiC regularisation weight are not stated in the paper. The code defaults above are not the paper's values unless the paper says so.

Lineage

• Builds on: Direct Preference Optimization, which KTO replaces with an implied-reward utility; the prospect-theory value function of Tversky and Kahneman (1992), cited in §3.1; PPO-Clip (Schulman et al., 2017), whose offline variant is a baseline §3.2.
• Compared against: DPO, offline PPO (an offline PPO-Clip variant), CSFT (Korbak et al.), SLiC (Zhao et al.) and ORPO (Hong et al.) §3.3 tab. 2 tab. 5.
• Superseded by: not stated in this wiki.