PPO: reinforcement learning for a GPT

Intro

0:00 Hello and welcome back to episode eight, where we're going to train our model with reinforcement learning: it writes stories, a reward scores them, and proximal policy optimisation makes the stories that score well more likely. Let's begin with a quick look at where we're going to be by the end of the episode.

0:17 After ten minutes of PPO, asked for kite, river and brave, our model uses all three, with brave again and again. Its stories use about forty-five percent of the words an instruction asks for, up from about twenty-seven, and we're going to see what it cost, next to DPO's.

0:36 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 reinforcement learning with PPO today. And each episode has a page here as well, with its chapters and its code.

0:55 Now, before we get to today's work, let's quickly go over what we did last time. Our fine-tuned model wrote four stories per instruction, and the one that used the most of the words asked for was preferred to the one that used the fewest. DPO learned from two thousand of those pairs, against a frozen copy of the model, with no reward model and no reinforcement learning. 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 from about two point three eight to two point eight seven, and fewer stories ended on their own. A longer story has more chances to use a word, so DPO learned to write longer stories, which repeat the words again and again. That's reward hacking: raising the score in ways that don't make the stories better.

Reinforcement learning

1:41 DPO learned from pairs that the fine-tuned model wrote once, before training. Today, our model is going to learn from its own stories as it writes them, which is reinforcement learning. It writes a batch of stories, and a reward scores each one, with the share of the words asked for that it uses, as last time, so we can compare the two methods. A critic learns what score to expect, so each token is credited for doing better or worse than expected, and proximal policy optimisation, PPO, makes the stories that did better more likely, in small steps. A penalty for drifting keeps the model's probabilities close to the fine-tuned model's, and then the model writes again. We're going to train it for ten minutes, twice, with a small penalty and a larger one, and compare both with DPO. All right, let's get started with where PPO comes from.

Where PPO comes from

2:32 PPO comes from Proximal Policy Optimization Algorithms, by John Schulman and others at OpenAI, from 2017. Its abstract describes policy gradient methods, which take turns between sampling data, such as our stories, and optimising an objective with gradient ascent. The standard methods make one gradient update per sample, and PPO's objective allows several passes of updates, in minibatches, over each batch of samples. It keeps some of the benefits of an earlier method, trust region policy optimisation, while being much simpler to implement, and the paper tests it on simulated robots and Atari games.

3:09 PPO is what trained InstructGPT, in Training language models to follow instructions with human feedback, by Long Ouyang and others at OpenAI, from 2022, which is the recipe behind the first ChatGPT. Its figure two shows the three steps of their method. First, supervised fine-tuning, on answers written by people, which is what we did in episode three. Then a reward model, trained on people's rankings of the model's answers, as people can't score each answer while the model trains. And last, PPO, which optimises the model against that reward. Our reward is going to be code that counts the words asked for, in place of a reward model. Equation two is the objective: the reward, minus beta times the log of how much more likely the new model makes the answer than the fine-tuned one, which is a penalty for drifting from it. There's also a term with the pretraining data, which we leave out.

4:04 In reinforcement learning, the model is called the policy, as it decides what to do, here which token to write next. The goal is to maximise the expected reward, the average score of the stories it would write for an instruction. The reward comes at the end of a story, from code that counts words, so we can't take its gradient the way we take the loss's. The policy gradient goes around that, by making each token of a story more likely in proportion to how well the story turned out. It follows the gradient of the token's log-probability, like fine-tuning does, weighted by A. A is the advantage: the reward the story got, minus V, the value, which is the reward a critic expected at that point in the story. A story that did better than expected has a positive advantage, and its tokens become more likely, and one that did worse, less likely. Subtracting what the critic expected leaves the gradient pointing where it did, on average, but makes it a lot less noisy, as a story that scores the way the critic expected teaches us very little.

PPO's clipped objective

5:05 PPO makes several updates from each batch of stories, so the model it updates drifts from the one that wrote them. Rho is the ratio of a token's probability under the model now to its probability when the story was written, and it starts at one. The objective multiplies the advantage by that ratio, so raising a good token's probability raises the objective. Clip keeps the ratio between one minus epsilon and one plus epsilon, with epsilon at zero point two, and the min takes the smaller of the two. So once a good token is twenty percent more likely than it was, pushing it further doesn't raise the objective any more, and for a bad token, pushing its probability down stops raising the objective once it's twenty percent less likely. The name comes from there: each update stays proximal, or close, to the model that wrote the stories.

The KL penalty

5:54 Last time, DPO's stories got longer and repeated the words, and beta was the single weight holding the policy near the reference. PPO keeps a frozen copy of the fine-tuned model too, as the reference, and r is the log of how much more likely the reference makes a token than the policy does. E to the r, minus r, minus one, estimates the K L divergence, which measures how far the policy's distribution has moved from the reference's. It's zero where the two agree, and above zero wherever they don't, as e to the r is at least one plus r. The loss has three parts: minus PPO's objective, the critic's squared error, weighted by a half, so that the critic learns to predict the rewards, and beta times the KL. With a small beta, raising the reward outweighs the penalty, and with a large one, the model stays close to where it started, so we're going to try two, zero point zero five and zero point five.

PPO in code

6:47 All of PPO goes in a new file, ppo dot py. It imports copy, to copy the model for the reference, along with os, sys, textwrap and time, NumPy, torch, nn and functional, and TensorBoard's writer. From our own files, it takes the checkpoints and runs folders, with loading and saving checkpoints, and the data folder. From dpo dot py, it takes the demo instruction, and words rate, the share of the words used on fifty held-out instructions. From evaluate dot py, it takes load instructions, and words used, which is going to be our reward. Then pick device, and from preferences dot py, the group size, the most tokens a story gets, and load asks, which reads the instructions by their length in tokens. And write, from sample dot py, for a sample along the way. Rollout is sixteen: groups of instructions per rollout, sixteen times eight, a hundred and twenty-eight stories. A rollout is a round of writing and scoring, before the model learns from what it wrote. Epochs is two, the passes over each rollout, and a minibatch is sixteen stories. Clip is zero point two, how far the probability of a token may move in a single update. The learning rate is one times ten to the minus five, as for DPO, and we evaluate each ten rollouts. The critic is a class of its own, value head: from each position's channels, a guess at the reward the story is going to get. It's one linear layer, a projection, from the embedding size to a single number. A linear layer starts with small random weights, which would make random guesses at first. Nn init zeros sets its weights and its bias to zero, so it starts by expecting a reward of zero everywhere, and its first guesses don't add noise to the advantages. Forward runs the layer, and squeeze drops the last dimension, of size one, for a guess per position. Token log probs gives the log-probability the model gave each token of a sequence. Position t predicts token t plus one, so the logits drop their last position and the tokens their first. Log softmax turns the logits into log-probabilities, and gather picks the one of each token in the sequence, as in DPO, but kept per token, without the sum, as PPO weighs each token on its own. Rollout writes the stories, with no gradients, as writing them trains no weights. The lengths it draws from are those with at least a group of instructions. For each of the sixteen groups, it picks a length at random, and eight different instructions of that length, so their prompts stack into one tensor with no padding, as in preferences dot py. Generate writes up to two hundred tokens for each, at a temperature of one, and over it, we write why there's no top k: PPO compares probabilities with the full distribution the tokens were sampled from. With top k, the tokens would come from a cut-down distribution, and the ratios PPO computes would be off. For each story, we cut it at its first end of text, which it keeps, so stopping is part of what it learns. Its reward is the share of the requested words it uses, between zero and one. Pack puts the prompts and their stories in one padded batch, for the update. The longest sets the width, the sequences start as padding, and the mask, one position shorter, marks the story. For each row, we copy in the prompt and the story. On the shifted frame, the positions that predict the story's tokens get a one in the mask, so the update skips the prompt and the padding. Advantages is where each token gets its share of the credit. The advantages start as zeros, in the shape of the rewards. Generalized advantage estimation, from the last token back: how much better each token turned out than the critic expected. Where the mask ends, so does the story, with no value after it. Running, the credit so far, starts at zero for each story. The loop goes through the positions in reverse, from the last token back. More is one where the story goes on at the next position, from the mask, and next value is the critic's value there, and both are zero past the end. Delta is the reward at a token, plus the value of the next token, if the story goes on, minus this token's value. Running adds up the deltas from the end, so each token's advantage is what came after it, against what the critic expected. With lambda at one, that comes to the story's reward minus the value at that token. A lambda below one would fade the credit for tokens far from the end, and at zero point nine five, the start of a two hundred token story would get almost none. It returns the advantages, along with the returns, the advantages plus the values, which is what the critic learns to predict. Masked mean averages over the story's tokens and leaves out the rest. Main takes the minutes and kl, which weighs the penalty for drifting from the reference, and each value gets a folder of its own. We seed PyTorch and NumPy's random generator, and load the fine-tuned model as the policy. Deep copy makes the reference, frozen, as in DPO. The policy goes into evaluation mode, and the comment gives the reason: dropout off, so the probabilities at the update match the rollout's. The critic is a value head on the policy's device. The critic reads the channels after the final norm, while the model returns the logits alone, so a hook keeps them in a dictionary, hidden, without changing model dot py. Register forward hook is PyTorch's way to call a function each time a module runs, with its inputs and its output, so this lambda stores the norm's output in hidden, on each forward pass. AdamW trains the policy's weights and the critic's together, from one list of their parameters, with no weight decay, as in DPO. Parameters lists a module's weights, and adding the two lists hands both to a single optimiser. Load asks reads the training instructions, as long as their prompt leaves room for two hundred tokens of story, and load instructions takes the fifty held-out ones that evaluate dot py uses. The TensorBoard writer logs to a folder named after the KL weight and the time. Evaluate measures the share of the words used on the fifty instructions, and writes a story of a hundred and twenty tokens for the demo instruction, rain, cake and jump. Add scalar logs the share to TensorBoard, add text logs the story, and flush writes them to disk. We print the share, and the story under it, indented, and save a checkpoint named after the rollout, so we can measure the model at each point later. The budget is the minutes, in seconds, and start is the time now. The loop starts with an evaluation at rollout zero, and runs until the time's up. Each rollout writes a hundred and twenty-eight stories, packs them, and keeps their rewards in a tensor. With no gradients, we compute what the update compares against. Old is the log-probability of each token under the policy that wrote it, values are the critic's guesses from the channels the hook kept, and ref is the reference's log-probability. Rewards start at zero, with a place for each token of each story. The reward goes on each story's last token, and the advantages spread it back over the story. Last finds the last position of each story's mask, as the arg max of the mask times the positions, and the story's reward goes there. Advantages gives the advantages and the returns, and we normalise the advantages to a mean of zero and a spread of one, over the story tokens, so the size of an update doesn't depend on how the rewards happen to be scaled. Then two epochs, two passes over the rollout, in minibatches of sixteen stories, from rand perm, a random order of the rows, split into sixteens. For each minibatch, we run the policy again, with gradients this time, and new is its log-probabilities. The ratio is e to the new minus the old, which is the ratio of the probabilities. It's about one on the first minibatch, and drifts from one as the updates go on. Then PPO's clipped objective, where moving a token's probability past the clip doesn't raise the objective. Surrogate is the smaller of the ratio times the advantage, and the ratio clamped between zero point eight and one point two, times the advantage, as in the formula. The policy loss is minus its mean over the story tokens, as the optimiser minimises. The value loss is the critic's squared error against the returns. We detach the channels, so the critic learns from the policy's channels without changing them. An earlier version let the critic's gradients reach the policy's channels, and the policy collapsed. Then the KL from the reference, estimated per token as e to the r, minus r, minus one, where r is the log of the reference's probability minus the policy's. The estimate can't go below zero, and it's zero where the two agree. So the KL loss is its mean over the story tokens. InstructGPT puts this penalty in the reward, per token. At our scale, its noise was larger than the spread of the rewards, and the policy learned little, so here it goes in the loss. The loss is the policy loss, plus half the value loss, plus the KL weight times the KL. And the step is fine-tuning's: zero the gradients, backward, clip them at one, and the optimiser's step. After the update, step, which counts the rollouts, goes up by one. Drift is the KL of each story from the reference: the sum of its tokens' log-ratios, from the rollout. We log the average reward and that KL to TensorBoard, print them with the minutes so far, and evaluate each ten rollouts. At the end, there's a last evaluation if the last rollout wasn't one, and the final model is saved in the run's folder, with a line for how many rollouts and minutes it took. Main runs with the minutes and the KL weight from the command line, ten and zero point zero five by default.

Ten minutes of PPO

16:30 Let's train for ten minutes with a KL weight of zero point zero five. At rollout zero, the policy is the fine-tuned model, and its stories use about twenty-seven percent of the words on the held-out instructions. Each rollout then takes about thirteen seconds on my laptop, to write the stories and make two passes over them, and prints the average reward and the KL. The KL is zero at first, as the policy starts as a copy of the reference, and it grows from there. Each ten rollouts, the demo story for rain, cake and jump shows the change: by rollout ten, all three words are in it, and later on, cake comes back again and again, as the reward counts the words and repeating one isn't penalised. By the end, after forty-eight rollouts, about sixty-one hundred stories in all, the held-out stories use about forty-eight percent of the words.

17:20 Then ten minutes with ten times the weight, zero point five. The KL stays down, at a few per story, and the reward moves a lot less, so the held-out stories go from about twenty-seven percent of the words to about thirty-two, in forty-nine rollouts.

17:54 On a chart, the reward, the share of the words in each rollout's stories, climbs with the small weight, from about thirty-four percent over the first five rollouts to about forty-nine over the last five, while the large weight takes it to about thirty-four. Each rollout scores its own instructions and stories, so the reward moves up and down by a few percent from one to the next. With the small weight, the KL from the reference grows rollout after rollout, to about twenty-two per story, while the large weight holds it at about three point five.

What it bought, and what it cost

18:23 Now evaluate dot py, on the fine-tuned model, DPO's, and the two PPO models. The words used go from twenty-seven percent for the fine-tuned model to forty-five with the small KL weight, and twenty-eight with the large one, against sixty-three for DPO. The loss on plain stories goes from two point three eight to two point five two and two point four zero, against two point eight seven for DPO, and the stories that end within two hundred tokens are forty-eight and fifty-two percent, against twenty for DPO. And after once upon a time, there was a little, the model with the small weight puts about seventy-two percent on girl, up from forty-three, so more of its probability goes to a single word, which is a sign its stories vary less.

19:07 Side by side, DPO still uses the most of the words, and PPO with the small weight gets about fifty percent of DPO's gain over the fine-tuned model. On plain stories, though, PPO's loss rises by about zero point one four, against about zero point four nine for DPO, and its stories end on their own as often as the fine-tuned model's, where DPO's ran on.

19:29 We give both PPO models the instruction from last time, rain, cake and jump, with the seed we used for the fine-tuned model and DPO's. With the small weight, all three words are in the story, cake comes back in almost each sentence, and the story ends on its own. With the large weight, the story uses cake, again and again, with no rain or jump, and it ends after a few lines.

PPO next to DPO

19:52 DPO learned from pairs that the fine-tuned model wrote once, so as the policy changed, its training data stayed the fine-tuned model's. PPO learns from the stories the policy writes now, so its training data changes as the policy does. That costs time, as each update waits for a rollout of a hundred and twenty-eight stories to be written and scored. And there's more to tune: the KL weight, the clip, the learning rate and the critic. With ten times the KL weight, the words used went from about twenty-seven percent to twenty-eight, against forty-five with the small weight.

What's real, and what's scaled down

20:26 The algorithm is the one InstructGPT used, with a policy, a critic, the clipped objective and a KL penalty, and what we scaled down is the reward, the model's size and the compute. Their reward came from a reward model, trained on rankings from forty contractors they hired, where ours is code that counts words, which is easy to game, as DPO showed us. Their largest model had a hundred and seventy-five billion parameters, and ours has fourteen point six million, with a context of two hundred and fifty-six tokens. And training their largest model with PPO took about sixty petaflop-days of compute, where each of our runs took ten minutes on my laptop.

21:04 Last, the docs get a page on reinforcement learning with PPO, with the reward, the loop of rollouts, credit and updates, and the results next to DPO's. And the README has the command to train with PPO, with a KL weight of zero point zero five.

Wrap-up

21:20 So, today our model learned from its own stories, with reinforcement learning. Each rollout wrote a hundred and twenty-eight stories, scored by the share of the words they use, and a critic learned what score to expect, so each token was credited for doing better or worse than expected. PPO's clipped objective kept each update close to the model that wrote the stories, and a KL penalty kept the model near the fine-tuned one. With a weight of zero point zero five, the words used went from about twenty-seven to forty-five percent, for a smaller rise in the loss on plain stories than DPO's, and its stories still end on their own. And with ten times the weight, the model stayed close to where it started, at about twenty-eight percent.

22:03 Let me know in the comments whether you'd pick PPO or DPO for a model like ours, and which parts of reinforcement learning 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.