Episode 7 of a series that builds mvGPT, a minimum viable GPT, from scratch and takes it through the training pipeline of a language model. Supervised fine-tuning shows a model the answers we want, and preference optimisation shows it two answers and which is better. Our preferences are scored by code: the fine-tuned model writes four stories for an instruction, and the one that uses the most of the words asked for is preferred to the one that uses the fewest. Then direct preference optimisation, DPO, learns from 2,000 pairs, against a frozen copy of the model, with no reward model and no reinforcement learning. We measure the share of the requested words its stories use, and what that costs on held-out stories, and see reward hacking: the model raising its score with longer stories, which have more chances to use the words, whether or not they're better.
Topics: Reinforcement learning
The code
Each step on screen is a tag in the code's repository. The code at the end of the episode is ep007/end.
Transcript
Intro
0:00 Hello and welcome back to episode seven, where we're going to teach our model from preferences, pairs of its own stories where we say which is better, with direct preference optimisation. Let's begin with a quick look at where we're going to be by the end of the episode.
0:15 After today, asked for a story with kite, river and brave, our model uses all three, again and again. Its stories use about sixty-three percent of the words an instruction asks for, up from about twenty-seven, and we're going to see what that cost.
0:33 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on preferences and DPO today. And each episode has a page here as well, with its chapters and its code.
Where DPO fits
0:53 Now, before we get to today's work, let's quickly go over where we are. Last time, we measured where our model spends its time, on the CPU and the GPU. Before that, we pretrained it to predict the next token, and fine-tuned it on instructions and the stories we want for them. The next stage in the pipeline behind assistants such as ChatGPT is preferences. Writing the perfect answer is hard, while saying which of two answers is better is a lot easier, so people compare pairs of the model's answers. The original recipe, RLHF, trains a reward model to score answers the way people did, then optimises the model against that reward with reinforcement learning. Direct preference optimisation, DPO, from a paper in 2023, learns from the pairs directly, with a loss of its own, and with no reward model and no reinforcement learning, and that's what we're going to build today.
1:42 The paper is Direct Preference Optimization, subtitled your language model is secretly a reward model, by Rafael Rafailov and others at Stanford, from 2023. Its abstract describes the usual recipe as complex and often unstable: fit a reward model to the preferences, then fine-tune the model with reinforcement learning, to maximise that reward without drifting too far from where it started. It shows that the best policy for that problem can be written in closed form, in terms of the reward, so the preferences can train the policy directly, with a simple classification loss. Equation seven is that loss, written with beta, the reference policy, and the chosen and rejected answers, y w and y l.
2:22 For a wider view of where this fits, Hands-On Large Language Models, by Jay Alammar and Maarten Grootendorst, is an illustrated book on how language models work and how they're used. Its chapter twelve, Fine-tuning Generation Models, goes from instruction tuning to preference tuning, and each chapter comes with a notebook on GitHub that you can run. All right, let's get started with where our preferences are going to come from.
Preferences, scored by code
2:47 We don't have people to compare our stories, so our preferences are going to be scored by code. For an instruction that asks for words, such as rain, cake and jump, the fine-tuned model writes several stories. The story that uses more of the words asked for is preferred, counted the way evaluate dot py counts them, which gave our fine-tuned model about twenty-seven percent in episode three. The share of the words is a stand-in for a person's preference that we can measure, and a narrow one, as it doesn't say whether the story is good, which we're going to come back to.
Pairs of stories
3:19 First, in evaluate dot py, the count of the requested words a story uses gets a function of its own, words used, as the preferences are going to use it too. It takes a story and the words asked for, and returns how many of them appear in it, each found with a regular expression, as a whole word, ignoring case. The comment moves up with it: a word counts in its plural, or with d, e d or i n g on the end. And instruction metrics now calls it, in place of the lines that moved. Then preferences dot py, which makes the pairs, and imports json, os, sys and time, NumPy and torch, along with our checkpoints, the data folder, words used, pick device, and from sft dot py, instruction and parse, to build the prompts like fine-tuning did. Samples is four, and as the comment says, the story that uses the most words is preferred to the one that uses the fewest. Group is eight, so the model writes for eight instructions of one length at once, in a batch of thirty-two stories. A story gets up to two hundred tokens, and the pairs go in a file in the dpo checkpoints folder. Load asks reads an instructions file, splits it into documents on end of text, skips the first ones if we ask it to, and parses each document into its fields. It keeps the documents that ask for words and have a story, grouped by their length in tokens, so a group stacks into one tensor, with no padding, as the comment says. Each keeps its prompt's token ids and its list of words, as long as the prompt leaves room for two hundred tokens of story. Make pairs takes the model, the tokenizer, the instructions by length, how many pairs we want, and a random number generator. The lengths it draws from are those with at least eight instructions, and the pairs and the ties start empty. Until there are enough pairs, it picks a length at random, and eight different instructions of that length, and stacks their prompts into one tensor. Repeat interleave repeats each prompt four times, one after the other, and generate writes the thirty-two stories at once, at a temperature of one, for more variety between the four. For each instruction, we take its four stories from the batch, with the prompt sliced off, and as the comment says, a story ends at its first end of text, which it keeps, so the model learns to stop too. Each story's score is the share of the words it uses. We sort them by score, so the first is the worst and the last the best. If the worst and the best scores are equal, all four stories scored alike, so there's no preference to learn from, and we count it as a tie and move on. Otherwise the pair keeps the prompt, the chosen story, the rejected one, and their scores. Make pairs returns the pairs with the number of ties, which tells us how often the four stories gave no preference at all. Load pairs reads the file back, which is how dpo dot py is going to get them. Main seeds both PyTorch, for the sampling, and a NumPy generator, for the instructions we draw, so the pairs can be made again. It loads the fine-tuned model, and a prompt may take up what's left of the context window after two hundred tokens of story. It makes two sets of pairs, in a loop over the two splits, each with its file, how many pairs it wants, and how many documents to skip: for training, from the training instructions, and held out, from validation instructions, skipping the first two thousand documents, as evaluate dot py uses the first fifty with words, and we don't want to measure on instructions the pairs came from. It prints how many pairs, how many ties, and the share of the words the chosen and the rejected stories used, with the minutes so far, as writing thousands of stories takes a while. Then it makes the folder if it's missing, and writes the pairs as JSON, with the training pairs and the held-out ones under their own names. By default, that's two thousand pairs to train on, and two hundred held out.
7:08 Let's make two thousand pairs for training, and two hundred held out, which takes about nine minutes. On the way, four hundred and eighty-seven instructions gave four stories that all scored alike, and were skipped. The chosen stories use about sixty percent of the words asked for, and the rejected ones about sixteen, so the pairs have a clear difference to learn from. Both come from the model as it is, sampled, so the model can already write the chosen stories, and DPO is going to make them more likely.
The DPO loss
7:38 DPO needs two models, and the first, the policy, is the model we train. Log pi of y given x is the log-probability it gives a story y, for an instruction x, which is the sum over the story's tokens, like the loss in fine-tuning. The reference is a frozen copy of the fine-tuned model, and the gain is how much more likely the policy makes the story than the reference does. The loss takes the gain of the chosen story, y w, for the winner, minus the gain of the rejected one, y l, for the loser, and pushes that margin up, through the log of a sigmoid. So the chosen story becomes more likely than it was, and the rejected one less. Beta, zero point one, scales the margin, and it comes from the recipe with a reward model and reinforcement learning, which the paper derives this loss from. There, beta weighs a penalty for drifting from the reference. So the smaller beta is, the further the policy may move to widen a margin.
8:33 Then dpo dot py, which imports copy, to copy the model, os, sys, textwrap and time, torch and functional, TensorBoard's writer, our checkpoints, the data folder, from evaluate dot py, load instructions and words used, pick device, load pairs, write and make batch. A batch is sixteen pairs, which is sixteen chosen stories and sixteen rejected ones. The learning rate is one times ten to the minus five, a twentieth of fine-tuning's, and beta, zero point one, sets how far the policy may move from the reference to widen a margin. We evaluate each two hundred steps, and the demo asks for rain, cake and jump. Log probs is the log pi from the formula, and its comment says what it computes: the log-probability of each sequence's story, the sum over its tokens, skipping the minus ones. Log softmax turns the logits into log-probabilities, gather picks the one of each target token, with the minus ones clamped to zero so they can be gathered, and the mask of targets that aren't minus one zeroes them out of the sum. We add up log-probabilities, as multiplying a hundred small probabilities together would round to zero in floating point. DPO loss builds one batch from the pairs, with the chosen stories in the first half, and the rejected ones in the second, each as the prompt and the story, with the prompt's length, and make batch from fine-tuning pads them and masks the instruction. With both halves in one batch, each model runs a single forward pass for the pairs, instead of two. The reference's log-probabilities are computed with no gradients, as we don't train it. The gain is how much more likely the policy makes each story than the reference does, in log-probability, as the comment above it says. The margin is beta times the chosen stories' gains minus the rejected ones', and the loss is minus the log sigmoid of it, averaged. It returns the margins too, as the training loop counts the pairs where the chosen story came out ahead. Held out runs the loss on the held-out pairs, thirty-two at a time, and returns the loss, the accuracy, which is the share of pairs where the chosen story gained more than the rejected one, and the average margin. Like words rate below it, it runs with no gradients, as no weights are trained there, which saves the memory they would take. Words rate writes a story for each of the fifty instructions evaluate dot py uses, and returns the share of their words the stories use. In main, a fixed seed makes the random pairs we draw repeatable from one run to the next, and the policy starts as the fine-tuned model. Deep copy makes the reference, eval turns its dropout off, and requires grad false freezes its weights. As the comment says, the reference is the fine-tuned model, frozen, and DPO measures each change against it. Dropout stays off in the policy too, so the two models' log-probabilities differ by training alone, and gradients still flow in evaluation mode. We load the pairs, the fifty instructions, AdamW with no weight decay, and the TensorBoard writer. Weight decay pulls each weight towards zero, which would move the policy away from the reference for a reason unrelated to the preferences. Evaluate measures the held-out pairs and the words used, and writes a story for the demo. It logs each number in TensorBoard under held out, and the story as text, then flushes them and prints them on two lines, the numbers and the story, indented. Then it saves a checkpoint, so we can measure each later. Writing fifty stories takes a while, so it runs once per two hundred steps. The budget is in minutes, as in pretraining, and we evaluate once before the first step, so the charts start from the fine-tuned model. The loop is fine-tuning's: sixteen random pairs, the loss, backward, the gradients clipped at one, and the optimiser's step, with the batch's loss and accuracy printed each twenty steps, until the time's up. After the loop, a last evaluation runs if the final step wasn't an evaluation step already, the writer closes, and the policy is saved as model dot p t, next to the checkpoints of the steps. Last, it prints how many steps the run took, and how many minutes.
Ten minutes of DPO
12:36 Let's train for ten minutes, and at step zero, the policy is the reference, so all the margins are zero, and no pair counts as ranked right yet. The batches' accuracy reaches a hundred percent within a few hundred steps, as the policy learns to tell the training pairs apart, while on the held-out pairs, it settles at about sixty-four percent. And the share of the words used, on the fifty instructions, goes from about twenty-seven percent to about sixty-seven, in twelve hundred steps.
13:10 Most of the gain comes early, and by step two hundred, the stories use about fifty-four percent of the words, and after that, it moves up and down by a few points between evaluations, as fifty instructions is a small sample. The held-out accuracy goes from zero to about sixty-four percent and stays there, while the training batches are ranked right almost all the time. Two thousand pairs, seen about nine times each in ten minutes, are memorised, which is a sign we're training for longer than the data supports.
What it bought, and what it cost
13:37 Now evaluate dot py, on the fine-tuned model and on DPO's. The words used go from twenty-seven percent to sixty-three, which is what we trained for. The rest of the table is what it cost. The loss on plain stories goes from two point three eight to two point eight seven, and the perplexity from about eleven to about eighteen. The stories that end within two hundred tokens fall from forty-eight percent to twenty, and the stories get longer.
14:05 Measured at the checkpoints along the way, the words used jump in the first two hundred steps, to about fifty-two percent, and more slowly after that. The loss on plain stories rises the whole time, from two point three eight to two point five one at step two hundred, and to two point eight seven at the end. So most of the gain in the words used came in the first two hundred steps, for a small rise in the loss, and the rest of the run added about eleven points, while the loss rose about three times as much again. Stopping around two hundred steps would keep the loss lower, and I'd expect a larger beta to do so too.
14:39 We give both models an instruction for rain, cake and jump, with one seed. The fine-tuned model writes about a red truck, with a jump, and no rain or cake. DPO's model writes about Lily and a cake, with a lot of jumping, and jump and cake come back again and again, while rain doesn't make it in. It runs to the end of its two hundred tokens without stopping, and its sentences are no better than before.
Reward hacking
15:05 Our preference rewarded stories for using the words, and DPO raised that score in ways we didn't ask for. A longer story has more chances to use a word, so the chosen stories tended to be longer, and DPO learned to keep writing, and to stop less. And the score counts a word once, however often it appears, so repeating it isn't penalised, and the stories repeat the words again and again. That's reward hacking, on a small scale: the score stands in for what we want, a good story that uses the words, and training raises the score by any route it can find, such as writing longer stories, whether or not they're better. So preferences for assistants come from people, who would prefer the better story, or from a reward model trained on their choices. And beta, the weight of the penalty for drifting from the reference, limits how far the policy moves to raise its score.
15:58 Last, the docs get a page on preferences and DPO, with how the pairs are made, the loss, the results, and what the preference rewarded. And the README has the commands to make the pairs, train and compare.
Wrap-up
16:12 So, today we taught our model from preferences, where the fine-tuned model wrote four stories per instruction, and the one that used the most words asked for was preferred to the one that used the fewest. DPO then made the chosen stories more likely and the rejected ones less, measured against a frozen reference. The share of the requested words its stories use went from about twenty-seven to sixty-three percent, while the loss on plain stories rose, and fewer stories end on their own. Most of the gain came in the first two hundred steps, and the rest of the run added about ten points while the loss on plain stories kept rising. And our preference, written in code, rewarded longer stories, which repeat the words, where a person, or a reward model trained on people's choices, would prefer the better story.
16:58 Let me know in the comments what you'd score our model's stories on, if you were writing the preferences, and which parts of DPO you'd like us to go deeper on. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.