DPO

DPO — Direct Preference Optimization

nnx.dpo_train_step_factory implements the Direct Preference Optimization objective from Rafailov et al. (2023) — a simpler alternative to PPO-based RLHF that fits in a single supervised-style training loop. Combined with nnx.NNPreferenceDataset and a GenerativeNNModel (the decoder-only LM path; see Language-Modeling), DPO becomes a drop-in train_step_fn= for the standard NNModel.train(...) call.

What DPO Does

Given a frozen reference policy π_ref (typically the SFT checkpoint the model was warm-started from) and a trainable policy π_θ, DPO directly fits the policy to a dataset of (prompt, chosen_response, rejected_response) triples by minimising:

L_DPO = −log σ(β · ( (log π_θ(y_w | x) − log π_ref(y_w | x))
                    − (log π_θ(y_l | x) − log π_ref(y_l | x)) ))

In words: maximise the policy's log-ratio margin (chosen over rejected) relative to the reference's. There is no separate reward model and no RL loop; the standard NNModel.train() machinery runs unchanged.

When DPO Beats SFT

SFT (supervised fine-tuning on (prompt, good_response) pairs) maximises the likelihood of "good" responses but says nothing about what's worse. When you have explicit preference data (A > B for the same prompt), DPO almost always outperforms SFT on the target preference metric because it directly optimises the gap between chosen and rejected.

Pick DPO over SFT when:

• You have preference pairs already (Anthropic HH, OpenAssistant, UltraFeedback, or annotator-labelled data from your own pipeline).
• The SFT model is already producing fluent output and you want to steer style / behaviour (helpfulness, harmlessness, refusal format, brevity, etc.).
• You don't have the infrastructure for a full PPO RLHF loop with a separate reward model.

Pick SFT (or layered SFT → DPO) when:

• You have lots of "good" examples but very few labelled comparison pairs.
• The base model is still learning the target task's basic format — preference data won't fix raw fluency.

Quickstart

import torch
from nnx import (
    Devices, GenerativeNNModel, Losses, Nets, NNModelParams,
    NNOptimParams, NNPreferenceDataset, NNSchedulerParams,
    NNTokenizerParams, NNTrainParams, NNTransformerParams, Optims,
    dpo_train_step_factory, set_seed, train_bpe,
)

set_seed(0)

# 1. Tokenizer (BPE for the demo; swap in a published tokenizer for real data).
tk = train_bpe(files=None, texts=["..."], vocab_size=8192,
               special_tokens=["<unk>", "<pad>", "<bos>", "<eos>"])
tokenizer = NNTokenizerParams.of(tokenizer=tk, path="artifacts/tok.json")

# 2. Build the policy (the trainable model) and load SFT weights.
net_params = NNTransformerParams(
    input_dim=tokenizer.vocab_size,
    output_dim=tokenizer.vocab_size,
    dropout_prob=0.0,
    vocab_size=tokenizer.vocab_size,
    n_layers=4, n_heads=4, d_model=128,
    ffn_mult=4, max_seq_len=128,
)
model_params = NNModelParams(net=Nets.TRANSFORMER, device=Devices.CPU,
                             loss=Losses.CROSS_ENTROPY)
policy = GenerativeNNModel(net_params=net_params, params=model_params,
                           tokenizer=tokenizer)
# policy.net.load_state_dict(torch.load("sft-checkpoint.pt"))   # in practice

# 3. Build the reference (a frozen copy of the SFT model).
ref_model = GenerativeNNModel(net_params=net_params, params=model_params,
                              tokenizer=tokenizer)
ref_model.net.load_state_dict(policy.net.state_dict())

# 4. Preference dataset — yields (prompt_ids, chosen_ids, rejected_ids).
preferences = NNPreferenceDataset(
    prompts=["..."],
    chosen=["..."],
    rejected=["..."],
    tokenizer=tokenizer,
    max_prompt_len=64,
    max_response_len=64,
    pad_token_id=1,  # "<pad>"
    batch_sizes=(8, 8, 8),
    seed=0,
)

# 5. DPO step — frozen reference + β temperature.
# pad_token_id matches the dataset's — padded positions are excluded
# from the log-prob sums.
step_fn = dpo_train_step_factory(ref_model, beta=0.1, pad_token_id=1)

# 6. Train. Callbacks, schedulers, checkpointing all work as usual.
policy.train(
    params=NNTrainParams(
        n_epochs=5,
        train_loader=preferences.train_loader,
        optim=NNOptimParams(name=Optims.ADAM, max_lr=5e-5,
                            momentum=(0.9, 0.999), weight_decay=0.0),
        scheduler=NNSchedulerParams(min_lr=1e-7, factor=0.5,
                                    patience=2, cooldown=1, threshold=1e-3),
    ),
    train_step_fn=step_fn,
)

dpo_train_step_factory

step_fn = dpo_train_step_factory(ref_model, beta=0.1, pad_token_id=1)

Argument	Default	Notes
ref_model	required	A GenerativeNNModel (or NNModel) whose weights serve as the frozen reference policy. Its parameters are not modified by the training loop.
beta	0.1	KL penalty coefficient — controls how sharply the policy may diverge from the reference. Rafailov et al. recommend 0.1; practical range is [0.01, 0.5].
pad_token_id	None	Token id used to pad responses in the dataset. When None, no positions are masked. When set, padded positions are excluded from the per-sequence log-prob sums.

Returns a TrainStepFn for NNModel.train(...). The step performs two forward passes through the policy and two through the reference per row. It reports loss (the DPO loss) and error (negated chosen−rejected log-prob gap — lower is better, same monotone direction as loss).

The beta Knob

beta (β) controls how much the policy is allowed to diverge from the reference:

• Higher β (e.g. 0.5): steeper implicit reward. The policy can drift further from the reference per gradient step. Risk: the policy stops being a language model and starts gaming the preference function.
• Lower β (e.g. 0.01): policy stays close to the reference. Slower convergence but safer.

If training diverges or the model collapses (output goes to gibberish or to a single fixed answer), lower beta first.

NNPreferenceDataset

A Dataset that tokenizes parallel lists of prompt/chosen/rejected strings, splits them into train/val/test, and exposes ready-made DataLoaders.

preferences = NNPreferenceDataset(
    prompts=[...],
    chosen=[...],
    rejected=[...],
    tokenizer=tokenizer,        # NNTokenizerParams
    max_prompt_len=64,          # truncate/pad prompt token ids
    max_response_len=64,        # truncate/pad response token ids
    pad_token_id=1,
    batch_sizes=(8, 8, 8),      # (train, val, test) batch sizes; None = no loader
    val_proportion=0.1,
    test_proportion=0.1,
    seed=0,                     # reproducible split
)

# Ready-made loaders:
preferences.train_loader
preferences.val_loader
preferences.test_loader

Each batch yields (prompt_ids, chosen_ids, rejected_ids) tensors.

Argument	Default	Notes
prompts	required	List of prompt strings.
chosen	required	List of chosen response strings. Must align with prompts.
rejected	required	List of rejected response strings. Must align with prompts.
tokenizer	required	NNTokenizerParams used to encode all strings.
max_prompt_len	64	Truncate/pad prompt token ids to this length.
max_response_len	64	Truncate/pad response token ids to this length.
pad_token_id	0	Id used for padding. Must match the value passed to dpo_train_step_factory.
batch_sizes	(None, None, None)	Per-split batch sizes. None = no loader for that split.
val_proportion	0.1	Fraction of data held out for validation.
test_proportion	0.1	Fraction held out for test.
seed	None	Random seed for the train/val/test split.

Pairing with HuggingFace Hub Preference Datasets

NNPreferenceDataset takes three parallel lists of strings — easy to fill from the datasets library:

from datasets import load_dataset

ds = load_dataset("Anthropic/hh-rlhf", split="train[:1000]")
preferences = NNPreferenceDataset(
    prompts=[row["chosen"].split("Assistant:", 1)[0] for row in ds],
    chosen=[row["chosen"].split("Assistant:", 1)[1] for row in ds],
    rejected=[row["rejected"].split("Assistant:", 1)[1] for row in ds],
    tokenizer=tokenizer,
    max_prompt_len=256,
    max_response_len=256,
    pad_token_id=1,
)

The (prompt, chosen, rejected) convention matches most HF Hub preference corpora directly.

How it Composes with the Rest of NNx

• train_step_fn= hook — dpo_train_step_factory returns a standard TrainStepFn; the same plug shape as KD, SimCLR, Mixup, CutMix, MoE, and the diffusion paradigm (see Training-Paradigms).
• NNRun content-addressed persistence — DPO runs hash the same way as any other run. The reference model is an opaque dependency; track its provenance separately (file hash, HF Hub revision, etc.).
• Callbacks — EarlyStopping, ModelCheckpoint, TensorBoardCallback, WandbCallback all work unchanged on the policy.
• Generation — after training, policy.generate(prompt=..., ...) produces text from the tuned policy with the standard sampling knobs (top_k, top_p, repetition penalty, seed).

Honest Scope

NNx's DPO is built for small-LM experimentation — the same TinyStories-class sub-30-minute-on-a-laptop scope as the LM path's GenerativeNNModel. It is not a production RLHF replacement. Specifically:

• The training step does two forward passes through the policy and two through the reference, per row. For a ~7B parameter model on a single GPU this dominates memory and time; production stacks cache reference log-probs offline or share the reference with the policy via LoRA.
• There's no IPO / cDPO / RPO variant, no offline reference log-prob cache, no PEFT integration (LoRA-DPO is the obvious next step but isn't wired in for v1).
• Mixed precision and gradient accumulation are explicitly not supported by the DPO step — finalize_step raises rather than silently dropping these knobs.
• No multi-GPU / multi-node sharding.

For production-scale preference tuning, use a dedicated stack (trl, axolotl, OpenRLHF, etc.) and treat NNx's DPO as the "how does this objective behave on my small LM?" experimentation path.

Reference

Symbol	Module	Notes
dpo_train_step_factory	nnx	(ref_model, beta=0.1, pad_token_id=None) → TrainStepFn
NNPreferenceDataset	nnx	Dataset for (prompt, chosen, rejected) triples. Exposes .train_loader, .val_loader, .test_loader.

See also

• Language-Modeling — GenerativeNNModel, NNTransformerParams, tokenizer pipeline
• Training-Paradigms — other TrainStepFn factories (KD, SimCLR, Mixup, CutMix, MoE)
• PEFT — LoRA-DPO is a natural next step
• Fine-Tuning — SFT (the typical warm-start before DPO)
• API reference · source