[Music] hello everyone welcome to this lecture in the build large language models from scratch series up till now in the attention mechanism part we have covered the following lectures first we covered the simplified self attention mechanism that was the first lecture in this series on attention mechanisms then we looked at self attention with training enable weights and in the last lecture we looked at causal attention now we have enough ammunition and and enough training to finally start understanding about the multi-head attention mechanism as I've mentioned previously if you directly start understanding the multi-head attention it will be very difficult to understand because it is fairly detailed mathematically as well as coding wise but now I believe that we have developed a strong Foundation if you watched the previous lectures I believe that with the understanding of simplified self attention self attention and causal attention we can use these as building blocks to truly understand multi-head attention in this series we are going to study two types of multi-head attention mechanisms the first type is basically by just concatenating the context Vector matrices which is obtained from different query key and value matrices and the second approach is a more unified approach which is more commonly implemented in modern llms and modern code bases so I'll be dividing this lecture into two parts in the first part we'll be looking at the first type of multi-head attention mechanism so let's get started before diving into multi-head attention I just want to briefly cover um what all we looked at in the causal attention mechanism if you have not seen the causal attention mechanism video I highly encourage you to go through that because multi-head attention is just an extension of the causal attention mechanism awesome so the way causal attention works is as follows we have inputs which basically are input embedding vectors corresponding to every single token which we have so if the tokens are your journey starts with one step we have threedimensional Vector embeddings for each of these tokens these Vector embeddings can be represented in a three-dimensional Vector space as can be shown here the goal of causal attention multi-head attention and any type of attention mechanism is to start with these input embedding vectors and convert them into context vectors context vectors are a more enriched form of input embedding vectors the input embedding vectors contain semantic meaning of the particular word but do not carry any information about how that word is related to the other words in the sequence so for example if you look at the input embedding for journey it does capture some semantic meaning of Journey but it does not really encode any information about when you look at Journey how much attention should be paid to the other words such as step such as your um such as with and such as one the context Vector encodes this information also and that's why it's an enriched representation as compared to the input embedding Vector okay so the first step is to basically take the inputs and multiply them with trainable weight matrices we have three trainable weight matrices one for the queries one for the key and one for the values so the if you look at the input Matrix and its Dimensions we have six tokens here and then the vector embedding size so 6x3 the trainable weight matrices have the First Dimension is fixed the First Dimension has to be equal to the vector embedding size but their second dimension is what will refer to in today's lecture also as the output Dimension and then when you multiply the inputs with these three trainable weight matrices for query key and value you get the queries the keys and the values Matrix so the dimensions of the queries key and the value Matrix will be the number of tokens multiplied by the output Dimension and the output Dimension is determined by this trainable weight matrices now note that when you initialize these trainable weight matrices their parameters are Rand om uh we will train these parameters and optimize them through back propagation in these attention series these four to five lectures we are not looking at back propagation we'll cover that later right now the whole goal is to understand different attention mechanisms like causal attention multi-head attention Etc the next step in the causal attention mechanism after getting the queries keys and values is to essentially compute the attention scores so the attention scores are computed by the multiplication of queries multiplied by the keys transpose so uh the way to interpret the attention score Matrix is you look at each individual row so you'll see that there are six rows and six columns right so if you look at the second row the second row corresponds to the second token which is Journey so your your journey begins with one step so the second row corresponds to journey and each element in the second row basically encodes information about how much that particular token relates to Journey or how much importance should we pay to that particular token or input embedding when we are looking at Journey so for example this third token this third value encodes the attention between journey and the third third token which is your Journey Begins so this is begins and journey this is Journey and with this is Journey and one and this is Journey and step so basically uh here you can see that uh the attention scores are essentially uh there for every query and then you multiply queries with the keys transpose and this value encodes information about when you're looking at a query how much information should be paid to a particular key right this is the meaning of the attention scores in the causal attention mechanism we do what we do is that we take these attention scores and we look at all the attention scores above the diagonal and we make these attention scores above the diagonal to be equal to zero so let me just show you how that looks like so let's say uh we have got the attention scores and let me show it here so let's say we have got the attention scores and they look something like something like this um actually let me show it here itself the attention scores which we have obtained looks something like this right but you can see that if you look at the second uh if you look at the second row which is the second row for Journey you will see that Journey has an attention score with all the other words now in causal attention what we say is that we should only look at Words which come before Journey so finally when we compute the attention weights when we compute the attention weights uh all the elements above the diagonal should be equal to zero let me show you what that means so so here is how the attention weights for a causal attention mechanism actually look like yeah like these and uh if you if you don't implement the causal attention mechanism then the attention scores look like something on the left attention weights look like something on the left which I'm showing here but when you implement the causal attention mechanism all the attention weights Above This diagonal will be essentially Switched Off why because if you're looking at Journey you should only pay attention to what comes before Journey your and journey if you look at starts you should pay attention to your journey and starts remember since we are predicting the next word Only The Words which come before a particular would really matter and that's what's implemented in the causal attention mechanism so we implement this and we get the attention weights like what I showed you below which are zeroed out above the diagonal and then we multiply the attention weights with the values uh when we multiply the attention weights with the values we get the context Vector Matrix which looks something like this so now this context Vector Matrix has six rows because one for every token and it has two columns which are dictated by the output Dimension which we have two in this case so the input embedding Vector for every token here which was mentioned over here is ultimately converted into a context Vector Matrix which looks like this that's the whole idea behind causal attention so the only difference between causal attention and normal self attention is that in causal attention these attention weights right all the elements above the diagonal of these attention weights will be set to zero and the rows of the attention weight Matrix will anyway sum up to one we will ensure that and these attention weights will then be multiplied with the values Matrix to get the context Vector Matrix that's essentially what is happening in the causal attention now there is a provision in the causal attention mechanism where we can do batch where we can process batches so when we input the first batch the first batch has six words your journey starts with one step right this is the inputs but the causal attention class which we developed in the previous lecture also can handle the second batch which also let's say has six tokens so for the first batch the output is the context Vector Matrix which is a 6 by2 Matrix as you can see here and for the second batch there is another context Vector Matrix which is 6x2 let me show you the the the causal attention class which we developed in the prev previous lecture so this was the causal attention class which we have let me quickly revise it right now the input X which is the input to the causal attention class has the shape of batch size comma number of tokens comma the input Dimensions so batch size because we can have two batches so if you look at these two batches which I have shown here right now um yeah so I've shown batch one and batch two over here and now let's look at the input if you look at the input for the first batch it's 6x2 right it's 6x3 whereas if you look at the input for the second batch that is also 6x3 so when you aggregate these batches together the input Dimensions will be um the input Dimensions will actually be 2 2 multiplied by 2 * by 6 * 3 2 * 6 * 3 that is actually the input Dimension which is also being mentioned over here so the input. shape is the number of batches which is if it's two number of tokens which is six and the input Dimension which is three then what we do is we multiply the input with the trainable key query and value Matrix and finally we get the keys queries and values so remember W key W query and W value are the weight matrices for the key query and value which we have to optimize and then after multiplication of the inputs with these weight matrices we get the keys queries and the value what we do then later is we get the attention scores by multiplying the queries with the keys transposed then we apply a mask so that we implement the causal attention remember we need to make sure that when we get the attention WS all the elements above the diagonal are zero I explained all of this in the previous lecture so I'm not going through the details right now and then finally we apply soft Max to the attention scores and get the attention weights we can also add another layer which is the Drpout layer towards the end but that's not strictly necessary and the last step is then multiplying the attention weights with the values to get the context Vector now uh remember the we had two batches right and the input size was 2A 6A 3 and ultimately when you let's say Implement when you create an instance of this class let's say I'm creating an instance of this class and storing it in C A so I have to specify D in which is three D out which is two the context length which is my number of tokens in this case that's six and the dropout rate that's zero and then we can get the context vectors remember we have two batches so we have two sets of context vectors the first set will be a tensor which is 6x2 why will it be 6x2 because that's exactly what we had seen on this white board also um you can take a look at this context Vector which we had obtained ultimately yeah the final context Vector which we obtained is 6x2 and there are two such context vectors right so ultimately the context Vector shape which we get out is 2 comma 6 comma 2 so that makes sense with what we had seen on the white board this is the causal attention class which we had implemented in the last lecture and a quick summary of how it works now what we are going to do is we are going to extend this causal attention class a bit into something which is called as the multi-head attention mechanism so the term multi-head essentially refers to dividing the attention mechanism itself into multiple heads um and each head will be operating independently to give you a understanding related to code when we get the attention scores here right now or the attention weights from one set of query key and value we say that this is one head this is one attention head right we are not decomposing into multiple query keys and values there is only one query Matrix one there is only one Keys Matrix one queries Matrix and one values Matrix so we have one attention head what multiple what multi-head attention does is that it extends the causal attention mechanism so that we have multiple heads and each of these heads will be operating independently um and then what we do in multi-head attention is we basically stack multiple single attention head layers together so what we'll simply do is that we will create multiple instances of the causal self attention mechanism each with its own weights and then combine their outputs it's actually very simple what the output which we had obtained before for one attention head we'll just combine them together and I'll show you how it can be done so as you might have expected this can be a bit computationally intensive but it makes llms powerful at complex pattern recognition tasks so researchers have found out that although when you stack multiple heads the computations which need to be made increase it really helps llm perform much better and when you look at Chad GPT right now it does have multiple attention heads it none of the modern llms work with a single attention head so let me show you what it means by stacking multiple attention heads together this diagram actually encapsulates it all earlier when the when I showed you the flowchart we just had one trainable Matrix for the query one trainable Matrix for the keys and one trainable Matrix for the values right so that was one attention head now when we look at multi-ad attention in this figure we have shown an example of having two attention heads so what we do here is that instead of one trainable weight Matrix for the query key and value we have two so here you see we have two trainable weight matrices for the query we have two trainable weight matrixes for the key and we have two trainable weight matrices for the values these will be multiplied by the input Vector input X and then we'll get two queries Matrix we'll get two keys Matrix and we'll get two values Matrix previously we just had one query key and value Matrix right then what we'll do is that we'll multiply the queries with the keys transpose and get the attention weights then we'll get the attention scores we'll multiply with the value Matrix and we'll get two sets of context vectors so this is the first set of context vector and this is the second set of context vector earlier we just had one context Vector right we had one context Vector Matrix now we get two context Vector matrices so if you look at this input here which has been highlighted this input 7. 2 and 1 let me highlight it once more yeah this input 7.
2 and one earlier we converted every threedimensional input input into a two dimensional context Vector right but now this 300 threedimensional input will be converted into two context vectors each of the output dimmension which in this case is two so then what we will do is that we'll get these two context vectors and we'll concatenate them together to give the ultimate context Vector Matrix so earlier if you see the context Vector Matrix was 6x2 so in this case the 6x2 and 6x2 will be aggregated together and the final context Vector Matrix size will actually be 6x4 it will be 6x4 okay so it actually looks something like this earlier when we applied the single head attention we just got one context Vector Matrix right but now since we have two attention heads we have two context Vector matrices and each context Vector Matrix has a dimension of 6x2 so when you stack these two and you stack them along their last Dimension which is two so the number of rows actually remain the same but the number of column colums increase so when you stack these two along the column the final con concatenated context weight Vector weight context Vector Matrix will be 6x4 that's all which is actually happening in multihead attention so let's go to our flowchart and see really what is going on okay so as we saw earlier we had this trainable query key and value matrices right now in multi-ad attention we have multiple of these matrices if we have five attention heads we have five trainable query matrices five trainable key matrices and five trainable value matri even here we had one one queries Matrix one Keys Matrix and one values Matrix right in multi-ad attention we have multiple of these similarly we have multiple sets of attention scores we have multiple sets of attention weights and ultimately when we multiply these with the values uh we have two set of context Vector matrices so if we have two attention heads if you look at the first context Vector Matrix this is from attention head number one which basically means it's from the first set of queries keys and values if you look at the second uh context Vector Matrix you'll see that this is from attention head two so the dimensions of the first context Vector Matrix is 6x2 and the dimensions of the second context Vector Matrix is 6x2 now these are concatenated along the columns and so we get a final context Vector Matrix which has the dimensions of 6x4 that is all what is being implemented in the multi-head attention that's why it's called multi- head because we are aggregating the output of multiple attention heads so if you find yourself getting confused just remember the main idea what we are doing is that we are just running the attention mechanism multiple times that's it that's the only thing to remember and this one figure which I've shown over here really summarizes it all we have multiple copies of the trainable query key and the value Matrix yeah here we have multiple copies of the trainable query key and value Matrix we have multiple copies of the queries keys and the value Matrix and we have multiple context Vector matrices which are stacked together now let us Implement multi-head attention in code for that you need to remember that we have already implemented the causal attention class and uh the output of this class is that it returns a context Vector correct so just remember this and let's move forward to Extended s extending single head attention to multihead attention so in Practical terms implementing multi-head attention involves creating multiple instances of the self attention mechanism each with its own weights and then combining their outputs so if you create one instance of the causal attention you get one context Vector so then you'll create another instance of the causal attention you'll get another context vector and you will merge them to that's exactly what we are going to implement in the code so to do that we we will actually Implement a multi-head attention rapper class that Stacks multiple instances of our previously uh implemented causal attention module so here's the multi-head attention rapper class now what is happening in this multi-ad attention rapper is actually pretty simple okay so we get the output from the causal attention mechanism which is written over here and we first Define number of attention heads so if the number of attention heads is five uh what we'll do is that we'll get outputs from uh we'll get five outputs and then we'll concatenate them together so first let's look at the forward method it's torch do cat which is which is concatenation and dimension equal to minus one why minus one because we are concatenating along the columns as we saw before and then here you can see we are looping over all the attention head in self. heads and what is self. heads so self.
heads is essentially it will create an instance of the causal attention class for how many number of heads so here if you see we are looping over the number of heads so if we specify the number of heads equal to two we will create two instances of the causal attention class and then the results of the two instances will be essentially uh stored in this function head of X okay and head is essentially in self.
