Introduction to Multi-Head Attention (MHA) | The attention family series - Lecture 8

0:00Hi, welcome to a new lecture from the attention family lecture series. In today's lecture we will learn about multi head attention. Multi head attention is implemented in almost all modern architectures and there is a very specific reason why we can benefit from having multiple attention heads looking at the same token sequence in parallel. We will first understand the logic of the implementation and then we will see mathematically how it is implemented. Specifically we will look at two techniques. In the first technique we will see how the attention heads can be processed in series using for loops and we will see that it is slow and in the second technique we will see the same thing can be implemented mathematically by using a simple matrix multiplication technique.

0:47So that is what we are going to do, what we are going to discuss in today's lecture. Before that, let us have a quick recap of what we saw in the previous lecture. So in lecture 7, first we discussed the need for having to mask some of the dot products between queries and keys, especially if query is going to look at the keys from the future. So, we saw that we have to construct an upper triangular mask like this in the attention weights matrix. So, what we did was because we are converting attention scores into attention weights using softmax, which involves calculation of exponents, we first convert all the upper triangular elements into minus infinity such that when you take softmax those numbers will become

1:31zeros because e raised to minus infinity is zero. So, the steps that were involved in this calculation was we first have input tokens, input tokens are converted into query key value space using query weights matrix, key weights matrix and value weights matrix and then we can calculate the attention scores by taking the dot product between query and key. So, this attention scores matrix is not yet scaled by dividing by square root of d. So, you scale by dividing the numbers by square root of d and we have discussed in detail in the past lectures why this scaling is needed and you assign minus infinity to the upper triangular mask and then you take softmax and once you do this you have the attention weights with masking applied and every row will add up to 1 and you will see

2:21that all the elements except the first element in the first row will be zeros, all the elements except the first two elements in the second row will be zeros and so on. After this we also discuss that to prevent over fitting during training we also implement dropout. And if dropout probability is p, we multiply or rather we scale the rest of the elements by a number 1 by 1 minus p, because if half of the elements are half the weights are unavailable to produce the output, we have to scale the output produced by the rest of the weights by a factor of 2. Because if 50 percent of the items are missing and if the output magnitude has to remain unchanged, the remaining weights will be scaled up by a factor of 2. And similarly, if in general we are

3:05using a dropout probability of p, we will multiply the remaining elements which are not dropped out by a number 1 by 1 minus p. So, all of this we discussed in detail in the last lecture. However, we were yet to discuss how all of this is implemented when it comes to multi-head attention or MHA. So, in this slide we are going back to the same figure that we have been looking at throughout the lectures and we will come back to this figure again and again and again. So, in the initial lectures we looked at how attention scores are calculated there were many different ways in which attention scores will be could be calculated. Initially we start started with additive attention or badano attention then we moved into dot product

3:47attention and then we looked at scale dot product attention and then we looked at self-attention with trainable weights and before that we also looked at self-attention without trainable weights in the last lecture we looked at mass multi-head attention sorry mass attention with single head not multi-head so that was about connectivity which tokens are allowed to look at which other tokens. In today's lecture for the first time we will look at the idea of can we do somethingto the attention heads themselves. So, right now we will be talking we will be discussing how we introduce multiple attention heads. Now, in this bucket there are a few more lectures that are upcoming they are not necessarily the next lecture immediately, but eventually all the lectures that we have

4:29in this series will fall into one of these 7 buckets. So, today we are going to look at the bucket number 3 for the first time. So, in the first part of today's lecture we will discuss about multi head attention and the lecture is actually divided into three parts. So looked at simplified self-attention without trainable weights. Then we looked at self-attention with trainable weights. Then masked or causal self-attention and now multi-head attention. So, before we jump into the implementation through loops, through four loops and also the implementation through the technique called weight split. We will first see what is multi-head attention and why do we introduce it? What is the advantage

5:30of having multiple attention heads, et cetera, from a qualitative perspective. And after that, we will look at what happens to the dimensionality of the input token sequence in both of these approaches. And finally, what happens to the final output, because if each attention head is producing an output we have to do something with that output to get back the context vectors so please keep in your mind the goal here is to produce context vectors corresponding to your input tokens that is the end goal and somehow if you have six input tokens you should have six context vectors that is the end goal so whether it is single head attention or multi-head attention that goal has to be met so we will see how that goal is met so firstly why multiple attention

6:11heads to understand the need for multiple attention heads look at this sentence the dog chased the ball but it could not catch it so here the last one it refers to actually it refers to ball right but the verb who is trying to act on this it it is dog so dog is trying to catch the ball if we have a single attention head at the output of the single attention head basically we have a set of attention weights and these attention weights individually multiply with the value vectors and then you sum up those value vectors and that's how you get the corresponding context vector now attention weight you can think of something that answers a question about that particular token so let's say the token is this it you can say that if you have multi-head attention

7:03each attention head answers one question about this token maybe the question that one of the attention head is asking and this is completely qualitative okay so when i say this when i say that attention head is trying to answer a question about the token it's completely vague uh it's a bit qualitative and that is the only way to kind of picture what may happen in attention head. Really what happens in attention in different attention heads is completely mathematical. It's very difficult to unless you have some sort of attention visualization map. It's very difficult to say this attention head is learning something. The second attention head is learning something else. Third attention head is learning something else, etc.

7:40But in general, we can it is safe to assume that there are multiple things to learn about this particular token which entirely cannot be possibly captured using one attention head. If you have a single attention head, you have one single context vector which is mapped to this particular token and that is because you have one single set of attention weights on the other hand when you have multiple attention heads each attention heads heads produce completely independent set of attention weights and using those attention weights you will get one context vector out of every single attention head so if you have three attention heads for example you will have three context vectors corresponding to this particular token.

8:26And you can hope that each of those attention head might be learning different things about this particular token. One attention head might be good at learning, okay, whether this is a noun or something else. Another attention head might be good at learning whether this is an animate object or inanimate object. Another attention head may learn some other feature while the examples that i give are very vague and over simplified versions effectively you can say each attention heads heads learn something specific about that token and when you have multiple attention heads acting in parallel the resulting output is equal same number of context vectors as there are number of attention heads and that collective output you can somehow make use of to produce a much

9:10richer representation of these tokens. So which cannot be accomplished using a single attention head. So that is the idea of using multiple attention heads. I know that this is a very vague way to explain this, but when you go through the mathematical steps, it will be very clear how exactly it is implemented. Because mathematically, how one attention head, how a single attention head operates, if you look at it in isolation, versus how any attention head in a multi-head attention operates, if you look at it in isolation, it is more or less the same thing. You are trying to produce context vector corresponding to tokens, right. So, at the end of the day, if you have five attention heads, you will have five

9:49tokens, rather you will have five context vectors corresponding to every single token. But finally, at the very end of all of this, you want one context vector associated with one token. You don't want five context vectors associated with each token. So the question is, do you accomplish that? How do you produce if you have five attention heads how do you produce five context vectors and then how do you combine that everything into a single context vector corresponding to one token? That is the question that we are going to answer now. So mathematically this is how multi-head attention works. So the simplest example for a multi head attention is two attention heads so you have an input x which is a six by six matrix so what is six by six you have six tokens

10:32and each token is living in a six dimensional space so this is token number one token number two three four five six and each token is living in a six dimensional space. So every row is basically one token. Now, one way to imagine what attention heads do is typically, and this is not always true, but you will see that in many cases, if you have n attention heads, so let's say the dimensionality of your embedding space is d, and if you have, let's say n attention heads, the dimensionality of the query keys and values in each attention head will be d by n. So, in this case, d is equal to 6 and number of attention heads that I am trying to show you here is 2, which means the dimensionality of the query key value space in each individual

11:18attention head will be 3. So, if that is the case, the context vector resulting from each of these individual attention heads will be of dimension 6 by 3. This is the context vector matrix by the way. Why is it 6 by 3? Because you have 6 tokens, that has not changed, you have 6 tokens, however the dimensionality of the space in which that attention head is operating is 6 by 2 equal to 3. So, the resulting context vectors will be three dimensional context vectors. So, this is context vector number 1, this is context vector number 2, etc. And where are these six context vectors emerging out of? They are emerging out of the first attention head. Similarly, from the second attention head also another set of six context vectors emerge and

12:12They are also in living in three-dimensional space now think about this your you have let's say this is your token number five Corresponding to token number five how many context vectors do you have right now you have two context vectors you have this first context vector Coming out of attention head number one you have the second context vector coming out of attention head number two so now for every token because you have two attention heads you have two context vectors however that is not the output that you want finally you somehow just want only one context vector per token so what do you do you concatenate this z1 which is the list of context vectors coming from attention head one and z2 in such a way that? You concatenate this Z1 which is the list of context vectors coming from attention head

12:451 and Z2 in such a way that you are concatenating them along the embedding dimension. So you have 6 by 3 concatenated along the embedding dimension with another 6 by 3 to get 6 by 6. So now you have 6 context vectors where each context vector is mapped to its corresponding token at the input and that context vectors are also living in 6 dimensional space just like how the input embeddings were. Your input embeddings were living in 6 dimensional space, your final context vectors are also living in 6 dimensional space, okay. So, this is simply how multi-head attention works. Now, you should have a list of questions here. What about WQ1, WK1, WV1? What are their dimensions? You can easily imagine what their dimensions should be. So the dimensionality of Q1, let's say this is the query of the belonging to the first attention head, K1, V1, etc. They all should be same as the dimensionality of z1,

13:47which means the WQ1 matrix will be something like 6 by 3. So, how do you get that? You have 6 by 6 incoming operated upon by WQ1 resulting in output which is 6 by 3. So, Q will be 6 by 3, K will be 6 by 3 and V will be 6 by 3. Why 6 by 3? You have 6 tokens, each of them living in 3 dimensional space in the query key value space within each individual attention head. Now, let us look at some of the details of this whole implementation. One more thing before we go into the details yes one more thing is after concatenation there is one small problem the problem is that here the output from every attention head is in a is arranged in a specific order but who says that this is the attention head number one and because

14:39of that that should go here and this is attention head number two and because of that that should go here like who decides what should come first right because i'm just labeling this as attention head number one i'm just labeling this as attention head number two at the end of the day one head does not get precedence over the other every head is doing its job they are all acting in parallel just because i i call this as attention head number one and just by because i call this as attention head number two at the output when you concatenate it should not be like a set of values are getting certain preference over the another set of when i say values a set of context vectors are getting higher preference over another set of context vectors that should not happen but that is

15:19happening here here if you look at the first three context vectors those not first three context vectors, if you look at the first three dimensions, those first three dimensions are exclusively coming from context vectors emerging from attention head number one. And the last three dimensions they are exclusively being contributed by the second attention head, but that should not happen like output should be sort of like it should be a mixed output meaning you should not be able to say that oh these three dimensions are completely contributed by this head and these three dimensions are completely contributed by the other head. So what happens in reality in practice when you implement is you also have an additional

16:00mixing step. So in this mixing step the output which is concatenated output It is passed through a mixing layer and mixing layers by the way should not change the dimensionality So if you have six context vectors each living in six dimensional space that should remain the same so six by six Should remain as six by six So what you do is the mixing matrix is also six by six So this happened to be six by six because your input is 6x6. If instead you had let's say 12 tokens living in 6 dimensional space. Your W output in that case would be again 6x6 and output would become 12x6. Because 12x6 multiplied by 6x 6 will become again 12 by 6. So the W out in this case is purely dependent on the embedding dimensionality.

16:53And W out is like a matrix which kind of jumbles the values in the sense like if you have a bunch of numbers here arranged in a certain order because you concatenated Z1 first followed by Z2. By the time you come to this output of the mixing step, you won't be able to say that, oh, the first three dimensions belong to head number one and second, the next three dimensions belong to head number two. You won't be able to say that. However, you can still track the context vector. So this is the input token, fifth input token. This is the context vector corresponding to fifth input token from the first attention head. This is the context vector corresponding to fifth input token from second attention head.

17:36And this highlighted row is the context vector corresponding to input token number one, input token number five after concatenation. But here you can clearly say that first three dimensions belong to head number one So the last three dimensions belong to head number two, but after mixing this highlighted row is the final context vector belong corresponding to fifth Token and here you won't be able to say that these dimensions belong to this head and these dimensions belong to that head you won't be able to say that these dimensions belong to this head and these dimensions belong to that head you won't be able to say that because of mixing and this is the actual way in which multi-head attention is implemented in practice now again there are a few more things we need to

18:16understand right when we implement batches because when you process tokens you are processing tokens in batches you have a token sequence and like that you have bunch of token sequences and to make sure that this whole operation happens efficiently we have to perform a certain you know matrix transpose operation so we will come to that in a few moments but just keep that in your mind so first is let's have a quick visualization of how the dimensionality of the multi head attention looks like. So let's say we have six tokens and let's say input dimensionality is 12 and we have four attention heads. In this case, what would be x? x is the input, x will be 6 by 12, 6 comma 12. Because you have six tokens, 12 is the embedding dimensionality.

19:03If we have four attention heads, the dimensionality of each within each head is 3, right? 12 by 4 equal to 3. So each head will be taken care of. Each head will look like, context vectors from each head will look like 6 comma 3. And then you will have 4 such items concatenated one after the other along the embedding dimensionality to get back 6 by 12. So just note this, so 12 is divided by 4, 12 is embedding dimensionality, 4 is the number of attention heads which means each attention head will be looking at 3 dimensions. By the way this is not completely necessary, you could define an attention head where every attention head is individually looking at you know 12 as the embedding dimensionality,

19:52you could do that, but then your compute will blow up. So in practice when multi head attention is implemented, it is implemented in such a way that each attention head looks at d by n, where d is the dimensionality of the embedding space of the inputs and n is the number of attention heads. Now, by the same logic, you can say GPT-2, it had 768 as the embedding dimensionality and it had 12 attention heads. And because of that, the dimensionality of the context vectors residing from the output of every attention head will be 64. 64 because 768 divided by 12 and in GPT-3 that number was 12288 that was the embedding dimensionality of the input and It had 96 attention heads. So each attention head the embedding space

20:42dimensionality of each individual attention head was 12288 divided by 96 which will be roughly 128. So, this is how you can calculate the embedding dimensionality living in each attention head. Now, let us look at the actual implementation of multi-head attention using an inefficient way inefficient way using for loops and then after that we will see how this is implemented using matrix multiplication. So, inefficient way, inefficient way using for loops and then after that we will see how this is implemented using matrix multiplication. So, inefficient way is very fairly straightforward just like how you implemented attention in the lecture number 7, you implemented causal or mass self-attention, whatever you did that

21:19step you repeat n number of times where n is the number of attention heads. So, first you calculate the context vectors emerging from attention head number one, then you calculate context vectors emerging from attention number two, repeat that for how many of our attention heads you have and once you have the output you concatenate all the outputs. This is fairly straightforward. However you can already imagine that this will be slow. What if you have thousand attention heads? This will become like a it may not become a bottleneck process but it will definitely be much slower than ideal the second approach is so instead of processing heads one by one first process attention at number one then go

21:56to attention at number two then go to attention at number three so this will result in z1 z2 etc up to let's say now we have 4000 attention heads only very few models have such large number of attention heads i'm not sure if any popular models have such very high number of attention heads but if we have let's say a very big model with 4000 attention heads you will have to calculate 4000 you know context vectors from these 4,000 heads and then you will have to concatenate and this will require a little bit of waiting as we will see in this Google collab notebook so I will show you the quick implementation of this and I will show you how much time it takes when you have two attention heads versus 20 versus 200 versus 4000 on the other hand we could

22:41have one single matrix multiplication helping us calculate all of this in one go and We will see the speed difference between these two when we are implementing this in Google Colab notebook. Okay So the first approach is something like this you have your inputs inputs are in this case six by six six tokens living in six dimensional space and let's say you are looking at attention head number one you have wq1 that's the query based matrix which is of dimension six by three resulting in queries which is again six by three you have six tokens living in three dimensional space uh six query tokens so that is q1 similarly you can calculate q2 so if you were to calculate uh sequentially going from attention head number 1 to number 2

23:26etcetera these are the two matrix multiplications involved assuming there are two heads. But the same calculation you can perform to get Q 1 and Q 2 separately using one single matrix multiplication. So, you can assume that W Q 1 and W Q2 are joined into one single matrix called WQ of dimension 6 by 6 ok. And the query weights are completely the same. So, whatever you have WQ1, whatever is the weights you have here in WQ1, they are exactly the same here in the first half. Similarly, whatever weights you have here in WQ2, they are exactly same here in the second half. But instead of treating WQ1 and WQ2 as two separate matrices, you just have them concatenated. Now, this concatenation is very different from concatenation of context vectors that we have discussed so far.

24:17Please don't get confused. I'm not talking about concatenation of the output here. I'm talking about utilizing one single matrix as the query weights matrix to be used as the equivalent of having multiple query weights matrices such that you have wq1 wq2 wq3 etc instead of having that just have one single query base matrix however there has to be similarity. Similarity is that if you had WQ1, if you had WQ2, etc., whatever would have been the output, that same thing you have to have here in such a way that you are concatenating WQ1 and WQ2 to produce WQ. So, let us say you had such a matrix, then you could produce something called query, which is just the normal query calculation, which is x multiplied by WQ. And query in this case will be 6 by 6

25:08now think about this mathematically this first half of this query will be exactly same as this Q1 and second half of this will be exactly same as Q2 but in this case we just performed one single matrix multiplication which GPU can process very fast. But if you had large number of matrix multiplications to be done separately, the time it takes for the GPU is much higher because GPU has to transfer all of that, then it has to load packages and it has to perform the calculation and send back the results But GPU can do the same thing so long as everything is packaged into one single matrix multiplication You will see this actually being implemented in Python in a few minutes But the basic idea is if your goal is to get q1 and q2 separately

26:00Why can't you produce one single queue and split them into two parts? You could do that right. So, you could kind of treat the attention, you can treat the query weights matrices as one single matrix and produce one single query and split that query into two parts and treat them as outputs of head 1 and head 2. This is completely valid and the numbers will look exactly the same as the approach above here now let's look at one more idea before we proceed ahead to any sort of implementation if you look at the actual implementation of let's say llm training or slm training you have the following dimensions first is you will have batches. So this 2, 6, 2, 3 is basically a set of dimensions corresponding to first

26:51is batches, second is sequence of tokens, third is number of attention heads and fourth is dimensionality within each attention head. So, here two means you have two batches that are being processed parallelly, six means in every batch you have a token sequence where there are six tokens in that sequence and you have two attention heads and dimensionality of each attention head obviously it will be three based on our previous discussion the number of the dimensions in each attention will be head will be 6 by 2 Now if this is the query, okay When we calculate the attention weights It is query dot key transpose and when we say key transpose what we are essentially doing is we are Transposing the last two dimensions. We are not transposing the in the transpose the batch dimension or sequence dimension

27:44They are not coming into picture. What is really coming into picture is the attention heads and dimensionality in naive case of transpose and you will see the mistake with that. So if we are naively transposing the last two dimensions, what will happen is the key transpose will be such that 2 comma 6 comma 3 comma 2. So this 2 is the batch, 6 is sequence of tokens, 3 is the dimensionality in each attention head and 2 will be the number of attention heads. Now, if you take the dot product between these two, the resulting dimensionality will be 2 comma 6 comma 2 comma 2 because when you when you multiply 2 by 3 with 3 by 2, the result is 2 by 2. 2 so you will get something like 2 comma 6 comma 2 comma 2 now what is the attention weights matrix looking like in this case

28:32it will look like a 2 by 2 matrix you will have 2 by 2 matrix and you will have such made 2 by 2 matrix for first batch and second batch separately. However, the actual attention weights matrix that you are expecting in this case is 6 by 6, because you have 6 tokens which means you have 6 queries and 6 keys and when you take query dot key transpose you have 6 by 6 right that is what you are expecting not 2 by 2. So, here there is a dimensionality mismatch when you calculate by naively taking query dot key transpose. What you have to do is you have to take a transpose between dimension number 1 and dimension number 2 basically these two dimensions. Why am I calling this as dimension number 1 because index number starts with 0. So,

29:20this is 0 1 2 3. So, when you transpose dimension number 1 and 2, what happens is instead of 2, 6, 2, 3. So, 2 is the batch size, 6 is the number of tokens in each sequence, 2 is the number of attention heads and 3 is the dimensionality of each attention head. You get something like 2 comma, 2 comma, 6 comma 3 you are you are swapping these two dimensions so that will be the key sorry that will be the query key also will be the same but key transpose will be 2 comma 2 comma 3 comma 6 because when you take a transpose key transpose you are basically taking transpose of the last two dimensions. So after you take transpose you have query as 2 comma 2 comma 6 comma 3 and key dot key transpose as 2 comma 2 comma 3 comma 6.

30:15Now what happens when you multiply 6 by 3 with 3 by 6 so this is 6 by 3 and 3 by 6 resulting will be 6 by 6 right and that is what this is. So, now this is the correct way to take the dot product between query and key and like this how many 6 by 6 matrices will you have in this case you will have two 6 by 6 matrices one corresponding to batch number 1 and another corresponding to batch number 2 because here you have two batches. But this is the right way to take dot product when you have when you consider the fact that you are not just individually considering just one batch you will you have an additional dimension called batch dimension and when you are implementing in practice in python you have to consider all of these practicalities ok.

30:58So, I know that this slide is a little bit crowded, but just trust me if if you walk through this one by one with me, you will understand everything. So I'll start from the very beginning and we'll walk out till the end where we do the merging and then we'll see the final resulting dimensionality. So first is the input. Input is 2, 6, 6. The first 2 is batch size. 6 is number of tokens in the sequence and the last 6 is dimensionality so batch size followed by token sequence length followed by dimensionality of each token so here we are assuming that tokens are living in a 6 dimensional space so if this is your input your final output that coming that comes out of the multi head attention mechanism should also be 2 comma 6 comma 6 because there also you should have two batches

31:51you should have six sequence length as six and embedding dimensionality also should be six and what are those actually they are context vectors so you will have six context vectors belonging to each batches okay so finally so after all the steps something that comes out of this should be 2 comma 6 comma 6 but in between there is a bit of dimensionality change and that is what we are going to understand now so now the second step is you have the input you are projecting the input into query key value space so how do you convert input into query i'll zoom in a little bit how do you convert input into query by I will zoom in a little bit. How do you convert input into query? By multiplying with query weights matrix.

32:26So, WQ is the query weights matrix. And let me just open the pen. So, 2 comma 6 comma 6 is multiplied by 6 by 6. What is the 6 by 6? 6 by 6 is WQ resulting in again 2 comma 6 comma 6. However, that will be splitting into two heads. So, here we are assuming that we have two attention heads, each attention head living in three-dimensional space, ok. So, the actual result of this operation is again 2 comma 6 comma 6 only, but the last 6 which is the six-dimensional embedding space is split into 2 comma 3 such that You have 2 attention heads and 3 dimensions for each attention heads. So 2 comma 6 comma 2 comma 3 Then you transpose the dimension number 1 and 2 which which means this 6 and this 2 gets exchanged

33:17So the dimensionality will become 2 comma 2 comma 6 comma 3. Now, what is this? This is query Key will be same in terms of dimensions but you are taking the transpose of key so when you take the transpose of key you have query as 2 comma 2 comma 6 comma 3 key transpose will be 2 comma 2 comma 3 comma 6 so 6 and 3 will be exchanged and the result will be 2 comma 2 comma 6 comma 6, right, 2 comma 2 comma 6 comma 6. Now, 6 by 6 is the attention weights matrix. Now, you do the rest of the things, you scale it, you apply masking, none of that will change the dimensionality. Does scaling change dimensionality? No, scaling means you are simply dividing by square root of d. Every element in your attention weights matrix will be divided by square root of d.

34:03Does masking change the dimensionality? No, masking means you are applying negative infinity mask and after that you are applying softmax, right? None of that changes the dimensionality. Now, you have to like once you have the attention weights, you multiply it with value. What is the dimensionality of value? Dimensionality of value is same as dimensionality of query 2, 2, 6, 3. So, this 2 comma, 2 comma, 6 comma, 6, which is the dimensionality of your attention weights matrix, that will be multiplied with your value matrix or value in this case tensor. Value will be 2 comma, 2 comma, 6 comma, 3. And what is the result? When you multiply 6 by 6 with 6 by 3, you get 6 by 3, by 3 right so the result is now 2 comma 2 comma

34:476 comma 3 now that you have you are done with your calculation the original swapping which you did so here you had taken a transpose between the dimension 1 and dimension 2 you apply that once more to reverse what you did earlier so you apply transpose 1 comma 2 once again so that the dimensionality instead of it being 2 comma 2 comma 6 comma 3 it will come back to 2 comma 6 comma 2 comma 3. So, this 6 the 6 and 2 will be swapped. So, here it was 2 comma 6 here it is 6 comma 2. So, now you are back into what you originally had and why are you doing this transpose because finally when you are merging the heads you want the numbers associated with different attention heads together right so here you have two attention heads and you have three items three dimensions belonging to each

35:35head so you can merge the last two dimensions so you here you merge the dimension number two and dimension number three assuming dimensionality starts with 0, then you will get 2 comma 6 comma 6 comma 6 ok. And then last is mixing, so mixing is after you after you do the merging you pass through a 6 by 6 matrix which performs mixing of the of this matrix and then the result is again 2 comma 6 comma 6 comma 6. 6 So this 6 by 6 and you multiply with this 6 by 6 the result is again 6 by 6 and you have 2 batches So you have batch size as 2. So what is the output 2 comma 6 comma 6 meaning you have that 2 batch size as 2 Token sequence length is 6 and embedding dimensionality is 6 So essentially you have 6 context vectors living in 6 dimensional space.

36:25And like that you have 2 batches. That is the final output, the meaning of the final output. Just go through these numbers, it is fairly easy to understand. But this is the whole process of transformation just from a dimensionality perspective. So this slide is a brief summary of the previous slide. So you have 6 tokens living in six dimensional space. If you consider query key value as one single matrix which you can split later, after projection you will have six by six. However, you are splitting them into two separate heads. So when you do that split, you have six by three in the head number one and six by three in the head number 1 and 6 by 3 in the head number 2. Then you calculate the attention

37:09scores in each head separately and after that the resulting context vectors will be 6 by 3. You have 6 tokens each of them living in three-dimensional space. When I say token I mean context vector. Then you merge them. So when you merge two 6 by 3, two set of context vectors of matrix dimension 6 by 3, you get result as 6 by 6. Then you pass them through mixing matrix which is also of dimension 6 by 6, the result is finally 6 by 6. So, you have 6 tokens living in 6 dimensional space converted to 6 context vectors living in 6 dimensional space ok. This is effectively what we did here, but what we saw earlier was a bunch of details that we have to realize before we actually implement this in practice in code. So, now it is a good time for us to

37:54walk through all of this in code, I have already written the blocks of code corresponding to today's lecture and it is definitely an extension of what we have been doing since the lecture number seven so i will walk you through block by block but my intention is not to teach you python coding or anything here it is just to show you the implementation in practice because i just want to remind you that this series is not about or not necessarily about python implementation of different attention mechanisms it is for for us to get a theoretical familiarity with respect to various attention mechanisms existing. There is group query attention, there is multi query attention, there is multi head latent

38:31attention, latent multi head attention, Kimi delta attention, there are many many variants. So, eventually we have to reach over there that is our goal. But I hope by end of today's lecture, you will have a fairly good idea about the way in which attention mechanism is implemented in LLMs or SLMs including GPT models and many other popular models ok. Now, let us look at the slide let us look at the Google Colab notebook in the. Let us walk through the code first what we start with is we have our 6 dimensional vectors. So, here you have 6 rows and 6 columns each row is 1 vector. So, the example that we were considering in the last lecture was dream big and work for it same thing here. So, this is the vector f corresponding to dream corresponding

39:19to big and so on ok. Now, we are just taking a a shortcut here we are just saying that here the batch dimension the number of batches the batch size here is two but what we are doing is instead of defining another set of six vectors because i'm lazy i'm just copy pasting this these set of inputs twice so the batch dimension is such that you are having this six vectors in batch number one in batch number two also you are having the same six vectors the batch is basically a stack of these set of inputs on top of each other okay so here now your final dimensionality when you run this block of code you will you can print the final dimensionality as 2 comma 6 comma 6 2 for the batch size 6 for the sequence length and the last 6 for the dimensionality

40:07now i could have you know we can add one more dimensionality here so let us say i do this our output will become 2 comma 6 comma 7 see this became 2 comma 6 comma 7 because i added an additional dimension so just understand what is the difference between this 6 and the next 6. We can also print the actual output by the way. We are not yet, we have the batch, I can simply print batch. So if I print the batch, you see this is the first item in our batch this is the second item in our batch so batch size is 2 but they are exact replicas of each other all right now this is the class that we had in the last lecture we had a causal attention class and without spending too much time I will give you a quick walkthrough. First is self.wquery, where it's a linear, you know, it's basically a linear transformation via matrix,

41:11where the input dimension is dn and output dimension is dout. So if we are projecting from a six dimensional space into three dimensional space, this will be six and this will be three. Same thing applies for w key and w value. So, the query weights matrix, key weights matrix and value weights matrix are all linear transformations. Basically, they are matrices which transform from input dimension to output dimension. So, these nn.linear means the total number of trainable parameters. So, this will introduce a bunch of trainable parameters. Total number of trainable parameters will be d in multiplied by d out. So, this will introduce a bunch of trainable parameters, total number of trainable parameters will be d in multiplied by d out.

41:46So, same thing applies for query key and value. Dropout we applied in the last lecture, we will drop out at us drop out certain parameters in the attention base matrix with a certain probability. And the other important thing to note here is now queries you can calculate by passing your input x into w query similarly input x you can pass into w key and w value to produce queries keys and values and attention scores are query dot product with key but here as you just understand why we are doing this transpose so So, if I had not explained earlier why we take the transpose you would not understand this this one step you will be confused not knowing why we are transposing dimension 1 and dimension 2. We are doing that so that it makes it easy for us

42:37it makes it possible for us to calculate an attention weights matrix of dimensionality 6 by 6 otherwise it will produce an attention weight matrix of dimensionality 2 by 6, otherwise it will produce an attention weight matrix of dimensionality 2 by 2 ok. So, that is the reason why we are transposing at dimension 1 and dimension 2. Please remember dimensionality starts with dimension 0 ok. Thenattention scores are now masked by putting upper triangular elements as minus infinity, then after that we calculate attention weights. But of course, before that, we do the scaling. So this is square root of d. And square root of d of what? Dimensionality of the key space. It can be called as dimensionality of key space or value space

43:15because all of that is same in this case. But yeah, in general, dimensionality of the key space. And then finally, what we return is attention weighted sum of value vectors. So this is attention weighted sum of value vectors. So, this is attention weighted sum of. So, by the way before we take the attention weighted sum of value vectors, the attention weights themselves are passed through dropout function. So, otherwise if you do not want to apply dropout, all you need to do is you can just you can simply delete this, you can just say weights multiplied by values that will be yourfinal context vectors, but in this case we are applying from home ok. And the only thingwe need to pass is so if this is the causal attention class what we

43:58are passing into that class here and this is the object one head is the object we are passing the input dimension output dimension and context length so here when i say output dimension it is not the output from the multi head attention okay or or even single head attention it is basically the dimensionality of the query key value space so please remember that so here our nomenclature happened to be such that the out is referring to the dimensionality of the query key value space so our embedding dimensionality is sixth query key value space dimensionality is three and uh context length is six okay now we can print anything so let me run this first we can uh okay something is missing dropout dot forward is missing input Yeah, so I have to pass weights here I forgot to add weight back here

44:50yeah, so now we are just printing the The size of this batch we can also just print the actual Batch itself not bash. We can pass the batch. into this one head object and we can print the resulting output. So, now this is the resulting output. So, what is this? This is so we our batch size is 2 right. So, this is the item number 1 in our batch this is the item number 2 in our batch they will both look identical because the inputs going in inputs that were part of the both items in the batch was the same we took the exact replica right in this step. So, here we took inputs comma inputs meaning we are stacking the same items to form a batch size of 2, but if you look at this one single let us say matrix what is that actually? They are 6 context vectors living

45:46in 3 dimensional space, 6 context vectors living in 3 dimensional space because you had 6 vectors at the input, but the input dimensionality was not 3 it was 6. Now in step number 2 we can concatenate the output from two different attention heads. So here what we are doing is, I mean we have defined a class called multi head attention wrapper. It just produces a wrapper so that the output from different attention heads can be concatenated. That is the only job this does and the only line of code that I want you to pay attention to in this case is just this. So here we are doing torch.cat which is torch.concatenate head of x is nothing but we are so head we are we have defined over here um where is head yeah actually head is not defined outside anywhere head is over here

46:42so here we are saying self dot heads equal to nn dot module list of causal attention of all these inputs for dash in range of number of heads. So, here we have a for loop. So, what we are doing is we are creating a list, but this list is different from the traditional list in Python. nn dot module list means the list will be curated such that it is useful for python to carry of calculations using gpu ok. So, here for dash in range of number of heads how many times will this for loop run this for loop will run as many times as there are number of heads in our case we will define the number of heads to be 2. So, this for loop will run 2 times which means this causal attention class is being or the object corresponding causal attention class is being defined two times in this nn.module list.

47:32Okay. So what that means is you can iterate through it two times. So now here in this torch.cat line of code, what we are doing is for head in self.heads. So self.heads have how many items? It has two items so head of x is the first attention head when you are in you know first item of the for loop head of x is the second attention head when you are in the second item of this for loop so this for loop will run two times this for loop will also run two times because here we are passing we are here we are explicitly saying the number of heads is equal to two so now this causal attention now this multi head attention is basically a concatenation of outputs from the two attention heads you concatenate it along which

48:17dimension along the last dimension and what is that last dimension the last dimension is the embedding dimension so the dimensions one 2 3 or rather 0 1 2 will belong to attention head number 1 the dimensions 3 4 5 will belong to attention head number 2 so this wrapper is the object corresponding to multi head attention wrapper class and you can pass into this object your batch so here this batch is defined in the first block itself so this was our batch this guy this 2 comma 6 comma 6 was our batch and that batch we are passing into the um object corresponding to our multi-head attention wrapper and out loop is nothing but the actual output coming out of our multi head attention. So here we can print either the shape of it or we can print the entire output.

49:07So we can also print out loop itself. Out loop simply means output produced as a result of looping. So here in this first example, I am not showing you the way of avoiding the loops, but rather here we are implementing using for loops because soon i can show you how this output will look like if the number of attention head instead of two if it was four thousand what would happen right so if i run this you can see the output so this is the output dimensionality two comma six comma six so batch size is two so this is batch number one and this is batch number two the item number two in in our batch and six comma six as you know you have six tokens living in six dimensional space so these are context vectors and why is the output

49:51dimension six and not three because we are showing the output here after concatenation so outputs from your your final multiple attention heads you are concatenating and that's what you are displaying here. Now let us see how having this kind of a loop based implementation of multi-head attention is not really ideal. So here is a very simple example and this also applies for inference right so not just for training when you are performing inference you are passing a bunch of tokens as a sequence and that you are using to produce the output. So now if the number of attention heads are very high, it will introduce tremendous amount of latency in your inference. So here I am picking two examples, two attention heads, 400 attention heads and 4000 attention

50:39heads and here I am just simply running these as loops. So in the first part of the loop, I will run with two attention heads, in the second part with 400 attention heads and third, 4,000. So the only part where the code changes is in this line of code, where number of attention heads we are passing, we'll pass variable number of attention heads. In all cases, input dimensionality is six, output dimensionality, the query key value space is living in a three dimensional space and number of heads uh is could be a small number or large number so of course this will uh break some of the things that we discussed earlier like usually the query key value space is such that the input dimension divided by number of attention head equal to the

51:21uh dimensionality in the query key value space that will not apply here okay that will apply for this one example but for 400 and 4000 that will not apply so now if we run the code you can see that two attention if you have two heads it performed uh in 0.003 seconds uh this is basically one forward pass through the attention heads 400 took more time an order of magnitude more time, 4000 took almost close to 1 second, 0.63 seconds, right. So this is very slow. And this is one of the reasons why we are doing the weight split approach. Weight split approach meaning we have, instead of having WQ1 and WQ2, we simply have one weights matrix, query weights matrix, and that we use for producing query.

52:03Same logic applies for keys and values so here if wq1 is a random matrix of 6 by 3 dimension wq2 is also a random matrix of 6 by 3 dimension the instead of trying to separately calculate them using two different using a for loop with two iterations we have one single query weights matrix which is a which is equivalently concatenating WQ1 and WQ2 along the along the column dimension. So, now you can see the output here when I print the query is basically inputs multiplied by WQ, otherwise what we would have had is Q1 is equal to inputs multiplied by WQ1. This would have been Q1. And Q2 would have been inputs multiplied by WQ2. And we can also print those. We can just see that print Q comma. Actually, we can print them separately print here we can print Q 1 and here we can print Q 2 so okay so

53:23this is the result inquiry where I am not calculating the Q 1 and Q 2 separately instead it is all one single matrix multiplication. On the other hand for calculating Q 1 I had to do one matrix multiplication and for Q 2 I had to make another matrix I had to perform another matrix multiplication. So, Q 1 is this and Q 2 is this and you can see that the final the big Q matrix if you look at the first three columns for first three columns are okay I am not able to select only the first three columns but if you look at these first three columns they are exactly same as this Q1 and if you look at the last three columns this last three columns are exactly same as Q2. So, what it shows is it does not matter if you split into two different matrix multiplication

54:14or if you just perform one single matrix multiplication and after you perform the matrix multiplication you can split. You can split the resulting Q into Q1, Q2, Q3, etc. But the key point is reducing the number of matrix multiplications that you need to do because one single matrix multiplication use using parallel compute gpus can do fast but if you have many matrix multiplications to be done that will slow down the total compute so it is always better to perform one single matrix multiplication so that is what we are trying to do here. Now, this is the case where we had two attention heads. What if we had more? What if we had 10? We will have to do 10 such multiplications. What if we had 4000? We will have to do 4000

54:55multiplications like this, right, which will make it very slow. So, now we can also look at the transpose idea which we discussed during the lecture. So we have first let me you know pick items one by one. I will put one more block of code here. So b is batch, number of tokens and d in. So, we can also print b comma actually we can just print batch. So, if we print batch we will see or we can print batch dot shape then we will see 6 comma or rather 2 comma 6 comma 6. So, the value of b is 2 number of tokens the number of tokens in our sequence is 6 input dimensionality is 6 we already know this then output dimensionality and number of tokens the number of tokens in our sequence is six input dimensionality is six we

55:45already know this then output dimensionality and number of heads so output dimensionality after everything is done we want it to be six basically what it means is the context vector resulting from six dimensional input tokens will also be six dimensional number of heads is two which means dimensionality in each individual head will be d out divided by number of heads. So d out since d out is 6 in this case 6 divided by 2 is equal to 3. So head dimension is 3 ok. Then let me just delete this. Then what we do is we define w query, w key and w value. So w query key and value are defined not based on t in comma head dimension so please note we are not doing this we are not having head dimension here if we were

56:32calculating query key and value separately for each heads using for loop we would have had separate query key value weight matrices for separate heads and input dimension would have been d in and output dimension would have been head dimension but that is not what we are doing here the output dimension is d out only okay and then we can say calculate queries is equal to w query of batch and same applies for keys and values so first what we do is we take we calculate queries so that will be w queries of batch and the dimensionality here would be we can we can print um yeah actually this statement is already printing when you print it the dimensionality will be 2 comma 6 comma 6 then you perform this view operation view operation is looking at the

57:18same tensor but you can you can view one of the dimensions by splitting it into two different dimensions so what we do here is instead of the last dimension being input dimension we split it into number of heads comma head dimension so that will be 2 comma 6 comma 2 comma 3 that is being printed here and then we transpose the dimension number one and two basically we will swap these two items so it will become so that the resulting dimension will become 2 comma 6 2 comma 2 comma 6 comma 3 if we don't do this the attention weights matrix will be 2 by 2 instead of 6 by 6 so and same same thing we do for keys and value so here whatever we wrote here in like four five six lines of code we are just writing in one line. So, we are initially calculating W, initially calculating keys by passing our batch which is the input into

58:11W key. Then we are viewing it such that the last dimension is split into number of heads and head dimension. Then whatever is resulting we are transposing the dimension number one and dimension number two we are transposing same thing we are applying for values. So, let us run this and now we can calculate attention scores attention scores is queries multiplied by keys transpose and by the way why is it keys transpose of 2 comma 3 because we are transposing the last two dimensions so when we transpose last two dimensions what we will have is 2 comma 2 comma 3 comma 6 instead of 6 comma 3 so instead of 6 comma 3 the key transpose will be 2 3 comma 6 and and the result will be 2 comma 2 comma 6 comma 6 so that is the that is

58:59what is written here but that can also be printed then we will apply masking so torch dot try u means upper triangular triangular triangle upper then upper triangular items are all converted into i mean here we are just defining the mass of the triangles basically we define an upper triangle where upper triangular items are zeros then that mask we can apply element wise such that whenever we have upper triangular mass as zeros we convert it into negative infinity in that in this case i think the upper triangular items are all if i'm not wrong upper triangle wave once i can i think instead of zeros let me just print it yeah upper triangles are all basically true so whenever you have a true item that true item will be replaced by negative infinity in this line okay so it's

59:53not zero or one it's basically true or false then uh after that you convert attention scores into attention weights by dividing the scores by square root of d and you can after this if you want you can also apply dropout ok. But here in this case we are not applying dropout just for a simplicity then context vectors are attention weighted sum of value vectors. So, weights multiplied by values. So, that is our context vector and the last thing that we have to do is we can print these context vectors. So, instead of printing something like this, we can just print context vectors. So, as you can see, the resulting dimension should be 2 comma 2 comma 6 comma 3. Let us just see how it is. First of all, what is this 2?

1:00:44This is the batch size. So, in the batch, we should have two items. What are the two items? This is the item number 1 and this is the item number 2. So, that is 2 comma, that is the first 2. What about the second 2? The second 2 is because of the transpose, right? So, we had taken a transpose and this 2 is corresponding to our head dimension. So, we had taken a transpose and this two is corresponding to our head dimension. So, we should have head 1 and head 2 where are the numbers corresponding to head 1 these are the numbers corresponding to head 1 in the batch item number 1 these are the numbers corresponding to head 2 in batch item number 1 these are the numbers corresponding to head 1 in batch item number 2 and these are the items corresponding to head 1 in batch item number 2 and these

1:01:25are the items corresponding to head number 2 in batch number 2 okay so this is corresponding to batch 2 this is corresponding to batch 1 this is corresponding to head 1 this is corresponding to head 1 this is corresponding to head 2 and this is corresponding to head two now what about six where is this six coming six is coming because the token sequence length is six you can see that you have six rows here one two three four five six so you have six rows in every single one of these items what about three what is this three three is the dimensionality of each individual attention head so you can see that there are three columns this is column number 1 column number 2 and column number 3 ok so this is how the final context vector looks like before you do concatenation before you do

1:02:19mixing ok now let us go to the next one so in the next one what you do you perform one more transpose you transpose the dimension number 1 and 2 this is to reverse what you did earlier so when you do this you can print print context what you will see is the dimension will become instead of 2263 it will become 2623 and let us actually move this to a new block and print it in isolation So, this is now and we can also print the shape print context.shape. Okay, while it is getting printed, let us just understand that the shape is basically 2623 right so where is 2623 here oh i think i somehow somehow my code is gone let me just get it back

1:03:54context equal to context for transpose of 1 comma 2 okay I know what happened what happened is I ran this block twice so because I ran it twice the transpose was applied once more that is what happened so I can I just need to run it once more that is all ok so now we are back so 2 6 2 3 see if I run this again the transpose will be applied again so it gets back to the original form if I run so you might wonder why do not we get same output when I run the same thing twice because I am I am transposing the same thing right so this is what we want so 2 6 two three where is the first batch item this is the batch item number one this is the batch item number two then where is the six we have six

1:04:34tokens right uh so this is token number one two three four five six and same thing applies for the batch item number two and where is this two two means you have outputs from two attention heads this is the output from attention head number one this is from head number two this is from head number one sorry this is from head number one this is from head number two and so on and what about three what is this it's the embedding dimensionality of that and of the individual attention heads so you have one first column second column third column okay now it's obvious you you know what is what we are going to do next what we are going to do next is we'll concatenate this with this so this last 2, 2 comma 3 will concatenate into one single item. So, when you concatenate, you have after transpose 2, 6, 2, 3, then you concatenate to make it 2, 6, 6.

1:05:33You can also print the actual item. You just print the actual context vector. And you will see that the only difference here now is here instead of you know this these numbers being in two rows currently they are in two rows minus 3.39 times something 2.11 times e to the power of minus 10 to the power of minus 1 etcetera. This will be one single row and that is what you can see here it became one single row. So, 2 comma 6 comma 6 y 2 you have batch size 2 y 6 because you have 6 items basically 6 input tokens and output is also 6 I mean the output dimensionality is also 6 which is why you have 6 columns 1 2 3 4 5 6 ok. So, this is the final output now all of this can be wrapped in one single class.

1:06:25So, I can just define a multi head attention instead of causal attention. The other class which we had here was causal attention class, this one causal attention. So just like how we have a causal attention class, we can also define multi head attention class. So, I will quickly walk you through that it is it is very simple we already implemented everything. So, first is define number of heads define the dimensionality in each individual heads define what should be the final output dimension final output dimension should be I mean unless we explicitly say otherwise mostly it is same as input dimension. Then w query here we have a linear transformation from d in to d out please note it is d in to d out it is not d in to head dimension

1:07:12we are not changing from input dimension to head dimension. Although we are internally imagining it that way, but when we are implementing it to avoid last number of matrix multiplications we just want to do this in one single matrix multiplication. That's the reason why here it is d in comma d out. Otherwise we will have to define w query number one, w query two, w query three, etc. Then we have w key, w value, then output projection. So this output projection is where you perform mixing. So once you have d out, you perform mixing and you also have a dropout in forward pass this is what you do you have queries keys and values calculated then you calculate attention scores i mean these are the raw scores before masking then you calculate mass course mass

1:07:56course you convert into attention rates then you have your context vectors uh then after wise you have context vectors uh i mean the resulting context vector uh the last two dimensions you concatenate so the resulting context vector dimensionality similar to two comma six comma six it will be the equivalent of that and before you return the final context vector you pass it through the projection the we are calling it as an output projection but essentially what it does is mixing uh the dimensionality of this will be same as it will be d out comma d out. So, we can see that here. So, output projection is d out comma d out. In this case, it happens to be equal to d in comma d in. Why? Because d in and d out are the same.

1:08:40d in is 6 dimensions. d out is also 6 dimensions. So, yeah, that is pretty much much it this also you can run and we can clearly see i mean we can write like a setup where we can compare whether the output is any different whether we run it in for loops versus where whether we calculate this in one single matrix multiplication there is absolutely no difference so I will share this Google collab notebook with you so please go through it take more time I did not intend to spend a lot of time on this Google collab notebook anyways but I am pretty sure that there will be some blocks or some lines of code where you will you would require clarification so what you can do is just like how I am using Gemini here you can simply use Gemini to understand the code you can use Gemini to modify the code

1:09:31it's pretty easy I mean none of us are coding by hand these days but it's always helpful to understand what the code is doing for for understanding internally what is happening so you can by all means try to write code using coding agents or gemini or whatever you prefer and but try to understand it also so that brings us almost to the end of today's lecture in summary we have finished now multi-head attention and now we are in a good shape to look at how position embedding is or how position information is injected into this whole setup. We have never discussed position embedding apart from discussing it in one word or something, we have never discussed it in detail. So, that is what we are going to do next in this lecture

1:10:15series. So, we have various position embedding mechanisms to discuss probably the most famous ones are learned position embeddings, rotary position embedding, sinusoidal position embedding and there are a few more. So, we will discuss them in the upcoming lectures. So, the main reason for position information being valuable is because if you look at this sentence all dogs are animals and if you look at this sentence all animals are dogs. If position information was not there the context vectors corresponding to or rather the attention scores corresponding to every token they will remain unchanged there will be no difference between attention scores corresponding to the those same tokens whether you are in sentence number one or sentence number two but clearly uh these two sentences are

1:11:03completely different right so something needs to change some position information is important because here although the same four tokens are present in both sentences they mean completely different things so how do you how do we teach that to a language model we teach that to a language model by teaching the language model to respect the positions of the tokens how do we do that we will discuss in the upcoming lecture thank you so much take care see you in the next lecture bye