Build a Small Language Model (SLM) From Scratch

Hello everyone. My name is Dr. Raj Gandekar. I graduated with a PhD in machine learning from the Massachusetts Institute of Technology and I'm one of the three co-founders of Vijuara AI Labs.

Today I'm incredibly excited to present this video to you in which we are going to build a small language model fully from scratch. I've been working on this video for the past 5 to 6 months and I'm truly excited right now to assemble all of it together and present it to all of you.

Like we do in all of our videos, we'll be showing you everything from the nuts and bolts. I'll not assume anything. I'll show you each and every calculation, each and every mathematical detail on the whiteboard and then we'll also do coding. Towards the end, we'll build a fully functional small language model which gives good quality responses. So let's get [Music] started.

All right. So let's get started on this journey. Today our main goal is that we are going to build a small language model fully from scratch. Let me first show you what we are going to build today. Here is the Google collab code file which we are going to run fully. I'll walk you through each and every line of code over here and towards the end as you can see over here we are going to give uh a sample sentence and we are going to generate certain stories out of it and I'm I'll make sure that the response of the um LLM or the response of the small language model which we build is good quality responses.

So you can think of this tutorial as a bit of a production level tutorial because I'll also teach you some things which we usually need to do to accelerate the training process. Some things which we usually need to do so that computations become memory uh memory efficient etc. So this is not just a toy project which we'll be assembling and the results which you get towards the end are do not make sense. But as you will see the sentences which our small language model can generate actually start to make a lot of sense. Okay. So this is what our goal is from this tutorial video.

First let me clarify to all of you what does a small language model actually mean. So if you look at language modeling itself, models just keep on getting larger and larger and larger. Right? GPT4 if you see over here was about 1 trillion parameters although it's not public information yet. GPT3 the largest model was around 175 billion parameters. And there's the scaling law which says that as the number of parameters increase the performance of the language model increases proportionately. So researchers went on increasing the size of language models to get better and better performance out of them.

However, people soon realize that what if let's say if there's a huge architecture for the language model which let me show here by this circle right and this huge architecture gets the job done it produces coherent responses what if I can get similar responses but with a smaller size language model what if I reduce the number of parameters to not as large as 1 trillion or 175 billion but what if I can take it down to a million parameters or a few million parameters. Can my model still do the task well or instead of a million what if first we try a few billion and see if the model does the task well. So if you see over here here I have presented some models which work relatively well and their sizes you can see here start from around 94 million 340 million 1.5 million etc. In this figure also you can see the Quen 2.5 billion model, the Quen 1.5 billion model, they perform very well and their size if you see 1.5 billion is about 100 times lesser than 175 billion right which is GPT3 size and it's about a,000 times lesser than GPT4 size and still they perform relatively well.

The appeal of small language models is that as the model sites get smaller, of course, it's more easier to run, it's more faster to implement. And if you connect the development of language models and apply it into production, you of course want things which work fast. So similar to the MOS law for the chip size, right? And the chip size reduced exponentially over time. I believe that the size of language models is also going to come down and within four to five years we may reach a stage where all of us will have a language model downloaded locally on our mobiles which we can interact with. This whole field of trying to reduce the size of language models essentially and to retain their performance is what's called small language modeling.

Now what is small that's not clearly defined but anything which is less than let's say 1 billion for example would be considered a small language model currently if you manage to have something which is of the size of 1 to 50 million that's awesome because that's definitely a small language model so today we are going to build such model um and we are going to use a GPU but I'm going to show you how you can get access to that GPU etc but we are not going to use a very very deeply computational traditionally heavy architecture with a reasonably light architecture we are going to get good results when we develop a small language model. So the language model which we are going to develop will have a size of around 15 million parameters and that's why it's safe to call it a small language model. So if you take this into context of GPT3 this is probably around uh 1,000 or 10 to 5 times lower than uh GPT3 because that's 175 billion and that's even around 10 to 6 or 10 to 7 times smaller than let's say GPT4 the size of the model which we are going to generate today. So that's the scale of uh the model size we are going to operate on in this tutorial. But you'll be surprised to know that at the specific task which we look at the model does relatively well with this context set.

Let's get started. We are going to divide this project into six parts. First we are going to go through our data set. Then we will do a bit of data prep-processing. Then we will assemble the entire model architecture. Then we'll set up the training pipeline of the small language model. Then we'll pre-train the entire small language model. And finally we'll do the inference. Don't worry if these terms or some of the terms I'm saying are not making sense right now. I'll try to go through everything in a lot of detail so that you understand every single aspect of this pipeline. These are the whiteboard notes which I'll be referring to throughout this project and the Google collab note or the Google collab code file which I've already shown you. We'll be going through this code file after we finish each and every part on the whiteboard. Okay.

So with that let's get started with the first part of this project which is the data set which we have. Every language modeling project begins with your data set right and as you might be knowing GPT3, GPT4 etc was trained on billions and billions of data set which we do not have right now or I don't have access to this right now and would have been very difficult to condense all of this in this small tutorial. So so I turned to a paper which became very popular in 2023 and this paper is called as tiny stories. You can check the number of citations which this paper have. So you'll see that it has around 219 citations, right? And what this paper tried to do, so they tried to answer the question that how small can language models be and still speak coherent English. And their main claim was that small language models were not working that well until the time this paper was published. And their main claim was that small language models are trained on this entire data set which is available all over the internet. Right? But they said that what if we reduce the size of this data set itself. What if we reduce the size of the data set so that the data set is for a specific task. So this is called as a specific task data set. And suppose let's suppose that we very intelligently curate this data set so that it captures all the nuances of the English language. Can we now train a small language model on this small data set?

Let me explain this further to you. In fact, the way these researchers constructed this data set is that they basically looked at stories which would be appropriate for 3 to four year old kids. They only looked at stories for 3 to four year old kids and that's why they called this data set as tiny stories. They constructed this data set. The way they constructed this data set is from GPT4. So they asked GPT4 to construct a huge amount of stories for 3 to four year old kids. And the smartness or the beauty of this approach is that this will be a very small data set, right? It won't why would we need all the information of the world for this? We only need stories for that age group. So those will only contain a specific words but they capture all the nuances of the English language because the stories have to be written in fluent English. The grammar needs to be correct. So if we are testing that can a language model essentially learn from all these data set and produce coherent English then ideally this is a great data set to have.

So what is our goal here? Our goal here is that can my model or can my 15 million can my 15 million model learn from this data set? learn from this data set. And what does it mean learning from the data set? It actually means that my model should be able to do two things. My model should understand the form of the English language. Sorry, my model should understand let's say the structure or the form of the English language which means it should generate grammatically correct sentences and second my model should also understand something about the meaning which means that my model should be able to come up with stories on its own which are suitable for 3 to four year old kids and if we could achieve this that would be truly incredible right because just from this small data set now I'm able to teach a language model about the English language The model starts from no knowledge. It has no information about language. It has no information about the fact that let's say uh subject verb object that is the structure followed in English language. It does not know this. But can we teach this? Can the model learn this from this data set? And if we are able to successfully create this model, which means that if we are able to successfully create a model which outputs coherent English stories, we are successful at our task of creating a small language model. And what is small? We are going to create only a 15 million parameter model. Okay.

So we are going to work with the tiny stories data set. This these researchers have made their data set public. So we are going to use this data set. If you scroll over to this, this data set is also available on hugging face. So you'll see that they split the data set into two parts. They split it into training and validation. So even although I say this is a tiny data set, it still has 2 million rows. And uh every row of this data set is a story meant for 3 to four year old kids. So for example, if you see this row, we see the first story, right? One day a little girl named Lily found a needle in her room. She knew it was difficult to play with it because it was sharp. Lily wanted to share the needle with her mom etc etc. And this whole story proceeds like this. Remember these stories have been generated from GPT4 but they constitute our data set like that. This is one story. Similarly, we have 2 million stories like this which are in our training data set and about 20,000 stories in our testing data set. So, we are going to work with this data. As I mentioned at the start of this video, this is not going to be a toy project video, but we are actually going to work with real data sets over here. Okay.

So, this is the data set and we'll use this data set which is specifically meant for 3 to four year old kids. This will help us reduce the size of our language model because now I'm not assembling all the data from all over the internet. I'm looking at a specific corpus of data meant for 3 to four year old kids. And the question is that can we construct can we construct a language model with just 10 to 15 million parameters and which produces coherent text. And if you think about this sounds like an incredibly difficult problem, right? Because you have to teach a model something about language otherwise how can the model come up with stories. We want the model to produce coherent stories towards the end of this. We want the model to produce stories which make sense on its own. So it's like imagine an alien coming to the earth and you want the alien to speak English. So instead of loading the alien with all the information available on Wikipedia, on Google, everywhere, you just give the alien stories for 3 to four year old because through those stories also the alien can learn the English language and then you feed all this data to the alien and then you hope that the alien will learn everything and come up with coherent stories on his or her own. So this sounds like a challenging problem but that's what's so exciting about this project.

Okay. So now I'm going to go to Google Collab and we are going to just start executing some code blocks as we uh finish understanding the steps on the whiteboard. So the part one is loading our data set. So I go over here and here you can see that you can um this data set is present on hugging face right. So you have to first install data sets uh which is a package through which you can load any data set on hugging phase. So you can run this command and it will take some time to install this. It took about 30 seconds on mine. Okay. After this point you have to import the load data set function and here you have to pass the URL of the uh here you have to pass the URL of your data set on hugging face. So if I go to hugging face right now, you'll see that this is the URL, right? Uh this is the URL which I have marked over here and that's the URL which I'm passing over here. So load data set and this URL. And when you run this, you'll see that the data set has been loaded now fully. And we see that the since the data had two parts already, right? We had a training data and validation data. We have two types of data which have already been loaded. The train split and the validation split. validation split. Okay. So the data set has fully been loaded. Now this loading of the data set might take some time for you because remember we are loading 2 million rows, right? We are loading 2 million stories in our training data. Um so as you can see here we are 2 million stories in our training data and we have about 20,000 stories in our validation data. So this does take some time but once this is done your data set is successfully loaded.

One more thing which I want to mention right at the start is that I'm changing the runtime here to T uh A100 GPU. So now if you are on the Google Collab free tire you can do T4 GPU but the training will take a huge amount of time on T4 GPU. You can run it overnight for 8 to 10 hours and then it will be complete. But if you switch to Google Collab Pro which is just $10 per month, you can run the training with an A100 GPU or the L4 GPU. Both are very strong and the training roughly takes just 25 to 30 minutes on an A100 GPU. But if you have a T4 GPU, it might easily take around 6 to 7 hours. It can be run on a free tire also. So no problem if you just have a T4 GPU. But keep in mind that it might take some time. And never run this code on CPU because it will fail. You will not be able to run this code on CPU. So it should be either of these GPUs which is provided by Google Collab. Okay.

So let's go to the whiteboard. Right now the step one is done where we have loaded our data set. I hope until this point all of you have understood the aim here. Our data set is only stories understood by 3 to four year old kids. We are going to use this data set to teach our language model about the English language. So that ultimately it produces stories which look like the ones in the data set. Okay. That's going to be the purpose of our small language model. Now let's move to part two which is data prep-processing.

All right. So data prep-processing is one of the most important parts of essentially training a large language model. And let me show you the whole training pipeline actually. So if you look at the pipeline of training a large language model, usually we have three blocks, right? We have the input block which I've shown over here. We have the processor block which I've marked over here and then we have the output block. We are going to assemble all three blocks today. But just keep in mind that in literature and even in uh everywhere everyone talks about the processor block because that's where cool stuff is happening. But input block does not receive that much attention although this is one of the most important steps. Here is where the data data processing actually happens for the input text which we have and uh that's why I'm going to go into a fair amount of detail for this part because without understanding how the data is processed in the input block you won't understand what happens to the data when it enters the processor or what happens to the data when it enters the transformer block. Okay.

So let's understand data prep-processing. So remember first that LLMs or essentially computer for that matter it cannot understand sentences right. It cannot understand sentences. It cannot understand text or words like humans do. It only understands numerical. It only understands numerical data. So what we have to do is that somehow we have to take our input data which has been present over here and we have to convert it into a numerical format. Now think about how would you actually do this right? Let's say let me take this story over here. Let me take a screenshot. And if you understand how this is done for one story, we'll of course understand how it can be done for all the stories. Right? So I've taken a screenshot over here and I want us to think how will you convert this entire paragraph into numbers. So the first thing which might come to your mind is that I'll do uh every word will be assigned some number. So I'll divide this paragraph into bunch of words and then I'll assign every word a number. Or the second thing which might come to your mind is that I'll divide this paragraph into characters and I'll assign every character a number. The first type of method in which every word is assigned a number is called as word based tokenization. And the second type of method where every character is essentially assigned a number or a token that's called as character based tokenization.

Now it turns out that to develop language models both of these tokenization forms are not that efficient. Why? Because when you have character level tokenization, every character is a token. So the input text will be get converted into a huge number of tokens and our transformer block which is the processor can only pay attention to a certain number of tokens at a time. So character based tokenization does not work at all. Um second word based tokenization right. So if every word is essentially one token, the issue happens when we have uh let's say words where spelling mistakes are there. If words have spelling mistakes, they will not be mapped to any token. That's the first problem. And the second problem is that whenever you do tokenization, you maintain something called as the vocabulary. You maintain something called as the vocabulary. This vocabulary consists of the list of tokens and their corresponding token ids. So if you do word based tokenization, these tokens will be all the words. And if you just look at the English language itself, we have about half a million or 500,000 words in the English language itself. So in the case of word based tokenization the vocabulary size becomes very large and that's a problem for us because vocabulary size affects our computations when we are going to predict the next token in a large language model. Usually the higher the vocabulary size the longer the computational the longer the computational time and the higher is the computational memory. So both of these are tokenization methods which we don't want.

So then what's the best solution? The best solution is something intermediate and that's that intermediate stage is called as subword tokenization. That intermediate stage is called subword tokenization. In subword tokenization what we do is that we maintain a vocabulary of tokens which can be characters which can be words and which can also be subwords. So words which are not very commonly occurring right like for example tokenization for example it is usually split into two words isation so isization becomes one subword that becomes a separate token and this token becomes another token so in token in subword based tokenization what we do is that characters are of course retained as individual tokens that much is fine all commonly occurring words are retained as tokens but all the words which are not that common they are broken down into subwords This solves our vocabulary issue because every word is now not a separate token. Some words can be broken down into subwords and this also solves the character based tokenization issue which breaks the text into huge number of tokens. So subword tokenization is that optimal scheme. So usually what is done in subboard tokenization is that there is an algorithm which is called as the bite pair encoding algorithm. Bite pair encoding algorithm and that's also called BPE. This is an algorithm which takes in a data set. This algorithm takes in a data set and it gives us the subwords or it gives us the list of tokens which the data set needs to be converted into. Right? And these tokens as I mentioned can be characters, these can be words or these can be subwords. All three are possible. Essentially what bite pair encoding does is that it just merges the commonly occurring bytes together until you reach a prescribed vocabulary size. I'm not going into too many details of bite pair encoding here for the purposes of this tutorial. It's okay if you have a higher level understanding of what is exactly done in bite pair encoding.

So what now we are going to do is that we are going to feed our entire data set into this bite pair encoding model and this is also called as the tokenizer. So the tokenizer what the tokenizer does is that the tokenizer takes our data set. The tokenizer takes our data set and then it converts it into a bunch of tokens. These tokens are each individual units and then every token will essentially be assigned a token ID. That's this whole process which needs to happen at the start before a data is fed to the transformer block. There are many other steps also but the first thing is you take the data set and then what you do is that you convert the data set into uh essentially a bunch of tokens and then you assign the tok each token a token ID and the way the data sets are converted into a bunch of tokens is through the tokenizer. Now we are going to use a tokenizer used by GPT2 which is the same as what I had mentioned earlier that's the bite pair encoding. Okay. And we are going to pass all of our stories over here as the data set. We are going to pass all of our stories over here. But instead of uh just passing every story and converting it into a bunch of tokens, there is uh uh there is a common way or there is a better way to do this which is good in terms of uh computational efficiency. And for this I have borrowed this part of the logic from Andre Karpati's nano GPT. This can be a good time to mention that Andre Karpath's nano GPT code has inspired a lot of portions of this small language model tutorial and I really want to thank um thank him for putting out such incredible content out there from which all of us are benefiting. So I've tried to take that repository and put it into a format which is much more easily digestible and when everything is exposed instead of directly importing some things from his library. So this is what we are going to do with respect to the uh tokenization. So let me show you what exactly we are going to do.

Okay. So now our data set looks like this right. Every row in our data set essentially corresponds to a story. Uh the first row corresponds to the first story. The second row corresponds to the second story. The third row corresponds to the third story etc. There are 2 million such rows for training and 20,000 such rows for validation. So 2 million stories for training and 20,000 stories for validation. We have already seen this in our data set over here. 2 million 2.12 million training rows and around 22,000 validation rows. All right. So now what we are going to do is that I'm going to take a small sample from our data set and I'm going to show you how we are going to do the tokenization. So I have pasted a screenshot of this sample over here. This is let's say one such story and our goal is to do two things right. We have to tokenize this data set which is we need to convert it into tokens and we need to store all the token ids in one place. So remember I mentioned to you that uh the every token will be assigned a token ID. Essentially what we need to do is that we need to take every single story and within every single story we need to convert every token into a token ID and merge all of these token ids together. That will be our training data set. And the way this is done is as follows. So first we are going to use this GPT2 subword tokenizer. We already saw about this that is the bite pair encoding tokenizer and we are going to apply this on every single story. So let's say when this bite pair encoding tokenizer is applied on this story we get certain tokens like these. Okay. So that can be 442 518 11 257 etc. We are going to store these tokens into uh let's say dictionary with ids. So these are my ids and my length is going to be 30. That's the length of the token ids in one story. So every story I'm going to store like this. This is the first story. This is the first story where I store the token ids and the length of how many token ids I have. So 1 2 3 4 5 6 7 8 9 10 11 12 13. So that's why this is 13. I'm similarly going to apply tokenization to my entire data set. So my first story will again have ids which are these and it will have this length. My second story will have a set of token ids which have this length. My third story will have again token ids and a length. Similarly my last story which is now 2.12 million will have a bunch of token ids and the length. So now what I want to do is that I don't want to store separate token IDs at separate places for each story. I want to merge all of these token ids together so that I can store them in one place. Right? And now we are going to see an efficient way of doing this merging.

So the goal is to store all the token ids in a single dot bin file. The reason we have bin file is because we'll store the file directly on the disk instead of storing it on the RAM. That avoids LA RAM overload on big data sets. So believe me, I've tried running this code without using the bin and it takes a huge amount of time to run this and the entire code becomes very slow. But if you just do this bin, if you create a dot bin file and if you store all the token in dotbin, you essentially store everything on the disk, right? Since everything is stored on the disk, um the data can be loaded very fast during the training. You just retrieve it from the disk. And secondly our RAM uh which is very precious for uh faster computations that does not get overloaded overloaded with huge amount of data espe especially since we have around 2 million rows right that is why we using dot bin and another advantage of storage on the disk is that we don't need to retokenize every time it is stored on the disk during multiple sessions. So once we tokenize that's it the token ids are all stored on the disk. every time you restart let's say you don't need to uh retokenize again so that's why we store everything in a dot bin file right uh currently we have something like this for every story we have a bunch of token ids and we have length of the number of tokens so that is for story one story two etc right so what we'll do is that we'll create a dot bin file as I mentioned so the dot bill file can be named something like trainbin we'll see this in the code also we'll create a dot bin file which is named at train.bin pin and we'll create a memory mapped array. Why do we create a memory mapped array? Uh because it means that the file is backed onto the disk. So as I mentioned over here, we need to do disk storage and for that we need to create an array which is memory mapped which means it will get stored on the disk. And uh so this train.bin file which is there we are going to store it on the disk. It looks like a numpy array. uh and when we use this np.mme map it means that the file is stored on the disk and we don't have to hold everything in the RAM essentially so we create a memory mapped array and then then what we do is that we split our entire data set into batches so let's say I have my 2.12 million stories right I have my 2.12 million stories I split these stories into 1024 batches so batch number one batch number two etc. Let's say badge number one has about 1,000 stories or 2,000 stories. Then what we do is that we take the token ids from all of those stories and put it into this train.bin file. Then I go to batch number two. I take all the token IDs from batch number two and put all of them into the train.bin file. Similarly, I go to the last batch batch 1024. I take all the token ids there and put it into the train.bin file. Ultimately, my train.bin bin file will have all uh we'll have all the token ids from all the different stories which and everything here will be saved to the disk.

So let me repeat what we are doing here. I hope until this part you have understood where every story is broken into a bunch of token ids and every token ID has some length which is associated with it. Ultimately what we have to do is that we have to store all of these token ids together in some place. That place is going to be this train.bin bin file and this file is going to be on the disk instead of storing it on the RAM. So in the code what you'll see is that you'll see the file name and then you'll see an array array which is initialized like np.mme map. So this is a memory mapped array. So this array eventually will contain all of my token ids which are stored on the disk right. um and then all and the file name where eventually this array is stored that's train.bin. So my array will contain all the token ids and this array will be stored in the file name train.bin. That's all which is done here. We uh we create a bin file. We create a memory mapped array. We split the data set into batches and we add everything to to the train.bin file. That's it. So my train.bin file now contains all of my token ids together. You might be wondering why do we do these batches. The reason we do these batches is for faster processing. Because if we don't do these batches uh then it takes a huge amount of time to serially go through my entire data set collect token ids and put it to my train.bin file. The reason I'm explaining all of this to you in detail is because that's exactly what we are going to see in the code right now. So if you go to the code and if you see step number two here is where we are going to tokenize the data set and ideally this should take a huge amount of time because we have 2.12 million stories. We have to essentially convert all of those stories into token ids and append all the token ids together. But this version of the code is actually very fast. The first thing which we are doing is we have we are getting encoding for GPT2. So this tick token is the library. If you see tick token GitHub, you'll see that this is a library which is open as library and we can get tokenizers from here and we can get them for different GPT2 model GPT models. So here we are getting the tokenizer for GPT2 and first thing what we are doing is that we are going through all of my stories and every story we are converting it into a dictionary of ids and length. So as you can see over here every story is getting converted into ids and length. We saw this over here also. That's the first step which is happening in this process. Right? Then what we are doing is that uh here you see we are creating a file name which will be created for training as well as for validation. For the training data we'll create a file called train.bin bin and we'll create an array which is a memory mapped array and we'll divide the entire data set into 1024 batches. Now what we'll do is that for each batch we'll collect the token ids and we will store them into this array. That is where so every batch will collect the token ids as has been shown in this part. Every batch will collect the token ids and store it into the array. This is what is happening in each iteration of the for loop here. And then we just go through different batches. And as we go through different batches, every batch will give us a bunch of token ids. All of those we'll put into this uh array. And finally, array.flush means that we store everything on the disk. That's the last command. Array. Which means we save everything to the disk ultimately. So now this array will be then uh stored in my file which is train.bin. File name is train.bin for the training data. And I have my validation data in validation.bin. So now after this after this point in the code, the train.bin file essentially it contains a huge list of token ids. It's a huge file because I have 22 million stories, right? And if you just imagine for now one word has one token. Let's imagine one word has one token, right? And if you see every word here, every every story here has at least 1, two, let's say 60 words. Every story at least has 50 words and have 2 million stories, right? So that means there are around 100 million words in the training data. So the trainbin might have around 100 million tokens, token ids rather. and my validation my validation is I think uh 22,000 right so that is around 100th of this so my validation will have around 1 million token ids which is still significant so the trainbin will have a list of all the token ids of my data and my validation dobbin will have all the list of token ids in my data set and again why are we doing this step because we need to convert my data which is in a text format into a numerical format and the first step of doing that is tokenization. So first we have to decide the tokenizer which we are going to use and that's going to be the bite pair encoding and we are going to use a subword tokenizer. So every subword will be one token ID. All the token ids will be appended together. That's what we are doing in a simple manner. But in this code we are just trying to improve the computational efficiency by storing everything on the disk. That's why we need a memory mapped array over here. If this part is a bit confusing to you, don't worry. Just remember that after this part, after this code is executed, we have tokenized the data set into token ids. We have created a file called train.bin. And this should be validation.bin. We have created a file called train.bin and uh validation.bin where all the token ids are stored. And we make sure that token ids are stored on a disk rather than on the RAM. for efficient computations. These three are the most important things to remember from step number two. So until this point we have our data set, we have converted it into token ids. We have a train.bin file and we have a validation.bin file. I hope all of us have understood until this point that is the step two which is data prep-processing.

Now we are coming to step number three which is creating input output pairs from the data set. The reason we need to create input and output pairs is because ultimately the loss function which we will define for training the language model will be based on these input and the output pairs. When we have regression tasks or classification tasks in machine learning, we are usually provided with the labels of the correct answer. Right? But in the case of language modeling, we are just provided with the data set which looks like this. But we are not essentially told that what is the correct answer, what is supposed to be the right answer generated by the language model etc. We have to create these input and output pairs from the data set. To understand how to create input and output pairs, first it's very important for all of us to understand what's the purpose of language modeling itself. So when large language models were uh constructed and when they became really popular, what they actually did was they took in the input as a sequence of tokens or a sequence of words and then the output which they generated was essentially the next token prediction task. So what is the next token prediction task? It's essentially just looking at the sequence of words which has come until now and then predicting the next token. So if you think of the language model as the engine, the engine receives the fuel and then the car moves forward. Right? Similar to that, the equivalent analogy to the language model is that the sequence of tokens which goes to the input that is the fuel and the next token prediction that is the movement which is the output which we want. What is so incredible and also beautiful about language modeling is that although we teach the model or train the model to just predict the next token at a time, it somehow understands the structure and the meaning of human language within that process almost as a byproduct of that process. And one more way to understand next token prediction is that if you go to chat GPT right and if you say something like uh make a travel plan for Italy you'll see that chat GPT response and the response is presented as a detailed paragraph but actually just one token which is predicted at a time and then it's appended to each other. So at every single instance only the next token or the next word is predicted. You'll notice that as the lecture goes along, I'm using the word token and the word in interchangeably. Although that's not the case. We are using the subword tokenization. So, usually one word is not one token, but we'll use that for the sake of simplicity. It's just better for explanation.

Now, let me tell you how these input and output pairs are actually created from the data set. Here I have taken one single story from the tiny stories and let me tell you how the input pairs are created. First to create the input pairs we have to essentially decide two things. First thing which we have to decide is the context size. The context size is the length of words or tokens which my language model is looking at at one time before it predicts the next token. Or put in another way, it is the maximum number of tokens my language model can look at before predicting the next token. And here I'm deciding the context size to be equal to four. What that does is that based on the context size, I will divide my entire data set into chunks and the length of those chunks will be equivalent to the context size. In this case, it is equal to four. So this is my first chunk. This is my second chunk x2. This is my third chunk which is x3. And this is my four chunk X4. Similarly, the whole data set will be broken down into these chunks. And you'll see that each chunk essentially consists of four words or four tokens because the context size is equal to four. Now although I'm showing words over here to do keep in mind that we have created train.bin and val.bin, right? validation.bin. So what the model actually sees is just a bunch of token ids like this. So it creates the inputs based on these token ids and not the words. I just wanted to clarify this because the model after this stage has no idea of the words. It just looks at token ids. Sometimes we we forget about that because as humans it's so natural for us to look at words. But just remember that when the model is creating the input output pairs it just looks at token ids. So the entire data set is broken down into these input pairs whose size or the chunk size is equal to four. The second quantity which I need to decide before I create these input output pairs is the batch size. Batch size is essentially when you're training language models, we don't process all the input and output data at one time. The data is usually processed in batches. So first you uh take the output of the first batch, calculate the loss, propagate it backwards, update up update the parameters. Then you calculate the output of the second batch, calculate the loss, propagate it backwards and update the parameters. The reason this is done is that if you have no batches, it would take a huge amount of time to update the parameters because you have to go through all of these all of the entire data before updating the parameters. So having batches makes it much more convenient to perform the training procedure. So here I'm deciding a batch size equal to five. Once you decide the context size and the batch size, you can create input matrix matrices such as this. So this is called as my first input batch. So this is now first this is my first input batch. The reason this batch has number of rows equal to four x1 x2 x3 x4 is because the batch size is equal to four. there will be four um there will be four sequences in each batch and in each sequence so for example if you look at X1 the number of columns is equal to four why because the context size is equal to four so what I'm just doing is I'm taking these four and stacking them on top of each other that's how we create the first input batch so similarly I can go through my entire data set and create multiple input batches that's how my input batches are created now let me talk a bit about how the output batches are created or how the output pairs of data are created. So let me rub this so that I'll show you how the output is created for the first input. Right? So if you take a look at this first input one day a little girl and I'm ignoring this comma here for the sake of simplicity. The output is just going to be the input shifted to the right hand side by one. So the output of this first input sequence is this day a little girl um yeah a day a little girl

Named that's going to be the output sequence. Um, and if you look at, for example, if I do it over here, actually here, the first input should be of a sequence size equal to four, right? So I should probably, so this is my first input. And, uh, my first output is going to be the first input shifted to the right-hand side by one. So this is going to be my first output. Similarly, if my second input, if my second input is, let's say this, which is x3, then the output corresponding to this is just x3 shifted to the right by one. So that is going to be my y3. Now, this is how the output pairs are created.

So if you look at every input pair, if x1 is 1, 11, 15, and 24, the output corresponding to that is just the input shifted to the right-hand side by one. So if you see 1, 11, 15, 24, 11, etc. Let's say this is how the token IDs are arranged, right? In my train.bin. If my first input is this, which is x1, my first output is just this input shifted to the right-hand side by one. So that will be, uh, 11, 15, 24, and 11. That's exactly what is shown over here. So you see they have an overlap of three over here. That is how every output is constructed from every input.

So essentially, once you have x1, x2, x3, x4, the output tensor or the output batch number one. This is my output batch one. It's quite easy to construct because it is just the input inputs x1, x2, x3, x4 shifted to the right-hand side by one. So imagine all the data set, imagine sliding over the data set in chunks of four. That will be my input, and just moving those chunks to the right-hand side by one and aggregating them together in batches. That's going to be my output.

Now, if you look at, let's look at the first input and the first output, right? So my x1 is 1, 11, 15, and 24, and my y1 is essentially 11, 15, 24, and 11. Okay. And let's assume some one day a little. So let's say this corresponds to one day a little, and this corresponds to day a little girl.

Now you might be thinking that if this is my input and this is the output, uh, what exactly are the input-output pairs over here? So it turns out, or what are we really predicting here? Remember I told you that language models are just next token prediction. So it turns out that in one input-output pair like this, there are multiple next token prediction tasks. So if "one" is the input, then "day" should be the output. If "one day" is the input, then "a" should be the output. If "one day a" is the input, then "little" should be the output. And if "one day a little" is the input, then "girl" should be the output. We train the model on all these four tasks. So when you look at one input-output pair, it's just not just one input and prediction. It's not just one prediction task. There are four prediction tasks which are happening here.

So let me write down this, um, these four tasks. So one is the input, and "day" should be the output. "One day" is the input. "a" should be the output. "One day a" is the input, "little" is the output, and "one day a little" is the input, and "girl" is the output. So remember, this column is my input, and this column is my target. So when my model is initially, in my model is initialized, its output which it gives will be something completely different than the target. What we have to do is that the model output has to be as close to the target as possible. That is what we have to ensure in the backpropagation.

So when you look at these input and target pairs, for every row, there are four input-output prediction tasks, right? And similarly, there are four such rows in every batch. So we are essentially training the model to predict the next token at every stage of the process. So if you start looking at this paragraph from now, from the start itself, what we are essentially doing by creating these input-output pairs is that when "one" is the input, "day" is the output, we are teaching the model that when "one day" is the input, "a" is the output, when "one day a" is the input, "little" is the output, when "one day a little" is the input, "girl" is the output. That's x1, y1 training. Then we go to x2, y2 training. So x2 is now "girl named lily found", right? So when "girl" is the input, "named" is the output. When "girl named" is the input, "lily" is the output. When "girl named lily" is the output, "found" is the output. So essentially, we're training the model at every stage to predict the next token, and we are doing that through the creation of this input and output batches.

I hope this part is very clear to you. It's very important to visualize this part before we go to the model architecture itself because here is where all the confusion in language modeling starts. You have to know how the input and target are created. And let me give you a short summary now of this whole process. Let's say we have this data set, right? Let's say we have this data set. The first thing which is done is that I decide my context size. If it's equal to four, I create input chunks which are like this. So I divide the whole data set into chunks of four. That will be x1, x2, x3, x4, etc. And then along with this, I essentially shift to shift this to the right-hand side by one. So that is how I create my output. That is how I create my output chunks, right? So I have my input chunks. I have my output chunks. Then I divide my input chunks into batches. I divide my output chunks into batches. If you look at one batch closely, there will be four in x1, y1, x2, y2, x3, y3, and x4, y4. If you analyze each row here, x1 and y1, there are essentially four input-output tasks here which we are training the model. And in each task, we train the model to predict the next token.

So when the whole training proceeds, the model knows that whatever is the input, let's say if this is the input, the model is trained to predict the next token as the output. So here you might be thinking that this is, is that, is that the only training objective in language models? And it turns out that the training objective is as simple as this. We are just predicting the next token at a time. And what's quite amazing is that how does the model really know language from this simple training? How does it understand? So if there's a sentence such as, let's say, "blue electrons eat fish." Let's say there is a sentence like this, right? Which is "blue electrons eat fish." How does the model actually know that it should not predict a sentence such as this? Because ultimately, we are just training on the next token prediction, right? We are not telling the model anything about the world around us. But through this simple task of the next token prediction, the model somehow understands that electrons cannot be blue, not just that electrons cannot eat fish. And so the model understands the form of the language and the meaning of the language, as we are showing over here. Okay.

So this is the main process of creating input-output pairs, and this is what we are going to exactly see in the code. Now, in the code, you will encounter a statement such as this: `ix = torch.randint(...)`. So what is exactly done in this statement is that we create an ID vector of length batch size. Um, so essentially, what this means is that if, for every batch, we sample from random places. So let's say, so here I've collected x1, x2, x3, x4 in together, right? It can be like this is x1, let's say this is x2, this is x3, and this is x4. That's also fine, and that's how we are going to do the code. So if this is x1, this is x2, x3, and x4, the corresponding y will be, of course, this shifted to the right-hand side by one. So whenever a batch is created, there we are going to define a function which is called as `get_batch`. Whenever the function `get_batch` is defined, what it will do is that it will create, so, and if the batch size is four, it will take from random places four vectors and create this input batch. It will shift the input batch to the right by one and create the output batch. So that is what you'll see happening in the code. That's the first thing.

The second thing which you'll see happening in the code is we are doing something like `X.pin_memory()`. So these input and output tensors or matrices which we have, we are going to pin, pin them to the memory. So what this does is that it locks the memory of the tensor in the RAM, and that allows faster transfer of the tensor to the GPU. Uh, now we are going to transfer this input and this output tensors to the GPU at some point for calculations in the transformer block. So having the pin memory locks the memory of the tensor in RAM, and we can allow faster transfer of the tensor to the GPU. So it's like telling our operating system that, "Hey, don't move this memory around, my GPU would want it very soon." And so the memory is kind of locked. The input and the output are locked in memory so they can be transferred to the GPU easily.

And then we are also doing this `non_blocking=True`. The reason `non_blocking=True` is done is because if you normally do `.to(device)`, it blocks the CPU until the copy to the GPU finishes. So `non_blocking=True` ensures that the CPU can continue to do other work like next batch preparation, etc. So the CPU does not have to be blocked while, uh, while we copy our data to the GPU. So you'll see this also in the code.

So let us jump to code right now. If you scroll down below to step number three, here is where we create the input and the output batches, right? So this is the `get_batch` function. Remember, every `get_batch` function is going to return one input matrix and one output matrix like this. Okay. So this is that first statement I mentioned. So this essentially samples four random integers, which essentially corresponds to, or four. I'm saying four because I'm assuming batch size equal to four. So it will correspond to randomly sampling from four different places in my data set. That's `ix`. What `x` essentially does is that it stacks, it stacks together x1, x2, x3, and x4. `y` stacks together y1, y2, y3, and y4. And if you look at x1 and y1, let's say if you see, `y` is just shifted to the right-hand side by one.

So what is being mentioned as `block_size` in here? That is the context size which we have in the code. And the batch size is the same as, so the `block_size` is what we discussed as context size on the whiteboard, and, uh, the batch size, the batch size is the same as the batch size which we had discussed. So here we are creating the input and the output batches: x1, x2, x3, x4, y1, y2, y3, y4, right? And now here we are just pinning the `x` and `y` to memory and `non_blocking=True` so that we can move to the GPU very easily. These statements are just written for computational efficiency. So within these small lines of code, there are actually a number of things which are happening. So if you just see this code block, you might be very confused by it and you won't understand what's going on here. But first, what we are doing is that we are doing these batches for both training data and validation data. That's why these first lines of code are there. Then here, what we do is that we sample random integers. So if the batch size is equal to four, we are going to sample four random integers. They correspond to four positions which we are going to look at in the data set. So for example, position one can be here, position two can be here, position three can be here, position four can be here. Wherever the position is, we get x1, x2, x3, x4. From there, we shift to the right by one and get y1, y2, y3, and y4. And then we just pin it to the memory. That's it. We are going to see this `get_batch` function when we run the pre-training loop. Every time we say `get_batch`, we are going to get one input and output batch. Then we are going to compute, we are going to run the input through the LLM, get the output, then compute the loss, backpropagate. Okay. So that's why `get_batch` serves a very important role.

Now that you have understood until here, we now move to the next part, which is probably the most important in this video, and that part is essentially defining the architecture of the small language model. So let's get started with this. The small language model architecture can appear a bit intimidating. It can appear a bit frightening because, as you can see over here, there are multiple different things which are happening in the model architecture, right? As we saw earlier, this whole architecture can be broadly divided into three parts: Part number one, Part number two, and Part number three. Part number one is the input block. Part number two is the processor, and Part number three is the output block. To truly understand the language model architecture, we have to understand all of these parts in detail. And I could easily spend 10 to 12 hours explaining every single aspect of this architecture. But my main aim in this tutorial is to give you enough understanding of this architecture so that you can understand what's going on in the code. I won't have the time to go into the details of each and every component here because that would make this tutorial very long. Uh, but I will give you enough intuition, enough writing on the whiteboard so that you understand step by step what is going on in each part, right? So let's get started.

The code associated for this part is in this, uh, code cell which is called "Defining the Language Model Architecture." And this is the longest code cell. So if someone looks at the code, they might not understand anything. And at that time, they might feel that language modeling is very difficult and it's not meant for me. But it's only when you start understanding the nuts and bolts of how this code is assembled, you'll see that everything really makes a lot of sense. Essentially, we are just assembling all of these components together so that when we get the input text, which is a sequence of words, at the end, the goal is to predict the next token, right? That's all we are going to see in the small language model architecture. What we are going to see is that when we have a sequence of words, which is the fuel to the language model, how does it predict the output, which is the next token? And once we get the next token prediction, we have already seen that we are going to find the loss between the output token produced by the language model and the target token for these input-output pairs, input-output batch pairs, which we have seen before. So the SLM model architecture is going to take some time to explain, but at the end of it, remember that we are just feeding in a sequence of words such as "one day a little," and we have to predict the next token, which ideally should be "girl." So if you see this sentence, "one day a little," if that is the input, ideally the next token should be "girl." So we are going to see how the LLM takes these sequence of tokens and produces the next token.

My approach here is that first, I'm going to explain to you everything on a whiteboard in a very sequential manner, and then what we are going to do is that then we are going to try to go to the code and we are going to map out every single thing which we have seen on the whiteboard to every single detailed aspect of the code, right? So let's get started. First, what we do is that we have already created this input and output batches, correct? So let's say we are looking at the first batch initially, and I'm looking at the first row of the first batch, which I have mentioned over here. That's my input right now. I'm looking at the first row of the first batch, and the goal is to take this input and to predict the next token ID, whatever it might be. So now, let us see the journey through which this input sequence of tokens actually go through before we have to predict the next token.

The first thing which we have to do is that although these tokens are in numerical format, so technically they can be processed by the computer or the language model, we convert every token ID into a higher dimensional vector, and that step is called as token embedding. The reason we do token embedding is because when we look at language, right? Let's say I take a screenshot of this, this story, and I paste it over here. Now, if you look at language and the way language manifests itself, every word carries some meaning, and that's what, that's what gives meaning to the whole paragraph, right? When you read "a little girl named Lily," something comes to your mind because "little" has some meaning, "girl" has some meaning, "cat" has some meaning, "dog" has some meaning. Both of them come under animals. So essentially, words have some sort of a semantic notion which is very important for the language model to capture.

So what can we do before we pass the input to the transformer block? Can we somehow take these token IDs and possibly represent them in higher dimensional vector spaces that capture some information about the meaning? So what we are going to do is that, let's say "cat," right? That has the token ID of 13, and let's say we have "dog," that has the token ID of 17. So I'm going to take these tokens and I'm going to project them into a higher dimensional vector space. For now, let's say I'm choosing a two-dimensional vector space, and "cat" comes over here, and "dog" comes over here. "Cat" and "dog" are animals, so they have some similarity in meaning, so they lie closer to each other. But both of these words, let's say, will be very different than a "chair," let's say. So "chair" will lie over here, but "chair" and "table" will lie very close to each other. So you see what we are doing here. We are projecting these words as vectors, and then we are hoping that vectors which are closer to each other would be similar in meaning. Vectors which are farther apart from each other would be different in meaning.

Why do we do this? Because we capture some notion about the meaning of words, and that is how we can exploit the properties of language when we look at convolutional neural networks. So I'm sure all of you must have seen this convolutional neural network explainer, right? Why do we take convolutional filters and why do we slide them over the images? Because we want to exploit properties in an image. The properties of translational invariance, the properties that two pixels will always slightly lie closer to each other, and even if a cup is present here or here, it will still be a cup. So there are some properties inherent in the images which are exploited when we use these feature maps or these convolutional layers. That's why CNNs work so well. So if you want to exploit language, which is our data right now, the best way to do it would be to capture meaning in vectors.

So that's why what's done as the next step is that every token ID is converted into a high-dimensional vector, and that's a token embedding vector. The way this is done is that we maintain something called as the token embedding matrix. We maintain something called as the token embedding matrix. And the token embedding matrix has a number of rows equal to our vocabulary size. So row number one, let me do this with a different color. Row number one, row number two, dot, dot, dot. If the vocabulary size is 50,000, then it will be 50,000. And each, uh, each token ID in my vocabulary is now encoded as a high-dimensional vector. So this can be, let's say, a 768. So I have to choose my embedding dimension, and that property stays with me throughout my architecture. I have to choose my embedding dimensions. So if I choose my embedding dimensions to be equal to 768, every token here is represented as a 768-dimensional vector. So this is also called as the embedding dimension. So we store this token embedding matrix where every token ID in my vocabulary is represented, let's say, as a 768-dimensional vector. Notice that the 768 is an experimental choice which we have to make at the start.

And now, once we have the token embedding matrix, how do you retrieve the embedding vectors for every token IDs? You essentially just use this as a lookup table. So if token IDs are 1, 11, 15, and 24, the way we retrieve their embedding vectors is that we just look at row number one, we look, we look at row number 11, we look at row number 11, we look at row number 15, and we look at row number 24, and we just get the vectors corresponding to these rows. So the token embedding matrix essentially serves as the lookup table. So that when we have token IDs, we can just look up this token embedding matrix and retrieve the corresponding token embedding vector.

So now this input which we have right now is transformed into this matrix over here. So I had 1, 11, 15, and 24, right? Now every token ID is represented, let's say, as a 768-dimensional matrix. So now if my input, let's say I'm going with the input as, let's say, "every effort moves you" or "one day a little," whatever. This will be a 768-dimensional vector. So let me change this. Actually, this was "one day a little," correct? So now "one" is a 768-dimensional vector, "day" is a 768-dimensional vector, "a" is a 768-dimensional vector, and "little" is a 768-dimensional vector. So now you may ask that how do we choose these vectors? Do we know the embeddings before? So are the embeddings such that they already capture meaning? No. All of these values are initialized randomly. And wherever I'm marking P, it means that these are trainable parameters of our model. So initially, this whole token embedding matrix is initialized as a random matrix. So all of these rows will be effectively normalized.

And what I was saying earlier was that, let's say if you have layer 1, layer 2, layer 3, layer 4, and you're backpropagating, right? Usually, what happens is that the gradients in every layer depend on the output of that layer. And if the output of the layer widely fluctuates between large and small values, then the gradient can either stagnate or the gradient can explode. It just becomes very unstable training dynamics. Also, if you don't do layer normalization, the distribution of inputs which is received at every layer, that distribution of inputs might change. For example, layer 1 in one iteration might have this input distribution, in another iteration, it might have a skewed distribution. So this training becomes hard as the distribution changes from one iteration to the next. Layer normalization prevents this problem, which is called internal covariate shift. It, it, it constrains the values of all inputs within a certain range, and it also makes sure that the distribution is kind of similar. So that just improves training. So wherever you see this layer normalization block, it is usually used to improve the training performance.

So the first step in the attention is the layer normalization. And if you go to the code, you will see that there is a separate block which we have written earlier. So there is a class called `LayerNorm`. Since layer normalization is utilized at many places in the transformer block, you see that it's introduced, it's introduced here. It's introduced here also. So it's there at two places, right? So it's better to define a class for it. And we just use the `layer_norm` method which is available through PyTorch. It just normalizes every single row, right? Okay.

So layer normalization is done. The next part is the place where magic happens. So magic happens here because the next step is multi-head attention. Now, what exactly happens in multi-head attention and why is the multi-head attention block needed? If you observe closely, what I have done until now, let's say we have the tokens, right? "one day a little." We have taken for each token, we have taken the token embeddings, we have added the positional embeddings, but until now, we have no information about how one token essentially relates to the other token, right? And that information is very important for us. Why? Because consider this sentence: "The dog, the dog chased the ball, it could not catch it." Let's say if we have this sentence, and this first "it" relates to the dog, right? But if you look at the second "it," the second "it" now relates to the ball. So if I want to do next token prediction, and if some "it" comes somewhere in the next sentence, I want to know whether it belongs to the dog or whether it's the ball, right? If it's the dog, I might predict something like "bark." If it's the ball, I might predict something like "fly," because the ball can fly. But if I don't know how one token relates to another token, I don't have the information about how to predict the next token. And that's what the attention mechanism actually does.

The attention mechanism augments the input embedding vector so that for every vector now, I can have information about the neighbors. So if I look at this input vector, I need to have information that I need to pay the maximum attention to "dog." If I look at this "it," I need to pay maximum attention to "ball." So the main purpose of the attention mechanism is that, let's say, uh, let's say we look at this "it," and currently we have the input embedding vector for this "it." It has no clue about all the other tokens, right? What the attention mechanism does is that it maintains the size of this "it," so that if it's 768, the size will remain 768, but it will convert it into a vector which is called as the context vector. Why is it called as the context vector? Because to get this context vector from the input embedding vector, we take into account how it, how this "it" relates to all the other tokens. So we somehow get the attention scores between "it" and all the other tokens, and we use those attention scores to go from the input embedding vector to the context vector for "it," and this has been captured here in this image very nicely. Uh, when we have one, one token, what we do is that through the attention mechanism, through the multi-head attention mechanism, we actually capture the attention score between one token and all the other tokens around it. So we capture the attention score between "next" and "the." We capture the attention score between "next" and "day." We capture the attention score between "next" and "is." And we capture the attention score between "next" and "bright." Why do we capture these attention scores? The reason we capture these attention scores is because, let's say I have "next," right? And currently we have an input vector for "next." Now I will use all these values: "is," I will use alpha 2 1, alpha 2 2, alpha 2 3, alpha 2 4, and alpha 2 5 to convert this input vector into a context vector. And the context vector, now remember, will be much richer. Why will it be much richer? Because it will have information about how my word "next" relates to all the other tokens surrounding it.

This attention mechanism, if it was not there, language models would not work the way they are working right now. And, uh, what the attention mechanism actually does is that it takes every token and converts them into these context vectors. So in this part is where we start understanding more context about the sentence. So words start understanding the meaning of their neighbors, and so next word prediction effectively becomes much better.

Now, what exactly happens in this multi-head attention block? There are several steps which happen in the multi-head attention block, but more importantly, the input is these input embeddings, right? So if you look at the multi-head attention block, the input are these same input embeddings, which are these, after layer normalization, we apply the masked multi-head attention, and the output is the dimension is the same. So if you see, before multi-head attention, each, each token has a dimension 768. After multi-head attention also, they have a dimension 768. But now each of these vectors are context vectors, which means they have enriched information of all their neighbors. The way we go from input vectors to context vectors is through these matrices which are called as queries, keys, and values. Now, I'm not going to go into detail of this because that will again, uh, tremendously increase the length of this video. But let me give you a quick recap of how this is actually done.

First, what we do is that, let's say I have my input embedding matrix now. So I'll bring this over here. Uh, let's say this, because it's easy to copy-paste. So I'll bring this over here. That's my input embedding matrix. Let's say what is done is that first we introduce some matrices which are called as the query matrix, the key matrix, and the value matrix. And the input right now is 4x768. Each of these will be of a size of 768x768. So when you multiply input with the trainable query matrix, it gives me the keys matrix. And that's again 4x768. When you multiply the input with the trainable keys matrix, sorry, this is the queries matrix. Then this is the keys matrix, and this is the values matrix. All of this must be sounding very weird because why do we introduce additional training matrices? The simple reason is that humans could not figure out a physical representation of how to capture context. So we pass all of that burden of figuring out how to capture context between tokens to these trainable weight matrices. So again, I'm marking these with P because these are the parameters which will be learned.

So let's see where all we have marked P so far. We have marked P over here in the token embedding, in the positional embedding. We'll mark P over here in the masked multi-head attention also, because these are the parameters which are actually learned. Once we have the queries and the key vectors, we'll take a dot product between the queries and the keys, and that will give me the attention score matrix. The attention score matrix will be a 4x4 matrix now because we have four tokens, right? Uh, and my four tokens are "every effort moves you" or, um, the "little girl" example which we had, "one day a little." So now, why do I have a 4x4 here? So if I say "one day a little" and "one day a little," the reason we have a 4x4 matrix here is because for every token, this matrix gives me how much attention should I pay to the nearby token. So if I look at "day," this value tells me how much attention is between "day" and "one." This value tells me between "day" and "day." This value tells me between "day" and "a." This value tells me between "day" and "little." This is my attention score matrix.

The attention score matrix is then converted into an attention weights matrix. And then the attention weights matrix is multiplied with my values. Um, so we have key, queries, and values, right? With my values. And that gives me a context vector matrix. And this context vector matrix is of the size of 4x768. And this is the output of the multi-head attention which we have mentioned over here. I know I went through this, this module a bit fast because usually I spend about four to five hours explaining attention mechanism itself. But here, the main purpose is for you to see a broader picture of how different blocks are assembled together so that you don't feel that this is something too complex. You should get a basic understanding of this, and later, if you want, you can dive into the multi-head attention mechanism a bit further also.

So if you see, what exactly happens in the multi-head attention mechanism is all these steps. We multiply with trainable queries, keys, values. We get the query matrix, the key matrix, the value matrix. We multiply queries with the keys transpose, and we get the attention scores. The attention scores are normalized by square root of keys dimensions, and plus we apply softmax, and we get the attention weights. We multiply with values, and we get the context vector matrix. One more thing to note here is that there is something called as causality here. Which means when we look at these, uh, attention scores, ideally, every token should have attention score between only that token and the tokens which come before it. It should not have attention scores between the token and what comes after it. So, for example, if I look at "day," we should not have the attention score between "day" and "a" because "a" comes after "day," right? And ideally, we are predicting the next token. So we should have no information about what comes after the next token. So essentially, what is done is that, so if you look at "one," it should not have information about these three. If you look at "day," it should not have information about these two. If you look at "a," it should not have information about this. So all the elements above the diagonal are essentially put to zero. And that is called as causal attention. So when you look at the code, you are going to see all of these steps exactly implemented the same way as I mentioned over here, and you will see the name "causal attention." But don't get scared or intimidated by this name. It's essentially because we cannot peek into the future. For one given token, we can only pay attention to that token or what comes before it, not to the tokens which come after it. So we put all the elements above the diagonal to be equal to zero.

If you want more information about the attention mechanism and if you want to dive deeper, you can have a look at the "Build LLM from Scratch" lecture series, which is a huge number of lectures. I think I have 5 hours of content on the attention mechanism. Okay. So this is the second block, which is the multi-head attention, and actually, this is the most important block of the transformer architecture. So if you look at the code, you'll see that this is that block which is called "causal self-attention." And remember, it is called causal because we cannot peek into the future, and it's called self-attention because we are looking at a given sentence. We are looking at a given sentence, and we are seeing computing the attention between tokens of that sentence with tokens of the same sentence. That's why it's called self-attention. Okay.

So now, if you see what is happening here. Um, so here we have the queries, keys, and the values matrix. Okay, the queries, keys, and the values. And then we multiply the queries with the keys transpose to get the attention scores, which has been shown over here. The queries will be multiplied with the keys transpose to get these attention scores. So these are also called causal attention scores. Causal attention scores. So here you can see, uh, here we get the causal attention scores, and then what we do here, you can see that the elements above the diagonal are replaced with minus infinity so that we can, when we take softmax, they will be equal to zero. So when we have the attention scores, we replace the elements above here with negative infinity so that when we take softmax, they'll be zero for the causal attention purposes. That gives us the attention weights. And remember, here we are also dividing by the square root of keys dimension. So we get the attention scores. We get the attention, the mechanism proceeds like this. We get the attention scores. We get the attention scores. We divide by the square root of keys dimension. This is done because the variance of the query is multiplied by the keys. The dot product, the variance of that dot product, should stay as low as possible. That's again to stabilize the training. So we scale by square root of keys dimension. Then after that, what we do is that we fill the elements above the diagonal. We fill the elements above the diagonal with negative infinity, and then we take the softmax. What this will ensure is that it will do two things. It will make sure that every row of the attention score sums up to one because we take softmax. And since we have negative infinity, and when we take that softmax, all elements above the diagonal will be zero. So then the attention scores, or the attention weight matrix, will be something like this. All this will be zero, and whatever is remaining, every row of whatever is remaining will essentially sum up to one. So then you can make quantitative statements like, "When this is my query, how much attention should I pay to each of the tokens?" So all of that is happening over here. We first take the queries multiplied by keys transpose, divide by square root of keys dimension, fill the elements above the diagonal with negative infinity, take softmax. Uh, and then there is also a layer of dropout, which we can apply over here. Again, dropout is applied to improve generalization performance. And the final step, as I've mentioned over here, is that we take the attention weights and we multiply it with the values, values matrix. The attention weights is multiplied with the values matrix, and we get the context vector matrix. That's it. Um, and then, uh, once we get the context vector matrix, then we can add one more layer of dropout after that, and we can also add an output projection layer. So a final neural network is added at the output. We'll come to that in just a moment. But for now, just remember that in the causal self-attention block, until here, we have the context vector matrix. There is one more output projection layer, which is a small neural network, is added at the end of the, so once we get this context vector matrix, a small neural network is added which kind of preserves the shape. So it still remains 4x768, and this is not usually required, but sometimes the neural network is there. So it's an optional neural network which can be added as an output projection layer. That's what's mentioned over here. And then one more layer of dropout is actually applied. So this is what happens in the causal attention block. And I know that I have gone a bit fast through the block, through this block. But if you have understood this much, it means you have understood the basics of the attention mechanism. Remember, what the attention mechanism is trying to do is that we take the input embedding matrix and convert it into a context vector matrix. That's it. And these trainable weight matrices are, uh, devised, or we have added them because we could not come up with a closed-form expression for how to capture context. So we offload that to the trainable weight matrices, and we hope that they will do the job for us. Okay.

So this is the masked multi-head attention. Until this point, we get the context vector for every single token. Then you see in the transformer block, if, let me rub this one so that we have a clearer view of what all steps are happening. Until now, what all have we covered in the transformer block? We have covered the layer normalization and multi-head attention. Now, after this point, there is one more layer of dropout. There is one more layer of dropout. So we randomly mask out some elements to zero again to improve generalization performance. So that is this dropout layer. And then you'll see this symbol over here, right? That symbol is a shortcut connection or a skip connection. And the reason we have shortcut connections is because we give another alternative path for the gradient to flow. Shortcut connections are added to prevent the vanishing gradient problem when multiple layers are chained together. So whenever we have a shortcut connection here, it means that there is a path which has been created. So if you see now, yeah, whenever you have this shortcut connection, it means that there is a shortcut between, uh, this path and this. Whenever you see there is the shortcut connection, there is a shortcut between this path and this, which means the gradient can also flow from here. What it means that there is a shortcut, it means that the input of this layer is added to the output here. So this input is added to this output, and that serves as the input to the next layer. Similarly, what happens here is that this input is added to this output, and that's the plus summation over here. So a simple operation like this prevents the vanishing gradient problem by giving the gradient an alternative path to flow. You can think of this as similar to the ResNet architecture. Okay.

So, there is a, there is a shortcut connection here, and there is also a shortcut connection at the end of the transformer block. There are two shortcut connections. There is one shortcut connection here, and there is one shortcut connection here. Okay. After this first shortcut connection, we have another layer of. After this first shortcut connection, we have another layer normalization. So then again, we ensure that for the context vectors which we have, the mean is equal to zero and the variance is equal to one. That is very important that we apply layer normalization once more. And after layer normalization is applied, there is one more crucial layer, and that's the feed-forward neural network.

Now, something very interesting happens in this feed-forward neural network. This feed-forward neural network actually looks something like this. So let's say we focus on one input, right? Now it has 768 dimensions, right? This input will go through my feed-forward neural network. So I have an input with 768 dimensions. The hidden layer has dimensions 4 * 768, and that I think is, uh, uh, 302, 3072. So the hidden layer has 3072 dimensions, and my output layer again has 768 dimensions. So I call this an expansion-compression neural network, where the neural network retains the size of the input. So every input which passes through the neural network, when it comes through the output, it will have the same size, 768. But remember that due to these expansion and contraction, and due to these additional trainable parameters, it allows the language model to explore a much richer space. That's what's done frequently in deep learning, right? If something does not work in a low-dimensional space, you project everything to a higher-dimensional space so that you can capture additional insights, you can capture more nonlinearities, your function can learn new things. So it turns out that without this neural network, the language model does not learn about the context patterns in the underlying data, and it cannot answer queries that well. But addition of this neural network completely changes the performance of the language model, and that's why this neural network is introduced at this stage of, uh, uh, this stage of the transformer block, where, uh, the input. So now we have four rows, right? So 4x768, when it passes through the neural network, and when it comes out, it's again 4x768. But through this process, we make the, uh, additional, we have additional parameters so that the language model can explore a much richer space.

Now, one more thing to notice is that usually the activation functions which we have in neural networks are ReLU, Tanh, etc. But the activation function here is GeLU, and that's a different function than ReLU because the ReLU is this, right? It's, it looks something like this. But a GeLU actually looks something like, it looks something like this. So at higher X's, when X becomes high, it becomes very similar to Y=X. So it approximates ReLU, but at negative X's, it does not directly become zero. But it's smooth. It's differentiable at X equal to zero, and it slowly becomes zero. So generally, this has shown to have good results. Again, this is an experimental choice. There is no real need to only stick with this. But as a hyperparameter, researchers have got very good results with language modeling if they use the GeLU activation function. Okay. So that's the activation function used in the feed-forward neural network. And once we come outside the feed-forward neural network, there is another dropout layer over here. So then the output of the feed-forward neural network, some of these outputs are again randomly turned off to zero. As you can see over here, some of these outputs are randomly turned off.

to zero to improve generalization performance. And then we have one more shortcut connection here. So the input of this is added to the output of this block. Again, the shortcut connection provides another path for the gradient to flow, thus preventing the vanishing gradient problem.

So now, actually, let us see the building blocks which we have assembled so far. We have the layer normalization block. We saw how the, um, causal attention block is there. There is one more class which we specially create for this feed-forward neural network. So here you can see we have a feed-forward neural network over here. Uh, this is the first layer, and that input goes into four times the embedding dimension. So we have an input with the embedding dimension, as we saw over here, that was 768. The input goes into a hidden layer with four times embedding dimension, that's the expansion layer. Then we have the JU activation function, which we saw, and then there is a, uh, there is a compression layer over here. So that is again projecting the 4 into 768 back into 768, followed with a dropout. So the neural network is again followed with a dropout, as we have seen over here. The neural network layer follows a dropout layer. So we call this class as multi-layer perceptron. Okay.

So here, if you see the code, we assemble three classes. We assemble the layer normalization class. We assemble the causal attention class, and we assemble the multi-layer perceptron class. Once we have these three classes, we can technically assemble the entire transformer block like this, because we have, let me rub this. Why can we assem- why can we assemble this entire block? Because now we have everything we need, right? We have the class for the layer normalization. So I'm marking this over here, one. We have a layer normalization class. We have a multi-attention class with dropout. And we have, uh, the feed-forward neural network with dropout, everything we have. So now we just assemble these different classes together. And that's why that's what's done in this class called block. So the class called block is my transformer block, is my transformer block. And if you see closely, what's happening here, we first have a layer normalization, as we saw here with number one. That's the layer normalization. Then we have the multi-head attention with dropout. So if you see the causal attention by default, towards the end, the dropout is already implemented. So this is number two, the multi-head attention with dropout, that's number two, that's done right now. Then we have, uh, uh, then we have another layer normalization. If you see over here, we have another layer normalization, that's again marked as one. And then finally, we have the feed-forward neural network with feed-forward neural network with the dropout. So that's marked as three over here. So that's how we have assembled the entire transformer block in just three lines of code or four lines of code, because we have defined classes earlier. The LN1 is the first layer normalization. LN2 is the second layer normalization. This is the causal attention block, and here we have the multi-layer perceptron block. Okay.

And if you see the shortcut connections are these. So the input to the, so if you see attention, right, uh, the self.attention attention. That's where the first, uh, shortcut connection exists. So the output of this is added to the input. That's what this line represents. And the second shortcut connection is after the feed-forward neural network over here. So the output of this is added to the input of this. That's what this, that's what this line actually represents. So this transformer block code here contains the entire code for the full transformer block. So these eight lines of code is where all the magic really happens. And remember that in GPT, there is not one transformer block. There are multiple transformer blocks, which we have to decide. So until now, in the code, if you see, we have reached until this stage, right, where we saw the token embedding plus the positional embedding. These were passed as an input to the transformer block. Okay. But now we have multiple transformer blocks. So if you see, uh, here, uh, yeah, h, right. So we have multiple transformer blocks, which is given by the number of layers. So the number of layers tells us how many transformer blocks we have. What our input does is that it has to pass through all of these transformer blocks. It has to pass through all of these transformer blocks. So whatever I have shown over here, whatever I have shown over here, right? There are actually multiple transformer blocks over here. There is not just one transformer block. So one transformer block has this entire architecture, right? These might be 12 or 24 such blocks. So my input has to pass through the first transformer block, then through the second, then the third, right up till the very end. And then, and only then, does it exit out of the transformer, and then it has to come to the output block.

So if you see over here, what we have written over here is that the input, which was the token plus the positional embedding, has to go through all of these multiple transformer blocks, and then it comes out of the processor. And this is that step in which we have the output defined. So I'll come to that in a moment. Um, all right. So the output has been defined, actually, the output has been defined over here, and also over here. So I'll come to that in a moment. But until this point, just remember that we have gone through the multiple transformer blocks, and then we have the output. What is the output size after going through multiple transformer blocks? After the first transformer block, the output size remains the same as the input. Right? When we went inside the transformer block, remember the size was four rows and 768 columns. Even after applying dropout, that's the same size. When we come out of multiple transformer blocks, the size remains the same. That's 4 * 768. Okay.

Now we have come out of the multiple transformer blocks, and finally, we'll go to this last third block, which is the output layer. Now, you see here, there are two things which happen. First, we have to pass through another layer of normalization, and then we have to pass through the output layer. So let me explain step by step what is happening over here. When we come out of the transformer block, we have this size, right, which is 4 * 768. Then we will do one more layer of normalization. And here, the same thing happens. Essentially, we subtract the mean, divide by the square root of variance. So now the mean is equal to zero, variance is equal to 1 for every row. Okay. And then we pass this 4 * 768 matrix into an output head. What this output head does is that this output head converts every vector is now of a size of 768. Right? It converts every vector into the vocabulary size. Now, it converts every vector into the vocabulary size. And the way it does is that by passing every vector through a neural network. So if this size is of 4x768, I will have a neural network of size 768 multiplied by my vocabulary size. So when this input passes through this neural network, it will be of a size of 4 into vocabulary size. So now, take a look here. Every token now will have 50257, or my vocabulary size, if my vocabulary size is something. Every token will have those many number of columns. Uh, so every token has these many number of columns, which are equal to the vocabulary size. This is called as the logit tensor. So if we just have one batch, the size of this logit tensor will be four multiplied by the vocabulary size. If we have two batches, it will be 2 mult* 4 multiplied by vocabulary size. Don't worry about multiple batches for now, because we are just looking at one batch currently. So every effort moves you. Let's say these are the four tokens, and we have the number of columns equal to my vocabulary size.

Now, this is exactly what is done in the code, right? So here we have come out of the multiple transformer blocks. What is done here is that if you see LNF, it is another layer of normalization. So we apply the normalization layer here. That is what is given in the, um, if you see the output block. Yeah, the, we have one more layer of normalization here, right? That's the first layer of normalization. And then, if you see what is done here is selfLM head. Now, this LM head, if you see, that's a neural network with the input of number of embeddings and the output is equal to my vocabulary size, right? Now, so that's my neural network with the number of rows equal to N embedding and the number of, uh, columns which is equal to my vocabulary size. All right, good. So this is the LM head, which is the output head now, which my u, which is the final output of my transformer block. So this is my logits matrix, and this logits matrix is now going to be used for the next token prediction task. How is my logits matrix going to be used for the next token prediction task? If you carefully look at the logits matrix right now, every logits matrix has size equal to the vocabulary, right? So, uh, and as I told you, every batch, every batch has multiple input target prediction tasks. So if you see this, let's take this same example, right? One day a little. When we come out of the entire transformer architecture, uh, when we come out of the entire transformer architecture, we have something like this: one day a little. And this is now 50 to 57, or my vocabulary size. Let's say all of these are now equal to my vocabulary size. All of these are now equal to my vocabulary size. Okay.

So the way the next token is predicted is that when one is the input, I look at that token index which has the highest probability. If that is equal to 40, I go to my vocabulary and I look for token ID 40, and I see the token associated with it. So ideally, if my model is trained, the output should be the same as the target. Now, the target, when one is the input, day is the output, right? That's the target. But maybe token 40 corresponds to something like random thing like the. Now, when one, when day, when I look at day, I again look at that thing with the maximum probability, the index with the maximum probability. When one day is the input, o is the output, but maybe I'll get something like flag. Then, oh, I look at the index which corresponds to the maximum probability. If that index is 60, maybe that corresponds responds to something like, uh, o. But my target is little. And when I look at little, and I look at the index with the maximum probability, maybe that corresponds to person. But ideally, my target is girl. So now I have my outputs and I have my targets, and then I'll calculate the loss function between the output and the target. This is how the logit tensor, the logits, the output logits matrix, which I get after coming out of the transformer block, this is how it is used for the next token prediction task.

So, so one, when one batch is passed, let's say one batch is passed, right? Which is one day a little. When it comes out of all of these three steps, when it comes out of the input processor and output, it has this logits. And within that, there are multiple input output tasks. There are four input output tasks which needs to be predicted in each batch. And this is exactly how these tasks are done. So now there are four output target pairs, right? We find a loss function based on all of these four, and it's this loss which eventually needs to be minimized in my training routine. So whenever one batch is passed, there is a loss which is an accumulation of, let's say, if this is L1, this is L2, this is L3, and L4 for every batch, my loss is L1 + L2 + L3 + L4. So I get my loss for all the batches, then I do the back propagation, etc., and that's how the training actually proceeds. Um, yeah. So until now, what I've shown you is exactly how the, uh, how the transformer architecture is assembled. And let's see some details in the code right now, which we might have missed, and we'll cover some of that, uh, in the future or in the rest of the modules as well.

So in the forward method, now we have token embeddings, positional embeddings, we add them together, we apply dropout, pass them through through the transformer block. In each transformer block, which we have, there are multiple things which are happening, and that's defined in this block. There is a layer normalization, a causal self-attention, another layer normalization, a feed-forward neural network. So we go through all of these blocks. Uh, we go through all of these in the, in one transformer block. And similarly, there are multiple such transformer blocks. When our input embeddings come out of this, they are now context embeddings. And when they come out of the transformer block, when they go through multiple transformer blocks, we pass them through one more, uh, layer normalization layer, and then we pass them through the output neural network, which is also called as the output head. That gives us the logits matrix. The logits matrix, when we look at the logit matrix for every row, we look at that token ID which corresponds to the next token prediction, because that has the maximum probability. And then I compare it with my actual targets. That's how the loss is calculated. So when I showed this L1, L2, um, when I showed this L1, L2, L3, and L4 here, actually that's the cross-entropy loss, which is calculated, and we'll see that in a lot more detail in, uh, in the next section, actually. So the next section after this point is going to be the loss function. So we are going to see that in a bit of detail. So until now, you would have understood everything in this forward method and in this class GPT. You have understood class GPT, class block, class MLP, and class, uh, causal attention.

There are just some initializations which are done at several places in this code, which I want to explain to you. So, for example, if you see, there are these initializations which are done. So different modules have different trainable parameters, which are kind of initialized in slightly different manner. So, for example, if you see over here, here is a list of all the trainable parameters which are usually there, right? So the token embedding layer has trainable parameters. The positional embedding layer has trainable parameters. Those are initialized as a Gaussian with normal distribution and standard deviation 0.02. Then the attention block has the query, keys, and the value matrix. They are also initialized as a Gaussian, each of these matrices. Then the attention block has an output projection. Remember I said that's optional, but you can keep that on. And that's initialized in a slightly different manner, just so that the training process becomes a bit more optimal. We divide it, uh, by square root of two into the number of transformer blocks which are there. Then this is the feed-forward neural network. Uh, the first layer, which goes from 768 to 4x768. We initialize the weights as a Gaussian. And the output layer, which is also the projection layer. It's called as the projection layer. Uh, that's initialized by a Gaussian. But we divide, uh, 0.02 by again, square root of 2 into the number of transformer blocks, which is, if it's 12, then it will be square root of 24. Then the final layer normalization, again, that has no, no initializations. And the final output head, which gets me my logits matrix, that is initialized with a Gaussian of mean zero and standard deviation of 0.02.

Now, whenever, now, when you run this code cell, when you run this whole code cell, ultimately we are going to only use class GPT, which is my model. And for that, we have to define a GPT configuration. So the configuration which we are going to use is vocabulary size of 50257. Um, then I'm going to use a block size, block size of 128, the embedding dimensions of around 384. Um, that's the embedding size which we are going to use, and, uh, the dropout rate of 0.1. So number of layers and number of heads, there is a distinction between these. Number of layers is the number of transformer blocks which I have. So remember here, I told you that transformer, a number of transformers will be, uh, added here. So here we have to decide the number of transformers, and that's equal to six. So here I'll just mention, uh, N layer equal to number of transformers. Okay, number of transformer blocks, and N heads is the number of attention heads in each block. So we did not talk about attention heads today. But essentially, when multi-head attention is computed, there are multiple attention heads, each computing a separate attention score, so that we can capture multiple perspectives. And that's equal to six. So you can think of this fact that in each transformer block, there is a multi-head attention, and the number of attention blocks which we have are now equal to six. And here also, we have six transformer blocks. Right? So let me replace this with actually six, since we have six transformer. So six transformer blocks, each have six attention heads. So overall, there will be 36 attention heads. So this is the configuration. The N embedding is the embedding dimension which we have, right? So at the start, I told you that we have to decide an embedding dimension. That embedding dimension is now going to be 384. Okay. So this is the embedding dimension which is there. Um, and then I just define my model as GPT. So this GPT, which is mentioned here, is the GPT class which we have defined over here. My GPT class, um, this GPT class, which takes my input embedding matrix and produces, um, a loss which is based on the logits output and my target outputs. So this is the model GPT, and I have to specify this configuration, as has been mentioned over here.

Now, what I, I'm going to do is that in the next section, I'm going to talk a bit more about the loss, because that's very important to understand. In this section four, the section four was the most important section, where we had to understand the different steps of how essentially the input, the input sequence goes through the entire transformer architecture, the input, the processor, and the output. And at the output, we get this logits, we get this logits matrix, and that can be used for next token prediction task and compare with the target next tokens. And that's how the loss function is determined. This is the whole workflow of how the LLM architecture operates. And I hope you have understood how it maps exactly to the code. I really encourage you to take a look at the code in in much more detail after going through this lecture, but I've given you enough starting point so that you'll understand many different things which are happening in the code.

Now, what we'll do is that we'll come to the next step, which is defining the loss function. It is very important for us to understand what exactly happens in the small language model loss function, because here is where we'll ultimately try to tell the language model that look, you're making a mistake, and you need to improve upon this. Right? We have already looked a bit at the loss in the previous part, where we saw that once the logits matrix is obtained, let's say here we get the logits matrix. Right? After this point, we look at every token, and we look at the token ID which corresponds to the highest probability, and we decode the output corresponding to that token ID. We compare it with the target, and that's how we get the loss function. But now we are going to look at this same step in a bit more detail, so that you can get an exact clarity of what exactly is happening in the loss function.

So let's say we have the batches which look something like this, right? So this is my first batch, and let's say that these four tokens are my input X1, and these are my target, which is equal to Y1. Okay. So what we'll first do is that we'll first take X1, which is this, and we'll pass it through the transformer model. We'll pass it through the input, the processor, and the output. So let me bring this over here, so that, um, you recollect what we have seen in the previous lecture. So let me copy-paste this and u, I'll bring it over here. Okay. So what we are going to do is that we have this input text, right, which is "I had always thought." "I had always thought." We'll pass it through first the input layer, then it will go through the transformer block, and then when it comes out of the transformer block, we essentially have this logits matrix. And try to think about the, uh, try to think about the size of this logits matrix, right? If you have four, if you have four tokens which go as an input, they are first converted into an embedding dimension of, let's say, 768, which we had seen before. But once they come out of the transformer block, the logits essentially look something like this. We still have these four tokens: "I had always thought." But now the dimensions of each of them are 50257, which can be the vocabulary size, or whatever vocabulary size we have. So if you take a look over here, this is what we have seen previously also. The size or the dimensions of every token are now equal to the vocabulary size. That's the first step which is, uh, which is done to calculate the loss.

Then what we do after this point here is that essentially we want to convert every row. So let's say this is row number one. Uh, let's say this is row number two. This is my row number three. And let's say this is my row number four. I want to convert every row into a vector of probabilities, which means that I want every element, I want each row to sum up to one. So if I consider all the elements of row one and I sum them up, I want their summation to be one, so that I can give a probabilistic meaning to every token. To do that, what I do is I apply the softmax to this logits matrix, and then essentially I get a vector of probabilities. Remember, what softmax essentially does is that it replaces every element with, let's say, e to x, which is that element, multiplied by e to x1 plus e to x2, etc., for a particular row. So it ensures that all the elements lie within 0 to 1 and they sum to one. So now every row can be interpreted in a probabilistic sense. So you can look at every row and try to find that token ID, token index, which has the maximum probability. So that will be the decoded next token. So for example, if this token ID has the maximum probability, and if this is token ID 35, that might correspond to a decoded subword of "ac". So we go back to the vocabulary. Now we look at that token ID and get the decoded token. Of course, this is my LLM output, and this is not the target output at all. The target output which I want is "had always thought jack." So this is my X1, and this is my Y1, right? So this is my target output. So for these inputs, ideally, I want the target output as "had always thought jack." Okay. And if I go to my vocabulary, I see that "had" has a token ID of 23. "Always" has a token ID of 3881. "Thought" has a token ID of 11223. And "Jack" has a token ID of 15.

So now let me go to this, uh, vector matrix of probabilities, and let me write this here again. "I had always thought." And let's look at "I" first. Okay, let's look at "I" first. This is that 50257 [Music] uh dimensional vector. And the maximum probability currently is at index number 35, let's say. Or let's say index number 35 is here. That's the maximum probability, and it corresponds to the token output of "ac". And what I want is that I want the next token to be "had" when "I" is the input. So I want the token ID corresponding to 23 to have the maximum probability. So ideally, what I want to do is that I want to make sure that for the token ID number 23, which is this, this probability is as close to one as possible. What that will ensure is that if this probability is as close to one, all others will be zero, and my next token which is predicted would be "had". So ideally, the loss function which I have, that is constructed only based on this token index, which is 23. So currently, let's say the token index 23 has a probability of around 0.1. Let's say in the current version, let's say token index number 23 has a probability of .1, which is not very good. I want this probability to be as close to one as possible, and that's how I will define my loss. I will look at the index which I actually want as the next token, and I will make sure that the probability corresponding to that index is as close to one as possible.

Let's do this exercise for the second row, which is "had." Now, and let me write down the u, vocabulary size length vector now for "had." Okay, let's say this is the vector, and I want to now look at "always," because when "had" is the input, I want "always" to be the output. So I want the token ID corresponding to 3881 to have the maximum probability. So here I'm going to look at the token ID corresponding to 3881, and I will look at its probability, and currently, let's say the probability is .02, which is not good. I want to take this probability as close to one as possible. So I'll define the loss between these two. So if you notice carefully, what is going on, we have these target indexes, right? Which play a very important role in getting my loss function. These are my target token IDs, and ultimately the loss function will be calculated based on these token IDs itself. And I want the probabilities which is associated with these token IDs to be as close to one as possible.

So what we do in the next step is that we get the output probability for the target token IDs. Which means that in the first row, we look at token number or ID number 23, and get the current probability. Let's say that is 0.1. In the token, uh, in the second token "had," we look at index number 3881, and get the probability corresponding to that. Let's say that's .02. Similarly, we get the probabilities corresponding to "always," whose true answer is "thought," and "thought," whose true answer is "Jack." Okay. And we get these probabilities, and they are P1, P2, P3, and P4. P1, P2, P3, P4 are probabilities corresponding to these target token IDs. And remember, we want these probabilities to be as close to one as possible. That's the main, uh, that's the main output which we are chasing. And the loss function which we are going to define is exactly based on this concept that we want all of these probabilities to be as close to one as possible.

And so we conduct, or and so we construct a loss which is called as the cross-entropy loss, which is based on the loss fun, which is based on the logarithm. So we take the logarithm of all of this and apply negative sign here. Since we want to maximize all of these probabilities, when we take the negative sign of here, essentially we want to minimize these quantities as much as possible. Okay. And since the probabilities always lie between 0 and 1, these, this entire quantity will be minimum when P1 is as close to one as possible, P2 is as close to one as possible, P3 is as close to one as possible, and P4 is as close to one as possible. This loss, which is the negative of the summation of the logarithms of these probabilities, is also called as the negative log likelihood, or it's also referred to as the cross-entropy loss. The cross-entropy loss is the main loss through which we are going to do back propagation for language models. And it's very important for you to visualize how we obtain this loss. Okay.

First, we go through this entire transformer block, the, the input, the processor, and the output. When we come out of the output block of the transformer, we have this logits matrix, right? From this logits matrix, essentially we have the tokens, and we have the token IDs, uh, which correspond to the decoded next token. So what we do is that we have the inputs, and we have the targets. So for every input, we have a set of target token IDs. Let's say for this input "I had always I had always thought," the set of target token IDs was "had always thought jack," that's 23, 3881, 123, and 15. We want the probabilities corresponding to these token IDs to be as close to one as possible, because that will be the decoded next token. So in the first row, I want the 23rd token to have the highest probability. In the second row, I want the 3881 token to have the highest probability. In the third row, I want the 12223 token to have the highest probability. And in the fourth row, I want the 15th token to have the highest probability. And these probabilities are currently P1, P2, P3, and P4. I want P1, P2, P3, P4 to be as close to one as possible. So I want to maximize P1 + P2 + P3 + P4. One more way to rewrite this same loss is if you want to minimize the average of log P1 plus log P2 plus log P3 plus log P4. Taking a logarithm is generally much more better, because if you want to maximize P1, P2, P3, and P4, you would construct your probability which you want to maximize as P1 into P2 into P3 into P4. Right? Um, and taking derivatives of products is not very easy. So you convert this into a summation by taking a logarithm of this, and taking a derivative of this quantity during back propagation is just easier, and we apply negative sign because it's common notation to say that, uh, loss needs to be reduced. So ideally, we want to reduce this quantity as much as possible, so that P1, P2, P3, and P4 will be as close to one as possible, and that's the negative log likelihood loss.

What I showed you right now is just for a single input, right? But when you are processing an entire batch together, what usually happens is that this logits matrix for multiple batches are aggregated together, and the target token IDs for multiple batches or inputs of the same batch are aggregated together. For example, if you look at one batch now, it has four inputs, right? X1, X2, X3, and X4. If the batch size is equal to four, these four inputs are aggregated together. Their logits are aggregated together, and their target token IDs are also aggregated together. And so here we have P1, P2, P3, and P4, right? So then we'll have 16 such probabilities which we'll need to maximize, because now there are four inputs, and there are four outputs. Uh, so let me repeat this once more. Let's say if you have X1, which is 1, 11, 15, and 23. You have X2, which is again these four. You have X3, which is these four, and you have X4, which is these four. That's my first batch, correct? And you have Y1, Y2, Y3, and Y4. The, what ideally you have is that you have this logits matrix for X1, which is now 4 by vocap size. You have logic logs for X2, which is 4 by vocap size. You have logits for X3, and you have logits for X4, which is the same size. So it actually looks something like this. This same vector, you, the same matrix, you can imagine, but now repeated four times. This is for the second, this is for the third, and this is for the fourth. Then what you do is that you collect all of these logits matrices together. That's called as the flattening operation. You flatten all of these together into one big logits matrix, right? And that will essentially now be, so if you have four, four such ones, right? So that will be 16 multiplied by my vocap size. So you have 16 rows over here. And this target, right? My target now, my target will also be 1, 2, 3, 4, that's corresponding to, uh, Y1. Then I will have 1, 2, 3, 4, that's corresponding to Y2. 1, 2, 3, 4, corresponding to Y3, and 1, 2, 3, 4, corresponding to Y4. Okay. And as we saw, for each Y, we have four probabilities, right? P1, P2, P3, P4. So now we have P1, P2, P3, P4, that's for Y1. Then P5, P6, P7, P8, that's for my Y2. Then P9, P10, P11, P12 for Y3, and P13, P14, P15, and P16, that's for my Y4. Now, this is my set of probabilities, and I want to make all these probabilities as close to one as possible. So we'll just again calculate a negative log likelihood for this entire batch. So that's how one batch is processed. Okay.

So when we see the code, you will see some terms such as flatten, etc. So whenever this flatten is used, it essentially makes it means that my multiple, uh, inputs are now merged together into one big, uh, matrix or tensor. So that's how the loss of the small language model is actually calculated. And you'll see that when you have this GPT class, right? When you have the class GPT, the loss is actually calculated within this class itself. If you look at the forward method, until now we have the logits matrix, right? So here you see what is happening. We have the cross-entropy loss again between the logits and between the targets. So this is that cross-entropy loss which I mentioned, um, over here. Essentially, that cross-entropy loss just means taking the negative log likelihood for the estimated tokens which we have. And then we have one more function which is called as the estimate loss. This estimate loss function essentially, uh, what it does is that it gets one batch and then it estimates the, uh, losses for that entire batch, actually. So here you can see that, uh, for a prescribed number of iterations, let's say we have 100, uh, we have 100 iterations, right? And and evaluations are done after every 100 iterations. So what happens is that for every iteration, we get one batch, and then u, essentially for every batch which we have, we get the, uh, we get the loss, and that's added to my losses, right? Now, okay. And then you see this out.outsplit equal to losses.mean. So for all of these iterations, I essentially will take a mean of mean of the losses. So this estimate loss, what it does is that I have to run the model for a prescribed, uh, I have to run the model for prescribed number of iterations. And when I do that, when I run the model, let's say for 100 iterations, for each iteration, I'm getting a different batch. When I get a batch for an iteration, let's say X and Y, I get a batch which is like what I've shown here. We get a batch for an iteration. We pass this batch through the model. We pass this batch through this entire model. We get the logits. We merge all of these logits together, and then we get the loss, which has been mentioned over here. We get the loss, the cross-entropy loss. That's the loss from one batch during one iteration. When we are doing another iteration, a separate batch is is given to us. Remember the get a batch function sources the inputs and the targets randomly. So if you remember what we had seen in get a batch, there is this equal to torch.random integer, right? So it takes batch batches randomly from the dataset. And once we have one batch, we can propagate it through the model and we can get the loss. That's what's happening here. Okay. Uh, we collect all of the losses which are happening through different iterations. And then what we do is that we essentially take the mean. We take the mean. So let's say in the iteration number one, the loss on the batch is this. Iteration number two, the loss on the batch is this. Iteration number three, the loss on the passive batch is this. For all of these iterations, the loss, the overall loss is the mean over, let's say, these 100 iterations, uh, the 100 iterations which we have. So that's the estimate loss function. This function will be very useful for us when we are, let's say, training the model, and we want to print out the loss after every evaluation iterations. So if you want to print out the loss after every 100 iterations, it it takes the mean of all of those 100 iterations and prints out the, and prints out the mean loss for those 100. That's what is happening in this loss function calculation.

So this is how we get the loss in a language model. We convert the dataset into batches. One batch is plucked out. That that's passed through the, that's passed through my model. Then I get the loss out of, uh, the target probabilities and my output probabilities for the entire batch. And then I go to the next iteration, etc. That's how the losses are estimated. Okay. So until now, what we have done is we have finished step number five, which is defining and calculating the language model loss function. And then we now go to step number six, which is coming to the training loop of the small language model. So until now, let's do a quick recap of what we have until now. What we can do is that we can take a batch. Let's say this is my first batch of input X1, which has, let's say, X1, X2, X3, and X4. And that's my first batch of output Y1, which is Y1, Y2, Y3, and Y4. I can take my whole batch. I can pass it through the language model, which consists of input, processor, and output. Then I get these logits. Then I have a loss function estimation, and I get my loss. So essentially, now we can do one full forward pass through my input. And once we can do this full forward pass, once we get the loss, then the only thing which needs to be done is that we need to propagate this loss backwards through my entire architecture, so that we can then update the parameters of my architecture. And that's when we start defining the training loop, okay, of the small language model.

Now, as I mentioned at the start of this tutorial, this is going to be a bit of a production-level tutorial, because we are going to do a bit more advanced things so that our iterations run faster. One such advanced thing which you will notice in the code is that we use something like torch.cast. The reason this is done is that it enables something which is called as automatic mixed precision. So precision corresponds to float 32, float 16, etc. Um, so usually by default, let's say if everything is in float 32, that takes the longest time to compute. When we use automatic mixed precision, the model converts these to float 16, so that the computations become faster. Now, when are the numbers converted to float 16? Let's say if we are doing something like matrix multiplications in the attention score calculation, or we are doing a dropout, let's say, then the automatic mixed precision would use something like a float 16, because it's safer to compute in lower precision over here. But let's say if we are doing softmax, where exponential operations are involved, cross-entropy loss, where where logarithms are involved, weight updates, etc., all of these which have some sensitive operations and which can lead to a blow-up of numbers if due to exponential ial or log operations, there it's better to go with float 32. So in these operations, we usually go ahead with float 32. But for operations such as matrix multiplications, dropout, etc., float 16 is also completely fine to use. And this automatic mixed precision helps us do that. So wherever possible, the model uses float 16, and, uh, if not, then the model sticks to float 32. So you'll see this in the code. So when we go to step number six, which is the training configuration, you'll see that here, this is where we are using the [Music] torch.cast. So this is important to understand exactly why we are, uh, why we are doing this, uh, um, torch.cast.

Then the second thing which is important to understand is that there is something which is called as the gradient accumulation. Okay. Now, when we go to the actual code, let's say we want to train with a batch size of 1024. Um, the problem is that 1024 input and output pairs will not fit in my GPU at one time. If only if my GPU can only accept, let's say, 32 samples. So what we can do is that we can set gradient accumulation steps equal to 32. So what that means is that we'll run 32 iterations. Um, and then we'll update the parameters essentially after every 32 steps. Instead of updating the parameters after every 1024 steps, what we do is that the parameters are updated after 32 steps. So we can run 32 iterations because the b, the GPU fits only batch size of 32, and then we do it, uh, then we do it once more, right? Um, so essentially, we are not doing the parameter update after every 1024 steps, but we do the parameter updates after 32 steps. This is done just so that if we are doing the parameter updates, let's say after 1024 steps, it means we'll have to compute the gradients of batch size equal to 1024, and we'll need to store all those gradients, but that is not possible to do, um, that is not possible to do in a GPU. If my GPU fits less number of, uh, samples, then we can only calculate those many gradients, and then we can essentially update the parameters. So here, what I'll do is that, let's say, uh, you do first forward pass. Let's say you do the, this entire forward pass, then you calculate the gradient. You do the forward pass once more. You calculate the gradients like this. You do it for 32 times, right? So there will be 32 gradients which which have accumulated. Then what you do is that the final gradient which is there, you then take the mean, you then, you then take the, you update the parameters based on that final gradient. So you don't update the parameters after every gradient. So think about what happens, right? In every step here, one batch is processed. So in the first step, one batch comes here, X1 and Y1, we do the forward pass. In the next step, the second, the second batch is the second batch is obtained, which is X2, Y2. Similarly, we do this until, let's say, gradient accumulation step is 32. We do this 32 times, and until now, remember, we have not updated the parameters. The parameters will be updated only after 32 steps are done, and then I go to the next, etc. So although there are 1024 batches, the parameter updates are not done after 1024 batches. They are done after 32 batches. Um, so one more thing, uh, okay, one more thing I want to mention is that gradient accumulation is typically done if the batch sizes are quite large, okay? So if the batch size is equal to 1024, that's quite large, right? Because in every batch which you have, X1 will have 1024, X2 will have 1024, etc. And ideally, the practice is to process one entire batch and then update the parameters, but that's not what's done over here. If the batch is used, then only 32 are done. The parameters are updated. Then we go to the next, the parameters are updated, etc. Um, okay, so let me clarify this once more, because I know that this part is a bit tricky. So I want to mention that, let's say, uh, let me do this once more. Let's say this is my one batch. Okay. And my one batch now has 1024 such samples. And this is my X1, and this is my Y1. Okay. What gradient accumulation actually does is that we first go through this first. It's, it's gone through the entire model. Then we update, uh, then we get the loss function. Then we, uh, get the gradients. Okay. Then we go through the second. We get the loss function. We get the gradients. Similarly, we go to 32. This is the gradient accumulation steps. And until then, we have got all the gradients. Once we have reached until 32, that's when we update the loss function, or rather, that's when we update my parameters. Okay. So, the parameter update is done after these gradient accumulation steps, not after I go through the entire batch. That's what, uh, is the main point of gradient accumulation. And you'll see gradient accumulation.

In the code also. So if you see here, gradient accumulation steps equal to 32, right? So you're accumulating gradients up till that point, and then you are updating the parameters. Essentially, that's the second thing which is important to note in this uh training loop.

The third thing which is very important to note is how we are doing the training process itself. Okay. So, uh, when we are doing the training process, there are optimizers which need to be used, and in this case, we are going to use the Adam optimizer, uh, with weight decay. So if you scroll down over here, you'll see that we are using the Adam MW optimizer. That just better so that we don't get stuck in a local minima, and then training becomes stable.

But one more thing which you will see in the code is that we use a learning rate which looks something like this. The learning rate warms up in a linear manner, and then it decays in sort of like a cosine manner. So the learning rate is a combination of a warm-up plus a decay regime that is used again for stabilizing the training. You'll see that we have a learning, a linear, a linear function for the warm-up, and then a cosine decay. So that's my learning rate which I'm going to use with my uh optimizer. Generally, this is seen to have the best uh performance when training language models such as this. So we are using this optimization routine.

And then finally, what we are going to do is that we are going to use the Adam MW optimizer, which I mentioned. Now, all of these steps are mentioned in step number six and step number seven. I already mentioned the autocast to you. So that is done. The gradient accumulation step. Uh, and this is the learning rate. Learning rate, a linear regime plus cosine regime followed with Adam optimizer.

One more crucial thing to note here is that these are the hyperparameters which actually I would encourage you to explore around a bit after I share this code file with you and note some of the things which which are there. We are using a batch size essentially of 1 32. Um, and uh, so batch size which we, the batch size which we are using is 32, and the gradient accumulation steps is equal to the batch size. So in this case, my update will be done after every batch size. But if you choose to have a larger batch size, you can stick with gradient accumulation steps of 32. My block size is equal to 128. Remember, that's my context size. If you go over to the earlier, uh, over here, when you create the input output pairs, the block size is equal to my context size. We have seen that in the earlier part of this video.

Um, then the batch size and block size are done. The evaluation iterations, which means I'm going to print out my losses after every 500 iterations. Remember here, what we did, if the evaluation iterations are 500, in each iteration, we are getting a batch and we are getting its losses, right? So this is essentially done for, uh, uh, this is done for 500 iterations, and then I take the mean of all those iterations, and I'm then, then I'm going to report the loss.

Then what else is the hyperparameter? We have the learning rate, minimum learning rate, 5 into 10 minus 4, the warm-up steps. These are the hyperparameters in the, in my learning rate routine, and my learning rate, 1 into, so 10 to minus 4, essentially, and the number of iterations which I have, that's going to be equal to 20,000. All of these are hyperparameters, and you can feel free to play around with this as you explore more with this code.

Now, once you define all of these hyperparameters, and once you have understood the different things such as, uh, automatic mixed precision, gradient accumulation, and, uh, the learning rate warm-up plus decay, then we can go ahead to pre-train the small language model. This is the last step in the pre-training part, and you'll see that this is the pre-training code.

So what exactly happens in the pre-training is that, uh, there are 1, 2, 3, 4, 5, 6, 7, and 8 steps which happen, and I think now you are fully equipped to understand all these eight steps. For each iteration, we have to choose an X and a Y. What this means is that for each iteration, what we have to do is that we have to essentially go through, we have to go through my data set, and remember there is this torch.rand. What this does is that this will essentially choose random indexes from my data set, and that will give me my input pair. So let's say this is the first index, second, third, and fourth. If the batch size is four, this will be x1, x2, x3, and x4. And, uh, my y1, y2, y3, and y4 will be these y1, y2, y3, and y4 just shifted to the right-hand side by one. So that's the first step, which is getting the batch. So I get my batch.

Then what I do is I pass my entire batch through the model to get the logits. This is the forward pass which we, we had seen. We take a batch, we pass it through the entire, uh, language model, and we get this logits. Once we have the logits matrix, then we calculate the loss between the logits and my targets. That's my loss function. So that I calculate in my step number, step number three.

Then what I do is we do the back propagation. Correct? We are going to do the back propagation, which I showed through this arrow. So that is the back propagation. And here, what we are going to do is that we are going to accumulate the gradients till we reach gradient accumulation number of steps. So essentially, we are going to do this gradient accumulation is equal to 32. We are going to accumulate the gradients for, uh, we are going to accumulate the gradients for 32 steps. So essentially, for 32 iterations, we are going to, uh, accumulate the gradient, and, uh, then what we are going to do is that once we accumulate the gradient for 32 steps, then and only then we are going to update the parameters. So let's say we do one step, we do back propagation. We are not going to update the gradients. Then we'll do one more, uh, iteration. We'll again do the back propagation. One more iteration. Again, do the back propagation. And then after we are done through 32 steps, then we are going to update the parameters which we have.

And once the, how are the parameters updated? The parameters are updated through the AdamW optimizer, which we have chosen. So the simplest vanilla version of the gradient descent looks something like this, minus alpha times this. Right? The Adam is a bit more complex. It takes into account the history of the curvature which we have in our loss landscape, and it also uses a, um, adaptive learning rate.

So once the parameters are updated, then we also update the learning rate, and then again, we evaluate the model after every eval iterations, which is after every 100 iterations, and this loop is repeated. So we do the forward pass. Once we do the forward pass, we back propagate the loss. Remember that, uh, remember that we are essentially accumulating, uh, we are accumulating the gradients till we reach the gradient accumulation number of steps, and only after that point, we are going to update the parameters. The parameters are updated through the AdamW optimizer, and then we update my learning rate. My learning rate is updated every time the parameters are updated because we have an adaptive learning rate in this case.

And then what we do is that after my learning rate is updated, I when I go through evaluation iterations, which in this case, let's say if it's 100 iterations, let's see how much it is. The evaluation iterations which we have is 500. So when I do 500 iterations, I'm going to evaluate my loss, which is the mean over all these 500 iterations, and I'm going to, uh, I'm going to print out my loss.

Okay. So that's the main thing which is, uh, happening in these eight steps. So if you see the pre-training loop for the small language model, you can start with this thing where we get the batch, we get the X and Y, we pass the X and Y to the GPU. Then as the next step, what we are doing is we pass the X and Y through the model, we get the logits, and we get the loss. We divide the loss with the gradient accumulation steps because remember, we are going to update the parameters only after we accumulate these many steps. So just so that, so the loss will also get accumulated right over these many gradient steps. So I'm going to just, if I have 32 steps, I'm going to divide the loss by 32, and then I'm going to do the backward pass. Here, the parameters will be updated only after the gradient. Uh, so here I'm going to do the backward pass. The parameters are not yet updated. Here, the parameters are updated. So scalar.step optimizer parameters are updated only after I reach the gradient accumulation number of steps. Uh, then my learning rate is also, my learning rate is also updated. Um, and then this whole thing basically, uh, repeats in a loop, right?

And here, what is happening is that once we have reached the eval iterations, we are going to print out the losses using the estimate loss function which we saw over here. This estimate loss function. So as you see now, many things for us are coming together when we are pre-training the large language model, right? We take the batch, we pass it through the language model, we get the output. That's what's happening until this stage. Uh, then what we do is we do the batch propagation, and we start accumulating the gradients. But remember that we are not, uh, um, we are not updating the parameters yet. Only when we reach the gradient accumulation steps, we update the parameters. Then we change the learning rate. And essentially, this process continues. And once we have 100 evaluation, 100 iterations, which are 500 iterations rather, which are done, we call the estimate loss function, and we print out the training loss, we print out the validation loss, and we store the training and the validation loss.

Here, what we are also doing is that we are going to save our model which is performing the best in this best models parameters.pt, so that we don't have to retrain again. Remember, this is also very important. Sometimes we forget to store the model parameters, right? And that's not usually a good thing. The good practice is to, whenever you're training a model, just save the parameters because training is a very intensive, computationally intensive procedure, right? And you don't want to be running the loop again and again. So we are going to save all the model parameters which ultimately result at the end of, or which give the best loss. So the best validation loss will be associated with a certain set of parameters. I'm going to store all of those parameters, essentially. Okay. And, uh, so that I can reload the model anytime in the future when it is needed. So this is my training loop. You can run this training loop now.

So once you start running this training loop, here is where the main iterations will start. And let me now explain to you how these iterations usually proceed. All right. So I hope all of you are excited to reach this stage because here is where the whole running of the loop actually happens, and this is the best part of this tutorial. Until now, you have spent around two, two and a half hours to go through the entire content so far, and I've tried my best to explain every single part of, um, this code as much as possible. So first, I started with data set, then data prep-processing, assembling the model architecture, setting up the training pipeline, pre-training the language model. So now here, what I've done is that I've chosen the runtime to be 800. So on a run, how many iterations? I've run 20,000 iterations over here, and it took me around 30 to 35 minutes to run all of these iterations. As I've mentioned to you earlier, you can even run these iterations on a T4 GPU. That's the free GPU. Um, but that will take around 6 to 8 hours, I think. And laptop crashes sometimes. I've also tested with that.

So if you look at the training loss and the validation loss, you see that the training loss continuously goes down. The validation loss is also continuously going down. That's a great sign that our training is proceeding very nicely. And towards the end, we have the training loss to be very low, validation loss to be very low. And more importantly, the training and the validation loss are very close to each other. That indicates that we are not overfitting over here. I very strongly think that if you run this for 40,000 iterations, you might even get better results. But here I'm sticking with just those number of iterations which are computationally feasible for everyone once you run these iterations.

So finally, I've got a training loss of 2.3919 and validation loss of 2.3918. You can plot these, and here you can see that both of them smoothly start going down. If we are doing overfitting, then what usually happens is that the training loss continues to go down, but the validation loss shoots up, which is clearly not happening in our case. This is a beautiful graph which shows that the training and validation loss are proceeding smoothly. So at this part, we have successfully run a small language model on our machine, and we all should be incredibly proud of it. This is a model which has around 15 to, I think needs to be estimated, but 15 to around 30 million parameters. So still, it's definitely a small language model, and it's probably around thousand times or 10,000 times lesser number of parameters than GPT3, GPT4, etc.

But the real test of this model is how well are we able to do in terms of generating new content, right? So remember, we started this, uh, lecture with saying that we want to, uh, we want to make a model which can essentially give, um, produce coherent text. Our model needs to produce coherent text. So it essentially needs to produce English stories, uh, which sound like a story, right? Uh, so that will be truly awesome because then it means that our model, our model has learned the English language, it has also learned how to form stories, etc.

So what I've done after this point, which is also going to be the last step in our, uh, video today, that is essentially running the inference. In the inference, what we are going to simply do is that we take the, uh, whenever the input is passed, it goes through the trained model. Now, similarly, we have the logits matrix. We get the index of the highest value. We get the next token ID, that's decoded back to text. We produce the next token. We append to the previous input. That again goes into the model. And then this loop essentially continues.

So let's say if you have certain input like this, what happens is that this input goes through my model. Okay, the input goes through my model, and then I have, uh, my model is fully trained now. So parameters don't update. Then I have the decoded token over here. My decoded token, that decoded token is appended to the input now. So now my input at the next iteration becomes something like this with the appended decoded token. It goes to my model, and then again, my model produces a new decoded token, which is again appended to the input. So this process continues until we have to prescribe how many new words we have to generate.

So if you see in ChatGPT, this is exactly what is happening, right? Let's take the same example which we took yesterday or in the previous part of this lecture, make a travel plan for Italy. I believe we had taken this example. So here you see, first, the first token is generated, then that will be again appended to the input, then the second token is generated, that will be again appended to the input, third token is generated, etc. That's how the inference process actually happens in language model tasks. We append tokens, then a new token is decoded, that's again appended to the input, that again passes through the model, new token is decoded, that's again appended to the input, etc., and this process actually continues. So this is the whole inference pipeline. Uh, this step-by-step process enables the model to generate text sequentially, building coherent phrases from the initial input context.

So what we are going to do now in this part of the code is that I have, uh, first, what I have done is that I have stored the model. Uh, let me see where that is. Yeah, actually, my model is already, the parameters are already stored in best model params.pt, which is the best, which is stored in a variable called best model params path. So what I have done here is that I'm getting this path over here, which is best model params path, and I'm loading the model here. So the reason I'm showing this part of the code is that even if you turn off your, uh, kernel or if you restart your instance, since this path still contains the best parameters, you can load your entire model using this function called load state dict. This is so useful, and you will thank me for sharing this because sometimes or many times it has happened with me also that I've trained a model but forgot to save the parameters. But it's a very simple thing, right? You just store all the learned parameters in this best model params, and then you just load it whenever it is needed.

So here I'm using load state dict, that loads the entire model parameters into the model, and then what I'm doing is that I'm starting with a sentence, that's my input sentence. Now, uh, I will encode the input sentence and then I will call this model.generate. So if you see, if you scroll above, my GPT class, which has been defined over here, my GPT class also has a function which is called as generate. What this generate function does is that it actually, the same logic which I mentioned over here, it applies this same logic, right? So first, what it will do is that it will generate this logit tensor. Okay. Then it will get the next ID based on the highest index of the, uh, last row of the logit tensor. That then it will append the ID to the current IDs, and then this whole process will repeat in the loop. That's exactly what I mentioned over here, right? Uh, we have the token ID, which is the input token ID, uh, uh, sorry, we have the, let me write this whole loop over here, actually. We have the current input sequence, right? And in my case, the input sequence, let's say something like, uh, "one day a little," "one day a little." That input sequence is passed through my model. That input sequence is passed through my model, and then I get my logit tensor. Now, remember, the logit tensor is "one day a little," and each of them is a 50 or vocabulary size vector, right? Now, since during the inference stage, we actually don't care about these first three at all, because we just have to know what comes after "little." So we only look at the last row over here, and we look at the token ID corresponding to the maximum probability. That's my decoded next token. If that is "girl," that's appended again, and then this whole appended sequence now goes through the model once more. That's how this loop actually proceeds.

So if you see what we do initially is that we pass this, uh, yeah, we get the, we get the logits now, and here you see f.softmax(logits). So then we get the vector of probabilities, and then we get the next ID based on the last row of this probability, and that is then appended to my current, uh, to my current IDs. So this is exactly, uh, this is exactly what happens in the generation part.

And you might be wondering what is this top k over here? So top k essentially means that right now we are looking at my whole set of probabilities, right? And then I'm looking at that token with the highest probability. Top k just restricts to, let's say, we look at the maximum five tokens, etc., and then we choose from these five tokens. So these are some inference strategies which are used. There are some other strategies also, which is called as temperature scaling. So you'll see one more factor here, which is temperature. So that temperature effectively controls the creativity of my output. The higher the temperature in the softmax, every token will have higher entropy. So it is more likely that sometimes some token will come up, sometimes some other token will come up, etc. Um, and in the top k, what is done here is that instead of distributing the probability among all the tokens, we essentially look at only the top tokens, and, uh, we are going to look at those top tokens and select the next token from among those top k, um, from among those top k tokens. Okay. And for that, we use this function which is called torch.topk. So you can search about this torch.topk. top k, and it returns the k largest elements of the given input tensor along the given dimension. Um, okay, so that's how this generation is done, the next token generation.

So you'll see that here towards the end, I pass my input sequence through this model.generate, and here it's 200. So what are these 200 over here? So this 200 are how many new tokens I want to generate. So if you see over here, uh, if you see over here, there is this max new tokens, right? So essentially, this loop continues for how many new tokens we are generating. So this entire loop continues for how many new tokens we want to generate. So I have said here that the maximum number of new tokens which I want to generate is 200. And then I can just print out my output.

So the input was "Once upon a time there was a pumpkin," and let's see what I have got as the output. It was very special. The pumpkin was very happy. The pumpkin were very happy. One day a kind lady came to the store to you. The delicious corn helped the world play games alone. She had two laser soldiers who lived in the forest. The lady felt ashamed. Her friend said, "Don't worry, I need a board." Now, first, what I'm amazed at looking at this is that this story is not there anywhere in the data set. Okay, that's the first thing we know. Second thing, the grammar of the sentences looks mostly correct. The meaning does not look quite correct, and that's probably fine. When we increase the number of iterations, I believe we'll get to a stage where the stories will even be coherent. But this is still awesome, right? It could have printed literally anything out there. But now our model has learned something about the model has learned something about the language. The model has learned something about how to generate stories. So you can see that the model is kind of making sense, right? The witch said, "So there is a witch in this story. I found it under the garden. It's important to help others. We need to be careful. The lady smiled and felt happy on herself." All right. So although you can see that all the sentences exactly don't quite merge together, but within just 15 to 20 million parameters, we have got the model to a stage where it's generating something which is coherent. It's not random tokens. It's not random words. It's trying to generate a story that "Once upon a time there was a pumpkin. It was very special. The pumpkin were very happy," etc. That's incredible in my opinion because we assembled this huge architecture. We assembled this huge architecture, and we never told anything to the model about language or grammar or anything. We just trained the model to predict the next token. And due to this training, the model has come up with a story on its own which is not present in the data set at all. And this story definitely needs to improve further. And I think it can be because this loss can further go down. The training and validation losses are both going down. So I encourage all of you to run for more iterations, and you can tell me what's the kind of stories which you are obtaining. What's the output which you are obtaining.

I did one more example. "A little girl went to the woods." So the model completed it like, "A little girl went to the woods and she found a basket for her. In one he enjoyed the great basket and the tasty treats doll. Then she decided to go home and the picnic. As they walked, the little girl was getting tired and didn't want to drink. She opened her pocket to her room and got out." This is good, right? So it kind of makes sense. Little girl went to the woods, bought a basket, then she walked back home. She got tired, didn't want to drink. Uh, so something is starting to make sense. Although everything is not making sense.

So we have literally constructed a language model out of nothing today and built it fully from scratch. And this is a language model which does not have billions and billions of parameters. It has parameters around 15, 20, 30 million parameters, I think. Now, one more thing which you can do once you have access to this code is that you can apply it to any application which you want, right? You can apply it to construct an SLM for healthcare applications where you can, let's say, develop a model which answers questions about chat diabetes patients, etc. But just make sure that the data set is small, the data set is constricted. The reason the small language model works so well in our case is because the data set is small. The data set is around 2 million rows, and it's restricted to stories for three to four-year-olds. That's what makes construction of the small language model possible. So when you get access to this code now, you can apply to so many sectors. You can apply to the healthcare sector to create a chatbot. You can create, you can apply it in the automotive sector where your, uh, language model answers questions about cars. Let's say you can build a language model therapist, a small language model therapist on your own, which when you interact with it, it replies with, uh, answers like a therapist would. But let's say if you're building this, right, the data set which you have should not be all over the internet because then your model will need to be huge. You need to intelligently construct a data set which, let's say, would be the most useful to a therapist, uh, similar to how we did for the present study.

So now this kind of, this video and this notebook which I'll share with you opens the door to you to do all kinds of experiments. But first, what you should do is that if you have free tier of Google Colab, just use that free tier. But if you have access to 800 GPU, that will significantly improve your life because it's 10 to 20 times faster than the T4 GPU. Right? So I would highly recommend going for this. Just go for this and then run this full code. And second of all, run it for 30,000, 40,000, 50,000 iterations. It might take two to three hours on A100, but I think it will be worth it because I think you will start to get awesome stories once you run for those many number of iterations. I have run for 60,000 iterations, and the stories are much more better than what I'm showing you right now.

So this is a fully functional small language model which we have built from scratch, and we went through these seven steps sequentially together. We started with the data set. Then we did the tokenization. We created input output pairs. We assembled the entire model architecture. Um, and the model architecture looks essentially something like this, what I've shown over here. And once we assemble the entire model architecture, we set up the training pipeline. Uh, we pre-train the LLM. We set up things such as gradient accumulation steps and adaptive learning rate.

Then earlier in the code, we have done some things to improve the performance, such as saving the data or saving the token IDs to a bin file for, so that, and we store it in a disk for easy retrieval. Then we have essentially, yeah, done something like x.pin_memory, y.pin_memory, etc., so that we lock the memory of a tensor in RAM for a faster transfer of tensor to GPU. After this, we have also done some things such as torch.cuda.amp, which is automatic mixed precision, so that we can automatically use float 16 wherever it's safe, so that the computational speed is increased and we are much more efficient. So all of these things we have done to improve the production quality of this application. This is not just a toy project. You can take this code and you can augment it to make it more industry ready.

And the last thing which we did in today's tutorial is to run inference and show you the quality of the stories which can be generated. Again, at the end of this video, I want to especially thank Andre Karpathy and the nanoGPT repository, which he had, which inspired a lot of sections over here. Many of the architecture sections, many of these, uh, efficiency-based calculations, Andre talks about a lot in his video also, and even these calculations later, which we have done with respect to gradient accumulation steps. Um, I've been inspired a lot from him, uh, as a whole. Thanks a lot, everyone. I hope you took a lot out of this tutorial. I'll be posting many such deep dive workshops, which is two to three hour workshops, which are not toy projects, but which are real projects which you can use and build on top of this to convert these projects into research publications. You can add those to your profile, and that will make you a truly strong ML and large language model engineer. Thanks so much, and I look forward to seeing you in one of these next workshops.