diffusion_kto_loss.py

instance_loss = (
    F.mse_loss(
        predictions.predicted_noise,
        diffused_latent_images.noise,
        reduction="none",
    )
    .flatten(start_dim=1)
    .mean(dim=1)
)

good_loss, bad_loss = instance_loss.chunk(2, dim=0)

with torch.no_grad():
    reference_instance_loss = (
        F.mse_loss(
            reference_predictions.predicted_noise,
            diffused_latent_images.noise,
            reduction="none",
        )
        .flatten(start_dim=1)
        .mean(dim=1)
    )

    reference_good_loss, reference_bad_loss = reference_instance_loss.chunk(
        2, dim=0
    )

    reference_weight = (
        (instance_loss.neg() - reference_instance_loss.neg())
        .mean()
        .clamp(min=0)
    )

beta_dpo = 5000
loss = (
    -F.sigmoid(
        beta_dpo * (good_loss.neg() - reference_good_loss.neg())
        - reference_weight
    ).mean()
    - F.sigmoid(
        -beta_dpo * (bad_loss.neg() - reference_bad_loss.neg())
        - reference_weight
    ).mean()
)

return loss