Lecture 14: Diffusion LLM Inference Pipeline

Key Takeaways

The last step which we need to do is the actual step of dnoising. What is dnoising? Well, dnoising is essentially starting from noise. Starting from complete noise like in the case of images, we start with noisy images and slowly our aim is to get to that probability distribution where the true image actually lives. If the training has been done correctly, then dn noising should actually help us to recover images in that exact or in the true probability distribution space or as close as possible to the true probability distribution space. What does it mean if training has been done correctly? Well, what it means is that if at each step of the training process, we are able to predict the noise which has been added. So at each step, yeah, at each step of the training process, if we are able to predict how much noise has been added in that step, essentially when we reverse when we reverse it or when we start from noise, we will be able to get back to the true image. That's the whole idea. So if we at each step if we are able to predict how much noise is added to make a noisy image if we start from the noisy image we will be able to go back in the reverse direction. So in the case of images it's fine right because in the case of images usually the dnoising might look something like this. Yeah like this we start from a completely noisy image and then we recover the clean image. So we start from a complete noisy image and then we recover Chinese characters as the dnoising proceeds. Right? But in the case of language, how does it happen? What does it mean by completely noisy image? We have already seen that a completely noisy or what does it mean by completely noisy text? It means that we start with a text in which everything is fully masked. So let's say this is the text in which everything is fully masked. We have a beginning of sequence and an end of sequence. Let's ignore that for a moment. But we have three masks over here. That's how the dnoising process actually starts. What we have to do is that now the model is trained. The parameters of the model will not change in the dnoising process. The model is fully trained. Um so all these parameters which we have seen over here, those parameters will not change. Okay. What happens in the dnoising is that we have to slowly uncover these masks and predict what was there here in the first place. Okay. So let's say we want to d noiseise from a fully noisy text which is this. The way the dnoising works in the case of language diffusion is through four steps. First, if we have in the forward pass, if we have time steps 1 2 3 4 or in this case, we had six time steps, right? In the forward pass, we had six time steps over here. If we had four time steps, for example, in the D noising process, we have to go backwards from four. So, we have to go 4 3 2 1, right? So, we have to start with four and then we have to sample the tokens for each of these masks. So we have to pass this sequence through the trained model and then see what the token prediction is at each of these masks at each of these positions. Then what we have to see is that which position is predicted most confidently and we have to uncover these masks sequentially. Right? So in the first iteration we only uncover one mask and keep all the other masks. Okay. Which mask is uncovered? the mass for which the prediction has the highest confidence of course. So let me let let's take an example. Okay. So if this is the input sequence we start at time equal to 4. So we start in the reverse manner. We pass this input sequence through my entire architecture which is now trained. When I say my entire architecture, this input sequence passed through is passed through this whole architecture. So let me bring it over here. This input sequence is passed um this input sequence is passed through this whole architecture. And if I clean this up a bit be easier for you. Yeah, this input sequence is passed through this whole architecture. And here we don't compute the loss because the training has already been done, right? But what we do compute here is what's the predicted uh logits matrix? What is the predicted logits matrix? Um so let's say we have position number we have three main positions which matter right because the beginning of sequence and end of sequence does not matter. We have position one position 2 and position three. So let's say the logits matrix is something like this. The logit's matrix for position one is 122 06.22 and 336. Um for position 2 it is.1 0 0 let's say and point 2 here. Uh for position three it is 0.1 let's say same 0 0. Now the way it works is that um let's see now for position number one this is position number one we see the maximum confidence is here and uh our dictionary actually is a uh our dictionary is let's see what our dictionary is I think the dictionary is beginning of sequence end of sequence sequence mask then we have a b and c right so let's actually take some different values of my let me refresh this um what I want to do here is that I want to take uh different um I want to take different values of my now it's working so let me rub this a bit okay and let's say the values here are such that it's 1 0 08 and.1 and this is also 0.100 08 and.1 okay so for this position for position number one the maximum confidence is for token C. So we'll predict C with this confidence for position number. So let's say this is.1 and8 and this is 1 and8. So for position two also it's C and for position three also it's C. Um so for position and actually the confidence let's say it's 336 here. So let's change this again. Sorry for that. So let's say this 336 and this 336. Let's say everything is the same for all. So I'll just uh take this and copy paste for all. Yeah. Yeah. So then I look at the token with the maximum confidence and that's C which is the last entry year for all my positions. Right. So after sampling my prediction should be CCC, right? But since we have just started the unmasking, we must end step three. So this first step we should end with three masks. That's the rule. So when we start unmask unmasking, first we'll end this step with three masks. Then we'll have two masks. Then we'll have one mask and then we'll have no mask. That's the idea. So we need three masks here. So again we'll mask all these three. So these three stay in masks. That's the output from the first dinoising step. Then we move to the second dinoising step which is time equal to three. Again let's say position number one from the logits matrix we have that it samples A with maximum confidence. Position two we have sample C with maximum confidence and position three we have sample C with maximum confidence. So after sampling we should get beginning of sequence A C and end of sequence. Now after the second D noising step we need to keep two masks. So the whole idea is that so I'll tell you here step step number four step three step two and step one at the end of this step we have to keep three masks at the end of this step we have to keep two masks at the end of this step which we have to keep one mask and at the end of the final step zero mask so that's my final thing which will appear on the screen this is exactly what's happening here we start with entire all masks and then we remove the masks one after one during the noising process I'm just showing you how what's the logic to be followed for demasking. So now here we have step three right and I have a c and c but I have to keep two masks. So which are the two masks I will keep I will keep the two masks with the highest um or I will keep so which what what will I unmask? So I have to keep two masks here right? So I have to unmask one thing. What's the thing I'll unmask? I'll unmask the thing with the highest confidence. Right? So either position two or three because position one has slightly lower confidence. So by default let's say unmask this position which is position number two. I could have unmasked position three also but let's say I unmask position two here. So then the input at the next step is Bos mask C mask and EOS. I pass it through the entire architecture. I get the logits and let's say again I sample A C and A. At this position we have a probab confidence of 3175. At this we have a confidence of 336 and at this we have a confidence of 3175. Now here the rule is that we have to unmask we have to keep one mask. So we have to unmask two which are the two I will unmask. Again I'll unmask those ones with the maximum confidence which are these two. Uh I could have unmasked this also but by default let's say I unmasked this. So I unmask position one and I unmask position two. Then at the last step I have this input sequence B ac and only one mask and then I pass it through the input architecture. I get that the after sampling we have BOS ACC and EOS. Um and this is the logits matrix right. So at first position I sample A with confidence of 3175. At second position I sample C with confidence 336. At third position I sample C with confidence.336. So here I have to not keep any mask. So I have to unmask everything. So the answer will now this everything will be unmasked. So the answer final answer final generated sequence is beginning of sequence A C C and end of sequence. That's the sequence which is generated by the diffusion language model. Now that's the end of it. We did the dnoising and that's the generated sequence produced by the dnoising process. This is what it looks like in action. So here you see we start with all the masks. So if you give a prompt the prompt along with these masks along with all the masks are passed to the input architecture and wherever there is mask that slowly unmasked right. So here we have 128 steps. So right now I just showed you four steps for 128 steps we slowly unmask everything one after another. Right? So just see how fast the inference happens here compared to the next token prediction. Right. So let's play this GIF if it opens. So see everything is masked right now. Right? Nothing is unmasked. We have 128 masks here. Now things are slowly getting unmasked here one after the other. Right? For some reason there is a weird black white contrast over here. So I can just bring this GIF over here and then just uh refresh it maybe. Yeah. See the masks are being removed one after the other. And uh it's not auto reggressive which means it's not one token at a time. Tokens masks can be unmasked here, masks can be unmasked at the top also. And everything happens very fast. That's the advantage of diffusion models. Right? This is the generation process or this is the dnoising process. This dinoising process is again if you have done the noising process correctly which means if you have obtained a very low loss in the noising process dn noising is kind of guaranteed to work even in diffusion language models. So even theoretically this is a safe method but it just as we'll see it takes a bit of time to train compared to auto reggressive language models. Okay, I hope all of you have understood the dnoising process, right? In dnoising, what happens is that let's say you fully trained the model as we have seen before. Then you start with a sequence of masks. Let's say if you want to predict a certain number of tokens, right? Uh you want to predict a certain number of tokens. All of them will be masked initially. Let's say you want to predict 50 tokens. All of them will be masked. You pass those 50 tokens into the architecture and you slowly go on unmasking the tokens one after the other. Uh that's the whole idea of unmasking. It's a very simple process but there is just one rule that uh you unmask those tokens which have the maximum confidence at every unmasking step. You unmask only those tokens which have maximum confidence and slowly you keep on unmasking one after the other. So this video or this GIF which Andre Karpati shared. Now if you look at the right hand side the diffusion you should have a much better understanding. We start with masks and what is shown on the right hand side is just unmasking one step at a time. Whereas on the left hand side this is auto reggressive. It's one token at a time. Whereas on the right hand side it's masked get it's masks getting unmasked one after the other. Okay. Now we'll see a summary of the three characteristics of the diffusion language models.