KTO Loss Shape

Problem

Compute the mean simplified KTO loss with an asymmetric sigmoid.

With delta = policy_logp - ref_logp, a desirable sample contributes 1 - sigmoid(beta * (delta - z0)) and an undesirable sample contributes 1 - sigmoid(beta * (z0 - delta)). Return the batch mean, or 0.0 when empty.

Hint

KTO treats desirable and undesirable examples asymmetrically, so the same delta gets different losses.

Test cases

case 1

in: [[0],[0],[true],1,0]

expected: 0.5

case 2

in: [[1],[0],[true],1,0]

expected: 0.2689414213699951

case 3

in: [[1],[0],[false],1,0]

expected: 0.7310585786300049

case 4

in: [[],[],[],1,0]

expected: 0

Starter codePython

import math
def kto_loss(policy_logps, ref_logps, desirable, beta, z0):
    # Your code here
    pass

4 test cases · run your code against them in the editor