Chapter 22.2 - Transformers & Attention Mechanism

0:08 >> So, I want to talk about one of the most exciting innovations in AI. You know, there have been so many interesting things that AI has done, but this one has really taken the world by storm, partly because with transformers and the attention mechanism and training that we've been able to do across the corpus of human knowledge, we've created agents now where we start thinking about their capability for intelligence or AGI. Right? So, that's how good they are. It starts to question the foundations of this thing being able to essentially be smarter than humans in interactions, at least from a technical or writing point of view.

0:49 and so, how does this actually happen? Well, these are the mechanisms allowing it to happen. And so, I want to walk through the basic architecture of a transformer and how the attention mechanism plays a role in starting to allow us to have this amazing learning paradigm, which allows us to do writing, reasoning, all kinds of the things that we think about associated with intelligence. So, first I want to go back to the last lecture and just make sure we understand that it still has this basic architecture, which is you have input, you have to have an encoding space that you build. This is going to be a tokenization procedure, and the tokenization procedure is not to, in fact, embed words as vectors, but in fact, to break words up into tokens. And then, that is what gets embedded in this space. It's still a numerical value, just like a vector for a word, but now you break it down into smaller chunks or tokens to go through. So, the encoder is here, and then you have to train a decoder to come from that latent space representation of that you're building to an output and producing the next word in the sequence, which then by the way could be fed back in to keep a rollout going. In other words, you give it a start of a a start of a sentence and you can tell it, you know, write 10 pages.

2:10 So, it's going to roll out 10 pages for you. Right? And it's just doing it through probability distributions of how the words are associated with each other. And the important thing is the transformer architecture with the attention mechanism has a lot of context cues built into it. So, it does a much better job than the standard language models with recurrent neural nets. Okay. So, the big key here is that the transformer and attention the transformer architecture and the attention mechanism essentially has an infinite look back period. So, it can basically look over the history of an entire book and preserve information or associations that were on page one even though you're on page 100. Right? So, it starts to track these things sort of in sort of a infinite look back period. It's not infinite, of course, it's finite, but it is large and essentially we call it infinite.

3:07 and this is very different than the standard recurrent neural nets, which have a finite look back period, which limits their capabilities significantly. So, here's the architecture of the transformer. You're going to come in with your embedding space or your encoder, which is you're going to basically take your input, which is going to be language, and you're going to have an input embedding. This is the tokenization idea. The tokenization, there's another key piece to this, which is the positional encoder. It starts to recognize the position of the tokens in the sequence, right? So, this positional encoding becomes very important before it goes into what's called the attention head or the multi attention head. We're going to dive down what's in this attention head because that is a critical piece of the architecture as well. Not only the positional encoding, the tokenization, now you actually go through, you do this attention mechanism, this multi-headed attention. Once you do this, you do some normalization procedures on the way out, a feed forward, and then essentially initial norm.

4:13 And this is your encoding space, which then proceeds over a few decoder. The decoder again is made up of a multi-attention head unit. And then when you want to come out to the output, you're going to do some you're going to basically add all these attention heads together. You're going to normalize it through a feed forward normal neural network, additional normalization again, a linear layer, a soft max to the output, so you make a prediction. And that prediction can then just be fed right back into the decoder. So you this is the roll out area. So this is how this works, and the key is to understand what's happening in that attention head right here. That's going to be the critical piece of this, and the attention mechanism has gotten it's comes from neuroscience, yes, or at least it's motivated by neuroscience, and it's a really important piece of how we learn how to encode information long-term.

5:07 Okay. So first of all, let's talk about the transformer architecture then. Now that we're in there, I want to talk about how do we do the attention mechanism itself. We're going to have a sequence there that's coming in. This is the word embedding sequence that we might have, or token embedding sequence that we might have. So that is what's being fed in, is a sequence, and our job is to predict the next token in that sequence, okay? And this is going to be done by building three key matrices, the Q, K, and V matrices. The Q matrix is the queries matrix, which I'll I'll a little bit more in a minute.

5:47 Keep the K matrix is the keys and the V is the value matrix. These are the three foundational matrices that we need to train the transformer and these are the ones that are responsible for the attention mechanism. So how do they work or at least let me tell you what they are. So first of all the query is essentially learning the word actually being queried in the sequence. So for instance if my dog is my the word is dog then that is the word I'm querying and I can ask questions like okay the Q here what it does it's associating your query word and it's association with many other words. If I have the word dog there's a very high probability or association work like with the word cat or the word Rottweiler or golden doodle.

6:41 Those are essentially high probability associations where if I have the word dog it's probably a low probability associated with scuba diving for instance. Okay? That's the way we can kind of think about this this matrix. The keys matrix are essentially the words that modify the keyword. It's sort of like adjectives and adverbs. Okay? So for instance if my word is dog okay? Then often words that are associated with dog are they're of high value.

7:19 Big, fluffy, red. That's often how we might describe a dog, right? Oh I have the I have a big dog, he's fluffy, he's brown or red or black, he wags his tail. Tail's another high value word. So those are all words that are high values in the key matrix. So they're kind of in some sense they're the modifiers of that a where dog might show up in a sense. The running dog dogs often run, that's a nice high value word.

7:50 But words like headphones, baking, skiing those are let's say low value words. We don't often have you know, aren't talking about headphones and dogs at the same time. I'm sure you you could, but it's a low probability event. So, in other words, it would be of low value. And remember, these matrices are going to learn how to encode value for these words and these word associations. That's how it's making predictions about the next word in the sequence. It's going to basically pick the word based upon probabilities.

8:28 Every word has a certain probability, and so you're going to roll the dice and see which one of the words gets selected as the next. The highest value words are most likely to get selected. And every time you do a roll out, you'll get a slightly different sequence. Why? Because it's choosing it by probability, okay? But again, often times it's going to make sense because it's choosing words of high value, in other words, that are often associated with this. How does it learn that? Well, it's learning it based upon a corpus of knowledge. In other words, you're going to train all of this with a massive data base of all human written knowledge, right? So, it's going to see every sentence written by everybody everywhere in the world that's been put up on the internet or written in a book, which is a PDF.

9:18 And so, what it's seen is all of these associations of the words through that language writing, and it's going to make these probability matrices, which are the K, Q, and V. Now, the V matrix, by the way, and by the way, words of high value to a keyword are said to attend to that word, okay? So, that's the attention mechanism. So, fluffy and dog, fluffy attends to dog. Because that is a word of high value and often is associated with dog. Okay?

9:50 All right. So, that's the attention mechanism. That's the key piece of this of the association. And not only that, we can look back through the entire history, let's say, of a book, of writing to get these attending behaviors or attending values. Unlike when we did the recurrent neural network, which had a finite lookback, this basically can scan over an entire all documents. In fact, it can see that fluffy and dog are associated in another in an entire another book to make the association that you have here in your current book, for instance. Okay?

10:27 And so, that is what we have. Now, the value matrix is the only one left, and this is the last matrix that we have. And the last matrix does a further assignment of the word in context. So, you It starts It the value matrix starts to understand the contextual usages of the words. How are you Why are you using these in the sentence? What are you actually talking about? So, that's going to be learned through the value matrix itself.

10:57 Okay. So, these matrices are all going to be learned in training. So, we're going to give it a corpus of human knowledge, and that's how we're going to train it. And the product of Q and K, right? This is the the value and key Sorry, the query and keys matrices against each other, this inner product is kind of in some sense how they're modifying each other. And so, this is the basis of what is called the attention mechanism. Here's the attention mechanism. It's a function of these three matrices. And you're going to take a softmax. The first thing you do is going to take product of the queries and the keys. Now, the queries are the words, the keys are things that modify the words. And so, when you take this inner product, you start seeing, "Oh, I have fluffy and dog often producing a high inner product piece." Right?

11:48 Remember, the inner product is the projection of one onto the other. Okay? And so, that this would produce a nice dot product, but when I say dog and scuba diving, it would produce a very low inner product. There's a normalization that has to happen after I've done all of this soft maxing of this product. or actually, I first take a normalization, soft max it, and then I take all of this through the value matrix at the very end.

12:14 So, that's the basic mechanism, but we're not going to just train one attention head or one Q and K and V matrix. We're going to actually train a bunch of them. Right? So, we're going to train a bunch of a head attention heads. This is multi-headed attention. And so, when we do this, since we're going to train them all jointly, at the end, we're going to concatenate all the results of these different heads and figure out how to weight them.

12:41 So, this is sort of in some sense, we're going to train an ensemble of attention heads. So, each one of them is have different K, Q, and V matrices that it's learning from con- from from the literature. They all were initially seeded randomly, right? So, they all learn different ways to embed and sequence model. And then, what we do is we weight them all together in this multi-headed attention, and then we decide how to listen to those that concatenation. So, it could be that one of the attention heads does a really good job. We would listen it to more strongly than one that does not.

13:16 But we don't know which one's going to do well ahead of time because remember, they're all initially randomized with random weights. Okay. So, here is how the multi-headed attention works. Right at the bottom here are the matrices query, keys, value matrices, and I'm training a bunch of them. As we go, I'm going to take the scale dot product of these here. That's KQ. After I do that, I scale it, I mask in a softmax, and that's matrix A. And then I hit that against the value matrix.

13:49 Then I do another concatenation, and then I do a linear output layer. So this is the transformer head, the multi-headed transformer head. I have taken in, and I'm training a bunch of these simultaneously, aggregating their results, and weighting those results against each other for the best prediction. So this is a massive scale you know, we're talking when we train modern language models, you know, we're hundreds of billions of parameters in the modern modern large language models.

14:23 And so that's all the weights that are sitting inside of these multi-headed attentions, as well as well as these encoder-decoder schemes. So that architecture of of a QKV matrices trained on the attention heads, as well as this architecture overall, is what our large language models are primarily built around, roughly this kind of architecture with modifications here and there to make them better. But a lot of the best modifications come on what's called post-training. So refinement tuning to make these things even better. And we're going to talk about that in two lectures from now. But all the training data you'd have here is, you know, all human knowledge to train hundreds of billions of parameters. So you have so much writing you can put in here. You can either pull it from the internet, you can pull all the books that have ever been written, right? In terms of that gives you this massive language model capability to to on that to just learn the probability distributions associated with the K, Q, and V matrices. That's what you're really going after to learn with all of this language that you're doing. But, remember, you're doing it not on the words themselves, but on the tokenized embeddings.

15:37 So, that's how this thing actually works is you've traded out words as the embedding to tokens as the embedding, and then you look at the sequence of tokens that you've learned across all of the all the human knowledge that you've been ingest into this thing to train this model. And just here, I just wanted to point this out. This is only goes up to GPT-3, right? So, this is even a little updated, but if you look at the number of parameters that we've gone from 125 million not that long ago all the way to GPT-3 has 175 million parameters. This is absolutely massive scale, right? So, this is why it takes just the electricity bill alone, right?

16:17 Is hundreds of millions of dollars, I think, to train these things just because you've got to actually gradient descent yourself or optimize, let's say it's not gradient descent, but it's some morally equivalent version of gradient descent, and you've got to figure out how to update these 175 billion parameters to get this thing to work at the at the performance levels that we currently have in our in our AI agents that we work with on a daily basis.

16:49 So, that's the the basic storyline. It's an architecture. Very few people could actually train these models, right? You need to be one of the big tech companies to be able to train something at that scale. So, that's potentially a frustration point because you have an architecture, but to truly unleash that architecture, you have to have thousands and thousands of GPUs, and you have to run these things over many, many months and have the kind of resources to actually train a model at that scale to perform at the levels that we're seeing today where essentially interacting with an agent, you often can't tell the difference if there was a human doing that or if it was the agent doing this the intelligence. So, this is kind of where we have achieved. It's it's just an incredible piece of science and engineering all based upon this idea of the transformer architecture with these attention mechanisms embedded in it and we're learning three key matrices, the keys, the values, and the queries, and that's how we build this whole infrastructure up.

17:54 >>