DPO (Preference Optimization)

DPO (Direct Preference Optimization) applies to chosen / rejected preference-pair data, training the model to favor the chosen response. On MinT it is implemented with forward_backward_custom plus a client-side custom Bradley-Terry loss.

Data shape

The training data consists of (prompt, chosen, rejected) triples, represented by PreferencePair. Each pair is flattened into two Datums, arranged as [chosen₀, rejected₀, chosen₁, rejected₁, …] — even indices are chosen, odd indices are rejected. The loss function relies on this ordering assumption.

Building a Datum

Prompt tokens get a loss weight of 0, and completion tokens get 1.0; the completion is encoded with a leading space and an appended eos_token_id; input takes all_tokens[:-1], target takes all_tokens[1:].

def build_datum(prompt_tokens, completion_text, tokenizer):
    completion_tokens = tokenizer.encode(f" {completion_text}", add_special_tokens=False)
    completion_tokens.append(tokenizer.eos_token_id)

    all_tokens = prompt_tokens + completion_tokens
    weights    = [0.0] * (len(prompt_tokens) - 1) + [1.0] * len(completion_tokens)

    return types.Datum(
        model_input=types.ModelInput.from_ints(tokens=all_tokens[:-1]),
        loss_fn_inputs={"target_tokens": all_tokens[1:], "weights": weights},
    )

Bradley-Terry Loss

forward_backward_custom() first passes the datums through the model to get logprobs, then calls the Python loss with (data, logprobs_list). Key point: keep logprobs as Tensors, don't convert them to a Python list early — otherwise the gradient breaks and backprop fails.

import torch, torch.nn.functional as F

def sequence_logprob(logprobs, weights):
    return torch.dot(logprobs.flatten().float(),
                      _to_float_tensor(weights))

def pairwise_preference_loss(data, logprobs_list):
    chosen_scores, rejected_scores = [], []
    for c_d, r_d, c_lp, r_lp in zip(
        data[::2], data[1::2], logprobs_list[::2], logprobs_list[1::2]
    ):
        chosen_scores.append(sequence_logprob(c_lp, c_d.loss_fn_inputs["weights"]))
        rejected_scores.append(sequence_logprob(r_lp, r_d.loss_fn_inputs["weights"]))
    margins = torch.stack(chosen_scores) - torch.stack(rejected_scores)
    loss    = -F.logsigmoid(margins).mean()
    metrics = {
        "loss":          float(loss.detach()),
        "pair_accuracy": float((margins > 0).float().mean().detach()),
        "mean_margin":   float(margins.mean().detach()),
    }
    return loss, metrics

Training loop

result = training_client.forward_backward_custom(
    data, pairwise_preference_loss
)
fb = result.result()
print(fb.metrics)
training_client.optim_step(types.AdamParams(learning_rate=1e-5)).result()

Parameters

Parameter	Default	Meaning
`MINT_BASE_MODEL`	`Qwen/Qwen3-0.6B`	Base model
`MINT_LORA_RANK`	`16`	LoRA rank
`MINT_DPO_STEPS`	`3`	Number of training steps
`MINT_DPO_LR`	`1e-5`	Adam learning rate

Full case study

For an end-to-end DPO experiment (with an eval-first data split and held-out pair_accuracy), see the chat-dpo case study.