Let's build GPT: from scratch, in code, spelled out.

Hi everyone. By now, you have probably heard of ChatGPT. It has taken the world and the AI community by storm. It is a system that allows you to interact with an AI and give it text-based tasks. For example, we can ask ChatGPT to write us a small haiku about how important it is that people understand AI, and then they can use it to improve the world and make it more prosperous. When we run this, the AI generates: "Knowledge brings prosperity for all to see. Embrace its power." Not bad, and you can see that ChatGPT generated all these words sequentially, from left to right. I had already asked it the exact same prompt a little earlier, and it generated a slightly different outcome: "AI's power to grow, ignorance holds us back, learn, prosperity waits." Both are pretty good and slightly different. You can see that ChatGPT is a probabilistic system, and for any one prompt, it can give us multiple answers, replying in different ways. This is just one example of what it can do. People have come up with many, many examples, and there are entire websites that index interactions with ChatGPT. Many of them are quite humorous. For instance: "Explain HTML to me, I'm a dog." "Write release notes for Chess 2." "Write a note about Elon Musk buying Twitter," and so on. As another example, you could ask it to write a breaking news article about a leaf falling from a tree. It might respond: "In a shocking turn of events, a leaf has fallen from a tree in the local park. Witnesses report that the leaf, which was previously attached to a branch of a tree, detached itself and fell to the ground." Very dramatic.

You can see that this is a pretty remarkable system, and it is what we call a language model, because it models the sequence of words, characters, or tokens more generally, and it knows how words follow each other in the English language. From its perspective, what it is doing is completing the sequence. I give it the start of a sequence, and it completes the sequence with the outcome. It's a language model in that sense.

Now, I would like to focus on the under-the-hood components of what makes ChatGPT work. What is the neural network under the hood that models the sequence of these words? That comes from a paper called "Attention Is All You Need" from 2017, a landmark paper in AI that proposed the Transformer architecture. GPT stands for Generatively Pre-trained Transformer. The Transformer is the neural net that does all the heavy lifting under the hood. It comes from this 2017 paper. If you read this paper, it reads as a pretty standard machine translation paper, and that's because I think the authors didn't fully anticipate the impact the Transformer would have on the field. The architecture they produced in the context of machine translation ended up taking over the rest of AI in the next five years, and with minor changes, it was copy-pasted into a huge number of applications in AI in more recent years. That includes, at its core, ChatGPT. Now, what I'd like to do is build something like ChatGPT, but of course, we are not going to be able to reproduce ChatGPT itself. It's a very serious, production-grade system. It is trained on a good chunk of the internet, with a lot of pre-training and fine-tuning stages, and it's very complicated.

What I'd like to focus on is just training a Transformer-based language model, and in our case, it's going to be a character-level language model. I still think that is very educational with respect to how these systems work.

I don't want to train on a chunk of the internet. We need a smaller dataset. In this case, I propose that we work with my favorite toy dataset, called Tiny Shakespeare. It is a concatenation of all the works of Shakespeare, in my understanding, and this is all of Shakespeare in a single file. This file is about one megabyte, and it's just all of Shakespeare. What we are going to do now is model how these characters follow each other. For example, given a chunk of these characters—given some context of characters in the past—the Transformer neural network will look at the characters I've highlighted and predict that "g" is likely to come next in the sequence. It will do that because we are going to train the Transformer on Shakespeare, and it will try to produce character sequences that look like this. In that process, it will model all the patterns inside this data. Once we've trained the system, just to give you a preview, we can generate infinite Shakespeare. Of course, it's a fake thing that looks kind of like Shakespeare. Apologies, there's some jank that I'm not able to resolve in here, but you can see how this is going character by character, and it's kind of predicting Shakespearean language. "Verily, my lord, the sights have left the—again, the king coming with my curses, with precious pale," and then Tranos says something else, etc. This is just coming out of the Transformer in a very similar manner as it would come out in ChatGPT—in our case, character by character. In ChatGPT, it's coming out at the token-by-token level. Tokens are these little subword pieces. They're not word-level; they're kind of word-chunk level.

Now, I've already written all the code to train these Transformers, and it's available in a GitHub repository called NanoGPT. NanoGPT is a repository you can find on my GitHub, and it's designed for training Transformers on any given text. What I find interesting about it is that, while there are many ways to train Transformers, this is a very simple implementation. It's just two files, each with about 300 lines of code. One file defines the GPT model—the Transformer itself—and the other file trains it on a given text dataset. Here, I'm showing that if you train it on an open web text dataset, which is a fairly large collection of web pages, you can reproduce the performance of GPT-2. GPT-2 is an early version of OpenAI's GPT from 2017, if I recall correctly. So far, I've only reproduced the smallest 124-million-parameter model, but this demonstrates that the codebase is correctly set up and that I'm able to load the neural network weights that OpenAI released later.

You can check out the finished code in NanoGPT, but what I'd like to do in this lecture is to write this repository from scratch. We're going to start with an empty file and define a Transformer piece by piece. We'll train it on the Tiny Shakespeare dataset and see how we can then generate infinite Shakespeare. Of course, you can copy and paste this approach to any arbitrary text dataset you want, but my goal here is to help you understand and appreciate how ChatGPT works under the hood. All that's really required is proficiency in Python and a basic understanding of calculus and statistics. It would also help if you've seen my previous videos on this YouTube channel, particularly my "Make More" series, where I define smaller and simpler neural network language models, such as multi-layer perceptrons. That series introduces the language modeling framework, and in this video, we're going to focus on the Transformer neural network itself.

I've created a new Google Colab Jupyter notebook, which will allow me to easily share the code we develop together. You can follow along, and I'll include this in the video description later. For now, I've just done some preliminaries: I downloaded the Tiny Shakespeare dataset from this URL, and you can see it's about a 1-megabyte file. I then opened the input.txt file and read all the text as a string. We're working with roughly 1 million characters, and if we print out the first 1,000 characters, they are what you'd expect—the beginning of the Tiny Shakespeare dataset, up to a certain point. So far, so good. Next, we're going to take this text—which is a sequence of characters in Python—and, by calling the set constructor on it, obtain the set of all unique characters that occur in the text. By converting this set to a list, we get an ordered collection, and then we sort it. This gives us all the unique characters in the entire dataset, sorted. The number of these characters will be our vocabulary size. These are the possible elements of our sequences. When I print out the characters, there are 65 in total: a space character, various special characters, and both uppercase and lowercase letters. That's our vocabulary—the possible characters the model can see or emit. Now, we need to develop a strategy to tokenize the input text. When people say "tokenize," they mean converting the raw text string into a sequence of integers according to some vocabulary of possible elements. In our example, we're building a character-level language model, so we'll simply translate individual characters into integers. Let me show you a chunk of code that does this. We're building both the encoder and the decoder, and I'll explain what's happening. When we encode an arbitrary text like "hi there," we get a list of integers that represents that string—for example, 46, 47, and so on. We also have the reverse mapping, so we can take this list and decode it to get back the exact same string. It's really just a translation to integers and back for any arbitrary string, and for us, it's done at the character level. The way this is achieved is by iterating over all the characters and creating a lookup table from each character to its corresponding integer, and vice versa. To encode a string, we simply translate all the characters individually; to decode, we use the reverse mapping and concatenate everything. This is just one of many possible encodings or tokenizers, and it's a very simple one, but there are many other schemes people use in practice. For example, Google uses SentencePiece, which also encodes text into integers but with a different scheme and vocabulary. SentencePiece is a subword tokenizer, meaning it doesn't encode entire words or individual characters, but rather subword units. That's usually what's adopted in practice. OpenAI, for example, has a library called TikToken, which uses a byte pair encoding tokenizer—that's what GPT uses. You can also encode words directly, such as turning "hello world" into a list of integers. As an example, I'm using the TikToken library here. I'm getting the encoding that was used for GPT-2. Instead of just having 65 possible characters or tokens, GPT-2 uses 50,000 tokens. When we encode the exact same string, "hi there," we get a list of three integers, but those integers are not between 0 and 64—they are between 0 and 52,256. You can trade off between the codebook size and the sequence lengths. You can have very long sequences of integers with very small vocabularies, or you can have short sequences of integers with very large vocabularies. Typically, people use these subword encodings in practice. However, I'd like to keep our tokenizer very simple. We're using a character-level tokenizer, which means we have very small codebooks and very simple encode and decode functions, but we end up with very long sequences as a result. That's the level at which we're going to stick for this lecture, because it's the simplest approach. Now that we have an encoder and a decoder—effectively, a tokenizer—we can tokenize the entire training set of Shakespeare. Here's a chunk of code that does that, and I'm going to start using the PyTorch library, specifically torch.Tensor from PyTorch. We're going to take all of the text in Tiny Shakespeare, encode it, and then wrap it into a torch tensor. To get the data tensor, here's what it looks like when I examine just the first 1,000 elements. We see that we have a massive sequence of integers, and this sequence is an identical translation of the first 10,000 characters of the text. I believe, for example, that zero is a newline character, and maybe one is a space. I'm not 100% sure, but from now on, the entire dataset of text is just stretched out as a single, very large sequence of integers. Let me do one more thing before we move on. I'd like to separate our dataset into a train and a validation split. In particular, we're going to take the first 90% of the dataset and consider that to be the training data for the Transformer. We'll withhold the last 10% at the end to be the validation data. This will help us understand to what extent our model is overfitting. We're going to keep the validation data on the side because we don't want just a perfect memorization of this exact Shakespeare text. We want a neural network that creates Shakespeare-like text, and it should be fairly likely for it to produce the actual, stowed-away, true Shakespeare text. We're going to use this to get a sense of overfitting. Now, we would like to start plugging these text sequences—or integer sequences—into the Transformer so that it can train and learn those patterns. The important thing to realize is that we're never going to feed the entire text into a Transformer all at once. That would be computationally very expensive and prohibitive. When we train a Transformer on these datasets, we only work with chunks of the dataset. During training, we sample random little chunks from the training set and train on just those chunks at a time. These chunks have some kind of length and a maximum length. The maximum length, typically—at least in the code I usually write—is called block size. You might find it under different names, such as context length. Let's start with a block size of just eight. Let me look at the first train data characters—the first block size plus one characters. I'll explain why it's plus one in a second. This is the first nine characters in the sequence in the training set. What I'd like to point out is that when you sample a chunk of data—let's say these nine characters from the training set—this chunk contains multiple examples packed into it, because all of these characters follow each other. When we plug this into a Transformer, we're going to simultaneously train it to make predictions at every one of these positions. In a chunk of nine characters, there are eight individual examples packed in. For example, in the context of 18, 47 likely comes next. In the context of 18 and 47, 56 comes next. In the context of 18, 47, and 56, 57 can come next, and so on. That's the eight individual examples. Let me spell it out with code. Here's a chunk of code to illustrate. X is the input to the Transformer—it will just be the first block size characters. Y will be the next block size characters, offset by one, because Y contains the targets for each position in the input. Here, I'm iterating over all the block size of eight, and the context is always all the characters in X up to and including position T. The target is always the Tth character, but in the targets array Y. Let me just run this, and it spells out what I said in words. These are the eight examples hidden in a chunk of nine characters that we sampled from the training set. I want to mention one more thing. We train on all eight examples here, with contexts ranging from one all the way up to the block size. We do this not just for computational reasons—because we happen to have the sequence already or for efficiency—but also to make the Transformer network accustomed to seeing contexts from as little as one all the way up to the block size. We'd like the Transformer to be used to seeing everything in between. That's going to be useful later during inference, because while we're sampling, we can start the generation with as little as one character of context. The Transformer knows how to predict the next character with just one character of context, and it can predict everything up to the block size. After the block size, we have to start truncating, because the Transformer will never receive more than block size inputs when it's predicting the next character. We've looked at the time dimension of the tensors that are going to be fed into the Transformer. There's one more dimension to care about, and that is the batch dimension. As we're sampling these chunks of text, every time we feed them into the Transformer, we're going to have many batches of multiple chunks of text, all stacked up in a single tensor. This is done for efficiency, so that we can keep the GPUs busy, since they are very good at parallel processing of data. We want to process multiple chunks all at the same time. However, those chunks are processed completely independently—they don't interact with each other. Let me generalize this and introduce a batch dimension. Here's a chunk of code. Let me run it, and then I'll explain what it does, because we're going to start sampling random locations in the dataset to pull chunks from. I am setting the seed in the random number generator so that the numbers I see here will be the same numbers you see later if you try to reproduce this. Now, the batch size here is how many independent sequences we are processing in every forward and backward pass of the Transformer. The block size, as I explained, is the maximum context length used to make those predictions. Let's say the batch size is four and the block size is eight. Here's how we get a batch for any arbitrary split: if the split is a training split, then we're going to look at the training data; otherwise, we look at the validation data. That gives us the data array. When I generate random positions to grab a chunk from, I generate a batch size number of random offsets. Because this is four, our X is going to be four numbers that are randomly generated between zero and the length of the data minus the block size. These are just random offsets into the training set. The X's, as I explained, are the first block size characters starting at each offset I. The Y's are those same chunks, but offset by one—just add one to each index. We're going to get those chunks for every one of the integers I in X and use torch. We use stack to take all those one-dimensional tensors, as we saw here, and stack them as rows so that they all become rows in a 4x8 tensor. Here, I'm printing the result. When I sample a batch, XB and YB, the inputs to the Transformer are now the input X, which is the 4x8 tensor—four rows of eight columns—and each one of these is a chunk of the training set. The targets are in the associated array Y, and they will be fed into the Transformer at the end to create the loss function. They provide the correct answer for every single position inside X. These are the four independent rows spelled out as we did before. This 4x8 array contains a total of 32 examples, and they're completely independent as far as the Transformer is concerned. For example, when the input is 24, the target is 43—meaning 43 in the Y array. When the input is 24, 43, the target is 58. When the input is 24, 43, 58, the target is 5, and so on. Or, when it is 52, 581, the target is 58. You can see this spelled out. These are the 32 independent examples packed into a single batch of the input X, and the desired targets are in Y. Now, this integer tensor X is going to feed into the Transformer, which will simultaneously process all these examples and look up the correct integers to predict at every one of these positions in the tensor Y. Now that we have our batch of input that we'd like to feed into a Transformer, let's start feeding this into neural networks. We're going to start off with the simplest possible neural network, which, in the case of language modeling, in my opinion, is the bigram language model. We've covered the bigram language model in my "Make More" series in a lot of depth, but here I'm going to go a bit faster. Let's just implement a PyTorch module directly that implements the bigram language model. I'm importing the PyTorch nn module for reproducibility. Here, I'm constructing a bigram language model, which is a subclass of nn.Module. Then, I'm calling it and passing it the inputs and the targets, and I'm just printing the results. When the inputs and targets come in, you see that I'm just taking the inputs X—here renamed to idx—and passing them into this token embedding table. What's going on here is that in the constructor, we are creating a token embedding table of size vocab_size by vocab_size, and we're using an embedding, which is a very thin wrapper around a tensor of shape vocab_size by vocab_size. When we pass idx, every single integer in our input refers to this embedding table and plucks out a row corresponding to its index. For example, 24 will go into the embedding table and we'll pluck out the 24th row; 43 will go in and pluck out the 43rd row, and so on. PyTorch will arrange all of this into a batch-by-time-by-channel tensor. In this case, batch is four, time is eight, and C (the channels) is vocab_size, which is 65. We're just going to pluck out all those rows and arrange them in a B by T by C tensor. Now, we're going to interpret this as the logits, which are the scores for the next character in the sequence. What's happening here is that we are predicting what comes next based on just the individual identity of a single token. You can do that because, currently, the tokens are not interacting with each other and are not seeing any context except for themselves. For example, "I'm an F, I'm token number five," and I can make pretty decent predictions about what comes next just by knowing that I'm token five, because some characters typically follow other characters in certain scenarios. We saw a lot of this in much more depth in the "Make More" series. Here, if I just run this, we get the predictions—the scores, the logits—for every one of the 4x8 positions. Now that we've made predictions about what comes next, we want to evaluate the loss function. In the "Make More" series, we saw that a good way to measure the loss, or the quality of the predictions, is to use the negative log likelihood loss, which is also implemented in PyTorch under the name cross-entropy. What we want to do here is compute the loss as the cross-entropy between the predictions and the targets. This measures the quality of the logits with respect to the targets. In other words, we have the identity of the next character—how well are we predicting the next character based on the logits? Intuitively, the correct dimension of the logits, depending on whatever the target is, should have a very high value, and all the other dimensions should have very low values. Now, the issue is that this is what we want—we want to output the logits and the loss—but unfortunately, this won't run. We get an error message, but intuitively, this is what we want to measure. When we go to the PyTorch cross-entropy documentation, we see that we're trying to call cross-entropy in its functional form. That means we don't have to create a module for it. However, when we look at the documentation, we have to pay attention to how PyTorch expects these inputs. The issue here is that PyTorch expects, if you have multi-dimensional input—which we do, because we have a B by T by C tensor—it really wants the channels to be the second dimension. It wants a B by C by T tensor instead of B by T by C. This is just a detail of how PyTorch treats these kinds of inputs, and we don't want to deal with that. What we're going to do instead is reshape our logits. Here's what I want to do: I want to give names to the dimensions—logits. The shape is B by T by C. Let's unpack those numbers. Let's say that logits equals logits.view, and we want it to be B * T by C—just a two-dimensional array. We're going to take all of these positions here, stretch them out into a one-dimensional sequence, and preserve the channel dimension as the second dimension. We're just stretching out the array to make it two-dimensional, and in that case, it will better conform to what PyTorch expects for its dimensions. Now, we have to do the same for targets. Currently, targets are of shape B by T, but we want it to be just B * T—one-dimensional. Alternatively, you could always use -1, because PyTorch will infer what this should be if you want to lay it out, but let me just be explicit and say B * T. Once we've reshaped this, it will match the cross-entropy case, and then we should be able to evaluate our loss. Now, we can do loss, and currently, we see that the loss is 4.87. Because we have 65 possible vocabulary elements, we can guess what the loss should be. In particular, we covered negative log likelihood in detail. We are expecting log, or ln, of 1 over 65, and the negative of that. We're expecting the loss to be about 4.17, but we're getting 4.87. That tells us the initial predictions are not super diffuse—they've got a little bit of entropy. We're guessing wrong, yes, but we are able to evaluate the loss. Now that we can evaluate the quality of the model on some data, we'd also like to be able to generate from the model. Let's do the generation now. I'm going to go a little bit faster here, because I covered all of this already in previous videos. Here's a generate function for the model. We take the same kind of input idx here, which is the current context of some characters in a batch. It's also B by T, and the job of generate is to take this B by T and extend it to B by T + 1, plus 2, plus 3, and so on. It continues the generation in all the batch dimensions and in the time dimension. That's its job, and it will do that for max_new_tokens. You can see here at the bottom that whatever is predicted is concatenated on top of the previous idx along the first dimension, which is the time dimension, to create a B by T + 1. That becomes the new idx. The job of generate is to take a B by T and make it a B by T plus 1, plus 2, plus 3—as many as we want, up to max_new_tokens. This is the generation from the model. Now, inside the generation, what are we doing? We're taking the current indices, getting the predictions, and those are in the logits. The loss here is going to be ignored, because we're not using it, and we have no ground truth targets to compare with. Once we get the logits, we focus only on the last step. Instead of a B by T by C, we're going to pluck out the -1, the last element in the time dimension, because those are the predictions for what comes next. That gives us the logits, which we then convert to probabilities via softmax. Then, we use torch.multinomial to sample from those probabilities, and we ask PyTorch to give us one sample. idx_next will become a B by 1, because in each batch dimension, we'll have a single prediction for what comes next. Setting num_samples to 1 makes this a one, and then we take those integers from the sampling process—according to the probability distribution given here—and concatenate them on top of the current running stream of integers. This gives us a B by T + 1, and then we can return that. One thing to note here is that when I call self(idx), it will end up going to the forward function. I'm not providing any targets currently, so this would give an error because targets is not given. Targets has to be optional—targets is None by default. If targets is None, then there's no loss to create; loss is just None. Otherwise, all of this happens and we can create a loss. If we have targets, we provide them and get a loss. If we have no targets, we just get the logits. This will generate from the model, so let's take that for a ride now. I have another code chunk here that will generate from the model. Okay, this is kind of crazy. Let me break this down. These are the idx. I'm creating a batch of just one, and time will be just one. I'm creating a little one-by-one tensor holding a zero, and the dtype is integer. Zero is going to be how we kick off the generation, and remember that zero is the element standing for a newline character. It's a reasonable thing to feed in as the very first character in a sequence—to be the newline. It's going to be idx, which we feed in here, and then we ask for 100 tokens. Generate will continue that now. Because generate works on the level of batches, we then have to index into the zeroth row to remove the single batch dimension, and that gives us the time steps—just a one-dimensional array of all the indices. We convert this to a simple Python list from the PyTorch tensor, which we can feed into our decode function to convert those integers into text. Let me bring this back. We're generating 100 tokens. Let's run it, and here's the generation that we achieved. Obviously, it's garbage, and the reason it's garbage is because this is a totally random model. Next up, we're going to want to train this model. One more thing I want to point out is that this function is written to be general, but it's a bit ridiculous right now because we're feeding in all this context, concatenating it, and always feeding the entire sequence into the model. That's unnecessary, because this is just a simple bigram model. For example, to make a prediction about "K," we only need "W," but what we actually feed into the model is the entire sequence, and then we only look at the very last piece to predict "K." The only reason I'm writing it this way is because, right now, this is a bigram model, but I want to keep this function fixed so it will work later when our characters look further back in the history. At the moment, the history is not used, so this looks silly, but eventually the history will be used, and that's why we're doing it this way. Just a quick comment on that. Now we see that this is random. Let's train the model so it becomes a bit less random. First, I'm going to create a PyTorch optimization object. Here, we are using the Adam optimizer. In the "Make More" series, we've only ever used stochastic gradient descent, the simplest possible optimizer, which you can get using SGD instead. But I want to use Adam, which is a much more advanced and popular optimizer, and it works extremely well. A typical good setting for the learning rate is roughly 3e-4, but for very, very small networks, as is the case here, you can get away with much higher learning rates—3, or even higher, probably. Let me create the optimizer object, which will take the gradients and update the parameters using those gradients. Up above, our batch size was only four. Let me use something bigger—let's say 32. Then, for some number of steps, we sample a new batch of data, evaluate the loss, zero out all the gradients from the previous step, get the gradients for all the parameters, and then use those gradients to update our parameters. This is a typical training loop, as we saw in the "Make More" series. Let me now run this for, say, 100 iterations, and let's see what kind of losses we get. We started around 4.7, and now we're getting down to 4.6, 4.5, and so on. The optimization is definitely happening. Let's try increasing the number of iterations and only print at the end, because we probably want to train for longer. We're down to 3.6, roughly, and then down to three. This is the most janky optimization, but it's working. Let's just do 10,000 iterations. From here, we want to copy this, and hopefully we're going to get something reasonable. Of course, it's not going to be Shakespeare from a bigram model, but at least we see that the loss is improving, and hopefully we're expecting something a bit more reasonable. We're down at about 2.5—let's see what we get. Okay, dramatic improvements, certainly compared to what we had before. Let me just increase the number of tokens. We see that we're starting to get something at least reasonable. It's certainly not Shakespeare, but the model is making progress. This is the simplest possible model. Now, what I'd like to do is—obviously, this is a very simple model, because the tokens are not talking to each other. Given the previous context of whatever was generated, we're only looking at the very last character to make predictions about what comes next. Now, these tokens have to start talking to each other and figuring out what is in the context so they can make better predictions for what comes next. This is how we're going to kick off the Transformer. Next, I took the code that we developed in this Jupyter notebook and converted it to a script. I'm doing this because I want to simplify our intermediate work into just the final product that we have at this point. At the top, I put all the hyperparameters that we need to define. I introduced a few, and I'll speak to that in a little bit. Otherwise, a lot of this should be recognizable: reproducibility, reading data, getting the encoder and decoder, creating the train and test splits, and using a data loader that gets a batch of inputs and targets. This is new, and I'll talk about it in a second. Now, this is the bigram language model that we developed. It can perform a forward pass, give us logits and loss, and it can generate text. Here, we are creating the optimizer, and this is the training loop. Everything here should look pretty familiar. Now, some of the small things that I added. Number one, I added the ability to run on a GPU if you have one. If you have a GPU, this will use CUDA instead of just the CPU, and everything will be a lot faster. When the device becomes CUDA, we need to make sure that when we load the data, we move it to the device. When we create the model, we want to move the model parameters to the device. As an example, here we have the nn.Embedding table, and it's got a weight inside it, which stores the lookup table that would be moved to the GPU. All the calculations here happen on the GPU, and they can be a lot faster. Finally, when I'm creating the context that feeds into generate, I have to make sure that I create it on the device. Number two, what I introduced is the fact that here in the training loop, I was just printing the loss. Inside the training loop, I was printing the loss, but this is a very noisy measurement of the current loss because each batch will be more or less lucky. What I usually want to do is have an estimate_loss function. The estimate_loss function is called up here, and it averages the loss over multiple batches. In particular, we're going to iterate eval_iter times, get our loss, and then compute the average loss for both splits. This will be a lot less noisy. When we call estimate_loss, we're going to report the much more accurate training and validation loss. Now, when we come back up, you'll notice a few things. Here, I'm setting the model to evaluation mode, and down here, I'm resetting it back to training mode. For our model as it is, this doesn't do anything because the only thing inside this model is this nn.Embedding, and this network would behave the same in both evaluation and training modes. We have no dropout layers, no batch normalization layers, etc. However, it is good practice to think through what mode your neural network is in, because some layers will have different behavior at inference time versus training time. There's also this context manager, torch.no_grad, which just tells PyTorch that everything that happens inside this function will not require calling backward. PyTorch can be much more efficient with its memory usage because it doesn't have to store all the intermediate variables, since we're never going to call backward. It can be a lot more memory efficient in that way. It's also good practice to tell PyTorch when we don't intend to do backpropagation. This script is about 120 lines of code, and that's our starter code. I'm calling it b.p, and I'll release it later. Running this script gives us output in the terminal that looks something like this. As I ran the code, it gave me the train loss and validation loss, and we see that we converge to somewhere around 2.5 with the current model. Then, here's the sample we produced at the end. We have everything packaged up in the script, and we're in a good position now to iterate on this. We are almost ready to start writing our very first self-attention block for processing these tokens. Before we get there, I want to introduce you to a mathematical trick that is used in self-attention inside a transformer, and is really at the heart of an efficient implementation of self-attention. I want to work with a toy example to get you used to this operation, which will make it much clearer once we get to it in the script. Let's create a B by T by C tensor, where B, T, and C are just 4, 8, and 2 in the toy example. These are the batch size, time steps, and channels. We have batches, the time component, and information at each point in the sequence. Now, what we want to do is allow these tokens—up to eight tokens in a batch—to communicate with each other. Currently, these eight tokens are not interacting, and we want them to talk to each other. We want to couple them, but in a very specific way. For example, the token at the fifth location should not communicate with tokens in the sixth, seventh, or eighth locations, because those are future tokens in the sequence. The token at the fifth location should only talk to the ones in the fourth, third, second, and first positions. Information only flows from previous context to the current time step, and we cannot get any information from the future because we are about to try to predict the future. What is the easiest way for tokens to communicate? The simplest way, if we're at the fifth token and want to communicate with the past, is to just take an average of all the preceding elements. For example, if I'm the fifth token, I would take the channels that make up my information at this step, but also the channels from the fourth, third, second, and first steps. I would average those, and that would become a feature vector that summarizes me in the context of my history. Of course, just doing a sum or an average is an extremely weak form of interaction. This communication is extremely lossy—we've lost a lot of information about the spatial arrangement of all those tokens. That's okay for now; we'll see how to bring that information back later. For now, what we want to do is, for every single batch element independently, and for every token in that sequence, calculate the average of all the vectors in all the previous tokens and also at this token. Let's write that out. I have a small snippet here, and instead of just fumbling around, let me copy and paste it and talk through it. In other words, we're going to create X, and B is short for "bag of words" because that's a term people use when you are just averaging things. This is just a bag of words—there's a word stored at each of these eight locations, and we're doing a bag of words by averaging. In the beginning, we're going to say that it's just initialized at zero, and then I'm doing a for loop here. We're not being efficient yet—that's coming later—but for now, we're just iterating over all the batch dimensions independently, iterating over time. The previous tokens are at this batch dimension, and everything up to and including the T-th token. When we slice out X in this way, X_prev becomes a shape with however many T elements there were in the past, and then, of course, C—all the two-dimensional information from these little tokens. That's the previous chunk of tokens from my current sequence, and then I'm just doing the average, or the mean, over the zero dimension. I'm averaging out the time here, and I'm just going to get a little C-dimensional one-dimensional vector, which I'm going to store in X_bag_of_words. I can run this, but it's not going to be very informative. Let's see—this is X of zero, this is the zeroth batch element, and then X_bow at zero. Now, you see how at the first location here, the two are equal, and that's because we're just doing an average of this one token. But here, this one is now an average of these two, and now this one is an average of these three, and so on. This last one is the average of all of these elements. The vertical average, just averaging up all the tokens, now gives this outcome here. This is all well and good, but it's very inefficient. Now, the trick is that we can be very, very efficient about doing this using matrix multiplication. That's the mathematical trick, and let me show you what I mean. Let's work with a toy example here. Let me run it, and I'll explain. I have a simple matrix here that is a 3x3 matrix of all ones, a matrix B of just random numbers that's 3x2, and a matrix C, which will be the result of multiplying the 3x3 by the 3x2, giving a 3x2 matrix. Here, we're just using matrix multiplication: A multiplied by B gives us C. How are these numbers in C achieved? This number in the top left is the first row of A dot-producted with the first column of B, and since the row of A is all ones, the dot product with this column of B just sums the column. For example, 2 + 6 + 6 is 14. The element here in the output C, which is also in the first column, is the first row of A multiplied with the second column of B: 7 + 4 + 5 is 16. Now, you see that there are repeating elements here. This 14 appears again because this row is also all ones, and it's multiplying the first column of B, so we get 14 again. This last number here is the last row's dot product with the last column. Now, the trick here is the following: this is just a boring array of all ones, but PyTorch has this function called tril, which is short for "triangular lower," and you can wrap it in torch.ones, and it will just return the lower triangular portion of this matrix. Now, it will zero out these elements here, and we just get the lower triangular part. What happens if we do that? Now, we'll have A as this lower triangular matrix, and B as before, and now what are we getting in C? What is this number? This is the first row times the first column, and because these are zeros, those elements are now ignored, so we just get a two. This number here is the first row times the second column, and because the zeros are ignored, it's just seven. This seven multiplies this one, but look what happened here: because this is one and then zeros, what ended up happening is we're just plucking out this row of B, and that's what we got. Now, here we have one, one, zero—so 1, 1, 0. The dot product with these two columns will now give us 2 + 6, which is 8, and 7 + 4, which is 11. Because this is 1, 1, 1, we end up with the addition of all of them. Depending on how many ones and zeros we have here, we're doing a sum of a variable number of these rows, and that gets deposited into C. Currently, we're doing sums because these are ones, but we can also do averages. You can start to see how we could do an average of the rows of B in an incremental fashion, because we can normalize these rows so that they sum to one, and then we're going to get an average. If we took A, and then we did a call to torch. Sum along the first dimension, and then let's keep them as true. Therefore, the broadcasting will work out. If I rerun this, you can see that these rows now sum to one. This row is one, this row is 0.5 and 0.5, and here we get one-third for each entry. Now, when we do A multiplied by B, what are we getting? We are just getting the first row here. Next, we get the average of the first two rows: 2 and 6, whose average is 4; and 4 and 7, whose average is 5.5. On the bottom, we are now getting the average of all three rows. The average of all the elements of B are now deposited here. You can see that by manipulating the elements of this multiplying matrix, and then multiplying it with any given matrix, we can perform these averages incrementally. We can manipulate that based on the elements of A, which is very convenient. Let's go back up here and see how we can vectorize this and make it much more efficient using what we've learned. In particular, we are going to produce an array A, but here I'm going to call it W, short for weights. This is our A, and it represents how much of every row we want to average. It's going to be an average because you can see that these rows sum to one. This is our A, and our B in this example is X. What's going to happen now is that we are going to have an expo2, and this expo2 is going to be W multiplied by RX. Let's think this through. W is T by T, and this is matrix multiplication in PyTorch: a B by T by C, and it's giving us—what shape? PyTorch will see that these shapes are not the same, so it will create a batch dimension here. This is a batched matrix multiply, and it will apply this matrix multiplication to all the batch elements in parallel and individually. For each batch element, there will be a T by T multiplying a T by C, exactly as we had below. This will now create a B by T by C, and expo2 will now become identical to expo. We can see that torch.allclose of expo and expo2 should be true. This kind of convinces us that these are, in fact, the same—expo and expo2. If I just print them—okay, we're not going to be able to just stare it down, but let me try expo at the zeroth element and expo2 at the zeroth element, just the first batch. We should see that these are identical, which they are. What happened here? The trick is that we were able to use batched matrix multiplication to do this aggregation, and it's a weighted aggregation. The weights are specified in this T by T array, and we're doing weighted sums according to the weights inside here. They take on a sort of triangular form, which means that a token at the t-th position will only get information from the tokens preceding it. That's exactly what we want. Finally, I would like to rewrite it in one more way, and we're going to see why that's useful. This is the third version, and it's also identical to the first and second, but let me talk through it. It uses softmax. "Tril" here is this matrix of lower triangular ones. W begins as all zeros. If I just print W in the beginning, it's all zeros. Then I use masked_fill. What this is doing is. Masked fill sets all elements to zero, and I'm saying that for all the elements where "tril" is equal to zero, make them negative infinity. All the elements where "tril" is zero will become negative infinity. Now, this is what we get, and then the final line here is softmax. If I take a softmax along every single dimension—specifically, dimension negative one, which is along every single row—what is that going to do? Well, softmax is also a normalization operation, and spoiler alert: you get the exact same matrix. Let me bring us back to softmax and recall that in softmax, we're going to exponentiate every single one of these elements and then divide by the sum. If we exponentiate every single element here, we're going to get a one, and everywhere else we're going to get zero. Then, when we normalize, we just get one. Here, we're going to get ones and then zeros, and softmax will again divide, giving us 0.5 and so on. This is also the same way to produce this mask. Now, the reason that this is a bit more interesting—and the reason we're going to end up using it in self-attention—is that these weights here begin with zero. You can think of this as an interaction strength or an affinity. It's telling us how much of each token from the past we want to aggregate and average. This line is saying tokens from the past cannot communicate. By setting them to negative infinity, we're saying that we will not aggregate anything from those tokens. This then goes through softmax and through the weighted aggregation, which is performed via matrix multiplication. What this means is that these zeros are currently just set by us to be zero, but a quick preview is that these affinities between the tokens are not going to be constant at zero. They're going to be data-dependent. The tokens are going to start looking at each other, and some tokens will find other tokens more or less interesting. Depending on what their values are, they're going to find each other interesting to different amounts. I'm going to call those affinities. Here, we are saying the future cannot communicate with the past. We're going to clamp them, and then when we normalize and sum, we're going to aggregate their values depending on how interesting they find each other. That's the preview for self-attention. Long story short, you can do weighted aggregations of your past elements by using matrix multiplication in a lower triangular fashion. The elements in the lower triangular part are telling you how much of each element fuses into this position. We're going to use this trick now to develop the self-attention block. First, let's get some quick preliminaries out of the way. The thing I'm kind of bothered by is that you see how we're passing in vocab size into the constructor. There's no need to do that, because vocab size is already defined up top as a global variable. There's no need to pass this stuff around. Next, what I want to do is create a level of indirection here, where we don't directly go to the embedding for the logits. Instead, we go through this intermediate phase, because we're going to start making that bigger. Let me introduce a new variable, n_embed—short for number of embedding dimensions. n_embed here will be, say, 32. That was a suggestion from GitHub Copilot, by the way. It also suggests 32, which is a good number. This is an embedding table with only 32-dimensional embeddings. Here, this is not going to give us logits directly. Instead, this is going to give us token embeddings, as I'm going to call them. To go from the token embeddings to the logits, we're going to need a linear layer. Let's call it "self.lm_head," short for language modeling head, which is a linear layer from n_embed up to vocab size. Then, when we swing over here, we're going to get the logits by exactly what Copilot suggests. Now, we have to be careful here, because this C and this C are not equal. This is n_embed C, and this is vocab size. Let's just say that n_embed is equal to C, and then this just creates one sparse layer of interaction through a linear layer. But this should run. We see that this runs, and this currently looks kind of sparse, but we're going to build on top of this now. So far, we've taken these indices and encoded them based on the identity of the tokens inside idx. The next thing that people very often do is not just encode the identity of these tokens, but also their position. We're going to have a second position embedding table here: self.position_embedding_table, which is an embedding of block size by n_embed. Each position from zero to block size minus one will also get its own embedding vector. First, let me decode B, T from idx.shape. Here, we're also going to have a pos embedding, which is the positional embedding. This will be just integers from zero to T minus one. All of those integers from zero to T minus one get embedded through the table to create a T by C matrix. Here, this gets renamed to just "x," and x will be the addition of the token embeddings with the positional embeddings. The broadcasting note will work out: B by T by C plus T by C. This gets aligned, a new dimension of one gets added, and it gets broadcasted across the batch. At this point, x holds not just the token identities but also the positions at which these tokens occur. This is currently not that useful because, of course, we just have a simple bag-of-words model. It doesn't matter if you're in the fifth position, the second position, or wherever—it's all translation invariant at this stage. This information currently wouldn't help, but as we work on the self-attention block, we'll see that this starts to matter. Now we get to the crux of self-attention. This is probably the most important part of this video to understand. We're going to implement a small self-attention mechanism for a single individual head, as they're called. We start off where we left off; all of this code should be familiar. Now I'm working with an example where I change the number of channels from 2 to 32. We have a 4x8 arrangement of tokens, and each token is currently 32-dimensional, but we're just working with random numbers. Previously, we saw that the code does a simple weighted average of all the past tokens and the current token. The previous information and current information are just being mixed together in an average, and that's what this code currently achieves. It does this by creating a lower triangular structure, which allows us to mask out the weight matrix that we create. We mask it out, and then we normalize it. Currently, when we initialize the affinities between all the different tokens—or nodes; I'm going to use those terms interchangeably—to be zero, we see that the weight matrix gives us a structure where every single row has uniform numbers. That's what, in this matrix multiplication, results in doing a simple average. Now, we don't want this to be all uniform because different tokens will find different other tokens more or less interesting, and we want that to be data-dependent. For example, if I'm a vowel, maybe I'm looking for consonants in my past, and I want to know what those consonants are. I want that information to flow to me, and I want to gather information from the past, but I want to do it in a data-dependent way. This is the problem that self-attention solves. The way self-attention solves this is as follows: every single node, or token at each position, will emit two vectors. It will emit a query and a key. The query vector, roughly speaking, represents "what am I looking for," and the key vector represents "what do I contain." To get affinities between these tokens in a sequence, we just do a dot product between the keys and the queries. My query is dot-multiplied with all the keys of all the other tokens, and that dot product now becomes the weight. If the key and the query are aligned, they will interact strongly, and I will get to learn more about that specific token as opposed to any other token in the sequence. Let's implement this. Now we're going to implement a single head of self-attention. This is just one head. There's a hyperparameter involved with these heads, which is the head size. Here, I'm initializing linear modules and using bias equals false. These are just going to apply a matrix multiplication with some fixed weights. Now, let me produce a key and a query—k and q—by forwarding these modules on x. The size of this will now become B by T by 16, because that is the head size, and the same here: B by T by 16, with this being the head size. You see here that when I forward this linear layer on top of x, all the tokens in all the positions in the B by T arrangement, all of them in parallel and independently, produce a key and a query. No communication has happened yet, but the communication comes now. All the queries will do a dot product with all the keys. What we want is the weight matrix—the affinities between these tokens—to be the query multiplied by the key, but we have to be careful. We can't just matrix-multiply this; we need to transpose k. We also have to be careful because, when you have the batch dimension in particular, we want to transpose the last two dimensions: dimension -2 and dimension -1. This matrix multiplication will do the following: a B by T by 16 matrix multiplies with a B by 16 by T matrix to give us B by T by T. For every row of B, we're now going to have a T squared matrix giving us the affinities, and these are now the weights. They're not zeros anymore; they are now coming from this dot product between the keys and the queries. This can now run. I can run this, and the weighted aggregation is now a function in a data-dependent manner between the keys and queries of these nodes. Just inspecting what happened here, the weight matrix takes on this form. You see that before, the weight matrix was just a constant applied in the same way to all the batch elements, but now every single batch element will have a different weight matrix because every single batch element contains different tokens at different positions, and this is now data-dependent. When we look at just the zeroth row, for example, in the input, these are the weights that came out, and you can see now that they're not exactly uniform. In particular, as an example, for the last row—this was the eighth token—the eighth token knows what content it has and at what position it's in. Now, the E token, based on that, creates a query: "Hey, I'm looking for this kind of stuff. I'm a vowel, I'm in the E position, I'm looking for any consonant at positions up to four." All the nodes emit keys, and maybe one of the channels could be, "I am a consonant, and I am in a position up to four," and that key would have a high number in that specific channel. That's how the query and the key, when they do a dot product, can find each other and create a high affinity. When they have a high affinity—say, this token was pretty interesting to this eighth token—then, through the softmax, I will end up aggregating a lot of its information into my position, and I'll get to learn a lot about it. Now, at this point, we're looking at the process after everything has already happened. Let me erase this operation as well—let me erase the masking and the softmax—just to show you the under-the-hood internals and how everything works without the masking and the softmax. What comes out here are the outputs of the dot products. These are the raw outputs, and they take on values from negative two to positive two, and so on. That's the raw interaction and raw affinity between all the nodes. But now, if I'm the fifth node, I do not want to aggregate anything from the sixth, seventh, or eighth nodes. We use the upper triangular masking so those nodes are not allowed to communicate. Now, we want to have a nice distribution—we don't want to aggregate, say, negative 0.11 of this node, because that's crazy. Instead, we exponentiate and normalize, and now we get a nice distribution that sums to one. This tells us, in a data-dependent manner, how much information to aggregate from any of these tokens in the past. In this way, the values are no longer just zeros, but are calculated based on this process. Now, there's one more part to a single self-attention head. When we do the aggregation, we don't aggregate the tokens themselves exactly. We produce one more value here, which we call the value. In the same way that we produced key and query, we're also going to create a value. Here, we don't aggregate x directly; instead, we calculate v, which is achieved by applying a linear transformation on top of x again, and then we output the weights multiplied by v. V represents the elements or vectors that we aggregate, instead of the raw x. Of course, this means that the output of this single head will be 16-dimensional, because that is the head size. You can think of x as kind of private information to this token. If you think about it that way, x is private to this token. For example, if I'm the fifth token, I have some identity, and my information is kept in vector x. Now, for the purposes of this single head, here's what I'm interested in, here's what I have, and if you find me interesting, here's what I will communicate to you—and that's stored in v. V is the thing that gets aggregated for the purposes of this single head between the different nodes, and that's the self-attention mechanism. This is what it does. There are a few notes I would like to make about attention. Number one: attention is a communication mechanism. You can really think of it as a communication mechanism where you have a number of nodes in a directed graph, with edges pointing between nodes. What happens is that every node has some vector of information, and it gets to aggregate information via a weighted sum from all of the nodes that point to it. This is done in a data-dependent manner, depending on whatever data is stored at each node at any point in time. Now, our graph doesn't look exactly like this; it has a different structure. We have eight nodes because the block size is eight, so there are always eight tokens. The first node is only pointed to by itself. The second node is pointed to by the first node and itself, and so on, all the way up to the eighth node, which is pointed to by all the previous nodes and itself. That's the structure our directed graph happens to have in an autoregressive scenario, such as language modeling. But in principle, attention can be applied to any arbitrary directed graph—it's just a communication mechanism between the nodes. The second note is that there is no notion of space in attention. Attention simply acts over a set of vectors in this graph, and by default, these nodes have no idea where they are positioned in space. That's why we need to encode them positionally and give them some information anchored to a specific position, so they know where they are. This is different, for example, from convolution. If you run a convolution operation over some input, there's a very specific layout of the information in space, and the convolutional filters act in space. Attention, on the other hand, just operates over a set of vectors out there in space—they communicate, and if you want them to have a notion of space, you need to specifically add it. That's what we've done when we calculated the positional encodings and added that information to the vectors. The next thing that I hope is very clear is that the elements across the batch dimension—which are independent examples—never talk to each other. They're always processed independently. This is a batched matrix multiplication that applies a matrix multiplication in parallel across the batch dimension. Maybe it would be more accurate to say that, in this analogy of a directed graph, because the batch size is four, we really have four separate pools of eight nodes. Those eight nodes only talk to each other, but in total there are 32 nodes being processed. However, you can think of it as four separate pools of eight nodes each. The next note is that, in the case of language modeling, we have this specific structure of a directed graph where future tokens do not communicate with past tokens. However, this does not necessarily have to be a constraint in the general case. In fact, in many cases, you may want all of the nodes to communicate with each other fully. For example, if you're doing sentiment analysis with a transformer, you might have a number of tokens and want them all to interact, because later you are predicting the sentiment of the sentence. In this scenario, it's acceptable for these nodes to talk to each other. In those cases, you would use an encoder block of self-attention. All it means for it to be an encoder block is that you would delete this line of code, allowing all the nodes to communicate freely. What we're implementing here is sometimes called a decoder block. It's called a decoder because it is used for decoding language, and it has this autoregressive format where you have to mask with a triangular matrix so that nodes from the future never communicate with the past, as that would give away the answer. In encoder blocks, you would delete this mask and allow all the nodes to communicate. In decoder blocks, this triangular structure is always present. Both are valid, and attention itself does not care; it supports arbitrary connectivity between nodes. The next thing I wanted to comment on is that you keep hearing me say "attention," "self-attention," and so on. There's also something called cross-attention. What is the difference? The reason this attention is called self-attention is because the keys, queries, and values all come from the same source, from x. The same source x produces the keys, queries, and values, so these nodes are self-attending. However, in principle, attention is much more general than that. For example, in encoder-decoder transformers, you can have a case where the queries are produced from x, but the keys and values come from a separate, external source—sometimes from encoder blocks that encode some context we want to condition on. The keys and values will come from this separate source, which are nodes on the side, and here we're just producing queries and reading information from the side. Cross-attention is used when there is a separate source of nodes from which we want to pull information into our nodes. It is self-attention if we just have nodes that want to look at and communicate with each other. The attention here happens to be self-attention, but in principle, attention is much more general. The last note at this stage is that, if we look at the "Attention Is All You Need" paper, we've already implemented attention. Given query, key, and value, we've multiplied the query and the key, applied softmax, and then aggregated the values. There's one more thing we're missing here, which is dividing by one over the square root of the head size. The dk here is the head size. Why are they doing this? They call it scaled attention, and it's an important normalization to have. The problem is, if you have unit Gaussian inputs—zero mean, unit variance—then k and q are also unit Gaussian. If you just do the multiplication naively, the variance will be on the order of the head size, which in our case is 16. But if you multiply by one over the square root of the head size, the variance will be one; it will be preserved. Why is this important? You'll notice that this value will feed into softmax, and it's important, especially at initialization, that the distribution be fairly diffuse. In our case here, we lucked out and had fairly diffuse numbers. The problem is that, because of softmax, if the weights take on very positive and very negative numbers, softmax will converge toward one-hot vectors. I can illustrate that here: say we apply softmax to a tensor of values that are very close to zero, then we'll get a diffuse output from softmax. But if I take the same tensor and start sharpening it by multiplying these numbers by eight, for example, you'll see that softmax will start to sharpen, and it will sharpen toward the maximum value—toward whichever number is the highest. We don't want these values to be too extreme, especially at initialization; otherwise, softmax will be too peaky, and each node will aggregate information from only a single other node. That's not what we want, especially at initialization, and the scaling is used to control the variance at initialization. Having said all that, let's now take our self-attention knowledge and put it to use. Here in the code, I created this head module, which implements a single head of self-attention. You give it a head size, and it creates the key, query, and value linear layers. Typically, people don't use biases in these layers; these are the linear projections that we apply to all of our nodes. Now, here I'm creating this "tril" variable. "Tril" is not a parameter of the module; in PyTorch naming conventions, this is called a buffer. It's not a parameter, and you have to assign it to the module using "register_buffer." That creates the "tril," the lower triangular matrix. When we're given the input x, this should look very familiar now. We calculate the keys and the queries, compute the attention scores inside "wei," normalize them using scaled attention, ensure that the future does not communicate with the past (which makes it a decoder block), apply softmax, aggregate the values, and output the result. Then here in the language model, I'm creating a head in the constructor and calling it a self-attention head. For now, I'm going to keep the head size the same as n_embed. Once we've encoded the information with the token embeddings and the position embeddings, we're simply going to feed it into the self-attention head. The output of that then goes into the decoder language modeling head to create the logits. This is the simplest way to plug in a self-attention component. Now, next up, in the "Attention Is All You Need" paper, there's something called multi-head attention. What is multi-head attention? It's just applying multiple attentions in parallel and concatenating their results. They have a diagram here—I'm not sure if it's super clear—but it's really just multiple attentions in parallel. Let's implement that. It's fairly straightforward. If we want multi-head attention, then we want multiple heads of self-attention running in parallel. In PyTorch, we can do this by simply creating as many heads as we want, and then deciding what the head size of each should be. We run all of them in parallel, collect them into a list, and simply concatenate all of the outputs over the channel dimension. The way this looks now is that we don't have just a single attention head with a head size of 32—because remember, n_head is 32. Instead of having one communication channel, we now have four communication channels in parallel, and each of these channels will typically be smaller. Because we have four communication channels, we want eight-dimensional self-attention from each channel. Each channel produces eight-dimensional vectors, and then we have four of them, which concatenate to give us 32, matching the original n_embed. This is similar to group convolution in convolutional neural networks: instead of having one large convolution, we do convolution in groups, and that's what multi-headed self-attention is. Here, we just use self-attention heads instead. Now, I ran it, and scrolling down, I ran the same thing. We now get the validation loss down to roughly 2.28. The output generation is still not amazing, but clearly the validation loss is improving, because we were at 2.4 just before. It helps to have multiple communication channels, because these tokens have a lot to communicate. They want to find the consonants, the vowels, or certain patterns from specific positions. They want to find all kinds of different things, and it helps to create multiple independent channels of communication, gather lots of different types of data, and then decode the output. Going back to the paper for a second, I haven't explained this figure in full detail, but we are starting to see some components of what we've already implemented. We have the positional encodings, the token encodings that are added, and we have the masked multi-headed attention implemented. Now, here's another multi-headed attention, which is a cross-attention to an encoder, but we aren't going to implement that in this case. I'll come back to that later. I want you to notice that there's a feedforward part here, and this is grouped into a block that gets repeated again and again. The feedforward part here is just a simple multi-layer perceptron. The position-wise feedforward networks are just simple little MLPs. I want to start in a similar fashion by also adding computation into the network, and this computation is on a per-node level. I've already implemented it, and you can see the diff highlighted on the left here where I've added or changed things. Before, we had the multi-headed self-attention that handled the communication, but we went too quickly to calculating the logits. The tokens looked at each other, but didn't really have a lot of time to process what they found from the other tokens. What I've implemented here is a small feedforward single layer. This layer is just a linear transformation followed by a ReLU nonlinearity, and that's it. It's just a little layer, which I call feedforward_n_embed. This feedforward layer is called sequentially after the self-attention. We self-attend, then we feedforward. You'll notice that the feedforward layer, when applying the linear transformation, operates on a per-token level. All the tokens do this independently. The self-attention is for communication, and once they've gathered all the data, they need to process that data individually, which is what the feedforward layer does. That's why I've added it here. Now, when I train this, the validation loss continues to go down, now to 2. 2.4, which is down from 2.28. The output still looks kind of terrible, but at least we've improved the situation. As a preview, we're now going to start interspersing the communication with the computation, and that's also what the Transformer does. It has blocks that communicate and then compute, groups them, and replicates them. Let me show you what we want to do. We want to do something like this: we have a block, and this block is this part here, except for the cross-attention. Now, the block intersperses communication and then computation. The communication is done using multi-headed self-attention, and then the computation is done using a feedforward network on all the tokens independently. What I've added here is that you'll notice this takes the number of embeddings in the embedding dimension and the number of heads that we want, which is kind of like group size in group convolution. I'm saying that the number of heads we want is four, and because this is 32, we calculate that the number of heads should be four, the head size should be eight, and everything works out channel-wise. This is how the Transformer structures the sizes. Typically, the head size will be eight, and this is how we want to intersperse them. Here, I'm trying to create blocks, which is just a sequential application of block, block, block, where we're interspersing communication and feedforward many, many times, and then finally we decode. Now, I tried to run this, and the problem is that this doesn't give a very good answer or a very good result. The reason for that is we're starting to get a pretty deep neural net, and deep neural nets suffer from optimization issues. I think that's what we're starting to run into. We need one more idea that we can borrow from the Transformer paper to resolve those difficulties. There are two optimizations that dramatically help with the depth of these networks and make sure that the networks remain optimizable. Let's talk about the first one. The first one in this diagram is, you see this arrow here, and then this arrow, and this arrow. Those are skip connections, or sometimes called residual connections. They come from the paper "Deep Residual Learning for Image Recognition" from about 2015, which introduced the concept. What this means is you transform data, but then you have a skip connection with addition from the previous features. The way I like to visualize it is the following: here, the computation happens from top to bottom, and you have this residual pathway. You are free to fork off from the residual pathway, perform some computation, and then project back to the residual pathway via addition. You go from the inputs to the targets only via plus and plus-plus. The reason this is useful is because during backpropagation—remember from our micrograd video earlier—addition distributes gradients equally to both of its branches that fed in as input. The supervision, or the gradients from the loss, hop through every addition node all the way to the input, and then also fork off into the residual blocks. You have this gradient superhighway that goes directly from the supervision all the way to the input unimpeded, and then these viral blocks are usually initialized in the beginning. They contribute very, very little, if anything, to the residual pathway. They are initialized that way in the beginning; they are almost not there. But during optimization, they come online over time and start to contribute. At least at initialization, you can go directly from supervision to the input, the gradient is unimpeded and just flows, and then the blocks over time kick in. That dramatically helps with optimization. Let's implement this. Coming back to our block here, what we want to do is we want to do x, x plus self-attention, and x, x plus self. Feedforward: this is x, and then we fork off and do some communication and come back, and we fork off and do some computation and come back. Those are residual connections. Swinging back up here, we also have to introduce this projection, nn.Linear, and this is going to be from after we concatenate. This is the projection and embedding. This is the output of the self-attention itself, but then we want to apply the projection, and that's the result. The projection is just a linear transformation of the outcome of this layer. That's the projection back into the residual pathway. In the feedforward, it's going to be the same thing. I could have a self.projection here as well, but let me just simplify it and couple it inside the same sequential container. This is the projection layer going back into the residual pathway, and that's it. Now we can train this. I implemented one more small change. When you look into the paper again, you see that the dimensionality of the input and output is 512 for them, and they're saying that the inner layer in the feedforward has a dimensionality of 2048. There's a multiplier of four, so the inner layer of the feedforward network should be four times larger in terms of channel size. I came here and multiplied four times embed for the feedforward, and then from four times embed, we come back down to embed when we return to the projection. This adds a bit of computation and grows the layer inside the residual block on the side of the residual pathway. Then I trained this, and we got the validation loss down to 2.08. We also see that the network is starting to get big enough that our training loss is getting ahead of the validation loss. We're starting to see a little bit of overfitting, and our generations here are still not amazing, but at least you can see that what we get now—this "grief syn"—is starting to almost look like English. We're really starting to get there. The second innovation that is very helpful for optimizing very deep neural networks is here. We have this addition now—that's the residual part—but this "norm" refers to something called layer norm. Layer norm is implemented in PyTorch. It's based on a paper that came out a while back, and layer norm is very similar to batch norm. Remember back to our Make More series, part three, where we implemented batch normalization. Batch normalization just made sure that, across the batch dimension, any individual neuron had a unit Gaussian distribution: zero mean and unit standard deviation. What I did here is copy-paste the BatchNorm1d that we developed in our Make More series. For example, we can initialize this module and have a batch of thirty-two 100-dimensional vectors feeding through the batch norm layer. What this does is guarantee that, when we look at just the zeroth column, it has zero mean and one standard deviation. It's normalizing every single column of this input. The rows are not going to be normalized by default because we're just normalizing columns. Now, let's implement layer norm. It's very straightforward. We come here, change this from zero to one, and instead of normalizing the columns, we normalize the rows. Now we've implemented layer norm. The columns are not going to be normalized, but the rows will be. For every individual example, its 100-dimensional vector is normalized in this way. Because our computation now does not span across examples, we can delete all of this buffer stuff. We can always apply this operation and don't need to maintain any running buffers. There's no distinction between training and test time, and we don't need these running buffers. We do keep gamma and beta, but we don't need momentum, and we don't care if it's training or not. This is now a layer norm, and it normalizes the rows instead of the columns. This here is identical to this here. Let's now implement layer norm in our transformer. Before I incorporate the layer norm, I just wanted to note that, as I said, very few details about the transformer have changed in the last five years, but this is something that slightly departs from the original paper. In the original, the add and norm is applied after the transformation, but now it is more common to apply the layer norm before the transformation. There's a reshuffling of the layer norms. This is called the pre-norm formulation, and that's the one we're going to implement as well. It's a slight deviation from the original paper. We need two layer norms. Layer norm one is nn.LayerNorm, and we tell it the embedding dimension. We need a second layer norm as well. Here, the layer norms are applied immediately on x itself: layer norm one applied on x, and layer norm two applied on x before it goes into self-attention and feedforward. The size of the layer norm here is embed, which is 32. When the layer norm is normalizing our features, the normalization happens over 32 numbers. The batch and the time both act as batch dimensions. This is a per-token transformation that just normalizes the features and makes them unit mean and unit variance at initialization. Of course, because these layer norms have gamma and beta training parameters, the layer norm will eventually create outputs that might not be unit Gaussian, but the optimization will determine that. For now, this incorporates the layer norms. I let it run, and we see that we get down to 2.06, which is better than the previous 2.08—a slight improvement by adding the layer norms. I'd expect that they help even more if we had a bigger and deeper network. One more thing I forgot to add is that there should also be a layer norm at the end of the transformer, typically before the final linear layer that decodes into the vocabulary. I added that as well. At this stage, we have a pretty complete transformer according to the original paper, and it's a decoder-only transformer. I'll talk about that in a second, but at this stage, the major pieces are in place. We can try to scale this up and see how well we can push this number. Now, in order to scale out the model, I had to make some cosmetic changes to improve the code. I introduced a variable called n_layer, which specifies how many layers of blocks we are going to have. I created a set of blocks, and we now have a new variable for the number of heads as well. I pulled out the layer normalization here, and this part remains identical. One thing I briefly changed is that I added Dropout. Dropout is something you can add before the residual connection, before the connection back into the residual pathway. We can apply Dropout at the layer here, at the end of the multi-headed attention as well, and also when we calculate the affinities—after the softmax, we can drop out some of those values. By randomly preventing some of the nodes from communicating, Dropout helps regularize the model. Dropout comes from a 2014 paper, and it works by randomly, on every forward and backward pass, shutting off some subset of neurons—randomly dropping them to zero and training without them. Because the mask of what's being dropped out changes every single forward and backward pass, it effectively trains an ensemble of sub-networks. Then, at test time, everything is fully enabled, and all of those sub-networks are merged into a single ensemble, if you want to think about it that way. I recommend reading the paper for full details. For now, we'll just stay at the level that this is a regularization technique, and I added it because I'm about to scale up the model significantly and was concerned about overfitting. When we scroll up to the top, you'll see that I changed a number of hyperparameters for our neural net. I made the batch size much larger—now it's 64. I changed the block size to 256; previously, it was just eight characters of context, but now it is 256 characters of context to predict the 257th. I lowered the learning rate a bit because the neural net is now much bigger. The embedding dimension is now 384, and there are six heads. 384 divided by 6 means that every head is 64-dimensional, which is standard. There are going to be six layers, and Dropout is set to 0.2. That means every forward and backward pass, 20% of all these intermediate calculations are disabled and dropped to zero. I already trained this and ran it. Drum roll—how well does it perform? Let me just scroll up here. We get a validation loss of 1.48, which is quite an improvement over what we had before, which I think was 2.07. It went from 2.07 all the way down to 1.48 just by scaling up this neural net with the code we have. Of course, this ran for a lot longer. I believe this trained for about 15 minutes on my A100 GPU. That's a pretty powerful GPU, and if you don't have a GPU, you won't be able to reproduce this on a CPU. I would not run this on a CPU or a MacBook or something similar—you'll have to reduce the number of layers, the embedding dimension, and so on. But in about 15 minutes, we can get this kind of result. I'm printing some of the Shakespeare output here, but I also printed 10,000 characters—much more—and wrote them to a file. Here we see some of the outputs. It's much more recognizable as the input text file. The input text file, just for reference, looked like this: there's always someone speaking in this manner, and our predictions now take on that form, except of course they're nonsensical when you read them. For example, "It is every crimp tap be a house, oh those prepation we give heed, Oho sent me you mighty Lord." Anyway, you can read through this—it's nonsensical, of course, but this is just a Transformer trained on a character level for 1 million characters from Shakespeare. It sort of blabbers on in a Shakespearean manner, but it doesn't make sense at this scale. Still, I think it's a pretty good demonstration of what's possible. Now, I think that concludes the programming section of this video. We did a pretty good job of implementing this Transformer, but the picture doesn't exactly match up to what we've done. What's going on with all these additional parts here? Let me finish explaining this architecture and why it looks unusual. What's happening here is that what we implemented is a decoder-only Transformer. There's no component here called the encoder, and there's no cross-attention block. Our block only has self-attention and the feed-forward layer. It is missing this third, in-between piece. That piece does cross-attention—we don't have it, and we don't have the encoder; we just have the decoder. The reason we have a decoder-only model is because we are just generating text, and it's unconditioned on anything. We're just generating sequences according to a given dataset. What makes it a decoder is that we are using the triangular mask in our Transformer. It has this autoregressive property where we can just go and sample from it. The fact that it's using the triangular mask to mask out the attention makes it a decoder, and it can be used for language modeling. Now, the reason the original paper had an encoder-decoder architecture is because it was a machine translation paper. It was concerned with a different setting. In particular, it expects some tokens that encode, for example, French, and then it is expected to decode the translation in English. Typically, these are special tokens. You are expected to read this in and condition on it, and then you start the generation with a special token called "start." This is a new special token that you introduce and always place at the beginning. The network is then expected to output "neural networks are awesome," followed by a special end token to finish the generation. This part here will be decoded exactly as we've done it. "Neural networks are awesome" will be identical to what we did, but unlike our approach, they want to condition the generation on some additional information. In this case, the additional information is the French sentence that should be translated. What they do now is bring in the encoder. The encoder reads this part here—we're only going to take the French portion—and we're going to create tokens from it exactly as we've seen in our video. We then put a Transformer on it, but there is no triangular mask, so all the tokens are allowed to attend to each other as much as they want. They are just encoding whatever is in the content of this French sentence. Once they've encoded it, the outputs come out at the top. Then, in our decoder, which does the language modeling, there's an additional connection to the outputs of the encoder, brought in through cross-attention. The queries are still generated from X, but now the keys and values come from the side—the top—generated by the nodes that came out of the encoder. Those keys and values feed in from the side into every single block of the decoder, and that's why there's an additional cross-attention. What it's really doing is conditioning the decoding not just on the past of the current decoding, but also on having seen the fully encoded French prompt. It's an encoder-decoder model, which is why we have those two Transformers—an additional block, and so on. We did not do this because we have nothing to encode; there's no conditioning. We just have a text file, and we want to imitate it. That's why we are using a decoder-only Transformer, exactly as done in GPT. Now, I want to do a very brief walkthrough of nanoGPT, which you can find on my GitHub. NanoGPT consists of two files of interest: train.py and model.py. Train.py contains all the boilerplate code for training the network. It includes all the stuff we've discussed here—the training loop—but it's a lot more complicated because we're saving and loading checkpoints and pre-trained weights, decaying the learning rate, compiling the model, and using distributed training across multiple nodes or GPUs. The train.py file gets a little bit more complicated; there are more options, etc. However, the model.py file should look very, very similar to what we've done here. In fact, the model is almost identical. First, we have the causal self-attention block, and all of this should look very familiar to you. We're producing queries, keys, and values, doing dot products, masking, applying softmax, optionally dropping out, and here we are pulling the values. What is different here is that in our code, I separated out the multi-headed attention into a single individual head, and then I have multiple heads and explicitly concatenate them. In nanoGPT, all of this is implemented in a batched manner inside a single causal self-attention, and we don't just have B, T, and C dimensions—we also end up with a fourth dimension, which is the heads. It gets a bit more complex because we have four-dimensional array tensors now, but it is mathematically equivalent. The exact same thing is happening as in our implementation; it's just a bit more efficient because all the heads are now treated as a batch dimension as well. Then we have the multi-layer perceptron, which uses the GELU nonlinearity, defined here, instead of ReLU. This is done simply because OpenAI used GELU, and I want to be able to load their checkpoints. The blocks of the Transformer are identical in the compute phase as we saw, and the GPT model itself is identical. We have the position encodings, token encodings, the blocks, the layer norm at the end, and the final linear layer. This should all look very familiar. There is a bit more here because I'm loading checkpoints and separating out the parameters into those that should be weight-decayed and those that shouldn't, but the generate function should also be very similar. A few details are different, but you should definitely be able to look at this file and understand the pieces. Now, let's bring things back to ChatGPT. What would it look like if we wanted to train ChatGPT ourselves, and how does it relate to what we learned today? To train ChatGPT, there are roughly two stages: first is the pre-training stage, and then the fine-tuning stage. In the pre-training stage, we train on a large chunk of the internet and try to get a first decoder-only Transformer to generate text. It's very similar to what we've done ourselves, except we've only done a tiny, basic pre-training step. In our case, for example, this is how you print the number of parameters—I printed it, and it's about 10 million. The Transformer I created here, the little Shakespeare Transformer, had about 10 million parameters. Our dataset is roughly 1 million characters, or about 1 million tokens. However, you have to remember that OpenAI uses a different vocabulary. They're not at the character level; they use subword chunks, and their vocabulary is about 50,000 elements. Their sequences are a bit more condensed. Our Shakespeare dataset would probably be around 300,000 tokens in the OpenAI vocabulary, roughly. So, we trained a model with about 10 million parameters on roughly 300,000 tokens. Now, if you look at the GPT-3 paper and the Transformers they trained, they trained a number of Transformers of different sizes, but the largest Transformer has 175 billion parameters. Ours, again, is 10 million. They used a certain number of layers in the Transformer—this is the n_embd, this is the number of heads, this is the head size, and then this is the batch size. Ours was 65, and the learning rate is similar. When they trained this Transformer, they trained on 300 billion tokens. Remember, ours is about 300,000 tokens—this is about a millionfold increase. By today's standards, that number wouldn't even be considered large; now you'd be going up to 1 trillion tokens and above. They are training a significantly larger model on a much larger portion of the internet, and that is the pre-training stage. Otherwise, these hyperparameters should be fairly recognizable to you, and the architecture is nearly identical to what we implemented ourselves. Of course, it's a massive infrastructure challenge to train this—you're talking about thousands of GPUs communicating with each other to train models of this size. That's just the pre-training stage. After you complete the pre-training stage, you don't get something that responds to your questions with helpful answers. Instead, you get a document completer. It generates text, but it doesn't generate Shakespeare—it generates internet text. It will create arbitrary news articles and documents, and it will try to complete documents because that's what it's trained for: completing sequences. When you give it a question, it might just generate more questions, or follow up with more questions. It will do whatever a similar document in the training data would do, and who knows—you get kind of undefined behavior. It might answer questions with other questions, ignore your question, or just try to complete a news article. It's totally unaligned, as we say. The second stage is fine-tuning to align it to be an assistant. This is the second stage, and the ChatGPT blog post from OpenAI discusses how this stage is achieved. There are roughly three steps to this stage. First, they start to collect training data that looks specifically like what an assistant would do. These are documents where the question is on top and the answer is below, and they have a large number of these—though probably not on the order of the entire internet. This is probably on the order of thousands of examples. They then fine-tune the model to focus only on documents that look like that, and you start to slowly align it. The model begins to expect a question at the top and to complete the answer. These very large models are very sample efficient during fine-tuning, so this approach works, but that's just step one—just fine-tuning. Then, there are more steps. The second step is to let the model respond, and then different raters look at the different responses and rank them according to their preference for which is better. They use that to train a reward model, which can predict, using a different network, how desirable any candidate response would be. Once they have a reward model, they run PPO, which is a form of policy gradient reinforcement learning optimizer, to fine-tune the sampling policy so that the answers ChatGPT generates are expected to score highly according to the reward model. There's a whole alignment stage here—a fine-tuning stage with multiple steps in between—and it takes the model from being a document completer to a question answerer. That's a whole separate stage. Much of this data is not available publicly; it is internal to OpenAI, and it's much harder to replicate this stage. That's roughly what would give you ChatGPT, and nanoGPT focuses on the pre-training stage. Okay, and that's everything that I wanted to cover today. We trained, to summarize, a decoder-only Transformer following this famous paper, "Attention is All You Need" from 2017, and that's a GPT. We trained it on Tiny Shakespeare and got sensible results. All of the training code is roughly 200 lines of code. I will be releasing this code base also, it comes with all the git log commits along the way as we built it up. In addition to this code, I'm going to release the notebook, of course, the Google Colab, and I hope that gave you a sense for how you can train these models, say GPT-3, that will be architecturally identical to what we have, but they are somewhere between 10,000 and 1 million times bigger, depending on how you count. That's all I have for now. We did not talk about any of the fine-tuning stages that would typically go on top of this. If you're interested in something that's not just language modeling, but you want to, say, perform tasks, or you want them to be aligned in a specific way, or you want to detect sentiment, or anything that, anytime you don't want something that's just a document completer, you have to complete further stages of fine-tuning, which we did not cover, and that could be simple supervised fine-tuning, or it can be something more fancy we see in ChatGPT, where we train a reward model and then do rounds of PPO to align it with respect to the reward model. There's a lot more that can be done on top of it. I think for now