Direct Preference Optimization: Data Preparation

In the previous chapter, we used SFT to teach the model to converse by following instructions. But SFT has a limitation: it can only tell the model how to answer, not what a better answer is.

For example, two answers may both be correct, but one is concise and the other wordy, or one is polite and the other blunt. SFT cannot tell them apart.

In practice, researchers found that telling the model which answer is more reasonable during fine-tuning teaches it to answer questions in a more accurate, honest, and safe way. This fine-tuning technique is called preference optimization (PO). It is usually stacked on top of SFT to further improve the model’s output.

Preference optimization currently comes in three main kinds:

Reinforcement learning from human feedback (RLHF): train a reward model and use a reinforcement learning algorithm to optimize the model;
Direct preference optimization (DPO): train the model directly on data that carries preferences;
Test-time preference optimization (TPO): adjust the model’s output at inference time with specific strategies and algorithms (such as guided decoding and rejection sampling).

Direct preference optimization¶

Direct preference optimization (DPO) teaches the model to compare: show the model two answers to the same prompt and tell it which one is better.

Data format¶

The data format of DPO is a triple:

Prompt: the user’s question or command;
Chosen: the more reasonable answer;
Rejected: the worse answer.

For example:

Prompt: ROMEO:
        Hello, Juliet!
Chosen: JULIET:
        Hello, Romeo! Nice to meet you!
Rejected: JULIET:
          Hi!

Data source¶

As with the SFT data, we again extract conversation samples from the test set of the tiny Shakespeare (tinyshakespeare) dataset, as the prompts and the chosen answers. The rejected answers use the same prompts and are continuations written by the model.

💡 Tip: To make teaching easier, the samples in this chapter, like those in Chapter 15, come from the Shakespeare script: the chosen answer is the real speech, and the rejected answer is a continuation by the pre-trained model, so no manual preference annotation is needed. In real projects, preference data is usually produced by humans or stronger models comparing and labeling multiple answers to the same prompt.

import pickle
import re
from itertools import pairwise

import numpy as np

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

np.random.seed(42)

Data extraction¶

Extracting the chosen answers¶

The same extraction function extract_turns as in the previous chapter extracts all the speeches from the test set of the Shakespeare script (the last 10% of the speech blocks, which lies inside the test set used in pre-training).

def extract_turns(filename, split=0.9, max_turn_chars=100):
    with open(filename, encoding="utf-8") as f:
        text = f.read()

    blocks = re.split(r"\n\s*\n", text.strip())
    turns = []

    split = int(len(blocks) * split)
    for block in blocks[split:]:
        lines = block.strip("\n").split("\n")
        if not lines:
            continue

        m = re.compile(r"^([A-Z][A-Za-z' ]{0,30}):\s*$").match(lines[0].strip())
        if not m:
            continue

        speaker = m.group(1).strip()
        content_lines = []
        length = 0
        for line in lines[1:]:
            line = line.strip()
            if not line:
                continue
            if length + len(line) + 1 > max_turn_chars:
                break
            content_lines.append(line)
            length += len(line) + 1
        if not content_lines:
            continue

        turns.append((speaker, "\n".join(content_lines)))

    return turns

Generating the rejected answers¶

The function generate_rejected uses the pre-trained model to generate the rejected answers:

Temperature sampling: set the sampling temperature to 0.8 to control the randomness; the larger the value, the more random;
Top-k sampling: set top-k to 20, sampling only from the 20 most probable tokens to avoid generating gibberish;
Length: the generated length equals that of the chosen answer.

# Use the answer generated by the model as the rejected answer
def generate_rejected(model, prompt, length, temperature=0.8, top_k=20):
    tokens = list(prompt)

    with Tensor.no_grad():
        for _ in range(length):
            # Inference
            logits = model(Tensor([tokens[-CONTEXT_SIZE:]]))
            last_logits = logits.data[0, -1] / temperature

            # Sampling
            k = min(top_k, last_logits.size)
            threshold = np.partition(last_logits, -k)[-k]
            last_logits = np.where(last_logits < threshold, -np.inf, last_logits)
            probs = np.exp(last_logits - np.max(last_logits))
            probs /= np.sum(probs)
            tokens.append(int(np.random.choice(len(probs), p=probs)))

    return tokens[len(prompt):]

Hyperparameters¶

Model parameters¶

All the hyperparameters used here stay the same as in pre-training.

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

Building the samples¶

Loading the model¶

The rejected answers are generated by the pre-trained model, so we load the pre-trained model here and switch it to evaluation mode.

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

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
trainer = GPTTrainer(model, None, None)
trainer.load(MODEL_FILE)
model.eval()

Extracting the samples¶

Build (prompt, chosen answer, rejected answer) triples from the test set of the dataset:

Pairing: as in SFT, extract adjacent pairs of speeches, as the prompt and the chosen answer;
Generating the rejected answers: use the pre-trained model to generate one output for each prompt as the rejected answer;
Filtering: remove the samples whose rejected answer is identical to the chosen answer.

pairs = list(pairwise(extract_turns(DATA_FILE)))
np.random.shuffle(pairs)
pairs = pairs[:512]

samples = []
for i, ((speaker_a, content_a), (speaker_b, content_b)) in enumerate(pairs):
    # Prompt
    prompt = dataset.encode(f"{speaker_a}:\n{content_a}\n\n")
    # Chosen: the real speech
    chosen = dataset.encode(f"{speaker_b}:\n{content_b}\n\n")
    # Rejected: the answer generated by the model
    rejected = generate_rejected(model, prompt, len(chosen))

    if rejected != chosen:
        samples.append((prompt, chosen, rejected))

    if (i + 1) % 64 == 0:
        print(f"generated {len(samples)}/{len(pairs)} DPO samples")

generated 64/512 DPO samples
generated 128/512 DPO samples
generated 192/512 DPO samples
generated 256/512 DPO samples
generated 320/512 DPO samples
generated 384/512 DPO samples
generated 448/512 DPO samples
generated 512/512 DPO samples

Saving the samples¶

Save them to a file with pickle.dump() as well.

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

with open(DPO_SAMPLES, "wb") as f:
    pickle.dump(samples, f)
print(f"saved {len(samples)} DPO samples to {DPO_SAMPLES}")

saved 512 DPO samples to ../../../../data/dpo-samples.pkl

Verifying the samples¶

Look at the first sample to confirm that the format is correct.

sample = samples[0]
print(f"{dataset.decode(sample[0])}\n{dataset.decode(sample[1])}\n{dataset.decode(sample[2])}")

SEBASTIAN:
Foul weather?


ANTONIO:
Very foul.


ISABELLA:
I hope, sir

In this chapter, we again relied on the test set of the tiny Shakespeare dataset and built 512 DPO training samples: each sample consists of a prompt, a chosen answer (the real speech), and a rejected answer (the model output).

In the next chapter, we will use these samples to train the SFT model with direct preference optimization, so that it learns to pick the better one among several answers.

Exercises¶

Try raising the temperature used to generate the rejected answers from 0.8 to 1.5, and observe how the rejected answers change. How would this affect the quality of the preference data?