Lecture 12: Diffusion LLM Noising Schedule

Key Takeaways

So what we aim to do is as follows. We have a true probability distribution, right? So let's say the meaningful sentences look like this in probability distribution space. We want to get a probability distribution or we want to predict a probability distribution which is as close to this true probability distribution. But we want to predict this not with an auto reggressive model but we want to use a diffusion model for getting this probability distribution of the predicted sentences. Right? So we will integrate the three main characteristics of diffusion models. First we need a noising process. Second we need to predict the noise using some sort of a model. And third we need a dnoising process. The noising process in the case of images was adding noise to an image using a noising schedule maybe a goshian schedule. The predicting noise was the ML model called unit in the case of images and the dnoising process which they had was just uh reversing uh the noising process backwards so that we can generate images from noise. Again I want to bring your attention to this visual which we saw in the dnoising process. What happens is we start with complete noise and then we bring it back to the original image. Now let's start implementing these characteristics step by step. So how do we first of all let's start with noising process. How do we integrate noising process in the text? Okay, this is where our understanding of the ARM model is going to play a very very crucial role. Right? So let's see how noising is included. What is the noising process in the text? Right? And for this the first thing which is going to be extremely crucial for us is that in the case of images we had an image right which we want to noise or which we want to subsequently add noise to. Um let's say the image was that of a Chinese character. That was an image. We can successfully go on adding noise to this image. We add noise. So this is for images. What about text? How do we go about adding noise to a piece of text? So how do we corrupt a piece of text? So the whole idea would be to take a bunch of text to corrupt it with noise to predict that noise and then when we start with noise, we can predict the original text just like how we did in images. So first is how do we add noise to text? So let's think about ARMs, right? What happens in ARMs? ARMs have this uh ARM have this architecture, right? And here I'm again it's the exact same architecture but just a different uh visual here. So what ARMs do is that you have this piece of text which is passed into this architecture. Right? Let's say I have this piece of text called the next day is bright. Okay, I have tokenized it. So let's say the token ids are 1 11 20,000 55 and 3,000. And then what I have done is that I have converted it into embedding vectors. That's the next step. So token embedding. So this is a 384 dimensional vector. This is a 384 dimensional vector. I'll just mute my discord. This is a 384dimensional vector. This is a 384dimensional vector. And this is again a 384dimensional vector. Right? This is how it works for auto reggressive models. Now the question which I'm asking to all of you is that I want to start with the first step of these three three characteristics which is the noising process. How do I corrupt the given piece of text by adding noise? What can I do over here? So that noise is added. Now what does noise mean? Noise essentially means that this piece of text right now let's say I have the actual distribution of true sentences. Right? This is the true true sentences probability distribution. Let me this is the true sentences probability distribution. This and this sentence the next day is bright belongs somewhere here. The next day is bright lies somewhere here. So the probability associated with this. So if this sentence is X, probability associated with X is very high because this is a good sentence. Now I want to corrupt this text which means I want to take it to spaces. So if this is my probability distribution space, I want to take it to spaces here where the probability is very low of it being a true sentence. How do I do that? How do I corrupt a sentence? The best way to corrupt a sentence is to just remove portions of it. Right? If I remove portion of this, if I remove a portion of this, this is no longer a true sentence. It becomes a cor, it becomes a corrupt text or it becomes a noisy text. Um, this is exactly how noising is done at the sentence level. So we have input sequences or we have sequences of tokens, right? In diffusion language models, we corrupt this text by introducing masks. What this means is that if I have the next day is bright, I will mask this input sequence. Which means I will let I will say that the input sequence is this. The mask day mask bright. This is basically corrupting the input sequence by adding masks. Right? Uh now how do I add these masks? The the way I add these masks is by defining something like a noising schedule. Right? So first I define a time step. Let's say I go from time equal to 1 till time equal to 6. And I define the number of tokens to be masked uh proportional to let's say my time let's say t divided by so if 6 is capital t divided by capital t so if my t is equal to 1 um I will approximately mask 16th of the tokens If my t is equal to 2, I will approximately mask 26 of the tokens. If my t is equal to 6, I will mask all of the tokens. So you see at t equal to0, I have my entire clean sequence which is not masked. So that's the correct sequence. Think of it like lying in the correct probability distribution space. As more and more noise gets added, I move out of this probability distribution space and start going haywire. So maybe initially when only one mask is added, I go to a lower probability. When two masks are added, I go to even lower probability. When all the masks are added, it's complete noise. So I go to a probability distribution or I go to a space in my probability distribution space which contains no meaningful text at all. I go to a random location basically. Uh so if you see diffusion noise schedule images and maybe let's see a gif. Noising schedule. Noising schedule diffusion. Yeah, take a look at this. This is a noising schedule, right? I have a dog image on the left. I keep on adding noise and then it becomes complete noise. The way it happens in text is can be represented by uh this kind of a thing. Yeah. So the way it happens in text is that uh at time equal to 1, you almost have an entire clean text with maybe some mask over here. But as time increases, as time becomes two, I have two masks. At time as time becomes three, I have three masks. As time increases almost everything becomes masked. So if I watch the full noising you'll see that you start with lower masks and then you increase the number of masks. That's basically what is meant by a noising schedule. In the case of diffusion LLMs it's also called as a masking schedule and the masking schedule is basically proportional to the time. The further along we are in the noising schedule the more masks there will be. Usually there is a probability associated with it. So it's a Bernoli distribution. So the probability of masking is way higher. Probability of masking all tokens is way higher as time is increased. But this animation suffices to understand how masking works for text. Okay. Um yeah. So what does it mean by masking essentially right? What it essentially means is that we just in our vocabulary. So if we have a vocabulary of 100,000 100,000 and the vocabulary is uh the boy etc. We add one more token to this vocabulary that's the mask token right. So now wherever there is a mask so here we have a mask right or here there is a mask and here there is a mask right or here there is a mask. So the original tokens were 1 11 20,000 55 and 3,000. Right? Now these are masked. So this is masked and this is masked. So this will be replaced with 1 0 0 0 1 which is the token ID for the mask. And this will also be replaced with 10 0 01 which is the token ID for the mask. And then this input sequence will go into the transformer architecture. Right? This input sequence goes into the transformer architecture or when I say transformer architecture actually this is the input sequence which goes over here which goes through the input the processor and the output. So we start with a noisy input with mask tokens whereas on the left hand side is the ARM. In the ARM we start with the entire sequence whereas in the case of diffusion language models we start with a noisy input with mask tokens. Okay. Everything else remains exactly the same. Now we'll see what are the changes when we go through this IPO block. What are the changes when we go through the I um I P and the O block in the next step. But the first thing which I wanted all of you to understand is how the noising process actually works. So to give a more concrete feel to it, let's say if we have four time steps, t= 1, t=2, uh here we have t= 3 and here we have t equal to 4. Right? Initially nothing is masked but as the time increases you'll see more and more masks to the input tokens. Um yeah so the first difference between the ARM architecture and the diffusion architecture if you compare them side by side is the input which goes into this input block the processor block and the output block this input itself is different here we take a bunch of token ids and we add masks right so when the forward pass is done first a random time step is chosen based on that time step we decide how much noise is added how do we know how much noise to add because we a noise schedule which is defined. So let's say if a time step equal to three is chosen. The time step is chosen randomly in each forward pass. So here what happens is that in ARM architecture I told you that a batch is selected right? batch is selected then we get the targets then we sorry we get the predictions then we get the loss then we update the parameters and then we choose a new batch this process stays the same for the diffusion LLM but it's just that whenever a new batch is selected a new time step is also randomly selected so if you have sim six time steps a new time step will be selected from 1 to six. So if time step number three is selected, if time step number three is se selected, roughly 50% of the input token sequence will be masked and they'll be fed into the architecture. The input, the processor and the output, the input, the processor and the output. Okay. So until now what we have seen is that what is the noising process in the case of diffusion language models. The next step which we are going to see is that how do we predict the noise. So in the case of uh images the noise is predicted by unit. But let's see in the case of diffusion language models how do we predict the noise? What kind of a model do we need to predict the noise? Um and the answer is much simpler than you think. So, let's get into that right