Lecture 13: Diffusion LLM Training Pipeline

Key Takeaways

Okay. So the next step in these key characteristics which we mentioned and probably that's the most important step is this noise prediction. Right. Basically we need a model which can help us predict the noise in the noising process because the whole idea is as follows. Right? If you take a look at this video, if you take a look at this video again, what we want to do is that when we start rotating the fluid slowly, we want to predict how much noise is added at each step. So that during dinoising, it will it will help us recover the original state. So at each step of the noising process, how can we predict the amount of noise which is added? The way this is done in in the case of images is that you take a noisy image. You take a noisy image. Let's say this is a noisy image of a panda which is the example which we have been seeing. You take a noisy image and then you use a noise prediction model which is a unit model to predict the noise. How does the unit model work? Well, it's a combination of CNN which is convolutional neural network upsampling and down sampling layer. So it can extract features from the image. So it can identify what's noise and what's not. We'll also we also know the actual noise right during training. So we can predict whether the predicted noise is actually close to the actual noise. So loss function is defined based on the predicted noise and the actual noise. So we need a loss function, right? We need to predict the noise and then we need to compare it with the actual noise. Essentially we we need a model which can predict the noise added during the noising process and compare it with some true value. What is such model in the case of text? I want you all to take a pause at this moment and think about it yourself. Let's say you have this noisy input text, right? which is uh uh the so the next day is bright and we are masking next and is so it's the mask day mask and bright now that's my input sequence right now okay and uh so this is my input sequence this is mask and uh this is mask. Okay. How do I predict the true values here? So what I want to do what does it mean prediction of noise? I somehow want to predict the true tokens at the places where we have mask. This is exactly what is done in images also. Right? What what does it mean to predict the noise from a noisy image? It just means that we need to predict how we go from the noisy image to the clean image. This is exactly what we are doing over here. We need to predict how to go from this noisy text to the clean text. So we need to predict what's the text which was here in the first place. Again, I invite you to pause this video at this moment and think about it. How do you know what the true token was at this position? And how do you get the loss function? Well, the answer is extremely simple, right? Because we we already know what the true tokens are. We know that here the true token is next and we know that here the true token is is, right? Then how do we get the loss function? We have exactly seen how to propagate this input sequence through my model and I'll get my predictions right. I'll get my predictions for the next I'll get my predictions for each of these. So I'll get the next token prediction for the I'll get the next token prediction for this mask. I'll get the next token prediction for day. I'll get the next token for mask and I'll get the next token prediction for bride. Right. So the next token which is predicted at this position I'll compare it with the true next token which should be next. The next token which is predicted at this position. I'll compare it with the true next token which is bright. And I'll use this to compute my loss function. So this is L1 and this is L2. So my loss function will then be L1 + L2 / 2. Think about this for a moment, right? We are corrupting the text. We are masking it. I'm just passing the whole input sequence with the masks through the architecture and I'm getting the predictions. Why? How do I know I get the predictions? Because this is exactly what we saw in ARMs. If you see, if you remember what we saw in ARMs, in ARMs, for every input, we are getting predictions, right? For the next, these are my predictions. These are my predictions. I'm just going to compare the predictions for my mask with the true tokens and I'm going to get the loss function. It's as simple as that. So I'll compare the predicted token at this position. I'll compare it with the true token at this position. I'll get the predicted token at this position. I'll compare it with the true token at this position. And then I'll get the loss function based on the predicted token and the true token. And then I'll take an average of these losses. That's the model which I'm going to use. So in simple way the model which predicts the this model which predicts the noise right which is the second um second characteristics the model which predicts the noise which was a unit in the case of diffusion for images. This model is nothing but a transformer. The model which predicts the noise is the or I should say the language model architecture itself. We are going to see this architecture now just in a moment. But this architecture which I'm showing on on this side that is the model which predicts the noise. Okay. So now I'm just going to uh take you through this entire model because you'll see that most of it actually remains exactly the same. Here we have the input in the ARM architecture. We saw that there is a input, processor and output. Here also there will be an input, processor and output. Here also there is a transformer block. Many things remain exactly the same. Let's see what the differences are. Right? In the ARM architecture, we remember we had the uh token embedding and position embedding. We added them for the input sequences. Right? In the case of uh uh diffusion models, the first change which happens is that in these token embeddings, one or more of these tokens are masked. So this is masked. That's what we have already seen. So token embeddings, it's masked. Position embeddings, these remain the same. Here we don't have anything masked but token embedding we have a mask token here and which tokens to mask are is decided probabilistically but how many tokens to mask is decided by that noise schedule which we have. So this this noise should so let's say we get the token embedding we get the position embedding these tokens are now masked we add the token embedding plus position embedding the one major difference in diffusion language models is that along with adding token embedding plus position embedding we also add time embeddings right so we have time steps for the noise schedule if those are from 1 to 4 we have a 768 embed 768 create dimensional vector for each of these time steps. And remember what I mentioned for diffusion models. In diffusion models, whenever we are doing a forward pass, we not only choose a batch, we also choose the time step as well. So if you have chosen let's say time step equal to two, the time embedding, we'll go to the time embedding matrix and we'll see the vector corresponding to time step equal to two and that same time embedding vector is added to all my tokens. So now the input embedding is not just the token embedding plus position embedding. My input embedding is token embedding plus position embedding plus time embedding. There are three embedding vectors which are added together. Why do we consider time embedding here? Because the time at which we are predicting the noise, it's also important. While dnoising, we need to know the time step at which we are denoising. Right? If we are way if it's extremely if it's extremely noisy image, it's further along in time. If the image is not that noisy, it's earlier in time. This is one major change. We add the time embedding also along with the token embedding plus the position embedding. That's the first change. So now if you see if you compare these block diagrams in the ARM architecture, we only had the token embedding and position embedding. In the uh diffusion model architecture, we have token embedding, we have position embedding, and we have time embedding. And all these are added to each other. This resulting input embedding is then passed through through the transformer block. Okay. Now the transformer block remains exactly the same. The transformer block has no change. We have the layer norm, we have the multi head attention, we have dropout, shortcut, layer norm, feed forward network and dropout. So the whole flow which we have seen layer norm, multi head attention, dropout, shortcut connection here. then another layer norm feed forward neural network dropout and another shortcut connection. This stays exactly the same. The difference is in the attention mechanism. Now think about it right here. We no longer have a next token prediction task. What we have to do is we just have to predict the masks right? We just have to predict the masks. So if we are looking at this mask, I can look at tokens before this mask and I can look at tokens after this mask also. I get information of tokens behind me and I get information of tokens in front of me also. This is very different than ARMs. In ARM, if I'm looking at this position, I only have access to tokens which come before this position. But in diffusion models, for one token, I have access to tokens before it and after it. Which means that when I'm calculating the attention so the mask the mask day mask bright right which means that if I'm calculating the attention score for any token I can calculate the attention score for tokens before that token and also after that token right so if this is my attention scores matrix this will be a 5x5 matrix 1 2 3 4 5 1 2 3 4 5 1 2 3 4 5 In the case of auto reggressive models I will not I will not have so these attention scores right which I'm marking right now these attention scores will not be there I will only have these attention scores for an auto reggressive model whereas for a diffusion model I can look at a token from all sides. So if I just rub this right now if I look at yeah if I look at diffusion models I essentially have this entire causal attention matrix which is available to me. I don't have to set these values to zero like in the ARM block. So one major difference takes place in this attention module. If you look here in the auto reggressive modeling the causal attention exists which means that all the attention scores above this diagonal are set to zero. Whereas this does not happen in the case of diffusion models. We don't have such a causal attention mask. We have a biirectional attention mechanism. So all of the attention scores remain valid. That's one major difference and I want all of you to keep this in mind. That's why we went into detail in the attention mechanism. That is the second major difference, right? Uh what's the third major difference? The third major difference comes in the output token prediction, right? Uh let's say here we have every effort moves you, right? and we have predictions and we have targets and we compute the loss for all of these and we add all of these losses together for diffusion based modeling it's very different for diffusion based models let's say this is my input sequence I have to only consider the prediction at this position and I have to consider the prediction at this position and I have to compare it with the actual value right so I just have to find the cross entropy loss for my mask one I have to find the cross entropy loss for my mask two and I have to take the average of these losses. That's my loss. I don't have to find the loss for the other positions at all. Even in the example which we have chosen over here um the mask day. So this example right the mask day mask and bright right. So here I only have to care. So let's say I have my prediction here and the prediction here the true value is next and that comes at ID equal to 3. I just have to take negative log of probability at this position and here the true value is let's say is comes at ID equal to 500. So then I just have to take negative log of probability of 500 and the loss and the loss is just then equal to negative of 1 by 2 log of P3 plus log of P500. So if this is let's say 0.1 this probability and this is2 my loss in this case will be - 1 by2 log of.1 plus log of 2. So let's see how much that is. -1 by 2 into log of.1 plus log of that's 84. So the loss which I've obtained in this case is um so the loss which I've obtained in this case is 84 and then I do back propagation. So the way this entire um noising process happens now I've tried to explain it not the entire noising process but the entire flow we first decide a time point let's say time equal to three right I asked 50% of the input sequences the masked input then goes through this whole architecture what are the changes in this architecture compared to the auto reggressive model we have a time embedding now which was not there auto reggressive model in multi attention block uh we don't have the causal attention now we can take the whole whole attention attention scores matrix because for each token we can look behind and we can look ahead also okay that's why generally b models are good at sentiment analysis right because bird is also biirectional in bird models we look at a token from behind and ahead so you can look at an entire sentence and predict the sentiment of the sentence right It's kind of similar here. The only difference from the BERT architecture is that in BERT the masking ratio is fixed. But here the masking ratio is different. Right? The masking ratio uh in this case actually depends on u the time step. More the time step, more is the masking. Less the time step, less the masking. That's the difference between BERT and uh language diffusion. So we have chosen time equal to three. 50% is masked that input sequence. we choose a batch. So if the batch size is 128, we choose a batch of 128 sequences. We mask 50% of those and we pass it through this entire architecture. So this is let's say the input sequence, it first goes through the token embedding. So every token is converted into an embedding vector. We add the positional embedding, we add the time embedding and uh so token embedding plus position embedding plus time embedding is my final input vector that goes through my transformer block. So I have my layer normalization multi head attention I have my dropout I have my skip connection I have second layer normalization I have my feed forward neural network I have dropout and another shortcut connection on the outset which means at the end I have my final normalization I have this logits matrix and I have the cross entropy loss which is only computed for the mask tokens that's the second major last major difference between language diffusion and auto reggressive model diffusion here the Cross entropy loss is only calculated for the mask tokens and for nothing else. So to summarize the entire noising process looks like this. We have the let's say we choose the batch size. We choose two things, right? We choose the batch size. We choose the time step, right? Let's say batch size is 128. We assemble an input sequence of 128 tokens uh and we choose a batch size. So we randomly mask let's say 50. We choose a time step. If the time step is half of the total time, we randomly mask 50% of these. So these will be replaced with mask. Okay. Now before now these this input sequence before going through the transformer architecture we add the token embeddings plus the position embeddings plus the time embeddings. So this vector which I have shown is token embedding plus position plus time. And this this input vectors then go into the transformer architecture. When they come out of the transformer architecture, we get the categorical cross entropy loss only on the mask tokens. Based on this loss, we find the partial derivative of the loss with respect to all the parameters. And remember we have seen what what do I mean by all the parameters? Uh we have seen everything marked P over here are all the parameters. We have seen that in one of the previous lectures. This remains the same. All trainable parameters remain the same in ARM and in the language diffusion. Right? So we get the partial derivative of loss with respect to all the parameters. Then we update then we update the parameters according to ADAM optimizer or any other optimizer which you want. And then this entire thing goes into a loop. This is how I pre-train. This is how I pre-train a language diffusion model. What is implicitly happening in this case is the better we pre-train a language diffusion model the better we are able to discover what noise is added and during the generation process or during the dnoising the better we'll be able to find the underlying probability distribution right so if you look at this analogy over here. Yeah. So what we are doing right now is the noising process, right? And we are training the model to recover the noise. And the better we train this model to recover the noise, the better we'll be at denoising. So the better we so in dnoising what we are going to do is that we are going to start with all masks and then we we are slowly going to uncover the masks and recover the actual tokens in that position. That's why we get gifs such as this. Start with all masks and then we uncover tokens correct tokens at that position. But dnoising will only work if the loss function during the pre-training or understanding the noise is good. Which means the loss function becomes lower and lower and lower. In this large language diffusion models paper, they have mentioned a theoretical result here which I just want to mention. So this is the loss function during pre-training. All which I have explained right now. Everything which I have explained on the screen this this loss function which I have explained on the screen until now is explained in one simple formula over here. That's why it's very difficult to understand research papers until you have a firm background of the foundations. So this is the loss function. What these people predict is that the probability of actually recovering the true true distribution or the chances of recovering the true distribution are directly dependent on the loss function. The lower the loss function, the better we can recover the true distribution during dinoising. This is what is meant by this second formula. So dnoising is very directly linked with this loss function over here. So there is a link between the noising process and the dinoising. In the noising process, we are adding masks and we're recovering the true value at those masks. The better we do this, the better will be the dnoising also. I'm not going into the mathematical details of this but they have mentioned that um this has been proved the loss function in equation three has been proved to be an upper bound on the negative log likelihood making it a principled objective for generative modeling. So what they are saying is that the lower the loss function is the better we'll be at this uh true distri better better we'll be at approximating this true distribution. So if the green is the true distribution and the orange is the predicted, the lower the loss function in pre-training, the closer these two will lie. So the dotted red line will lie much closer to the orange. If the loss function during pre-training is lower, what is the loss function? It's exactly what we saw saw over here. The loss function essentially quantifies what is there at the masked masked positions versus what should be there at the masked positions. Okay. So, until now we have seen the first two steps of diffusion based language models. First is what is the noising process and second is how do we predict the noise? We predict the noise through the language model architecture itself which I've depicted using this diagram. So, anytime you are confused, you can take a look at yeah this diagram. This is the diagram which is essentially used for predicting the noise in the case of diffusion models. And I hope all of you have seen the differences between the diffusion architecture and the uh ARM or the reg auto reggressive architecture. Time embedding is the first difference. Then in multi attention the causality is the second difference. Then the cross entropy only for the mask tokens that is the third difference. Noisy input here that's the four difference. There are four major differences. Everything else essentially remains the same. That's why I mentioned earlier that understanding the auto reggressive model architecture is so crucial for understanding the language diffusion model architecture. Okay. Now when I'm going to share this resources with you, there is also a website in which I have explained all of these calculations from scratch. So if you scroll down below, I have explained there is a section called vector flow. So what we do over here is that we take a toy model configuration where we assume vocabulary size is equal to six sequence length is equal to 5. Uh we this is a we construct a token embedding matrix. We construct a position embedding matrix and we construct a time embedding matrix also. And I show you how everything is done. So at each position we add the token embedding, position embedding and the time embedding. Then we do layer normalization vector wise. So if you want to go into the mathematical details of how the dimensions are retained, I'll share this link with you where even I've shown everything from the attention mechanism to the shortcut connections which are also called residuals to the second layer normalization feed forward neural network final layer normalization final logits matrix and the final cross entropy loss. I have already shown this to you on the whiteboard but just for the sake of simplicity I've added this hole over here. So the entire forward pass pipeline for diffusion language model can be summarized here. We have token ids convert them into token embedding vectors add position embedding vectors add the time embedding vectors since you have to decide a time step at the start of the forward pass. Then we go into the transformer architecture. We pass through the first layer norm, attention mechanism, shortcut connection, second layer norm, feed forward neural network, another shortcut connection. Final then we come to the output header. We have the final layer normalization. We have the logits matrix and we have the cross entropy loss on the mask positions. All right, cool. And now we are ready to move to the final step which is essentially the dnoising process. So let's get to that right now.