Transformers & Diffusion LLMs: What's the connection?

Key Takeaways

This is a diffusion-based large language model. It starts from gibberish and iteratively turns it into a coherent response. Patterns emerge from noise. >> I'll always look at it in coding. >> I I don't even see the code. All I see is bong, brunette, redhead. This is in stark contrast to traditional or auto reggressive models like GPT which produce one token at a time without ever revising it. The fusionbased LLMs promise faster inference, better quality outputs, and more flexible prompting. I covered this in my previous video together with an intuition on how they were inspired by physical diffusion. But one question kept coming up in the comments. How are they different from transformers, the famous architecture behind traditional LLMs? Well, we'll see that that's a false dichotomy. Transformers actually power the fusionbased LLMs. In this video, we'll unpack how the transformer went from a humble architecture for machine translation to becoming the backbone of the entire natural language processing field. If you're particularly interested in the fusion, I suggest you skip to the bird chapter. Otherwise, stick around because it's a really interesting origin story. The transformer is a model architecture introduced by Google researchers in 2017. Soon after its publication, it replaced recurrent neural networks as the status quo architecture for natural language processing. Conceptually, machine learning connects an input X to an output Y through a function with parameters W. W are also known as the model weights or the numeric values discovered through training. F is where the human intervention happens. It's the model architecture designed by researchers. And the transformer is one convenient way of defining F. If we had infinite resources, we wouldn't need any human intervention. Given a set of neurons, we would define f in the least constraining way possible by connecting all pairs. But this quadratic is intractable. It gets really big really fast. The role of a model architecture is to prune these connections and rearrange the neurons in an informed way. It needs to strike the right balance between cost and expressivity. So, how does the transformer pull that off? Contrary to common belief, the original transformer was not intended as a general language model. Instead, it focused on machine translation, mapping text from a source language to another target language. But despite its humble beginnings in machine translation about 8 years ago, the transformer is now dominating the entire machine learning world. In this video, we'll talk about its usage in the language world, where among others, it enables both auto reggressive and diffusion-based large language models. In the next video, we'll see how the transformer expanded to the visual world, where it's now used for tasks like image recognition or image and video generation. Machine translation was an interesting and somewhat lucky start for the transformer. On the surface, translation rewrites text from one language to another, but ultimately it's a sequencetosequence problem or sequence transduction if you want to be pretentious. That's a very general interface shared by many other tasks. The core challenge in machine translation is alignment. Words in different languages don't line up one to one in the same order. For instance, in French, pronouns that are indirect objects like Louie go before the conjugated verb don, but in English, the order is reversed. These correspondences are never spelled out in the training data, so the model needs to uncover them by itself. When generating the next target word, the model has to keep track of what's already been translated and decide which source words to tackle next. The most common mechanism for handling this is attention, which assigns a weight to each source word. Here, the pronoun I gets no attention because it's already translated. The pronoun him elicits the strongest attention, so it gets translated into Louie. Attention was actually not introduced by the transformer. It was published three years prior in a paper that jointly learned to align and translate. That meant attention was a mere appendix to the main translation engine, a recurrent neural network. What the transformer did was to promote the attention mechanism from a supporting role to the main building block of the translation network. Hence the now famous title, attention is all you need. The need for attention is actually what connects all of these tasks. For instance, during image recognition, when classifying dogs versus cats, you need to pay attention to key distinguishing features like the eye collar or the shape of the nose and ignore distractors like the floor. Similarly, when classifying text, maybe to analyze sentiment on my YouTube comments, the model needs to pay particular attention to strongly emotional words like repulsed or amazing and disregard UX labels like reply. We'll now open up the hood of the transformer to understand what makes it so versatile. Tiny bit of housekeeping first. Throughout this video, I'll jump back and forth between training and inference. And I'll use these tabs at the top to clarify which one I'm currently talking about. I'll also cross reference multiple models. And to avoid confusion, I'll show their names in the bottom right corner. Now, back to the transformer. This is a famous figure from the original paper. In this video, I won't exhaustively open up all the boxes. Instead, I'll strictly focus on those aspects that can help us understand the implementation differences between auto reggressive and diffusion LLMs. For a more complete coverage, I'll link in the description an older video I made back in my days at Google. To get familiar with the transformer, let's walk through the inference process for a machine translation example. We'll start high level and then keep digging deeper. So the transformer has two main components, an encoder and a decoder. The encoder on the left processes the English input and feeds it into the decoder. The decoder on the right produces the translation one word at a time. It pays attention both to the English sentence coming from the encoder and to the French words it has predicted so far. And the loop continues until the model outputs the special end of sequence token. Now let's zoom into the encoder. Each source word is looked up in a static embedding dictionary and converted to a continuous input embedding. At this point, token embeddings are independent of each other. For instance, the article 'the' gets mapped to the same representation regardless of the noun it refers to. The encoder combines these embeddings to make them more contextual. For instance, the article the is now a weighted sum of all the original word embeddings, presumably putting a high weight on book. This mechanism is called full self attention. Full because every token is expressed in terms of all the other tokens and self because the input only looks at itself. Now let's shift our attention to the decoder which produces a translation word by word. After predicting Jelui, the output tokens are looked up again in the dictionary and fed back into the decoder. You're probably aware that transformers have a fixed context window. Regardless of where they are in the generation process, they're hardcoded to operate on this many tokens. In this story example, only three of the six tokens in the context window are present. So, we'll fill the rest with a special padding token. The original transformer had a context window of 512 and current LLMs go up to hundreds of thousands or millions. Now, we need to build an embedding for the next token to be translated. But this time, we can only leverage the tokens to the left. Everything to the right is just useless padding. This mechanism is called causal attention. It's what enables the model to look at its past. For translation, of course, we have to also consider the input sequence processed by the encoder. We'll compute a new next token embedding which is a function of the source tokens as well as the history. This mechanism is called cross attention because it crosses between languages. And we're almost there. We'll pass the final embedding through a linear layer which expands its size to one dimension per vocabulary token. The softmax layer turns it into a probability distribution from where we can sample the next token. I know we covered a lot of ground, but there's only two key things to remember. First is that transformers started in machine translation which is a task the very general API sequence in sequence out and second transformers have two main components an encoder which processes the input with full attention and a decoder which produces the translation one token at a time using causal attention. Now here's the interesting thing. Researchers realize that these two components are a little bit redundant. They're actually very similar to each other. So, we don't need both of them at the same time. A lot of subsequent transformer-based models use either one or the other. When using them separately, cross attention as a concept disappears. When looking at these components side by side, we can see they're actually not that different. Ultimately, they both take input embeddings and contextualize them. The main difference is the type of attention. In the encoder, everything attends to everything, while in the decoder, tokens only attend to their left. We'll now look at how subsequent models borrowed one of the two components to extend the transformer to other tasks. And we'll start with auto reggressive LLMs. On the surface, language modeling has the same interface as machine translation. Prompt sequence in response sequence out, same old sequence transduction. But while machine translation used both an encoder and a decoder, auto reggressive LLMs only reuse the decoder. So what happened? What broke this equivalence between translation and language modeling? Well, in 2018, OpenAI released their first GPT model, which stands for generative pre-training. At the time, the bottleneck for language modeling was data. There simply weren't enough response prompt pairs to be mined on the internet. In contrast, free form text like Wikipedia articles was and still is a lot more abundant. Plus, it stores a huge amount of world knowledge and linguistic patterns. To make room for it, GPTE split the training process into two stages, pre-training and fine-tuning. So, to leverage free form data during pre-training, we need to gamify it or come up with a learning objective that implicitly coerces the model to extract valuable information from the data. And that's what next token prediction is. A madeup artificial yet useful task. Given the start of an article, the model needs to predict the next word. And the same challenge repeats for every single word in the article. The next token prediction task is the reason why auto reggressive LLMs are picking the decoder side. Causal attention is a good fit by nature. When predicting the next token, it only looks at the previous ones. It's worth taking a closer look at GPT's training process in order to understand the trade-offs it makes compared to the Fusion LLMs. Given a pre-training instance like a Wikipedia article, GPT ingests the entire input sequence at once. So, each token in the article gets its own embedding. For each position, the decoder computes an embedding for the next token. For simplicity, let's focus on one position at a time. Say the model is expected to predict dog. Then the next token embedding attends to everything up to breed of retriever. Say it wrongfully predicts the word cat. Since this doesn't match the ground truth that means the original article the model's weights will be nudged to downweight cat in the future. The same mechanism applies for the next position but this time the model prediction is correct. So the word of is reinforced in the weights. Notably all predictions can be parallelized. With a single pass through the model we're able to learn from every single token. As we'll see later, this is a lot more efficient than training the Fusion LLMs. But it does come with a drawback. The token history always comes from the training data, never from the model's past predictions. Remember, we previously predicted cat instead of dog. Yet, in the current step, we're attending to dog. Using a clean history that comes from the original article is very useful at the beginning of training when initial guesses are bad and the model would have a hard time learning anything. However, this introduces a discrepancy at inference time. All of a sudden, the model is expected to ingest its own predictions. A single wild output can derail the entire response because the model never learned how to deal with outliers. This issue called sampling drift is something that the fusion LLMs avoid altogether. So pre-training is helpful for acquiring world knowledge. But in regular applications, we want more than just word autocomplete. The second training stage called fine-tuning leverages prompt response pairs to induce a conversational behavior. The two elements are concatenated together with a special beginning of sequence token in between. The rest of the process is similar to pre-training. The only notable difference is that tokens are predicted only for the response, not for the question. Once the model is fine-tuned, inference happens auto reggressively, one token at a time, as we already discussed. Now that we covered auto reggressive LLMs, we can make our way to diffusion LLMs. But on the way there, we need to talk about text classification, which is perhaps unexpectedly a stepping stone towards diffusion. Text classification comes in many forms, but a simple application is sentiment analysis. Given a product review, decide whether it's positive or negative. Compared to machine translation, its interface is simpler. The output is a single value because the input to text classification is available right from the start. Full self attention and therefore the encoder is a natural fit. All tokens can attend to all the others. But pre-training gets a bit trickier with an encoder. The next word prediction task is no longer supported since tokens peak into the future. The solution came from BERT, a pre-trained text classifier published by Google in 2018. BERT came up with a new task for pre-training, masked language modeling. Exactly 15% of the tokens are replaced with a special mask token and the model needs to recover them all in one go. Let's see how to pre-train the transformer encoder with the help of the masked token prediction test. As usual, the encoder produces a contextual embedding for every input token, including the masks. To turn these mask embeddings into token predictions, we can borrow the same trick we used in the decoder, a linear layer followed by softmax. Note that mask predictions can be paralyzed since they're independent from each other. The model predictions are then compared against the original article and the feedback is back propagated through the model weights. Right before fine-tuning, the additional token prediction layers are thrown away. They've served their purpose. We're now left with an encoder that is very good at processing text inputs in a contextual way. What's left is to teach the model to judge sentiment using supervised pairs of reviews and sentiment labels. BR has one more trick up its sleeve to make this happen. It appends a special classification token at the end of the text input. Due to the full attention in the encoder, its final embedding will compress information from the entire input. So feeding it through a new linear layer different from the one added for pre-training can now yield a sentiment label that is checked against the ground truth label. With BERT under the belts, we can now talk about diffusion LLMs. Text diffusion is still a very young field. As I covered in my previous video, there are many formulations, some more math heavy than others. And while there isn't a de facto standard yet, two things are clear. One is that all diffusion MLMs under one form or another are using the transformer. and two is that there's a clear convergence towards masked diffusion models, the ones inspired by BERT. Most likely, it's because they're easier to understand and easier to implement. To make it concrete, we'll walk through the implementation of Lada, a masked language model that came out earlier this year. But keep in mind the field is constantly changing and Lada will probably not be the one to establish the status quo. When designing the Fusion LLMs, there are a few important hyperparameters. Just like in traditional LLMs, there's a context window. I'm arbitrarily setting it to 35 here, but in practice, it's hundreds of thousands or millions. At the start of inference, we'll populate it with the prompt followed by the special beginning of sequence token. For masked diffusion models in particular, the rest is filled with masks. In other words, the current noise level is 100% since the entire response is full of noise. A second hyperparameter is the number of diffusion steps. That is how many times we'll refine the response draft. I'm setting it to three because I'm lazy, but in practice, it's somewhere in the tens or low hundreds. This acts as a knob between speed and quality. And finally, we need a noise schedule, which dictates how many masks we'll uncover at each diffusion step. It's common for language models to use a simple linearly decreasing schedule, like I did here. So, at step one, all masks are uncovered in one go. Because it's so early in the process, most guesses are bad, perhaps with the exception of the token 42. So, we'll throw the dice again, this time for only 60% of the tokens. The tokens to be remasked can be chosen at random or using some huristic like selecting those tokens that were assigned low confidence by the LLM. The next reveal improves the response quite a bit. Three tokens are now on the right track. It's also interesting that this time the special end of sequence token was generated. This signals that the intended model response goes all the way to hollow planet and everything after it should be ignored. We'll continue with 30% noise mask and unmask. This was our last step. So the final answer is 42 give or take. We'll simply ignore everything else. Now let's see how diffusion LLMs are trained. As I alluded to earlier, BERT pre-training, which uses a masked language modeling objective, is the basic building block for diffusion pre-training. However, BERT adopted a fixed masking rate of 15%. chosen empirically to maximize its capabilities as a text classifier. But a diffusion model needs to handle multiple levels of noise. Remember, my simplistic noise schedule had three thresholds, 100, 60, and 30%. Of course, if we increased the number of diffusion steps, the noise schedule would also be more granular. Going back to BERT, this is what pre-training with a fixed noise level looked like. Lada makes a simple modification. When choosing a training instance, it also picks a random noise level from the schedule. So for instance, for the parakeet article, it might choose to mask 30% of the tokens. For the cat article, it might arbitrarily pick 60%. Now, choosing a single noise level for each training instance might seem like an odd design decision at first. Why not apply all noise levels to every article just like at inference time? Well, we could do that, but implementation is a little bit easier this way. And anyway, when we train for millions of steps, chances are the same article will be selected multiple times with different levels of noise. Now, I do want to point out that diffusion training is overall somewhat inefficient. During a single training step, only some of the tokens are predicted. In contrast, GPD was able to fully leverage a training instance by predicting all the next tokens in a single pass. Fine-tuning a diffusion LLM is similar to the pre-training process. For a supervised prompt response bearer, we'll choose an arbitrary noise level, but this time we'll only mask the response part. There's no point in ever offiscating the prompt. And the rest is identical to pre-training. And that's how diffusion LLMs extend BERT with random masking to implement a whole new paradigm for text generation. In this video, we discussed how the transformer came to dominate language tasks. In the next video, we'll see how it expanded to the vision world as well, replacing convolutional neural networks as the go-to architecture. I know this stuff is a little dry, so thanks for bearing with me until the end. As usual, you'll find the slide deck and reading list for this video on Patreon for free. Thanks for watching, and I'll see you next time.