Blowing up Transformer Decoder architecture

Key Takeaways

hello everyone and welcome to another episode of Code Emporium where in this video we're going to blow up the architecture of the decoder and continue our Transformer architecture discussion in probably the last video or two I've described the architecture of a transformer neural network that you see over here where it consists of an encoder and a decoder in my last video we also looked at how we can blow up this encoder to kind of look something like this where it's a wonderfully amazingly complex diagram that you all definitely understood so we could take for example our inputs that we see yeah it's my name is a j if you kind of zoom in there my name is Ajay that would be inputs over here and then the outputs when it passes through this entire network is gonna be something that looks like let's just go all the way scroll scroll with all this wonderful attention going on and then it's going to look like this where we have my name is a j I thought that this was a decent picture of a first stab of trying to draw the entire network but I realized that there was a lot of confusion about how you would add the batch Dimension because when you're coding things out you typically add another dimension to all of this and it just became increasingly complex to draw and so what I've done is well in this video I'm going to do the decoder architecture where I actually talk about that batch Dimension and you'll see like all of these are much bigger squares you'll see oh wow there's like three dimensional shapes over here which I thought was very amazing and so we're gonna actually talk about the batch intervention and exactly how you would kind of code this out but explaining it completely in architecture first so there's not going to be too much code in this video but this is a good starting point of just understanding like what's really going on behind the scenes when you look at like this decoder part so let's get to that so first off when we're looking at this Transformer encoder here right we first just have like English words that we pass in and throughout the course of the network the idea is to use this attention mechanism so that each word in the sentence will get the context of other words and let's say for now that we are trying to build a system that will convert English to some other language called kanada so Canada is a regional language within India in case you didn't know I am Brown I look kind of brown I also do speak kannada a little bit so I thought this is an easy way to like validate things when I'm also coding it out which I will do in the future and so this encoder is essentially just going to take the English sentences and try to make every single word at least the vector representation of every single word that you see here it's going to try to make sure that it kind of understands each other so my for example will actually be the VAC the representation of my which is very contextually aware after passing through the trained encoder architecture the same thing is name is Anna J and the reason why all of these vectors are so contextually aware is because of the wonderfully drawn attention mechanism that I've Illustrated over here now despite the decoder having very similar components let's actually take a huge pass at how these new vectors that we saw here actually play into the decoder's hands when we're actually you know processing the Canada sentences so here's a decoder architecture and if we zoom in here let's just go from the very top of like what is actually going on so to the decoder architecture you need to pass in the sentence that you are going to be translating to in this case it's Canada and let's say that since all of these are this is like a word translator it's going to translate like one word at a time and that's why you see like one word on different lines so the first token that we will pass in is going to be a start token this is just going to initiate the start of the sentence and then you'll see like there's there's a thing that says which is my name is a j and then there's an N token to signify that the sentence has ended followed by a bunch of padding tokens and so the reason why we add padding tokens is because we want a fixed size input for any kind of sentence that we add in in any language sentences have variable length but if we want to fix size input then we basically finish the sentence and then just add padding tokens and so the dimensions of this are going to be the batch size which is the number of elements and the number of samples at a time that we can pass through the network then cross the max sequence length which is the maximum number of words in the sentence across some some 512 I've put here and this is the dimensions of each and every single word so every word is encoded into a 512 dimensional Vector the reason we need batch size is because well to perform something called mini batch gradient descent when we pass a sentence or sentences through the network we only every time we pass a single sentence if we were to perform you know forward propagation calculate the loss perform back propagation you'll be updating millions of parameters that is going to take a lot of time if we only do one sentence at a time but if we kind of take multiple sentences at a time let's say like 30 or 50 or 100 and put 100 sentences all together in in parallel we can actually learn a lot faster and all because the weight updates are only going to happen once every let's say 30 elements or once every 50 elements that are seen and so because parameters are updated less frequently it happens a lot faster but the network is also seeing all of that data and hence mini batch gradient descent or some variant of like batch gradient descent is used next we look at positional encoding so because all of these words are passed in simultaneously into the network we need a way to make sure that some ordering is preserved because in English ordering matters but clearly here it doesn't so in order to encode ordering we use the sine and cosine functions and we add it to every single element over here so it's basically the same dimensions and we get another tensor of well the same dimensions itself but what's different about this is now it is positionally encoded now at this point each word is represented just by a 512 dimensional tensor but we'll actually make that word be represented by three more vectors that is a query Vector a key vector and a value Vector so you can think about this like in some intuitive that the query Vector is basically what do I have the key Vector is what can I offer and the value Vector is going to be what is actually being offered when Computing the attention mechanism that you'll probably see later here I've seen in many implementations that this same exact Vector that you see here can just be repeated three times to create the query key and value Vector but in other implementations I've also seen there it being like a feed forward layer which is also pretty good because it adds other parameters to the to the entire network so that there's an additional case there to learn if required and this can also potentially improve performance down the road so you would basically have this tensor you would have like a one is to three linear layer and you would map it to this big tensor I've also written it in this way because in practice the query key and value vectors are stacked on top of each other and so we'll end up with batch size cross math sequence length cross well 512 times 3 which is 1536. at this point we are going to now perform the multi-headed self-attention over here and you'll see that for the number of attention heads is going to be eight which is basically eight parallel processes that go on and that's why I've kind of stacked eight of these rectangles on top of each other all of them have the same process going on and for each of these eight heads what we'll do is divide each of the query key and value vectors into eight parts so in this query tensor for example we're going to divide it along this vertical Dimension over here that is the embedding Dimension that first part will be like the first head for the query Vector similarly the key we're going to divide into eight parts vertically where this first part will just be the key vector and then the same thing goes with the value Vector we just divide into eight parts itself and the first part is like the first head of the value Vector now because we've divided each of these into eight parts that's basically dividing this entire embedding of 1536 Dimension size into 24 parts right that's eight times three and if you do the math of like 1536 divide that by 24 you'll get like 64. right and because you know the way that we pass in these vectors together is a qkv concatenated so you'll end up with a qkv of like 64 64 64 which is adding those three together you get 192. so for one attention head you're gonna get essentially you're passing in this one tensor of well of size well this and like this there's going to be eight of these attention heads all happening in parallel now what happens here is the Crux of the attention mechanism where we will take the query Vector which I said which would be the batch size cross Max sequence links cross 64. multiply by key tensors this is a typo and once we actually do this multiplication you're going to get um what we what the essential is initial parts of the attention Matrix is which is going to be a batch size cross Max sequence length cross Mass sequence length since this is batch size cross Max sequence length cross 64 and this is batch size cross 64 cross matte sequence length and so multiplying these two together you get this tensor of this shape why I say that this is a rudimentary attention Matrix at this point is because like intuitively what this is doing is for a single word word we have a query Vector a key vector and a value vector and so every words query Vector is trying to compare itself to the key Vector of all of the other words in that sentence and if the multiplication of those two is high it means that they're very similar to each other but if it is not high then they're very dissimilar to each other and we want the network and self-attention we basically wanted to pay more attention to the cases that it is more similar to each other and so you would end up with a self-attention matrix typically where the diagonal is going to be like very high Affinity because uh you know it's attending on itself so there's going to be a very high attention value but there also would be attention values for other cases as well because like in the at the end of the day the vectors are going to be learned parameters for each of these words and they're going to represent well two words that are close to each other would actually have vectors that are very similar to each other over over time and so close to the numbers higher the affinity and that's kind of the Affinity that we're trying to capture in this like multiplication operation over here and hence this is like an unnormalized attention tensor you can call it now what we do here is actually pretty interesting too so for one there is an implementation where we would actually scale these values and in fact we do we do this mostly because the variance of some of these values may be pretty large and they can lead to erratic behavior of your network and so kind of scaling these values would help all of these numbers have like a mean zero standard deviation of one for like stable training now because this is the decoder architecture and something that we did not see in the encoder we apply some masking over here and you can see like the way I've designed the mask is like it's like a triangular Matrix with the bottoms or zeros and the upper triangle is a bunch of negative Infinities now regardless of why this negative negative Infinity right now you could say is this is massed and we have zeros over here which means that this is not masked just take it for a face value for now and the reason that we do this is because we don't want the decoder to cheat during inference time during inference time we're not going to have like we're not actually going to pass in everything as we do over here let me just go back here like we're not gonna actually pass in all of this stuff we're gonna probably just start with a start token and then everything is just going to be like a padding token because like we don't have any information at the moment and so if we try to perform like all this qkv attention stuff over here it's going to try to understand attention to words that it doesn't even have yet and so that's considered cheating and what we don't what we want to make sure is that the training and inference happen very similarly to each other so that we don't lead to something called like data leakage at all and so what we would do is take this Matrix we would now add it to our initial attention tensor now what this will do is like for all for every word it will only be able to attend on words that come before it right that's kind of why we see a triangular lower triangular Matrix here and all the upper triangular matrices will just be like negative Infinity we can't do anything about it why we apply negative Infinity is because we want to perform a soft Max when we perform a soft Max on this we would essentially it's essentially like taking e to the power of something right and e to the power of negative Infinity has the effect of turning it into zeros which is like masking so basically it's saying that this first word up here it can only pay attention to like the actual first word which is probably the start token but it cannot pay attention to any other words that come after it which is exactly what we want so that you know we don't cheat like we already mentioned that there was like a query and a key Vector that we used before now we have the value tensor over here so what we do is we multiply this value tensor with the attention Matrix and this is going to be a batch size cross Max sequence like cross Max sequence length this is going to be a batch size cross Max sequence length cross 64. and so you'll end up with a batch size cross Max sequence length cross 64 dimensional tensor now what this will actually give us is 64 to dimensional vectors for every single word and these 64 dimensional vectors are now context aware because of attention so they're really much more higher quality tensors so to speak and remember that we did this we create this for eight attention heads there's like eight parallel ones that go around here and so what we will do now is concatenate all of these eight together and when we concatenate them we will get batch size cross Max sequence length cross 64 times 8 which is 512 and so we ended up with a concatenated value tensor so this is now the full tensor each each word represented by 512 dimensional Vector that's very context aware and now what we're doing is well let me go back to the figure here so that we have now a general view we have outputs that have been shifted right which is like using the start token and so they've been kind of shifted towards the right direction we've created embeddings we have the positional encoding we created query key and value vectors We performed this masked multi-head attention and now we're going to perform add and layer normalization by using this residual connection so you'll see like a lot of these residual connections and they're required to propagate signal throughout the network and also ensure that back propagation does not fail otherwise you with very very deep networks you will learn you will eventually run into situations where the gradients will become so small that the neural network stops learning this is kind of like akin to what we call the vanishing gradient problem and so to ensure that we continue to propagate the signal much stronger and the C signals like over here are better felt as we go further back in the network we use residual connections and so coming back over here you can see like I've drawn an arrow here and this actually come this is a residual connection that is actually coming all the way from the very beginning um essentially what is what we see in the figure so we have that this is our entire multi-cell Mast multi-head self attention that we just discussed and now we have a residual tensor which is going to be the same exact shape that we're adding to this and then we're going to perform some layer normalization and what layer normalization does is it's going to try to normalize values along the layer Dimension so for the 512 Dimension Vector for every word we're going to essentially compute a mean standard deviation and try to scale these values such that they become more stable during training and in this process too we might be introducing some parameters called gamma and beta which are essentially going to be like learnable parameters that we can multiply and add to these tensors so that all these tensors will be comparable to each other so we're essentially going to take layer normalization over here and we have this tensor of all of these words right now but now what we're going to perform is multi-head cross attention so just to go back to the figure you can see we're going to do this other orange this is a multi-head attention technically this is cross attention because it's going to take two arrows here from the encoder and there's one Arrow coming here from the decoder so this is going to be the query vector and these are going to be the key and value vectors from the encoder and if you want to blow that up you'll see that well essentially this tensor is itself we can use it as the query vector or you can introduce another you know linear layer with like the same number of parameters itself and we get the query tensor and then you see from here you can kind of see this is like the encoder over here right so when you have the encoder I didn't label it but this is the encoder from the encoder what we're going we already see like batch size cross Max sequence link cross 512 we will now map it to uh from 512 we will now represent each with a key vector and a value vector and so it becomes match sequence length or sorry batch size cross Max sequence length cross 1024. and so that's like 512 and 512 over here and now that we have the query key and value we do exactly the same process as we did with self-attention where we now pass it through we'll create eight attention heads hence we have like eight of these in parallel it's the same Dimension shape if you look at all of these together so we have now just looking at one attention head it'll be the query key and value each of size 64. the query coming from the decoder but the key in value coming from the encoder and really batch size crossback sequence link cross 192. now we perform a normal attention where we're doing like the query that is for every single word in the kannada sentence we are trying to see its affinity for every single word in the English sentence and once you do a multiplication this is not an addition once you do that multiplication here you're going to end up with an attention vector or a tension tensor where every word in the content sentence is attended to every word in the English translation sentence and then after performing the scaling and the soft Max you'll have an actual attention Matrix now note here that when we perform softmax these values are now all going to be in like akin to like probability values between zero and one but note that we're actually not performing any masking here this is because that during inference time we will always have the each kannada word will always have access to every English word in the sentence from beginning to end so we don't need to perform this Mass detention and that's why we don't do it here and we have we make use of all the information that we can next we now have a value tensor here which is going to be from our our English sentence and this essentially is now going to be like a tensor where for every Mac sequence we're going to have like every single a row is going to be like an English word and each is going to be represented by a 64 dimensional Vector with some contextual awareness about itself and also about the Canada language we now then like concatenate all these so that 64 Dimensions here we are going to concatenate all of them together so it's 64 times 8 because we have eight attention heads we'll end up with a concatenated value tensor over here and before then we'll perform the add in layer normalization so we'll have like a residual addition that comes all the way over here where should we add and then we perform layer normalization like we mentioned before we now have like a feed forward Network for which we have now the final tensor but in order to understand the complexities of the English and the kannada language sometimes we're going to have to we need some more complex units and that's why we will take this tensor and kind of like put it back in the flow so it goes back all the way over to the beginning of this blue blob over here where it's the same batch size cross Max sequence like cross 512 and goes to the flow again so this entire decoder layer if I zoom out a little bit I've written in blue and I've just said we repeated like times n because that's kind of how it looks over here if we repeat it end time so we always come in go out repeat it about end times and that way we can understand the complexities and intricacies of English and kannada and then once we're ready we can basically once we have our final tensor over here we can now expand this such that each word is now corresponding to it's going to be mapped to some kind of word and so we'll use a feed forward layer in order to map it to the larger Canada vocabulary or larger or smaller depending on how big your vocabulary is and then we will represent them with probability values of like what is the probable next word that we're going to generate so in in the case of training it's actually very interesting how Transformers are trained and inferred in the case of training though every single one of these rows is going to be like a probability distribution and so the first row is supposed to predict the word this is going to be like let's say none now was supposed to be like this you know in the vocabulary it comes here and it's like 0.6 which is let's say it's pretty decent right so based on that we will compute the loss but we can also compute the loss in parallel for every single word this should have been hesaru this should have been a j this should be end this should be pad so this is like the raw you could say it's like imagine if this was like a one hot encoded tensor of the same Dimension and these are the probability values we can basically actually draw conclusions about like the loss function itself which if we use like a standard cross entropy loss function and we ignore all of the padding because padding is just useless we don't really care about that we just take you know these four losses that we would get here and we combine them together and then we can back propagate this loss through the entire network so that way that's one network update but in doing so we were able to feed an entire sentence at a time in fact not just an entire sentence we can do this for an entire batch of sentences and this is how like despite the Transformer being in extremely complex architecture with like millions of parameters training can happen pretty fast because if you know because of taking advantage of batches and also the fact that we can literally pipe in an entire sentence and get the labels for an entire sentence at the same time like we do so here now that's going to do for like the training procedure but if you want to infer this like I said before you would only pass in like a start token at the set at the start of all of this there would only be a start token followed by a lot of padding tokens here we then let it propagate through the network do some self-attention over here and then from the encoder we still get like all the sentences and we can now begin to encode all of these values over here perform this cross attention and then when we actually perform the vocabulary we would not we would only infer just the first part because we don't have information for the second row just yet we only have information to predict the next word which is in this case the first word and it should just predict and then we'll take Nana pass it again into the network from the beginning as the second word so we start and then then we try to predict hesaru and then we take whatever prediction we have put it back into the network and in that way the Transformer decoder during the generation phase or inference phase is going to generate one word at a time and so I hope that this decoder architecture kind of makes more sense here and I'm going to be pairing this with some code eventually later on that I've been I've actually been writing out so I'll be explaining all this code in the next video just to see how it relates to exactly this entire architecture but for now thank you all so much for watching and I hope that you kind of understood like what's going on here at least at a very high level if you have any questions please do leave them in the comments below please do give this video a like if you think I deserve it thank you all so much and I'll see you next time bye