Reformer: The Efficient Transformer

1 INTRODUCTION . The Transformer architecture ( Vaswani et al. , 2017 ) is widely used in natural language processing and yields state-of-the-art results on a number of tasks . To obtain these results , researchers have resorted to training ever larger Transformer models . The number of parameters exceeds 0.5B per layer in the largest configuration reported in ( Shazeer et al. , 2018 ) while the number of layers goes up to 64 in ( Al-Rfou et al. , 2018 ) . Transformer models are also used on increasingly long sequences . Up to 11 thousand tokens of text in a single example were processed in ( Liu et al. , 2018 ) and when processing other modalities , like music ( Huang et al. , 2018 ) and images ( Parmar et al. , 2018 ) , even longer sequences are commonplace . These large-scale long-sequence models yield great results but strain resources to the point where some argue that this trend is breaking NLP research1 . Many large Transformer models can only realistically be trained in large industrial research laboratories and such models trained with model parallelism can not even be fine-tuned on a single GPU as their memory requirements demand a multi-accelerator hardware setup even for a single training step . Do large Transformer models fundamentally require such huge resources or are they simply inefficient ? Consider the following calculation : the 0.5B parameters used in the largest reported Transformer layer account for 2GB of memory . Activations for 64K tokens with embedding size 1024 and batch size 8 account for 64K × 1K × 8 = 0.5B floats , requiring another 2GB of memory . If our memory use was only per-layer , then we should fairly easily fit a large Transformer even on sequences of length 64K on a single accelerator . Further , the whole corpus used to train BERT only requires 17GB to store . Why is it then that we can not even fine-tune these models on single machines ? The above estimate includes only per-layer memory and input activations cost and does not take into account the following major sources of memory use in the Transformer . • Memory in a model with N layers is N -times larger than in a single-layer model due to the fact that activations need to be stored for back-propagation . • Since the depth dff of intermediate feed-forward layers is often much larger than the depth dmodel of attention activations , it accounts for a large fraction of memory use . • Attention on sequences of length L is O ( L2 ) in both computational and memory complexity , so even for a single sequence of 64K tokens can exhaust accelerator memory . ∗Equal Contribution 1https : //hackingsemantics.xyz/2019/leaderboards/ We introduce the Reformer model which solves these problems using the following techniques : • Reversible layers , first introduced in Gomez et al . ( 2017 ) , enable storing only a single copy of activations in the whole model , so the N factor disappears . • Splitting activations inside feed-forward layers and processing them in chunks removes the dff factor and saves memory inside feed-forward layers . • Approximate attention computation based on locality-sensitive hashing replaces the O ( L2 ) factor in attention layers with O ( L logL ) and so allows operating on long sequences . We study these techniques and show that they have negligible impact on the training process compared to the standard Transformer . Splitting activations in fact only affects the implementation ; it is numerically identical to the layers used in the Transformer . Applying reversible residuals instead of the standard ones does change the model but has a negligible effect on training in all configurations we experimented with . Finally , locality-sensitive hashing in attention is a more major change that can influence the training dynamics , depending on the number of concurrent hashes used . We study this parameter and find a value which is both efficient to use and yields results very close to full attention . We experiment on a synthetic task , a text task ( enwik8 ) with sequences of length 64K and an image generation task ( imagenet-64 generation ) with sequences of length 12K . In both cases we show that Reformer matches the results obtained with full Transformer but runs much faster , especially on the text task , and with orders of magnitude better memory efficiency . 2 LOCALITY-SENSITIVE HASHING ATTENTION . Dot-product attention . The standard attention used in the Transformer is the scaled dot-product attention ( Vaswani et al. , 2017 ) . The input consists of queries and keys of dimension dk , and values of dimension dv . The dot products of the query with all keys are computed , scaled by √ dk , and a softmax function is applied to obtain the weights on the values . In practice , the attention function on a set of queries is computed simultaneously , packed together into a matrix Q . Assuming the keys and values are also packed together into matrices K and V , the matrix of outputs is defined as : Attention ( Q , K , V ) = softmax ( QKT√ dk ) V ( 1 ) Multi-head attention . In the Transformer , instead of performing a single attention function with dmodel-dimensional keys , values and queries , one linearly projects the queries , keys and values h times with different , learned linear projections to dk , dk and dv dimensions , respectively . Attention is applied to each of these projected versions of queries , keys and values in parallel , yielding dvdimensional output values . These are concatenated and once again projected , resulting in the final values . This mechanism is known as multi-head attention . Memory-efficient attention . To calculate the memory use of the attention mechanism , let us focus on the attention computation from Equation 1 . Let us assume that Q , K and V all have the shape [ batch size , length , dmodel ] . The main issue is the term QKT , which has the shape [ batch size , length , length ] . In the experimental section we train a model on sequences of length 64K – in this case , even at batch-size of 1 , this is a 64K×64K matrix , which in 32-bit floats would take 16GB of memory . This is impractical and has hindered the use of the Transformer for long sequences . But it is important to note that the QKT matrix does not need to be fully materialized in memory . The attention can indeed be computed for each query qi separately , only calculating softmax ( qiK T √ dk ) V once in memory , and then re-computing it on the backward pass when needed for gradients . This way of computing attention may be less efficient but it only uses memory proportional to length . We use this memory-efficient implementation of attention to run the full-attention baselines presented in the experimental section . Where do Q , K , V come from ? The multi-head attention described above operates on keys , queries and values , but usually we are only given a single tensor of activations A of the shape [ batch size , length , dmodel ] – e.g. , coming from embedding the tokens in a sentence into vectors . To build Q , K and V from A , the Transformer uses 3 different linear layers projecting A into Q , K and V with different parameters . For models with LSH attention , we want queries and keys ( Q and K ) to be identical . This is easily achieved by using the same linear layer to go from A to Q and K , and a separate one for V. We call a model that behaves like this a shared-QK Transformer . It turns out that sharing QK does not affect the performance of Transformer , even if we additionally normalize the length of the keys K , as we show in the experimental Section 5 . Hashing attention . For the LSH attention , we start with two tensors , Q=K and V of the shape [ batch size , length , dmodel ] . We keep the multi-head mechanism intact and focus on the attention computation from Equation 1 . As already mentioned , the main issue is the term QKT , which has the shape [ batch size , length , length ] . But note that we are actually only interested in softmax ( QKT ) . Since softmax is dominated by the largest elements , for each query qi we only need to focus on the keys in K that are closest to qi . For example , if K is of length 64K , for each qi we could only consider a small subset of , say , the 32 or 64 closest keys . That is much more efficient , but how can we find the nearest neighbors among the keys ? Locality sensitive hashing . The problem of finding nearest neighbors quickly in high-dimensional spaces can be solved by locality-sensitive hashing ( LSH ) . A hashing scheme that assigns each vector x to a hash h ( x ) is called locality-sensitive if nearby vectors get the same hash with high probability and distant ones do not . In our case , we actually only require that nearby vectors get the same hash with high probability and that hash-buckets are of similar size with high probability . We achieve this by employing random projections as follows ( see Figure 1 ) . To get b hashes , we first fix a random matrix R of size [ dk , b/2 ] . We then define h ( x ) = argmax ( [ xR ; −xR ] ) where [ u ; v ] denotes the concatenation of two vectors . This method is a known LSH scheme ( Andoni et al. , 2015 ) and is easy to implement and apply to batches of vectors . LSH attention . Knowing our LSH scheme and the general idea of hashing attention , we will now formalize the LSH attention we use in this paper . We first rewrite the equation for normal attention , ( 1 ) , for a single query position i at a time : oi = ∑ j∈Pi exp ( qi · kj − z ( i , Pi ) ) vj where Pi = { j : i ≥ j } ( 2 ) We introduce the notation Pi to represent the set that the query at position i attends to , and z to denote the partition function ( i.e . the normalizing term in the softmax ) . For clarity , we also omit scaling by √ dk . For batching purposes we typically perform attention over a larger set P̃i = { 0 , 1 , . . . , l } ⊇ Pi while masking out elements not in Pi : oi = ∑ j∈P̃i exp ( qi · kj −m ( j , Pi ) − z ( i , Pi ) ) vj where m ( j , Pi ) = { ∞ if j /∈ Pi 0 otherwise ( 3 ) Now we turn to LSH attention , which we can think of in terms of restricting the set Pi of target items a query position i can attend to , by only allowing attention within a single hash bucket . Pi = { j : h ( qi ) = h ( kj ) } ( 4 ) Figure 2 ( a-b ) shows a schematic comparison of full-attention with a hashed variant . Part ( a ) depicts that the attention matrix for full attention is typically sparse , but the computation does not take advantage of this sparsity . In ( b ) , the queries and keys have been sorted according to their hash bucket . Since similar items fall in the same bucket with high probability , the full attention pattern can be approximated by only allowing attention within each bucket . Hash buckets in this formulation tend to be uneven in size , which makes it difficult to batch across buckets . Moreover , the number of queries and the number of keys within a bucket may be unequal – in fact , it is possible for a bucket to contain many queries but no keys . To alleviate these issues , we first ensure that h ( kj ) = h ( qj ) by setting kj = qj ‖qj‖ . Next , we sort the queries by bucket number and , within each bucket , by sequence position ; this defines a permutation where i 7→ si after sorting . In the sorted attention matrix , pairs from the same bucket will cluster near the diagonal ( as depicted in Figure 2c ) . We can follow a batching approach where chunks of m consecutive queries ( after sorting ) attend to each other , and one chunk back ( Figure 2d ) . Following our earlier notation , this corresponds to setting : P̃i = { j : ⌊ si m ⌋ − 1 ≤ ⌊sj m ⌋ ≤ ⌊ si m ⌋ } ( 5 ) If maxi |Pi| < m , then Pi ⊆ P̃i . In practice we set m = 2lnbuckets ( where l is the sequence length ) . The average bucket size is lnbuckets , and we assume that the probability of a bucket growing to twice that size is sufficiently low . The overall process of LSH attention is summarized in Figure 2 . Multi-round LSH attention . With hashing , there is always a small probability that similar items nevertheless fall in different buckets . This probability can be reduced by doing multiple rounds of hashing with nrounds distinct hash functions { h ( 1 ) , h ( 2 ) , . . . } , such that : Pi = nrounds⋃ r=1 P ( r ) i where P ( r ) i = { j : h ( r ) ( qi ) = h ( r ) ( qj ) } ( 6 ) The multi-round case essentially involves performing LSH attention nrounds times in parallel ; the details of the procedure are described in in Appendix A. Causal masking for shared-QK attention . In a Transformer decoder , masking ( denoted by m ( j , Pi ) in Equation 3 ) is used to prevent positions from attending into the future . To implement masking in LSH attention , we associate every query/key vector with a position index , re-order the position indices using the same permutations used to sort the query/key vectors , and then use a comparison operation to compute the mask . While attention to the future is not allowed , typical implementations of the Transformer do allow a position to attend to itself . Such behavior is undesirable in a shared-QK formulation because the dot-product of a query vector with itself will almost always be greater than the dot product of a query vector with a vector at another position . We therefore modify the masking to forbid a token from attending to itself , except in situations where a token has no other valid attention targets ( e.g . the first token in a sequence ) .

Summary: This manuscript presents a number of algorithmic techniques to reduce the computational and space complexity of Transformer, a powerful and very popular deep learning model for natural language processing (NLP). Although Transformer has revolutionized the field of NLP, many small groups cannot make a full use of it due to lack of necessary computational resources. As such, it is very important to improve the space and computational complexity of this popular deep model. The techniques presented in this manuscript seem to be very reasonable and the experimental results also indicate that they are effective. My major concern is that the authors shall present more detailed experimental results. In addition to bits per dim, it will also better if the authors can evaluate the performance in terms of other metrics.