Reformer: The Efficient Transformer

Key Takeaways

this video will explore the new reformer efficient transformer model from Google AI this model is really interesting because it presents a way to approximate full attention and tension over long sequences as well as ideas to reduce the memory requirements of training transformer models this idea of increasing the attention span is really exciting in a similar evolution from a simple recurrent neural network to a long short-term memory network and to the transformer the idea of increasing attention span allows the model to incorporate more context from the past when making current decisions a lot of research has come up in the space of expanding the attention span such as Facebook's adaptive attention span and open AI sparse transformers most of these looking at how to take advantage of the sparsity and attention how usually if you have these long sequences due to the query key transpose pass-through softmax only a small amount of the words or you know the pixels whatever context you're using attention on is going to actually contribute to the output anyways so this paper uses a locality sensitive hashing to increase the attention span by approximating full attention and only focusing on the dot product between keys most similar to the query that are actually going to contribute to the attention in the end even if you did the full attention this paper also integrates reversible layers from the revenant architecture to avoid storing the intermediate activations and reduce memory cost of training transformer models at the expense of longer training speed this video will explore the reformer efficient transformer model from Google ai this is a new paper presenting a few new techniques to dramatically increase the attention span in transformers the length of the sequence of which the transformer attends on and reduce the memory cost particularly the memory costs associated with attending over long sequences as well as storing the intermediate activations through these multiple layers in these transformer networks improvements in the performer model compared to the standard transformer architecture is an improvement in memory efficiency reducing the memory from the order l-square required to have this Curie times the transpose key matrix of order l squared to order L log L by having this kind of an LS H attention mechanism while also reducing the amount of memory you need to do to store the intermediate activations by using these reversible layers so you're increasing your memory efficiency but you're trading this off for copy and speed so the key idea in the reformer and a lot of these papers like adaptive attention span from Facebook and sparse transformers from open AI is they're trying to take advantage of the sparsity of attention to attend to longer sequences so even if you say extend the sequence length that you're attending on to 64,000 you still would probably only look at 32 or 64 words to actually incorporate when you do that query terms of key transpose divided by this normalizing factor and you put it through that softmax function you're only gonna have so many words that you're really attending to the softmax function is going to blow up large values and really shrink small values so the idea is that you're taking advantage of the sparsity of attention by using this approximate attention based LSH for the nearest neighbors so this locality-sensitive hashing technique is going to help you to you know approximate the way that you do the query times the key matrix to only do the only multiply it by the keys that are similar to the query because the way you do this dot product is that you know the only the similar keys are even going to come out with a large value anyway so this is a technique to reduce this computation dramatically and make it so you can attend over longer sequences without having to look at all the keys then they're gonna introduce this reversible layer so as in another paper titled red Nets they're gonna introduce this way of structuring the output of the layers so that you can break the Y output into y1 y2 and you can sort of use this way to backtrack and get to X 1 from Y 2 and do this kind of a calculation in order to only have to store the activations on the last layer and then you can propagate backwards during training and then during updating the network the idea is that instead of having to store all the intermediate activations for back propagation you would have more computation through back prop because you got to recompute the outputs but you save that memory cost of storing the activations we'll start looking at the reformer by looking at the reversible layers because it's a bit easier to understand than the LSH attention approximation the reversible layers also presented in red net is this idea of having to avoid store in the in activations in the intermediate layers so these transformer networks they are composed of these blocks where you pass the query key and value matrices into the self attention layer and then a feed-forward layer the idea is that you repeat this block like six to twelve times forming this deep neural network and then naturally the memory requirements of that kind of an architecture come with respect to the number of layers sandy here is that you want to get rid of this bottleneck in the memory by having to store all these intermediate activations of each layer in order to do back propagation say idea behind reversible layers is that you decompose the output into y1 and y2 such that you can reconstruct the input so that this way you only store the activations in the last layer and then you will you know do this computation as you going backwards in order to recover the intermediate activations and the benefit of that is even though it definitely increases your computing time and will make this training take longer you don't have to worry about running out of memory with this the idea here is that so you break it your output into y1 y2 you can then deconstruct say x2 by taking y2 and subtracting the output of the feed-forward Network when you give it y1 so the y2 is a result of applying this feed-forward network to y1 and you still have access to this feed forward Network and you can run in France during back propagation as well so what you do is you take this separate component of the y1 and you pass through the feed-forward network and then you just do y2 minus whatever this is and that's how you get x2 and then you can do the same thing with the attention layer on x2 in order to get back X 1 and then X 1 X 2 would become y1 y2 and you keep going backwards doing this kind of a thing now will get how you go from the order l squared memory bottleneck to order L with respect to the reformer efficient transformer Network so the first idea is that the order l squared bottleneck comes from this query times the key transpose operation with respect to the attention layer city here is that the query matrix has the dimension length by D sub K D sub K being the embedding dimension for the queries and then the key is also going to have the same D sub K by length so this matrix multiplication is going to come out as length by length or order L squared in the memory requirement to store this matrix so the first idea is that instead of doing this multiplying all the queries by the keys at once we're going to index the queries individually with the Q sub I so we have these individual vectors of dimension 1 by D sub K so when you do 1 by D sub K times D sub K by length it'll come out one by length and then so you avoid having to intermediately store these length by length matrices so this individual indexing is the first beat up to speed up but the first you know memory reduction in order to reduce the memory requirement of this now we'll look at the locality sensitive hashing approximation to attention so the motivation is that you're doing this dot product attention where you're taking the queries and you're multiplying them by the keys so what you're doing is you're doing this element-wise multiplication and then summing it up so it's really a measure of similarity so say X 1 is 0 y 1 is 85 X 2 is 85 and Y 2 is 0 even though it has high magnitude at certain spots you're gonna get 0 as the output of this dot product compared to this kind of x1 is 10 y 1 is 10 and they both have similar values so they're gonna multiply and be 100 plus 100 is 200 this kind of dot product attention being approximated as a measure of similarity and we're gonna look at how the locality-sensitive hashing does this kind of a similarity metric in order to sort the sequence so you have this sequence of say 64,000 words it's gonna sort it such that you only do the query times say like the most 16 similar keys and then you're gonna kind of sort it along this way the idea of LSH attention is to approximate full attention by only multiplying the queries by the keys they're the most similar and are thus gonna dominate the full attention dot product anyways especially as it gets passed through that softmax function so locality sensitive hashing starts off by framing the similarity problem as an intersection of sets this is really interesting algorithm that's applicable to things like plagiarism detection and documents collaborative filtering where you're filtering between users picking out different products or watching different movies things like this and also applications like nearest neighbor image search where you have maybe these little dimensional vector representations where images are passed through convolutional networks and then you want to compare these vector representations to find the most similar images the idea is to frame similarity as an intersection of sets problem so this idea of Jaccard similarity is you take the intersection divided by the union so you just see this random example of having a as a set of 1 2 3 4 b 1 & 2 and so the similarity Jaccard similarity between a and b will be these two over the union which is 4 and then the Jaccard similarity of a and c is they have this two in common and the union is going to be 8 so this is the high-level idea of locality sensitive hashing as your framing similarity as an intersection of sets problem a great description of locality sensitive hashing is found in this book massive data sets and this video series it's linked in description of this video and this image particularly came from finding similar sets again the videos linked in description so this pipeline is describing a motivating example of finding plagiarism by comparing the similarity between documents and the text that appears in the document so the first stage of this pipeline is shingling which is like engrams where you're extracting the strings of length K that appear in the document and you're just sliding that k sized window across the document to form this set of all the strings that appear in the document of this length so naturally this set is gonna be gigantic especially if K is you know say 10 or some longer lengths like that so then the idea is to reduce the size of the set by doing min hashing so min hashing is gonna take say a string of length K and it's gonna hash it into a bucket and the bucket is going to be defined with less bytes than the strings would so say you use 4 bytes to encode the bucket number and then you use like 9 or 10 to encode the strings so now you have these buckets that are the purpose of the bucket is again just to reduce the storage cost of the strings of length K so now what you're doing is the locality-sensitive hashing where you're looking for the similarity based on set membership so you have all these documents like document 1 document 2 and then they're mapped into these buckets that you know they contain the strings of length K that are mapped to the same bucket so what you do with locality sensitive hashing is you randomly permutate the order of the documents and then you look for the first appearance of each set and then you use this in order to find similarity so probably make a little more sense by reading this textbook chapter so just another idea of the LSH algorithm is that you're gonna randomly use different hash functions and repeat this in order to get a better estimate of similarity so in the paper they show this diagram of doing the rotation between vectors and showing how this results in different bucket assignments so this is kind of the idea of the buckets that we're talking about you would in this first case you rotate the vectors and then this one ends up in buckets 0 this one ends up in bucket 3 so you're not gonna align this as being similar but when you rotate it this way it's 2 & 2 so the high-level idea of this is that you can see how doing the different hash functions and this kind of analogy was like a you know sphere in two-dimensional space is showing vectors that are nearby are resulting in similar bucket placement with the hash functions compared to vectors that are farther away so to tie this back into the main motivating idea of this the LSH attention is we want to group the keys together based on this kind of LSH bucketing so we take the comparison between the queries and the keys we do the LS h bucketing in order to get the similarity between the keys and then you chunk them and do the attention on the similar keys so another interesting characteristic of this is that you attend on the same bucket and the previous bucket as well so the blue would be the most similar and then you have the second where you have probably a mix of you know k most similar and then say like tier 2 similar and you're also going to attend backwards which is kind of interesting as well an interesting comment i saw when i was looking at the reddit discussion of this paper is somebody wondering about if you have these randomly initialized weight matrices for the queries and the keys won't this kind of random assignment of similarity be permanently stuck well i think that kind of the way that these nearest neighbors get updated is that similar to like this paper exploring randomly wired neural networks where you're searching for the edges within neural network that make up a sparse network within the overall set of possible edges you could use you have this kind of differentiation where even though in the case of the exploring randomly wired networks you're going to explicitly increase the probability of edges that weren't even originally included in this case I think you're just looking at the top end of that or you're decreasing the probability of having the similarity because you still passes through that weight matrix that assigns an input into the similarity so that matrix is going to have a partial derivative that will maybe you know perturb it so it's less similar as you go on through the training these are two other papers that are interesting with respect to the area of increasing the attention span or increasing the length of the sequences L with respect to the attention layers so this paper sparse transformers from open AI it looks like what this does is it breaks up the attention into the multiple heads so each of the heads of the attention sort of like the stack of these attention layers is going to look at a different subset of this sequence so that way they can fit the larger sequences into the attention layer by just having it explicitly passed into each of the heads tension layer so in this other paper from Facebook the adaptive attention span has this parameter over the attention span of each transformer head so basically they have different attention spans in each of the layers of the network it's a bit of a different idea than the reformer where you know each of these layers would have a high attention span that wouldn't be like a variance of attention span between layers thanks for watching this overview of the reformer efficient transformer architecture from Google AI I hope from this explanation you're able to get a good sense of how the reversible layers help you to avoid storing the intermediate activations as well as to get a general sense of what this locality sensitive hashing is about and sort of the high-level idea about how this might reduce the memory cost of storing these long sequences and attending over long sequences thanks for watching and please subscribe to Henry Air labs for more deep learning in AI videos