Direct Preference Optimization: Model Training

In the previous chapter, we built the DPO training data: each sample contains a prompt, a chosen answer, and a rejected answer. Now we use this data to perform direct preference optimization on the SFT model.

Direct preference optimization¶

The core idea of direct preference optimization: make the model give a higher score to the chosen answer and a lower score to the rejected answer.

Sequence log probability¶

How do we measure the score the model gives to an answer? We use the sequence log probability:

\log P(y|x) = \sum_{t=1}^{T} \log P(y_t | x, y_{<t})

(1)

Where:

$x$ : the prompt;
$y$ : the answer, where $y_t$ is the $t$ -th token of the answer and $y_{<t}$ are all the tokens before it;
$T$ : the number of tokens in the answer.

That is, given the prompt $x$ , it is the sum of the log probabilities of every token in the answer generated by the model. The larger this value, the more the model “likes” the answer.

DPO loss function¶

The DPO loss function is based on the log probability ratios of the policy model and the reference model:

\mathcal{L} = -\log \sigma\left(\beta \left(\log \frac{\pi_\theta(y_c|x)}{\pi_{ref}(y_c|x)} - \log \frac{\pi_\theta(y_r|x)}{\pi_{ref}(y_r|x)} \right) \right)

(2)

Where:

$\pi_\theta$ : the policy model being trained;
$\pi_{ref}$ : the reference model, usually the model after SFT, whose parameters are frozen during training;
$y_c$ : the chosen answer;
$y_r$ : the rejected answer;
$\beta$ : a coefficient that controls how far the model may deviate from the reference model;
$\sigma$ : the Sigmoid function.

The goal of the loss function: make $\pi_\theta$ assign $y_c$ a higher probability than $\pi_{ref}$ does, and assign $y_r$ a lower probability than $\pi_{ref}$ does.

Reference model¶

Without a reference model, the policy model may go to extremes: assigning the chosen answer an infinitely high probability and the rejected answer an infinitely low one, which degrades its ability to generate text. The reference model acts as an anchor that keeps the model from drifting too far.

import os
import pickle

import numpy as np

from src.core import *
from src.gpt import *

np.random.seed(42)

Datasets¶

SFT dataset¶

The same SFT dataset SFTDataset as in Chapter 16. The DPO dataset extends it.

class SFTDataset(Dataset):

    def __init__(self, filename, context_size=64, split=0.9):
        self.filename = filename
        self.context_size = context_size
        self.split = split
        super().__init__(1)

    def load(self):
        with open(self.filename, "rb") as f:
            examples = pickle.load(f)

        split = int(len(examples) * self.split)
        self.train_data = self._pack(examples[:split])
        self.test_data = self._pack(examples[split:])

    def _pack(self, examples):
        xs, ys, masks = [], [], []
        for prompt, response in examples:
            x, y, mask = self._build_example(prompt, response)
            xs.append(x)
            ys.append(y)
            masks.append(mask)
        return xs, ys, masks

    def _build_example(self, prompt, response):
        tokens = list(prompt) + list(response)
        if len(tokens) > self.context_size + 1:
            overflow = len(tokens) - (self.context_size + 1)
            prompt = prompt[overflow:] if overflow < len(prompt) else []
            tokens = list(prompt) + list(response)
            tokens = tokens[-(self.context_size + 1):]

        response_start = len(prompt)

        x = np.array(tokens[:-1], dtype=np.int64)
        y = np.array(tokens[1:], dtype=np.int64)
        mask = np.arange(len(y)) + 1 >= response_start
        return x, y, mask

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y, mask = self.data
        return Tensor(x[s]), Tensor(y[s]), Tensor(mask[s])

DPO dataset¶

The DPO dataset DPODataset inherits from the SFT dataset and processes (prompt, chosen answer, rejected answer) triples:

split each triple into two: (prompt, chosen answer) and (prompt, rejected answer);
encode each of the two groups into features and labels, and add a mask;
return a six-element tuple: (chosen features, chosen labels, chosen mask, rejected features, rejected labels, rejected mask).

class DPODataset(SFTDataset):

    # Packing function: convert (prompt, chosen, rejected) into two groups, (prompt, chosen) and (prompt, rejected)
    def _pack(self, examples):
        # Chosen features, chosen labels, chosen mask, rejected features, rejected labels, rejected mask
        cxs, cys, cmasks, rxs, rys, rmasks = [], [], [], [], [], []
        for prompt, chosen, rejected in examples:
            x, y, m = self._build_example(prompt, chosen)
            cxs.append(x)
            cys.append(y)
            cmasks.append(m)
            x, y, m = self._build_example(prompt, rejected)
            rxs.append(x)
            rys.append(y)
            rmasks.append(m)
        return cxs, cys, cmasks, rxs, rys, rmasks

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        cx, cy, cmask, rx, ry, rmask = self.data
        return (Tensor(cx[s]), Tensor(cy[s]), Tensor(cmask[s]),
                Tensor(rx[s]), Tensor(ry[s]), Tensor(rmask[s]))

Loss function¶

DPO binary cross-entropy loss function¶

The DPO loss function DPOLoss implements the loss above, which is the binary cross-entropy after applying the Sigmoid to the margin:

Compute the differences:

p_{diff} = policy_{chosen} - policy_{rejected}

(3)

r_{diff} = reference_{chosen} - reference_{rejected}

(4)

Compute the margin:

margin = \beta \cdot (p_{diff} - r_{diff})

(5)

Sigmoid:

sigmoid = \frac{1}{1 + e^{- margin}}

(6)

Loss: return the mean of $-\log(sigmoid)$ .

Where:

$policy_{chosen}$ , $policy_{rejected}$ : the sequence log probabilities that the policy model gives to the chosen and rejected answers;
$reference_{chosen}$ , $reference_{rejected}$ : the sequence log probabilities that the reference model gives to the chosen and rejected answers;
$\beta$ : a coefficient, corresponding to beta in the code, with a default value of 0.1.

class DPOLoss(Loss):

    def __init__(self, beta=0.1):
        self.beta = beta

    def loss(self, pc: Tensor, pr: Tensor, rc, rr):
        # Log probability difference between chosen and rejected for the policy model
        p_diff = pc.data - pr.data
        # Log probability difference between chosen and rejected for the reference model
        r_diff = np.asarray(rc) - np.asarray(rr)
        # Margin: how much more the policy model prefers the chosen answer than the reference model does
        margin = self.beta * (p_diff - r_diff)
        # Create the loss tensor
        sig = 1.0 / (1.0 + np.exp(-margin))
        e = Tensor(-np.sum(np.log(np.clip(sig, 1e-10, 1.0))) / pc.shape[0])

        # Gradient (closure) function of the loss tensor
        def gradient_fn():
            common = (sig - 1.0) / pc.shape[0] * self.beta * e.grad
            pc.grad += common
            pr.grad += -common

        # Set the gradient (closure) function and the parent list of the loss tensor
        return e.attach(gradient_fn, parents=[pc, pr])

Trainer¶

DPO trainer¶

The DPO trainer DPOTrainer also inherits from the GPT trainer GPTTrainer. The main differences are:

Two models: besides the model being trained, model, the DPO trainer also has a reference model, reference. The parameters of the reference model are frozen and do not take part in training;
_sequence_log_prob (sequence log probability function): computes the model’s log probability for each (features, labels, mask) group, counting only the positions where the mask is 1;
train (training function): computes the scores that the policy model and the reference model give to the chosen and rejected answers, and passes them to the DPO loss function.

💡 Tip: The reference model and the policy model have identical parameters at the start of training. During training, the policy model’s parameters keep being updated while the reference model’s parameters stay unchanged, serving as a control group. Note that the policy model is in training mode (dropout enabled) while the reference model is in inference mode (dropout disabled), so even with identical parameters, the scores they give are not exactly the same.

class DPOTrainer(GPTTrainer):

    def __init__(self, model, loss_fn, optimizer, reference):
        super().__init__(model, loss_fn, optimizer)
        self.reference = reference
        # Put the reference model in inference mode
        self.reference.eval()

    def train(self, dataset, epochs, scheduler=None, filename=None):
        self.model.train()

        steps = 0
        for epoch in range(epochs):
            order = list(range(len(dataset)))
            np.random.shuffle(order)

            total_loss = 0.0
            for step, i in enumerate(order):
                if scheduler is not None:
                    self.optimizer.lr = scheduler.step(steps)

                cx, cy, cm, rx, ry, rm = dataset[i]
                # Sequence log probabilities of the policy model
                policy_chosen = self._sequence_log_prob(self.model(cx), cy, cm)
                policy_rejected = self._sequence_log_prob(self.model(rx), ry, rm)
                # Sequence log probabilities of the reference model
                ref_chosen = self._sequence_log_prob(self.reference(cx), cy, cm).data
                ref_rejected = self._sequence_log_prob(self.reference(rx), ry, rm).data
                # Compute the loss
                loss = self.loss_fn(policy_chosen, policy_rejected, ref_chosen, ref_rejected)

                self.optimizer.zero_grad()
                loss.backward()
                total_loss += float(loss.data)
                self.optimizer.clip_grad_norm()
                self.optimizer.step()
                steps += 1

                if (step + 1) % 100 == 0:
                    lr = f" lr {self.optimizer.lr:.6f}" if scheduler is not None else ""
                    print(f"epoch {epoch + 1} step {step + 1}/{len(dataset)} loss {(total_loss / 100):.4f}{lr}")
                    total_loss = 0.0

            if filename is not None:
                self.save(filename)
                print(f"epoch {epoch + 1} saved DPO model to {filename}")

    # Test function: no testing during fine-tuning
    def test(self, dataset):
        return None

    # Sequence log probability function
    @staticmethod
    def _sequence_log_prob(x, y, mask):
        # Built-in Softmax
        exp = np.exp(x.data - np.max(x.data, axis=-1, keepdims=True))
        softmax = exp / np.sum(exp, axis=-1, keepdims=True)

        y = np.asarray(y.data, dtype=np.int64)
        mask = np.asarray(mask.data, dtype=softmax.dtype)
        # Create the probability tensor
        picked = np.clip(np.take_along_axis(softmax, y[..., None], axis=-1).squeeze(-1), 1e-10, 1)
        p = Tensor(np.sum(np.log(picked) * mask, axis=-1))

        # Gradient (closure) function of the probability tensor
        def gradient_fn():
            grad = -softmax.copy()
            target_grad = np.take_along_axis(grad, y[..., None], axis=-1) + 1.0
            np.put_along_axis(grad, y[..., None], target_grad, axis=-1)
            grad *= mask[..., None]
            x.grad += p.grad[:, None, None] * grad

        # Set the gradient (closure) function and the parent list of the probability tensor
        return p.attach(gradient_fn, {x})

    # Load the reference model parameters and zero the gradients
    def load_reference(self, filename):
        if os.path.isfile(filename):
            data = np.load(filename, allow_pickle=False)
            for i, p in enumerate(self.reference.parameters):
                p.data = data[f"param_{i}"]
                p.grad = np.zeros_like(p.data)

Hyperparameters¶

Learning rate¶

The preference optimization stage usually uses a smaller learning rate than the SFT stage. This chapter uses 0.000001, one hundredth of that in the SFT stage.

LEARNING_RATE = 0.000001

Model parameters¶

The other model parameters must stay the same as in the pre-training stage.

BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6

Modeling¶

Loading the models¶

Create two models:

model: the policy model;
reference: the reference model.

Both models load the SFT model parameters.

DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-sft.npz"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
# Policy model
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
# Reference model
reference = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = DPOLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = DPOTrainer(model, loss_fn, optimizer, reference)
# Load the policy model parameters
trainer.load(MODEL_FILE)
# Load the reference model parameters
trainer.load_reference(MODEL_FILE)

Loading the DPO data¶

DPO_SAMPLES = "../../../../data/dpo-samples.pkl"

dpo_dataset = DPODataset(DPO_SAMPLES, CONTEXT_SIZE)

Training¶

DPO_MODEL = "../../../../data/tinyshakespeare-dpo.npz"

scheduler = WarmupCosineScheduler(LEARNING_RATE, len(dpo_dataset), 50, LEARNING_RATE / 10)
trainer.train(dpo_dataset, 1, scheduler, DPO_MODEL)

epoch 1 step 100/460 loss 0.6913 lr 0.000001
epoch 1 step 200/460 loss 0.6926 lr 0.000001
epoch 1 step 300/460 loss 0.6761 lr 0.000000
epoch 1 step 400/460 loss 0.6767 lr 0.000000
epoch 1 saved DPO model to ../../../../data/tinyshakespeare-dpo.npz

Inference¶

print(trainer.generate(dataset, prompt="ROMEO:"))

ROMEO:

HORTENSIO:
Villor, will I come: did a hear,
Not our drow supposessiance, that you married haste,
Like a joy a Jack for joy her.

First Servingman:
I do go of year.

PITONES:
Hereford thy country.

PRUTUS:
Call now.

CoRLIFFORD:

SICINIAN:

BRUTUS:

IS:
N

In this chapter, we trained the SFT model with direct preference optimization on the DPO dataset: the model is made to give a higher score to the chosen answer and a lower score to the rejected answer, and the reference model acts as an anchor that keeps the model from drifting too far.

The training loss fluctuated between 0.68 and 0.72 without a clear decrease. When the policy model and the reference model give exactly the same scores, the loss is exactly $\ln 2 \approx 0.693$ , so this indicates that the trained model changed little relative to the reference model. Our focus is to understand the DPO workflow; with such a small teaching model, we should not expect a clear effect.

With this, we have completed model fine-tuning, step by step teaching the pre-trained model to converse effectively with users:

SFT (supervised fine-tuning): teaches the model to answer questions;
DPO (direct preference optimization): teaches the model to answer questions reasonably.

Exercises¶

Try changing the beta parameter of the DPO loss function from 0.1 to 0.5 or 0.01 and observe how the training result changes. What does the beta parameter control?
Try raising the learning rate to 0.00001 or 0.0001 and observe how the training loss changes.