Alignment: RL 2

Where this lecture sits

0:05 Let's get started. So before proceeding with today's lecture, I just wanted to kind of orient since we're us a little bit since we're at the end of the quarter now. So this is the final lecture that either I or Tatsu will be giving. So now let's begin today's lecture. So last time Tatsu gave an overview of RL from verifiable rewards and he talked about policy gradient algorithms such as PPO and GRPO. This lecture we're going to deep dive into the mechanics of policy gradient GRPO. So we won't necessarily cover new material but just to take the existing material and go a bit deeper show you code and some math as well. Okay. So just to kind of remind us what we're doing. So we're doing reinforcement learning for language models and in reinforcement learning you have to define the states and actions and rewards and so on. And just to be very clear I think this hopefully should be obvious by now but just to state the obvious.

RL, restated for language models

1:13 So the state is the prompt plus the response that you've generated so far and each action is generating a particular token. The next token. Okay. The reward is essentially how good the response is. And for the purpose of this class, we're going to focus on outcome rewards where the reward is a function of the entire response. We're going to focus on rewards that are verifiable. So the computation is deterministic. Some function just returns a reward as opposed to having to ask a human. And in this setting, you know, typically in reinforcement learning, you see notions of, you know, discounting and bootstrapping where it's multiple turns where you're getting and that really makes sense when you have rewards in the middle. If this is more of kind of a you just generate a bunch of actions and you see if your answer is correct. So that obviously introduces a bunch of difficulties because it's very sparse and delayed reward but it also simplifies I think some of the conceptual thinking around this. So, so for example, you might have a question math problem and then there's a you know chain of thought you think for a while and then the finally the language model might generate therefore the answer is three miles and then the reward would have to extract the three miles and see if that matches some ground truth. So in reinforcement learning typically people talk about the transition dynamics. It's very important. In this case in the language model

2:49 Setting it's very simple. It's just the plus operation. You just append the action into the state. This has a number of important implications because you can because you understand the dynamics of the world you can do planning aka you know test time compute. And this is something that you know roboticists would only dream of. If they had a transition dynamics of the world, they could simulate and they could do all these things. In language modeling, we can do that, but not true in general reinforcement learning. And the other thing to kind of realize is that the notion of state is really made up, right? If you think about reinforcement learning in context of robotics, the state is the joint angles, the position, maybe an image or something. It's really much grounded in the real world. Whereas language models the state is just tokens that the language model happens to generate and this introduces a lot of flexibility. It introduces a lot of degrees of freedom where the language model can kind of make up its own scratch pad in order to derive at the final answer. Whereas often a robot is very much limited by the external world in terms of what states are even achievable. So the issue here isn't about reaching any state which is often the a challenge in you know robotics and control. You can't even reach the state here. You can reach any state you want by just writing down tokens. But the challenge is making sure that these tokens actually lead to some ground truth answer that's you know correct.

4:26 So I think it's worth reflecting on it's all reinforcement learning but the intuitions and what's hard and what's easy kind of changes in the language modeling setting. Okay. So the policy in this case it's just a language model action condition on the state which is the prompt and response given so far. This is often fine-tuned from an existing pre-trained language model. And then there's multiple names for this rollout episode trajectory. You in this outcome reward setting you have start with a state you generate a bunch of actions aka tokens you get a single number which is a reward and then your objective is to maximize the expected reward where the expectation is taken over prompts which the environment gives you and also the response tokens which your policy pi that you're optimizing gives you. Okay, so hopefully that's you know makes everything kind of very grounded and you know precise. Okay, let me know if you have any

Policy gradient: SFT with a scalar in front

5:30 Questions about this the setup. So let's now talk about policy gradient. Which is a whole class of methods for learning the policy which essentially start with a policy and try to improve it using gradient methods. So for notational simplicity, we're just going to let a denote the entire response instead of a single action. Because we're in the outcome reward setting, it doesn't really matter. You can just think about the actions as being all generated at once by the language models. And then you get a single reward at the end. Okay. So remember, we want to maximize the expected reward with respect to the policy pi. So just expected reward is integral over the distribution over states times policy of action sequence given state times reward. So this is pretty straightforward and policy gradient says do the obvious thing you want to optimize something let's take the gradient. So you start taking the gradient so the gradient only touches the policy and this is the you know key trick I think it's like the policy gradient theorem some people call it but it's really just the chain rule for log so this basically gets rewritten as a pi and then the gradient of the log and I'm going to rewrite it in terms of expectations. So expectation is hiding basically the distribution over s and a here. Okay. So this form should be very familiar. So this is the gradient

7:13 Of the expected reward. Okay. So naive policy gradient basically says I'm just going to sample a prompt from P sample a response from A and then I'm just going to update the parameters based on whatever's inside this expectation. Okay, so think about if you're doing SGD, you look at the, you know, the expected loss and you just sample a term and you just evaluate and you know go with that. So interpretation this is essentially the same as SFT except for everything is weighted by the reward. Okay, so in SFT a human being writes down A for you and you just mimic A and update the parameters to maximize the probability of that action and now A is generated by the model and now but you get to hit it with a reward. Okay, so just as a kind of an intuition building. So suppose that the reward is either zero or one. Is it correct or wrong? So in this case the naive policy gradient is just going to only update on the correct responses. So r is going to be zero for the actions that lead to incorrect responses. So those are going to be completely ignored. And then on the correct responses r is one. So you just essentially sample you select the responses that are correct and you just do SFT on this. The only difference is that the data set is changing over time because after you do an update now you have a new policy and you generate the data set based on that policy and you iterate and so on. But this is very natural.

8:54 And you know the only challenge here is not the only challenge the main challenge is that people talk about policy gradient as being high noise or high variance. So in stochastic gradient descent if you're doing supervised learning this is also can be the case that your gradient is very noisy and typically if you use larger batch sizes and just kind of keep on going then it sort of empirically seems to work out. Reinforcement learning is a kind of

Sparse rewards and the variance problem

9:27 Whole another level in terms of noise and variance that you have to deal with. So in particular if you look at this setting where you have zero or one as a reward think about solving a hard math problem right so this in this setting is a sparse rewards the rewards are sparse which means that very few responses are going to get reward one most are going to get zero. So if you have a bad policy and you rollout and you generate actions, most of your rewards are going to be zero, which means if you look at the gradient, you're basically making a zero gradient update. So if you're in a setting where your policy is so bad that you're not able to get reward at all, then guess what? You're not making any gradient updates and you're you're sort of stuck. So that's a pretty bad setting to be in. At least in supervised learning, you're at least making an update towards something. You're guaranteed to make an update towards something. But in reinforcement learning, in sparse rewards, that's not true anymore. Question, why is the data set changing? So the question is why is the data set changing over time? So if you think about this is your objective function. So you're going to optimize. So you iterate each iteration you're going to compute the gradient and then update the parameters. So the gradient computing a gradient involves you know sampling a response and then updating. Okay. So now your

11:08 Parameters are different. So the next time you sample response, it's going to be might be different now because it's coming from the new model and if you sample multiple responses, those responses that's going to be a different data set from a different distribution than where you started. So your data set hopefully is going to get better over time. So intuitively even if you have this sparse reward as long as you can get some of the easy questions correct and you can get some reward you can update your model and you hope that the model can generalize beyond those examples so that next iteration when you're sampling and solving a slightly more difficult questions you're going to get a bunch of responses that happen to get the right answer and now your data set will have higher and higher reward you over time. Okay, so this is in contrast to reinforcement learning with human feedback where you use pairwise preferences, you estimate a reward model and the reward model is generally a bit more you know continuous. It's going to assign essentially a numerical value to every response how good it is. This is quite different than the sort of verifiable reward setting where you it's sort of like correct or incorrect. And notice that I mean I'm sort of emphasizing this point that it's all reinforcement learning. It's all even kind of within the scope of language models. Sometimes you even use the same algorithms like PPO but the intuitions for and for maybe

12:45 How you set the hyperparameters will and how you tune the model will be kind of different. Yeah. So I notic that the rewards are sparse, right? So you don't update the parameters if the reward is zero. Is that correct? Then should we just like set it to likeative one so it goes the other direction so we can use those rewards? Yeah. So the question is well in this case if the reward is zero then you don't update the parameters. Why not make the reward like minus one? So at least you're kind of pushing the model away from those incorrect responses. This is something that will happen automatically when we look at you know baselines because there this is just a simple naive policy gradient where you're updating against the rewards. We're going to show that there's much better things than just using the reward. Yeah. Question I think yeah question and then I think there was a question over there. Yeah, we know the reward only when we what's the whole response the last word of the response. So the reward is a function of the prompt and the response the whole response. Yeah. Yeah. And we know the response is correct or not only when we generated it completely. So we have only on the last token or so yeah. So you're getting at sometimes you have reward before you even finish and so the other type of reward is called process you know supervision or process rewards which allow you to basically keep track of how well you're doing. There's things you can do there and generally

14:23 That's better if your process reward is good. For simplicity right now I'm just assuming it's a outcome reward but absolutely if you have reward earlier you should always evaluate the reward as early as you can if you trust it. The problem with process reward in language models is that you just don't it's really hard to get a good process reward because unless you it's unless you're completely going off the rails it's hard to tell like how well you're you're doing. Is there another question same question like what's the reward model here that we are working with so how are we getting zero and one for that so question is what is a reward model I'll show you some examples of this yeah but you can think about the reward model is which I think for your assignment you're implementing basically the language model generates some stuff and then you parse the answer out and you check the answer against your answer key and you return one if it matches and zero Otherwise, okay, so maximize expected reward policy

Baselines, and why they cost nothing

15:32 Gradient sort of does the obvious thing. Naive policy gradient sorry does the obvious thing. But you have this high noise and variance problem. So the a key idea in you know policy gradient methods is using what is called a baseline. And so recall that this is the expected reward. If you could actually compute this expected reward, you'd be pretty good because then you could just compute that. But you can't. So you instead use unbiased estimate which usually means ripping off the expectation sampling at s and a and ripping off the expectation. This is unbiased estimate. But its variance can be pretty high. Okay. So just to give a toy example to appreciate what I mean by kind of high variance. So suppose you have just two states. So two prompts if you want S1 and S2. And from S1 you have two possible actions. A1 gives you reward 11. A2 gives you reward nine. And you from state S2, A1 gives you reward zero and A2 gives you reward two. Okay. So S1 fundamentally gives you higher reward and S2 just gives you lower reward. And of course the best policy would be in S1 you want to choose A1 and S2 you want to choose A2, right? So you don't want S1 to and going to A2. But you notice that this nine is actually bad larger than this two. So you can end up in situations

17:19 Where each state think about this as like a really s1 as being a really easy question and s2 is maybe a harder question. So you just because something gets high reward doesn't mean you absolutely should be you know updating aggressively against it. So if nine is greater than two, so if you're doing naive gradient updates, you might pick S1 going to A2 and it's like, wow, this is great. I'm going to get nine times the gradient. And then I look at S2 A2 and I say, well, that's only before two. So I'm pushing up that less. When in fact you know if you look at the whole picture it's clear that's maybe just locally not doesn't really make sense. I mean I would stress that in expectation if you look at overall states and actions it all works out the math all works out but you have to look at kind of everything to you know get that. Okay, does this example make sense? Yeah. So why would you choose A2 with the first screen? Because A1 already has 11. So the question is why do you choose A2? So you're going to choose whatever your policy decides, right? So in your policy when you're starting out, you don't know which is higher reward. So you're just going to maybe choose something random, right? So if you happen to choose A2, then you might upgrade that and get reward nine. I mean, it's actually worse because you say, "Wow, reward nine, this is great." and you the more you choose it, the more you're updating against it and then you

19:02 Sort of get stuck in this local optimum where you can't even choose there's no probability of mass left on A1 because you've updated against A2 so much. Okay. So, so here's a key idea behind baselines is that you're just going to maximize the baseline reward. Instead of expected reward, we're just going to subtract off some b(s). Okay, call the baseline, you know, function. So this is just the expected reward shifted down by a constant expectation of b(s). And the beauty of this is that this quantity doesn't depend on the policy pi. So even if you expand it out you'll have a x you know integral over p of s times pi a given s but integral over a given s is just one. So this is just you know it doesn't depend on the policy at all right and so of course optimizing expectation r is identical to optimizing expectation of r minus this because it's just a constant. Okay. So that's the whole idea of you know baselining it preserves the sort of optimization problem. And now if you do sort of naive policy gradient on this new objective then basically you sample a s you sample a and then you're just looking at you know this difference. Okay. So notice that this is true. This works for any b(s) as long as b(s) doesn't

20:48 Depend on A which you know it's not in arguments or can't depend on A. Then this is valid. So now the question is well over all possible functions b(s) which one should we use? And let's just revisit this two-state example just to provide some intuition. I'm going to make some kind of assumptions. You assume you your initial policy is a uniform distribution over all s and a and your gradient of your policy initially is one just to kind of ignore that for now. So if you look at the variance over your gradients you'll see that some of the time you get 11 sometimes you get nine and sometimes you get zero and sometimes you get two and the variance is 5.3. Okay. So if you define a baseline which is 10 for S1 and one for S2 then let's see what happens. So you subtract off the baseline from these reward values. So for S1 you subtract 10. For S2 you subtract one. So now the baseline variance is a lot less. It's only 1.1. So we reduce the variance from five to one which is pretty good by just subtracting a baseline which again doesn't affect your you know end results in expectation. Yeah. So I think there's the assumption that like this expectation of b(s) doesn't

22:27 Depend on the policy pie. So I was wondering like if it is okay still for us to initialize like you know B as let's say like the initial the very original frozen policy. Yeah. Yeah. So good question. So b(s) is not supposed to depend on policy. But your question is can I just take an initial policy and set b(s) to absolutely. So it's one of these things with reinforcement learning you have to be a bit careful of is what's constant and what's not. So you can always take a frozen model and say that's a constant and as long as your expression doesn't depend on that's actually treated as a constant like you're not differentiating through it then it's all fine which means also that if you are sampling and you're performing the update that you have to kind of keep a old copy of that policy because if old if that thing is changing now you could get into these weird settings where you're Not actually it's not actually a constant anymore. We like what if we initialize as the policy from the last step the last reading step instead of like a very variable. Yeah. So you can you can do this. So typically what you do is you essentially have a policy. Kind of a an old policy or you know anchor for a while or if you're doing PPO you have a value function and then you kind of hold it for a while and then you do some updates and then you kind of update it on so I mean in practice is what you do part of this like it can't depend on a

24:07 Is really like for the theory kind of to work out but in you know practice there's you know there's all sorts of things that people do. Okay. So what we showed is that at least for this example anecdotally you can define some baseline that reduces the variance. So this is kind of giving you hopefully a concrete sense of the power of a baseline and the lower variance you have the faster the convergence. Okay so you might wonder what is the

The optimal baseline, the heuristic, the advantage

24:41 Optimal baseline to use and if you do some math realize there's actually a very kind of nice closed form you know solution for this is written for one parameter models if it's multi-dimensional I think you have like some covariances and this is just expectation of this derivative squar time r divided normalized by the same expression without the r. Okay. So, so this is in general kind of annoying to compute especially if you have high dimensions and because this thing is going to be some sort of like you know you know covariance so typically what people do is you just ignore this stuff and you say well I'm just going to try to I'm going to set baseline equals expected reward divider given Okay, so this is not optimal but it's kind of like a you know good heuristic and even this is you can't compute this right because anything with expectation you can't compute you can only estimate but this is now I think a good guiding principle for what you want your baseline to be it's the expected reward given your state okay so just to recap we said we have the naive policy gradient we can introduce this b baseline it's valid for any b(s) and now we've come to a conclusion that maybe this is like a good choice if we can try to get a

26:21 Handle on expected reward given s okay so this particular choice has connections to advantage function so in reinforcement learning you typically have a value function which is the expected you know reward given a state. And then you have a Q function which is expected reward given a state and I've taken a particular action. And note here that Q and R are the same because we assume outcome rewards and A has all the actions. In general, R is going to be like the returns from that state rather than just a single reward. Those two are same for now. And the advantage function is just defined as the difference between the reward I would get if I took action A and sort of if I you know average over possible actions A. Okay. In particular it's asking it's measuring how much better is action A than the expected value that I get from S. What is the advantage of taking A as opposed to just going along with a baseline? Oh no, sorry baseline is ambiguous. We're going with a policy. Okay. So then if you look at B equals this sort of heuristic form that we talked about earlier then the baseline reward where you just subtract off B of is identical to the advantage. So you can interpret you know the

28:01 Baselining if you have the advantage sorry if you have this sort of this choice of baseline is identical to optimizing the advantage. All right. So, so that's just to give you some intuition behind what this baseline is doing. In general, we are probably going to deviate from that because we can't estimate any of we can't compute these expectations exactly. So, we're going to do something else. So, in particular, we're just going to I'm just going to let delta denote whatever thing we put in. To the estimate. So all of the methods are going to basically take the gradient of the log policy and times something that's going to be based on the reward. Okay. Okay. So all the algorithms you see you PPO GRPO have this sort of form where you're basically you rolled out and got an action and now the question is how much do you want to move towards that action and it's going to depend on you know the reward. Okay. So you can put the reward here then that's not your policy gradient. You can do reward minus some estimate of the expectation and you can also do divide by the standard deviation which is what GRPO does which we'll see later okay but this is a sort of the general form and so the exact form of GRPO isn't that

29:39 Important because you know next year there's going to be GRPO2 or whatever and but the all this like policy gradient stuff is I'll be giving the same lecture next year. Okay. All right. Speaking of lecture, you can check out Chelsea's deep RL class for a few more you know derivations and intuitions about RO in general. Okay. So any questions about policy gradient and baselines and advantages? Yeah, just for the baseline like a baseline reform. That's what so what's the definition of a baseline? Baseline is any function that b(s) which only depends on s and not a that you're essentially sticking into a policy gradient algorithm. You can think about it that way. So

Why GRPO is a language-model-shaped algorithm

30:48 Basically you update policy gradient will update based on this and baseline is whatever function you choose there. So those of you have seen kind of control variates this is basically the same idea from you know statistics. Okay. So now let's try to dig a little bit you know deeper. We're going to actually define a simple model, a simple task, and actually walk through some code that may or may not look like assignment 5 and we'll try to hopefully do some runs. Okay, so a lot of this exposition is based on the GRPO, you know, algorithm. So this slide I might have to update next year. But it's sort of trying to generalize it beyond the particular choices here. So GRPO remember what Tatsu said last time the it's sort of kind of interesting lineage where you had PPO and then people did DPO and GRPO is a kind of you know simplification of PPO as opposed to kind of algorithms you start simple and then kind of grow in complexity but really I think GRPO exists in 2024 as opposed to you know in 2017 or something when PPO was introduced because the language model setting provides some kind of natural structure that motivates GRPO. Can't really do GRPO naturally I think in you know classical RL settings with like the robot trying to walk and the particular thing is that in the

32:34 Language modeling setting you have a prompt you can generate a bunch of responses from that prompt and this gives you a natural group structure to compute the essentially the baseline which turns out to be the average over the rewards right the whole point is that you're trying kind of minimize the variance and the group is a kind of a natural comparison set. That's why this is kind of relative. It's relative to the group. Whereas if every rollout looks different, there's no natural group structure. Then GRPO doesn't really make sense and you have to find something else and that something else generally tends to be a value function which aggregates over all the possible states and actions. And here this case you kind of have an empirical estimate over all the possible responses without considering everything because of the LM setting. Okay. Okay. So here's the pseudo code. We're going to you know look over it in actual Python in a in just a second but I'll just flash it up

The sorting task and its two rewards

33:37 Here. Okay so let's define a simple task. So you know I wanted something simple enough where I can actually train a very simple model on my laptop without and running a language model transformer wasn't going to I mean I don't have a fancy laptop so that wasn't going to happen. So here's a simple task sorting n numbers. So the prompt is n numbers and the response is n numbers that are hopefully sorted. So now you have to define a reward to complete the definition of the environment and the reward should capture how close to the sorted ground truth the response is. Okay. So immediately you actually see one of the interesting things in reinforcement learning is that you can define a reward in you know different ways. So, for example, you can define a reward, which is like one if it's sorted and zero otherwise, which is presumably what you care about, but you're going to remember the sparse reward problem. If you start out with a policy that doesn't do anything, you're mostly going to get zero reward and you're not going to really make much progress. So, let's define a reward that gives partial credit. So, we're gonna define reward that returns the number of positions where the response matches the ground truth. Okay, so you basically the code looks like this. You look at the ground truth which is the sorted of prompt and you just count the number of positions where the response matches the ground truth.

35:15 Okay, so the reward for the correct answer is four because four positions match. If you have this thing that doesn't really solve the task then reward is one because I think you got lucky and you put a two in the you know second position and you have this thing which also is not solving the problem but you know you kind of look at this and say well it's it's really trying and again you get a one because the zero is in the right position. Okay, so that's one reward you could use. Maybe you're kind of unhappy with this because, you know, this is pretty far off and this is closer, but yet they're all getting reward one. So let's define alternative reward that gives more partial credit. So what I'm going to do here is I'm going to give one point for each token in the prompt that shows up in the response. Okay, so at least you'll you know output the tokens that are in the input. Not too much to ask, right? So for this particular example, I get inclusion reward of four and then I'm going to get one point for each adjacent pair in the response that's sorted. So that's a three because there's three pairs adjacent pairs and the reward is the sum. So that's seven. So, if you're kind of paying attention, you realize that this reward is actually has a loophole in it. But I'll let you

36:57 Figure that one out. Okay. So, now this particular very wrong answer gets a reward of three and this one gets a reward of six because it's really trying its best and we give a lot of partial credit for our algorithms. Okay. So, so that's the task and you know in general we're going to go with the second reward function. Okay, so the task is clear that's an environment. So now let's define a model to solve this task. And so we're going to define

A deliberately tiny model, and the shapes

37:29 A really simple model. We're going to assume the prompt and the response length are fixed and of the same length. When I was trying to iterate on examples, there was one I was doing some sort of, you know, dynamic thing and it was just like annoying to do. So I've simplified it. I'm going to try to capture positional information with separate per position parameters because obviously position is the kind of the name of the game here or position matters for sorting obviously. And one thing I'm going to do is decode each position in the response independently. So this is not an autoregressive model. I'm just going to encode the sequence and I'm just going to output all the tokens separately. Not what you want to do generally in real life except for I guess in some speculative decoding. It's can be helpful to do this in limited doses. But it makes the code a lot cleaner because this autoregressive generation that's a part which was getting really hairy so if I simplify it okay so here's the model so I have embedding matrix I have for each position I have a matrix for basically transforming the for decoding and a matrix for sorry encoding and then decoding okay So that'll be easier to explain as I look at the forward pass. Okay. So I start with a prompt. I'm going to use this notation in general to denote the dimensions of the of the tensor here. So there's a batch dimension which is the number of examples and number of

39:13 Prompts and then a position dimension which is obviously the position I'm in. So I'm going to generate responses. So this is inference. So let's walk through this. So just some you know terminology. So I have a batch size which is the number of prompts here. Each prompt has a particular length three. Number of responses per prompt. I'm going to generate in this case two. And the response length is four in the sorting task. You know this response length is going to be identical to the prompt length. Okay. So, so first I'm going to pass the prompts through the model to get logits which is going to give me for each prompt and each position a distribution over tokens. So what this looks like is I'm going to embed the prompts into this dimensionality. Now I'm going to transform using the per prompt position matrix and I'm going to collapse it. So using einsum I have batch position dimension one and then this weight is per every position I have this dim1 to dim2 transformation and I'm going to end up with batch dim2. Okay, so I've kind of summed out the position and dim1 and then given that vector for each prompt I'm going to for each position apply this matrix and get back a

40:57 Vector and then finally given a vector for each position in prompt I'm going to convert that into a logit over my vocabulary. Okay. So, you know, pretty I mean this is kind of a you know simple model I made up. I don't know if it is actually any useful but anything useful but it fit this sort of did this toy example. Okay. So now I have logits and this is just kind of shenanigans to deal with the fact that torch.multinomial needs something flattened. So I'm going to essentially basically flatten this into a batch times position by vocab. I'm going to sample num_responses and that's going to give me batch position number of responses which I'm going to call trial is that dimension and I'm going to you know basically rearrange things to make it so that it's for every prompt for every trial I'm going to get a sequence of examples. Okay. So maybe I'll pause just to check that everyone's, you know, with me. So this is this is nice. Be this is usually where you would call out to vLLM, right? And do the inference. I'm just doing this here because it's like a very simple model and I can just everything is just like matrix operations. Yeah, your responses are going to be very close by, right? You're kind of sampling from the same logs. Responses are close by. The responses are going to be close by. So

42:39 We're sampling num_responses independently from that distribution. So they're definitely going to be similar. But that would be the same as if you sample from a language model as well. The only difference between a language model is autoregressive and this is independent. The point was shall we change the temperature basically? Yeah. So you could change the temperature. You can do various things. I'm trying to be kind of faithful to the algorithm. Okay. So I get my responses. I can

Deltas: four ways to turn rewards into updates

43:14 Compute a reward over all these responses. And this basically goes through and for every prompt and every response I just call this reward function which we looked at earlier. Nothing too interesting there. And sorry, this is, you know, a little bit long. Oh no. Okay. Well, this part is going to be far less interesting because I thought I fixed the random seed. Just generally, you wanted a little bit more variance in the rewards. Anyway, so the rewards all happen to be three. So, this is going to make some of this other stuff kind of less interesting. But nonetheless I'll go through the computation here. So there's this function called compute_deltas which takes the rewards and converts them into the delta that I'm going to use for my updates. And the framework is there's many ways many choices here which some of which you'll explore in your assignment. So for example if you can just return delta equals the rewards you can you can center the rewards. So you can compute the mean over all the responses for each prompt. So this is what this is doing and then you can just subtract off the mean that gives you center rewards. And to answer a question from earlier, this is where if your

44:55 Rewards look like one and a bunch of zeros, your average is going to be like point, let's say 0.1, which means that if you subtract that, now you're going to end up with a bunch of minus ones and something that's slightly below one. And therefore, you're going to perform updates on all those negative examples as well. Okay, so that's centered rewards. And then you can do this thing in GRPO where you normalize. So you center the rewards and then you divide by the standard deviation and add something small so you don't divide by zero just in case and you return that. Okay. So one thing to notice is that what happens if all your rewards are five for example. So rewards is just going to be five. So you're going to make a bunch of updates. Centered rewards is going to be all zero, right? So that means you're not going to make any updates. So this is kind of interesting to think about. If you have an example where your policy is going generating everything is the same reward, then you're actually not going to make any updates on that example. And the intuition, you know, it might or may not make sense depending on what you're you're trying to do, but the intuition is that there's no, you know, gradients. You're not there's in terms of relative comparison between that group, there's no reason to favor any particular example more than other one. So, you're just going to abstain and hope that some

46:33 Other examples give you some signal. And the normalization means that if all your rewards are let's say 10 versus all the rewards are you know if you multiply all rewards by 100 then the normalized rewards is kind of invariant to the scale of the reward. So I think that's the you know good thing about that. Just for fun there's another thing which the GRPO paper didn't do but you know you might think is kind of a cool thing to do. So you can zero out any reward that isn't the maximum for each batch as well. So the idea here is that often in reinforcement learning sometimes you really get into these local optimum where if you assign too much partial credit then the policy is going to sort of latch on to that and go after the low-hanging fruit instead of the kind of you know the thing that would give you more reward. So you could try to kind of just say well if it's not getting you know the maximum reward then I just don't want that. It's kind of all or nothing thing. Okay. So far what we've done we generate a bunch of responses we compute rewards and these deltas which allow me to provide the update. Now I need to look at the other piece which is the log probabilities of the responses because remember in the gradient policy gradient I have delta times the gradient of the log policy.

Log probs, the naive loss, and freezing p_old

48:12 So this part is also fairly you know straightforward. I pass this through the model. Convert the logits into log probs. And then there's some shenanigans to because a prompt is doesn't is only batch pos and responses has a extra dimension in there. So I'm just going to kind of upcast the these log probs to that higher you know tensor dimensional tensor and I'm going to index you know log probs using the responses. So the responses basically for every batch trial position specifies the index into log probs that I need to get. So this is batch trial pos vocab and this is batch trial pos with an index and that will give me batch trial pos. Okay. So this gives me so logits that are returned by the model give me a distribution over all the vocabs for every position and prompt and now I have log probs for the particular responses. So these are log probs for these particular responses I got. Okay. So now we want to finally compute the loss and the loss is just going to be and here there's another set of choices here. So I'm trying to present this as kind of a menu of sorts. So the most naive thing is this is naive policy gradient. You just

49:54 Take the deltas and the log probs and you just multiply them together in average. Okay. So this is log probs. Remember batch trial pos and deltas is batch trial. And this is because it's outcome reward. You only get there's no position dependence on the deltas. So basically every position gets the deltas get broadcasted to every position. If you had process reward then this deltas would be kind of per position as well. Okay, so that's the naive thing. Just a kind of interlude to make sure that we kind of into this idea of freezing parameters. So later we'll see in GRPO you see these kind of ratios right? This kind of relates to this idea that you have a probability of action given s and you're dividing by some kind of you know other model and it's really important that you kind of remember to freeze and not differentiate through p old. So just give a kind of toy example. Suppose I have a single parameter W and that drives both the probability think about this as a policy and also there's like a old value which in this particular iteration is just the same thing. And now what happens if I compute the ratio and do a backward pass? What would the gradient be? So I'm dividing P by p_old. So the ratio is going to be one and the gradient of one is

51:36 Zero. So where are the okay so it's it's zero which is obviously not very helpful. So the way you do this properly is that you have your P as usual and then anything you want to kind of treat as a constant you put in wrap in this no_grad and then this allows you to get the value of but not the whole computation graph and you just divide by that and now the gradient is not zero. Okay. So just like a kind of be careful when you're because in RL as opposed to when you're doing pre-training there's many different values that all you know depend on different either old or not old parameters. Yeah discount like an outcome when we go back to the secret I guess you probably have a hypothesis that like the last 100 tokens first. So the question is why don't you discount do some sort of discounting when you have outcome rewards. So and maybe the last 100 tokens matter more than the first 100 tokens. I think that is less not so clear. I mean certainly like the last tokens where you're generating the literally generating the answer is important but I think the problem is that in reinforcement learning in general you have to do credit assignment right there's certainly decisions early on that matter in fact they even matter more than what you get because if you basically have the right strategy then the stuff

The clipped loss and the KL penalty

53:17 That happens at the end maybe just follows out and the hard part in reinforcement learning is that you don't know where to do the credit assignment. Especially when you have a sparse reward. So I think in this setting we're just going to smear the credit all over or the blame all over the place. Okay. So with that kind of interlude so this is important for computing the GRPO loss. So imagine we have an old model okay with old so this is an old checkpoint and using this old model we're going to compute the log probabilities of the responses so that gives us the old log probabilities and now let me just skip the unclipped version so this is the GRPO loss okay so let's step through it so it's the clipped version. So you look at the log probs and you divide by the old log probs, which remember should be a constant. Actually I sort of didn't follow my own advice. I don't think I wrapped the old log probs in with no_grad, but anyway, you know how to do that. So you look at this ratio and then that ratio gets multiplied by delta. Okay. And so if you contrast with the naive thing, you're just multiplying log probs times delta. So this is basically the same thing except for I've divided by old log probs which is constant. So you know the I'm just

55:00 Scaling the gradient. But now there's the other part which is I'm clipping. And the reason I compute this ratios in the first place is that I'm going to clip the gradients between 1 minus epsilon and 1 plus epsilon. So this is to make sure that the ratios don't go too out of whack. And I'm going to also multiply these clip ratios by delta. And I'm going to take the minimum and you'll return. And there's a negative sign because I'm talking about rewards and convert that into a loss. Okay. So, so this is just following the kind of the math in the you know GRPO paper and in the assignment you'll see this kind of explanation of why this makes sense. So if your updates are within the you know the clipping range then you just get the kind of the standard you know update but if you sort of your ratio is outside then your the magnitude of your updates are going to be bounded by 1 minus epsilon and 1 plus epsilon. Okay. So there's also this KL penalty which provides additional regularization. And so the just a kind of a quick you know little math exercise here. So remember the definition of KL between two distributions. It's your sample from the first left distribution and it's log

56:40 Over the left over the right. Rewrite that as you know q/p with a negative sign and then so if you're doing you know KL penalty you could just you know use this directly right but there's this and if you want to estimate the KL you can just take strip out the expectation and you get a unbiased estimate turns out that this particular estimate that can have you know lower variance where you just do q/p minus the log minus one. So it looks a little bit weird but you can sort of check that this is identical. Everything's equal because expectation q/p is just one. So this part is just one and you subtract one which is zero. So you're just left with this original term. So the math works out and the inside just gives you a better estimate of the KL than doing this kind of naive thing. So a lot of cases when you're trying to find these estimates of the expectation you're just trying to find an unbiased or sometimes even biased estimate that has lower variance than just doing the naive thing. So you kind of try to rewrite the inside to so that it reduces the variance. Okay. So just a quick implementation of the KL penalty. This is pretty you know straightforward. The only thing you have to pay attention as you're summing over the last dimension but taking the mean. So you're summing over the vocab

58:22 But you're taking a mean over batch trial and position. Okay. So just to kind of summarize the so far the components that are needed to introduce the actual algorithm you generate the responses given a particular fixed model from the responses you compute the rewards and you know the deltas and in the case of GRPO and friends these don't depend on the model it's just like a function of the you know responses and in you know if you think about kind of architecting this system sometimes the reward is very simple it's just like a oneliner exact match sometimes the reward will involve like an agent like executing an environment and doing all sorts of crazy things so this could be you know expensive or cheap operation but it's independent of the actual model unless you have like LM as a judge which case that's another model simultaneously you can compute the log

The full loop: three models, two loops

59:31 Props of the responses. And then you compute the loss from the log probs and you know the deltas that you got here. And you combine them either clip or not clip or you normalize or you sorry that's that's up here. So you either clip or don't clip and then you take a gradient you know step. Okay. So let's try to put this all together. So this is the full algorithm which is essentially a implementation of what's in here in code. So there's some boiler plate I'll just skip over. So I'm going to define you know some trivial data set here. Three prompts of length three each. And then define the model optimizer and then I have this there's multiple loops here and you know GRPO PPO so there's an outer epoch loop and here if I'm doing a sometimes if I'm using the doing a KL penalty then I'm going to kind of freeze the model and have that as a reference so I can regularize towards that. And then otherwise I sample the responses. So we talked about this you compute the rewards you compute the deltas depending on you know what variant flavor you like and if you're doing the you know the KL penalty then you also need

61:19 You know to compute the log probs under the reference model and if you're doing not the if you're doing some sort of clipping then you need to compute also the log probs under you know the old you know the current model and then you take a number of steps given the responses. So the idea here is that you know often response remember inference is expensive you generate and you want to do a bunch of gradient steps with respect to the responses. Okay, so this inner loop is basically taking a bunch of steps with respect to the same set of responses. And to do that, you compute the log probs, you compute the loss, if you want, you can add a KL penalty, which sometimes helps, sometimes doesn't really make a difference. And then you just optimize. Okay. So, so that is you know maybe I'll stop and ask answer any questions people have about the setup here. So the big picture is that you have the outer iteration which you ignoring the KL penalty because that's there's another wrinkle I'll mention later. So you generate a bunch of responses and then you quickly do a bunch of iterations on them and you generate a bunch of responses and do some

62:58 Iterations. That's kind of the basic idea. Now there's an additional thing which if you have the KL penalty then you actually regularize towards a reference model which is changing even slower than this outside model. So the way this is written is actually slightly different than the GRPO code. In the code, you actually have three loops. I only have two because I'm sort of just updating the reference model. You know you know a little bit kind of infrequently and I'm also not doing you know this part. Okay. Okay, so there's a there's a few implementation details, but hopefully you could kind of get the rough shape of what this looks like. So you actually have kind of three models. You have the reference model that you're doing KL regularization against and then you have the old model which you're using to compute these like kind of importance ratios between the current model and the old model which you're also freezing and then you're you know kind of updating you know the current model online. So one kind of annoying thing about the reference model is that you actually have to store a whole copy of the of the model which is can be you know just doubles your memory. The old model is actually a little bit, you know, better because you just remember if in the code you actually don't have to copy the old model, you can just generate a bunch of responses and

64:41 Compute the log probs on those responses and just store the log probabilities because those are going to be you're dealing with the same responses in an inner loop. So you don't have to store a whole copy of the model. You just have to store the log probs. Okay. Question. Why don't we do KL divergence between old policy and the latest policy instead of So question is why don't you do the KL penalty against the old instead of the reference. The way to think about this is that you let's see so remember how we said we define objective function in reinforcement learning and we optimize that. It's sort of a little bit of a lie because really you just define this objective function so you can compute a gradient and the gradient you compute and you update for a while but then your kind of objective function changes right so really you can think about let's say so this point you fix a reference model and then you write down this nice objective function which is let's see if I don't have the full expression here but you know the expected reward with all the you know the expected reward and then the KL you know penalty so if you change that too

66:24 Quickly then it's I guess the more you change that then the less you're actually like optimizing that objective function. So that's I guess maybe one kind of reason. I mean honestly I think you know empirically you know you it's sort of a gradient because depending on how you how many steps you take here and how often you can make it either more like pi_old or less like pi_old. I mean you sort of have two kind of mechanisms for regularization. One is the clipping thing which is applied to pi_old and then the KL which is applied to pi_ref. Okay. Yeah. Question for the first iteration of for the most inner loop isn't pi_theta_old equal to pi_theta is it like small to do that? So the first iteration pi_old is equal to pi_theta. So you're essentially going to well you're still performing an update.

The runs, and why the loss curve lies

67:33 You're just scaling by some remember pi_old is from the point of view of the update is just a constant. And I guess the point only thing that matters is that you're not going to get clipped because the ratio is one. So you're just going to guarantee to make an update. So yeah, it's fine. Okay, so maybe just kind of round this out a bit. So actually so I'm run some experiments. So you run the same piece of code that I presented. I'm going to do 100 epochs. Number of steps for epoch is 10. And we see this kind of learning curve. So just to maybe walk through what this is doing. So on epoch 0 we have this prompt the responses it generates a bunch of stuff. You score them you compute the reward and then you compute the you know the delta. In this case it's identical to the reward. So I'm going to update towards the threes a little bit more than you know the twos. And then you see, you know, this prompt, you know, kind of the same thing. This one, I'm still going to update against all these, towards these, responses even though they all have reward three. And then let me scroll down and see what happens. You know eventually

69:13 You see that I'm getting reward three but you know this is not you know looking I mean it's not really sorting is a short answer it's it's producing numbers which are sort of more or less in the prompt but not really sorting let's see if it how it does at the So it's it's not terribly great. I actually haven't looked at the other examples, so it'll be exciting to see what happens here. So here we're using the centered rewards. And see the mean reward is getting higher. So it's like, you know, past three instead of two, which is promising. So let's look at this run. So notice here that I generate some examples reward is two and three but because I'm centering I'm sort of giving pushing the model towards the ones with reward three and away from the ones that have reward two. So this is what centering is doing. And if they're, you know, all the same, then centering doesn't really, is actually going to be a noop. Whereas before I was performing updates, which is which is kind of silly, I think, because if you think about it, I'm sorting one, two, three, and here all the responses are, you know, equally bad. So updating towards this, these don't really maybe make as much sense.

70:54 Let's see towards the end how this does. So still not I guess it's not great, but it's got a little bit higher reward. I probably need to tune this a little bit more. But at least it's getting you know, reward four here for getting the right numbers. And 01 is, you know, it's not the most unsorted. Here I'm I think it's sort of getting the permutation right. And here I think it just gets stuck. And at this point you see all the deltas are zero. So like more training isn't really going to help on this batch. But if you regenerate, maybe you'll get some fresh samples and be able to update. So you can see that it's kind of like the partial reward is sort of double-edged sword, right? Because I'm giving you four points. So it's like it's, you know, it's good reward. But I still didn't solve the problem, you know, completely. So I think you have to be careful with you know reward dishing out reward. And then finally the normalization here didn't really do all that much so I'm going to skip it. One thing to I think note is that if you look at these loss curves, you think, well, this looks kind of good, but why is this like looks pretty bad even though the reward is going up? And this is goes back to kind of my statement that minimizing the loss is a little bit of a lie here. It's not like we were min

72:33 Minimizing one loss function that you can measure over time because the set of responses is changing over time. You don't have really kind of a reference point for the same loss. I mean you can measure loss on some maybe some validation set. But if you only have rewards then you're sort of back to only being able to rely on reward because even think about it even if you if you have a bunch of examples where you're the model is just happy generating them. It's sort of like the loss is only respect to what the model is generating. So it's a kind of a self circular thing. So monitoring the loss is not really kind of a so such a meaningful thing to do. Okay. So maybe just to wrap up here. So you know reinforcement learning just popping back up. I think it's really exciting because it's really I think the key to making models that actually you know surpass you know human abilities in some sense because label data is only going to get you as far as mimicking the behavior that's presented in the label data. You know I think the model is if you can measure it you can optimize it if you have a reward. I think there's you know optimization as we've seen even with this sorting examples I mean is can be challenging in itself but there's also a perhaps a bigger problem is you know how do you design rewards that are not hackable in environments that are you know generalizable I think that's a pretty

74:14 Active area investigation for the optimization part I mean you know the sorting isn't that hard I think I just ran it for like 20 seconds on my laptop so you don't you shouldn't expect greatness. I think you know just review the policy gradient framework I think is conceptually clear. You write down the expected reward. You subtract a baseline. It's the same thing and then you just you know take an you try to find a you know estimate of that reward or advantage. And that's one thing that we didn't talk about which I think is actually important and maybe next year I should try to do a bit more of this is that building our systems and scaling is actually much more complicated than you know pre-training and the reason is that you have to do inference and inference is its own can of worms you also have to manage multiple models so you have the policy you're training you have if you have PPO then you have the critic if you have you know the inference and you have maybe a inference worker then you need to ship the model over there to run the inference. If you have like an agent environment you have like the environment which requires spinning up and then on top of that you also have like the pi_old and pi_ref. So you have multiple models that you need to keep around and on top of that you have to do everything kind of in parallel and distributed and make it all sing and dance. So there's a lot to be done there which you know we're not really going to have time to talk about in this class but if you're interested you know there's definitely pointers I can send you to. All right

75:59 So that's it and see you next time.