KTO (Kahneman-Tversky) Loss

Background

KTO (Kahneman-Tversky Optimisation) is an alignment loss that — unlike DPO — does not require paired preferences. Each example is labelled independently as desirable or undesirable. The objective is informed by prospect theory's gain/loss asymmetry: pull desirables above a reference threshold, push undesirables below.

Problem statement

Implement kto_loss(rewards, is_desirable, kl_ref, beta=0.1) returning the mean loss:

$$\mathcal{L} \;=\; \frac{1}{N}\sum_i \begin{cases} -\log\sigma\bigl(\beta(r_i - z)\bigr) & \text{if desirable}_i \\ -\log\sigma\bigl(\beta(z - r_i)\bigr) & \text{otherwise} \end{cases}$$

where $z$ is kl_ref (the reference KL term). Use the numerically stable log-sigmoid identity -logaddexp(0, -x).

Input

• rewards — 1-D array of floats, length $N$. Per-example reward = log_policy − log_ref.
• is_desirable — 1-D array of bool, length $N$.
• kl_ref — float, the reference KL value $z$.
• beta — float, temperature; default $0.1$.

Output

Returns a Python float.

Examples

Example 1 — desirable with high reward → low loss

Input:  rewards=[10.0], is_desirable=[True], beta=0.1
Output: small positive value (sigmoid(1.0) ≈ 0.73 → -log ≈ 0.31)

Example 2 — desirable with low reward → higher loss

Input:  rewards=[-10.0], is_desirable=[True]
Output: significantly higher than Example 1

Example 3 — undesirable mirror case

Undesirable + very-positive reward gets penalised hard.

Constraints

• Length mismatch raises ValueError.
• is_desirable should be boolean (or 0/1).
• Returns a plain Python float.

Notes

• No pairs needed. DPO requires (chosen, rejected) pairs from the same prompt; KTO needs only a single per-example label. Makes data collection much easier.
• The role of $z$. $z$ acts as a learned (or fixed) reference threshold. Setting $z$ near the typical reward keeps desirable and undesirable gradients balanced.