Linear Complexity in Attention Mechanism: A step-by-step implementation in PyTorch

hello everyone in our last video we explored eight distinct algorithms that were aimed at improving the efficiency of the attention mechanism by minimizing its memory and arithmetic complexity in this video we'll be presenting a pytorch implementation of four of those algorithms that have linear complexity with respect to the sequence length this demonstration is intended to help you get familiar with the underlying Concepts so we'll be working with this Jupiter notebook um and we'll use pytorch to build uh those attention uh different attention algorithms so at first we will import these two modules that are required uh for this Hands-On exercise and at first let's also look at the theory so we assume that we are given input X which has dimensionality n by D so that is our input tensor and we also also have these projection matrices WQ w k and WV and for Simplicity we assume that the dimensionalities of these DQ DK DV are all equal to D which is the uh input X representation uh after we have uh once we have these four matrices then we can build our q k and V query key and value with these projections so q k and V are obtained with the projection of x with this WQ WK and WV so um so let's uh Define our parameters so since we want to U work with long sequences we assume that n which is the sequence length is 1,000 and the dimensionality of the representation vectors is 256 and as I said for Simplicity DQ DK and DV are also equal so now let's build uh our input X so X is equal to torch do Rand and the dimensionality is n by D and I have put some printed statements already in this notebook so that uh we can see the dimensionality of these different tensors uh so after X we then build WQ also using random so torch. Rand and dimensionality is d by DQ so we do the same thing for WK and WV and then we print their shapes so X is 1,000 by D or 1, by 256 so that is n by D and these are d by DK d by d q DK and D V so now we are ready to build our query key and values so Q is X so that is dot product between X and WQ so we use matm x. matm with w q and similarly for K we say x. matm WK and for V we say a x. matm WV yes so q k and V are also n byd now uh let's go over the original uh scale do product attention which has a quadratic complexity and uh and we can use this equation which is the soft Max of qk transpose the dot product between q and K and then we escape scale that by a square root of DK and then we multiply the entire product with v so we do that in these Three Steps step one we compute the attention scores so for attention scores that is defined as q k transpose so we say Q do matm K transpose for transpose we use this T so that that gives us the attention scores and in step two uh we need to do normalization using softmax so normalization and scaling so we do those two here so we call them attention weights equals uh so that these attention scores um so we use torch. softmax before that attention scores then we scale that using math.sqrt of DK uh and we also need to specify the dimensionality like we want to apply softmax over the rows right so we say de equals one so we can say dm equals one we have to be careful uh with the dimensionalities so in this case our queries keys and values they are all two- dimensional so we don't have any batch Dimension so one is fine but if we had the batch Dimension then we should have used U uh uh two or if we use negative -1 uh that is U safer because negative 1 would work for both cases and we do have a sanity checking just for uh since we know that softmax um uh like the output of softmax will all sum to one so let's do that real quick so we say Su over rows so we use torch. sum to calculate the sum over the rows of attention weights and we say dim equals uh the same thing as before we want to do that over the row so we say minus one so we build the sum over the row tensor here and then we uh compare that with tor. one's like that is is built here and we expect the result to be true the result of torch. all close and then in a step three we compute attention output so all we need to do now is to multiply these attention rates with v so we say attention rat dot matm with v and and that will give us uh uh the final output so let's run this cell um y uh so okay I have some over the rows but here some off rows so some over rows so let's run that again yeah so as you can see attention scores is n byn attention weights is also n by n um the sci check is true and the final output is also n by D so the fact that attention scores is n byn that is the qk transpose um so this will make it both in terms of memory and arithmetic complexity to be um quadratic um so in the rest of this video we will see other algorithms that will reduce this comp uh complexity uh but before we do that uh we also we should note that um when we are dealing with um deep learning Frameworks we always have u a batch Dimension but in this exercise we didn't have any batch Dimension so let's for for practice let's also build the batch version of these query keys and values so to do that we use uh q. UNS squeeze and we add one dimension uh from the left so that is batch Q we also be B build batch k and batch V all right so now we have the batch versions of query key and value and now uh to compute the attention scores um that is and that qk transpose so now we cannot use uh the mat Mo we should use torch. bmm that is the batched matrix multiplication and we provide the two matrices that are involved in this uh dot product as arguments so the first one is batch q and the second one is batch K but uh K must be transposed and here is another thing so we cannot use T here uh we should use do transpose because uh batch K is has three dimensions uh so which Dimensions you want to transpose that is uh Dimension One and dimension two so that gives us the attention scores and then we do normalization so attention weights that is torch. softmax of attention is scores divided by math. square root of DK and we specify the dimensionality uh which dimension we want to compute the softmax so we want to do that over the rows and here we can use either two or we can use netive 1 so negative 1 means the last Dimension and for Sci checking we still uh do that so sum over rows equals torch. sum of these attention weights and dim equals same as here -1 and this one should be some over rows and now um in a step three we can compute the attention output so that these attention weights dot matm with oh here is uh what I'm doing wrong so since we have batch versions We should use torch. bmm of attention weights comma um batch V and that will do it for the batched versions so now we can see batch Q is one um the First Dimension is one but that is the batch Dimension so b n and D attention scores is B NN and attention rate is also the same so the error that I'm getting here um yeah it looks like I have another some off rows so some over the rows and then we can see the attention output is also B and D so now let's move to the first method that reduces the complexity and it is called called efficient attention so we use this equation um the key idea here is that instead of computing qk transpose they first compute uh this uh Sigma call of K transpose time V so by uh changing the order of computations they reduce the uh both uh memory and computational complexity and the sigma here stands for soft Max when it's Sigma row it means it is uh applied to rows and when it's Sigma call means it is applied to the columns so for the rest of this video we will be using just uh the normal q k and V not the bashed versions all right so now we do this in these three steps first we apply Sigma to q and K so we call them Sigma Q is uh torch. softmax of Q and we specify dm equals 1 so Q is has two Dimensions um and we want to apply that to the rows uh so that is one and sigma K is torch. soft Max of K and we apply that to the columns so for that we specify D equals z so that is it for step one now in step two we want to do this computation in this parenthesis so we call the result of this step two result and we do this Sigma K and we here we have to transpose this Sigma K so we use that t and then we do mat Mo with v ma with be and then uh in a step three we calculate the final output so that is attention output equals U this multiplication Sigma row of Q times this step two result so that is Sigma Q do mat Mo with step two results so let's run this cell and uh we can see Sigma Q Sigma K have the same dimensionality as q and K uh but the step two result um is d by D which is much smaller than n byn assuming that n is large um so this way by switching the order instead of computing Q * K transpose they did K transpose time V uh and uh this way they were able to reduce both memory and arithmetic complexity so method two is also very similar uh so instead of using Sigma for Sigma that is applied to q and K they Define this uh feature function f and this is like a kernelized version of attention and when you do this then you can change the order from uh instead of computing uh F of Q * 5 of K transpose you can compute this quantity first that is of K transpose time V and then you you bring in the F of Q from the left and what feature function they use this feature function uh e of x uh which is exponential linear unit uh the curve is uh the curve for E of X alone is this blue curve but as you can see it is negative when X is um less than zero so uh since for car attention for for any kernel function we uh we should not have any negative feature function then they added this plus one so by this plus one it shifts the curve toward positive regions so let's build our feature function X feature function f uh so that is uh at first we build U using torch. nn. elu and then we build our f feature function as e of x + one so now uh we do this uh in these three steps at first we apply F to q and K so we say five of Q equals F of Q Yore k equals 5 of K and then we compute we move to a step two we compute the step two result that is a step two result equals so for step two this is our step two F of K transpose time V so we say Fiore K do T for transpose and then matm with v and the final out output is uh F of Q dot mat Mo with a step two result right so this was very similar except that um instead of using the sigma in method one they Ed this feature function file so now moving to the third method uh this method uh is using uh first order tailor series um approximation and the final equation is given here so uh 1 subn stands for uh a vector of all all ones and Q Prime and K Prime are the normalized version of Q and K and this normalization is based on the L2 Norm so as a prerequisite first we need to build this one subn so we do that by um torch. ons and this will build uh this Vector uh which has U 1,000 ons in it now we do this also in these three steps so in a step one we normalize q and K so that is this Q Prime and K Prime so first we compute Norm of q that is the L2 Norm of Q so we use torch. Norm applied to q and for L2 we use p = 2 and we use de equals 1 and uh we maintain the dimensionality using keep di equals two we do the same thing for K Norm of k equals also tor. Norm k p equals 2 de one and keep them also true so after we build is Norm of Q and Norm of K then uh we compute Q Prime as Q ided by Norm of Q and K Prime is K ided by Norm of K so that is it for step one in a step two now we build uh or we compute these three terms so if you look at this equation we have four terms two in the numerator and two in the denominator but this one is just a constant that is the sequence length it's not constant but uh um because the sequence length could vary but uh we have that as n right so we Define n um as a scalar value so we have term one we transpose. 1 subn and then term two is this term three is this one so term one is V transpose so we say V.T do matm with one subn so that is for this term then term two is Q Prime K Prime transpose V so at first we build uh or we compute uh the quantity inside this parenthesis so term two is uh so that is K Prime dot transpose. T then do matm with v and now we bring Q Prime from the left so we say let's put this into parenthesis and then we say Q Prime do matm this and finally term three term three is also very similar to term two except that so we can copy that and instead of V we use 1 subn and then finally the attention output is term one plus term two divided by n which is a scalar Value Plus term three but there is a problem here um the dimensionalities of the numerator and denominator do not match so we will get an error so uh but that we'll fix that later so Norm of Q uh is n by one um so that's uh it is still has two Dimensions because we used keep them equals through otherwise it would have been just n term one is a vector of size d term two is n by D term three is n a vector of size n uh but this uh division uh did not succeed because of uh this dimensionality mismatch so all we have to do we just need to add one dimension uh from the right side of this term Tre so we can say uh UNS squeeze negative one and let's run that again and yes that works so um as you can see this was another method for um for uh linear complexity of attention and yeah now moving to Method four uh this method is called L forers attention and um it is using low rank approximation and uh they they bring in two new random projection matrices E and F uh which has size R by n and R is the um assumed rank of the attention Matrix so once we compute like uh build this E and F then we compute this uh projection matrices EK and FV um which have dimensionality R by DK and R by DV after that we build this context mapping Matrix p p as shown here so this is similar to the original attention equation scale do product attention except that instead of K they are using e k and finally the the final output is p * FV so for the prerequisite of this we assume that R is 100 and because we assume that sequence length is 1,000 uh so R must be something uh much smaller than 1,000 so we we say like 100 is a good assumption so then uh we build this E and F so e is torch. Rand and the dimensionality is R by n so R is defined here and we say Tor Rand RN and F is torch. rant also r n and now uh we do this in these three steps in a step one we compute EK that is e do matm with K and FV is f do matm with v then in a step two we compute this context mapping M Matrix p so that is torch. softmax of uh q. matmo with e do T and then we divide that by a square root of DC and we run that over uh the rows so that is equals 1 and finally the attention output is p do mat Mo with FV so now notice that EK and FV have dimensionality R by D so uh it is reducing the dimensionality of K and V from 1,000 by 250 6 to 100 to 256 and then the context mapping Matrix p is n by R so 1,00 by 100 so previously in the original scale. product attention this was 1,00 by 1,000 So This Is How They reduced uh the complexity so that brings us to the end of this video which is uh the first handson video in PML Studio Channel I hope that the detailed work through of the implementation has helped clarify the concepts you can find the links to this notebook and the solution uh in the description thanks for watching and happy learning