Reformer: The Efficient Transformer
Authors: Nikita Kitaev
Lukasz Kaiser
Anselm Levskaya
Venue:  (2020)
Keywords: attention
locality sensitive hashing
reversible layers

Abstract: Large Transformer models routinely achieve state-of-the-art results on
a number of tasks but training these models can be prohibitively costly,
especially on long sequences. We introduce two techniques to improve
the efficiency of Transformers. For one, we replace dot-product attention
by one that uses locality-sensitive hashing, changing its complexity
from O($L^2$) to O($L \log L$), where $L$ is the length of the sequence.
Furthermore, we use reversible residual layers instead of the standard
residuals, which allows storing activations only once in the training
process instead of N times, where N is the number of layers.
The resulting model, the Reformer, performs on par with Transformer models
while being much more memory-efficient and much faster on long sequences.

Published as a conference paper at ICLR 2020
REFORMER: THE EFFICIENT TRANSFORMER
Nikita Kitaev U.C. Berkeley & Google Research k kitaevêcs berkelerkyyy... edu
Lukasz Kaiser"
Anselm Levskaya
Google Research
Google Research
( Lewskayaj@@oog com
ABSTRACT
Large Transformer models routinely achieve statc-of-the-art results on number tasks but training these models can be prohibitively costly, especially on long sequences. We introduce ers. For one, we replace dot-product attention by one that uses locality-sensine hashing, changing its complexity from O(I2) to O(L L), where L is the length of the sequence. Furthermore, we use reversible residual layers instead of the standard residuals, which allows storing actívations only once in the training pro- cess instead of N times, where N is the number of layers. The resulting model, the Reformer, performs on par with Transformer models while being much more memory-efficien and much faster on long sequences.
techniques to improve the efficiency of Transform-
INTRODUCTION
The Transformer architecture (Vaswani al. 2017) is widely used in natural language processing and yields state-of-the ta-thh--atttt results on number of tasks. To obtain these results, researchers have resorted to training ever larger Transformer models. The number parameters exceeds 0.5B per layer in the largest config configuration reported in (Shazeer et al. 2018) while the number of layers goes up to 64 (Al-Rfou et al. 2018). Transformer models are also used on increasingly long sequences. Up 11 thousand tokens of text in single example were processed in (Liu et al. (018)) 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-scquence models yield great results but strain resources to the point where some argue that this trend is breaking NLP researct Many only realistically be trained in large industrial research laboratories
large Transformer models and such models trained with model parallelism cannot even be fine-tuned on a single GPU as their memory requirements demand multi-accelerator hardware setup even for single training step.
Do large Transformer models fundamentally require such huge resources or are they simply ineffi- cient? Consider the following calculation: the 0.5B parameters used in the largest reported Trans- former layer account for 2GB of memory. Activations for 64K tokens with embedding size 1024 and batch size 8 account for 64K x 1K 8 0.5B floats, requiring another 2GB of memory. our memory use was only per-layer, then we should fairly easily fit large Transformer even on sequences of length 64K on single accelerator. Further, the whole corpus used to train BERT only requires 17GB to store. Why is it then that we cannot 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 single-layer model due the fact that activations need to be stored for back-propagation.
Since the depth dys of intermediate feed-forward layers is often much larger than the depth dyoode of attention activations, accounts for large fraction of memory use.
Attention on sequences of length L is O(L in both computational and memory complex- ity, so even for single sequence of 64K tokens can exhaust accelerator memory.
"Equal Contribution
https:://hackiiggeemmaa wwt.sssx.....9.0

---

Published as a conference paper at ICLR 2020
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 single copy of activations in the whole model, so the N factor disappears.
Splitting activations inside feedwrd layers and processing them in chunks removes the d factor and saves memory inside feed-foree forward layers.
Approximate attention computation based locality -eedstvvee hashing replaces the O(L2) factor in attention layers with O(L.ogg L) and so allows operating on long sequences.
We study these techniques and show that they have negligible impact on the training process com- pared 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 negligible effect training in all configurations we experimented with. Finally, locality-sensitiv hashing in attention is a more major change that can influence the training dynamics, depending the number of concurrent hashes used. We study this parameter and find a value which both efficient to use and yields results very close to full aftention.
We experiment on 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 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.
LOCALITY-SENSITIV 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 de, and values of dimension dy The dot products of the query with all keys are computed, scaled by ddg, 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:
QKT
Attention(Q. K, K) softmax
V VaL
Multi-head attention. In the Transformer, instead of performing a single attention function with lnnodel-dimensionnn keys, values and queries, one linearly projects the queries, keys and values times with different, learned linear projections to d2, de and du dimensions, respectively. Attention applied each of these projected versions of queries, keys and values in parallel, yielding d,- dimensional output values. These are concatenated and once again projected, resulting in the final values. This mechanism 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 [ Let us assume that Q, K and all have the shape [batch_size, length, dmodee]. The main issue is the term QKT, which has the shape Ibatch_.sie,, length, length]. In the experimental section we train a model on sequences of length 64K in this case, take 16GB of memory. This is impractical and has hindered the use of the Transformer for long sequences. But it is important note that the QKT matrix does not need be fully materialized in memory. The attention can indeed be computed for each query @ separately, only calculating the backward pass when needed for gradients. This way of computing attention may be less efficient but it only uses memory propor- tional to length. We use this memory-effictee implementation of attention run the full-attention baselines presented in the experimental section.
at batch-size of 1, this is 64K x 64K matrix, which in 32-bit floats would
vdn " softmax "KK"" V once in memory, and then re-computing
Where do Q, K, V come from? The multi-head attention deseribed above operates on keys, queries and values, but usually we are only given a single tensor of activations of the shape (batch size, length, dmodet] e.g., coming from embedding the tokens in sentence into vectors.
2

---

Published as a conference paper at ICLR 2020
Sphare Projected Points
Random Rotation 0
Random Rotation
Random Rotatioo 2
X: 0
B,
G
2
y: 320
X:
y: 021
Figure 1: An angular locality sensitive hash uses random rotations of spherically projected points to establish buckets by an argmax over signed axes projections. In this highly simplified 2D depiction, two points and are unlikely to share the same hash buckets (above) for the three different angular hashes unless their spherical projections are close to one another (below).
To build Q, K and V from A, the Transformer uses 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 We call model that behaves like this a shared-QK Transformer. 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, [batch_size, length, dmooee]. We keep the multi-head mechanism intact and focus on the atten- tion computation from Equation [] As already mentioned, the main issue is the term QKT which has the shape ibatch size, length, length]. But note that we are actually only interested in softmax(Q (T).. Since softmax is dominated by the largest elements, for each query q we only need to focus on the keys in that are closest to For example, if K is of length 64K, for each @ could only consider small subset of, say, the 32 64 closest keys. That is much more efficient, but how can we find the nearest neighbors among the keys?
start with two tensors, Q-K and V of the shape
Locality sensitive hashing. The problem of finding nearest neighbors quickly in high-dimensional spaces can be solved by locality-sensiine sensitive hashing (LSH). A hashing scheme that assigns each vector a hash h(z) is called locality-sensiinn 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 follows (see Figure D To get hashes, we first fix random matrix R of size [d,, b/2]. We then define ((x))= = arg max([= where Iu: ul 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, [I for single query position i time:
0: = Σ
exp (g: ky - P.))
where P. = (j:121)
JEP,
We introduce the notation P; to represent the set that the query at position í attends to, and & to denote the partition function (i.e. the normalizing term in the softmax). For clarity, also omit scaling by Vdg.
For batching purposes we typically perform attention over larger set P. = {0, 1, ...,, 2 P: while masking out elements not P::
0g = Σ exp (a: kg - -m(j. P.) P.)) where m(j, P.)
80 if
(3)
otherwise
JEP
3

---

Published as a conference paper at ICLR 2020
9:
9,
%
9
9 9.
9, 9,
9
Sequance que koskkeys
buevating
Sort by LSH bucket
Chunk sorted paralletize
(el Kokmmo
ON Buriorted
9
9, 9. 9
Actend within some bucket in chunk previcus
(61 Chunked
Figure 2: Simplified depiction of LSH Attention showing the hash-bucketing, sorting, and chunking steps and the resulting causal attentions. (a-d) Attention matrices for these varieties of attention.
Now we turn to LSH attention, which we can think of in terms of restricting the set P, of target items a query position can attend to, by only allowing attention within a single hash bucket.
P: = {j (g:)) = h(ky))
(4)
Figure (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 be uneven size, which makes it difficult batch across buckets. Moreover, the number of queries and the number of keys within a bucket may be unequal- fact, it is possible for bucket to contain many queries but no keys. To alleviate these issues, we first ensure that h(k;,) h(g;) by setting kj- T Next, we sort the queries by bucket number and, within each bucket, by sequence position; this defines permutation where r 8; after sorting. In the sorted attention matrix, pairs from the same bucket will cluster near the diagonal (as depicted in Figure 22)) We can follow batching approach where chunks of consecutive queries (after sorting) attend to each other, and one chunk back (Figure 2t). Following our earlier notation, this corresponds to setting:
4
P:
8j
m
m
m
(where is the sequence length). The average bucket size is Pbbackets' and assume that the probability of a bucket growing twice that size is sufficiently low. The overall process of LSH attention is summarized in Figure
max; [P&l < m, then P. S P.. In practice we set m
Multi-round LSH attention. With hashing, there is always a small probability that similar items nevertheless fall in different buckets. This probability be reduced by doing multiple rounds of hashing with Rrounds distinct hash functions (1), (2(2) ....), such that:
P. = U p(r)
Thoands
where p())= =(:)){j;00) - h))
=
(6)
The multi-round essentially involves performing LSH attention Nyounds times in parallel; the details of the procedure are described in in Appendix A
Causal masking for shared-QK attention. In Transformer decoder, masking (denoted by m(j. Pa) in Equation 3) is used to prevent positions from attending into the future. To implement masking in LSH attention, associate every query/key vector with position index, re-order the position indices using the same permutations used to sort the query/key vectors, and then use a comparison operation compute the mask.
4

---

Published as a conference paper at ICLR 2020
Table 1: Memory and time complexity of attention variants. We write for length, for batch size, ny for the number of heads, ne for the number of LSH chunks, 7, for the number of hash repetitions.
Attention Type
Memory Complexity
Time Complexity
Scaled Dot-Product Mefory-Efficient LSH Attention
max brulde, e bnal?)
max by ddk. br (2)
max (bmgldys brg max (bn Inn, bnk (4l (n)
max (bn ldy, bn max(bn, nn
Table 2: Accuracies the duplication task of 1-layer Transformer model with full attention and with lecality-sensitive hashing attention using different number of parallel hashes.
Eval
Full Attention
LSH-8
LSH-4
LSH-2
LSH-1
Train
Full Attention
100%
94.8%
92.5%
76.9%
52.5%
LSH-4
0.8%
100%
99.9%
99.4%
91.9%
LSH-2
0.8%
100%
99.9%
98.1%
86.8%
LSH-1
0.8%
99.9%
99.6%
94.8%
77.9%
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 shared-QK formulation because the dot-product of query vector with itself will almost always be greater than the dot product of query vector with a vector at another position. We therefore modify the masking to forbid token from attending itself, except situations where token has no other valid attention targets (e.g. the first token in a sequence).
2.1 ANALYSIS ON A SYNTHETIC TASK
To verify the performance of LSH attention and study its behavior, we start with the following synthetic task: duplicate sequence of symbols. In this task, each training and testing example has the form OwOw where w E {1....,N]*
sequence symbols ranging from to (we use N 127 in experiments). An example with the word of length 3 is given below.
Example:
0
19
113
72
19
113
72
To study LSH attention, train language model on examples of the above form where each is of length 511 (so the whole input OwOw is of length 1024). As this is a language modeling task, always predict the next symbol given all the previous ones, but we mask the loss and accuracy to only consider positions in the second half of the input, i.c., those that can actually be predicted.
The above task can be solved perfectly (to accuracy 100% and loss 0) by I-layer Transformer model. Note though, that it requires non-local attention lookups, so it cannot be solved by any model relying on sparse attention with a limited span. To make easy and fast train but similar
to models used in NLP, we use 1-layer Transformer with dmodel = dys 256 and 4 heads. We train it for 150K steps different settings: with full attention, LSH attention with Brounds 1,
Nrounds 2 and Trounds 4.
From the results summarized Table see that model trained with full attention can be imme- diately used with LSH attention, but at some loss of accuracy. When trained from scratch with LSH attention, the model trained with 4 hashes achieves almost perfect accuracy as well. Interestingly, the accuracy becomes perfect when evaluated with 8 hashes. It goes down when evaluated with 2 or hashes. Models trained with less hashes show worse results but even the model trained with just hash performs almost perfectly when evaluated with 8 hashes.
5

---

Published as a conference paper at ICLR 2020
3 REVERSIBLE TRANSFORMER
As the above section shows, the complexity of attention can be reduced from square in length to linear, provided an approximation is acceptable. But it is clear from Table [ that each field starts with a b- n4, . term: the b n,,, I dg. or alternatively b dmadel cost cannot be avoided. Indeed, the activations before each layer are already of the size b-l dmodet, so the memory use of the whole model with ny layers is least b- dmoddd ni, Even worse: inside the
layers of Transformer this goes up to b. 1 dys Ty. In a big Transformer it is usual to set dsf = 4K and ny 16 so with = 64K this again would an impractical 16GB of memory
In this section, we show how to reduce this cost by first dealing with the ny part of the term using reversible layers and then showing how chunking can allow handle the dyf problem. The effects of each of these approaches on memory and time complexity are summarized in Table
RevNets. Reversible residual networks were introduced by Gomez et al.) (2017) where it was shown that they can replace ResNets for image classification. The main idea is to allow the activations at any given layer to be recovered from the activations the following layer, using only the model parameters. Rather than having to checkpoint intermediate values for use in the backward pass, layers can be reversed one-by-one as back propaaraaagttion proceeds from the output of the network to its input. Whereas normal residual layer performs a function x + that operates on a single input and produces a single output and has the form y = x + F(x), a reversible layer works on pairs of inputs/outputs: (21, 12), and follows the equations:
yl x F(22)
y/2 x2 + G(yi)
A layer can be reversed by subtracting (rather than adding) the residuals:
T2 = 3/2 G(y11)
1 = y1 F(22)
(8)
Reversible Transformer. We apply the RevNet idea to the Transformer by combining the attention and feed-forward layers inside the revnet block. In the notation above, F becomes attention layer while G becomes the feed-forward layer. Note that Layer Normalization (Ba et al. 2016) is moved inside the residual blocks.
X1 + Autention(X., (X2)
1/2 = X2 + FeedForward(Y.))
The reversible Transformer does not need to store activations in each layer and so gets rid of the ny term. In Section [ we show that it performs the same as the normal Transformer when using the same number of parameters; we achieve this by having both 21 and x2 have size dmodel-
Chunking. While reversibility covers the ny term, the thicker layers still use lot of memory- The ferd-forward layer in particular can use intermediate vectors of dimensionality djf = 4K or higher. However, computations in feed-forward layers are completely independent across positions in sequence, so the computation can be split into c chunks:
Y2
Y
(c)
[x FeedForward(( ()::: : FeedForward((((0) (10)
This layer is typically batched by performing operations for all positions in parallel, but operating on one chunk at a time can reduce memory. The reverse computation in (8) and the backward pass are also chunked. In addition the feed-forward layers, for models with large vocabulary (more than dmodet word types) we also chunk the log-probabilities at the output and calculate the loss for sections of the sequence at time.
Chunking, large batches and parameter reuse. With chunking and reversible layers the memory we use for activations the whole network is independent the number of layers. The same is not true for parameters though as their number grows with the number of layers. This problem is remedied though because we can swap layer parameters to and from CPU memory when this layer is not computing. In a standard Transformer this would be inefficient because memory transfer to CPU is slow. The batch size multiplied by length in Reformer is much larger though and therefore the amount of compute done with the parameters amortizes the cost of their transfer.

---

Published as a conference paper at ICLR 2020
Table Memory and time complexity of Transformer variants. We write dmodet and d for model depth and assume d 2 dmodel; b stands for batch size, / for length, ny for the number of layers. We assume n, = 1/32 so 41/n.- = 128 and we write e = 128?.
Model Type
Memory Complexity
Time Complexity
Transformer
max(bldy (bldss, b ymu max (bld bmy
(bldgs + tral))yn
Reversible Transformer
(bmnldgs + bng 1 P)n
Chunked Reversible Transformer
max bld model br,(?)
(bngldss + bn, )
LSH Transformer Reformer
max (bldss f,, dn, C)m max moded, bn,
(bldgs + nc)m, (bldgs + n,)
RELATED WORK
The Transformer model introduced in (Vaswani et al., 2017) has been used widely in natural lan- guage tasks and further extended to model diverse data such as music scores (Huang et al. 2018)), and images (Parmar et al. 20188] |Ramachandran et al. 2019). Most notably, this model class has been applied successfully in the self-supervised training of extremely large language models (Devlin ct al. 2018. Radford et al., 2019)..
Given the enormous computational requirements of state of the art sequence models, there has been increasing interest finding methods to reduce the memory footprint and computational require- ments of Transformer models. In addition to standard methods such as precision reduction and gradient checkpointing (Sohoni et al.] 2209)), more efficient versions the Transformer model's self-attention mechanism (Sukhbaatar et al. 2019a b) have also recently been explored.
In particular, leveraging sparsity the attention layers has proved fruitful. OpenAl introduced the sparse Transformer (Child et al.) 2019) which exploits factorized sparse representation of atten- tion. Using product-key attention to increase the key space has also been used to reduce memory requirements in the feed-forward layers with no loss in performance (Lample et al. 2019).
Locality-sensitive hashing (LSH) has, to our knowledge, not been directly applied to Transformer attention layers before. But previous work using external memory with neural networks has dealt with memories of large sizes. The original implementation of memory networks (Weston et al.] 2014) and later work on scaling it (Bordes al. 2015 Chandar et al. (2166) used memory with size in the millions. The cost of doing so is that the memory must be fixed prior to training. Moreover, since during the beginning training the model unlikely to query the memory correctly, strong supervision is used to encourage the model to query memory locations that are useful. These hints are either given as additional supervising information by the task or determined heuristically Hill et al. (2015). The requirement that the memory be fixed before has been removed in [aantor et al. (2016) at the cost of memory size and later alleviated by Rae et al. (2016). The last paper considered memory lookups with approximate nearest neighbors including both LSH and random kd-trees, but only for lookups in external memory.
in
5 EXPERIMENTS
In this section we present experimental results demonstrating the techniques described above. We analyze the techniques one-by-one to make clear which combinations have impact on performance. We start by showing that reversible layers and shared query-key spaces do not impact performance, then proceed to analyze hashing attention and finally the full Reformer model.
We ran our experiments on the imagenet64 and enwik8-64K tasks, where the latter is variant of enwik8 that is chunked into subsequences of 216 = 64K tokens. We use 3-layer models for our ablations so
make it tractable to compare with the regular Transformer, which has high memory usage and performs full O(l2) attention. All experiments have dmmodet = 1024, d = 4096, Theads 8, and total batch size of 8 sequences. We used the Adafactor optimizer (Shazeer & Stern 2018) for training these models. We also evaluate on the WMT 2014 English-to- -orermmnn translation task, following the hyperparameters of Vaswani et al. (017)). Training for all experiments
7

---

Published as a conference paper at ICLR 2020
Sharing Query-Key enwik8
Reversibility
7.2
7.2
separace q
regular
3.8
5.9
bpp
2.9
1.4
1.4
20K
Sharing Query-Key Space imagenet61
Reversibility imagenet64
separate q shared
regular reversible
4.0
4.0
3.5
3.6
3.5
50K
100K
150K
200K
50K
100K
150K
200K
steps
Figure 3: Effect of shared query-key space (left) and reversibility (right) on performance on enwik8 and imagenet64 training. The curves show bits per dim held-out data.
Table BLEU scores on newstest2014 for WMT English- German (En-De). We additionally report detokenized BLEU scores as computed by sacreBLEU (Post 2018).
sacreBLEU Uncased
Model
BLEU
Cased
Vaswani et al. (201)) , base model
27.3
Vaswani al. (201)) big
28,4
Ott et al. 2018 9, big
29.3
Reversible Transformer (base, 100K steps)
27.6
27.4
26.9
Reversible Transformer (base. 500K steps, no weight sharing)
28.0
27.9
27.4
Reversible Transformer (big, 300K steps, no weight sharing)
29.1
28.9
28.4
was parallelized across 8 devices (8 GPUs 8 TPU v3 cores). Code for training models is made publicly available
Effeet of sharing QK. We first consider the effect of shared-QK attention regular Transformer 4i model. Shared-Qu QK attention sets hn = ||:; || and prevents tokens from attending to themselves (except when no other context is available). In the left part of Figure 3 we plot perplexity curves for both regular and shared-QK attention. shared query-key space does not perform worse than regular attention; in fact, for enwik8 it appears to train slightly faster. In other words, we are not sacrificing accuracy by switching to shared-QK attention.
Effect of reversible layers. In the two plots on the right in Figure we compare regular Trans- former per Vaswani al. (2017) with the reversible one describe in Section The two models have identical parameter counts, and the learning curves likewise appear to be nearly the same. These results show that the memory savings the reversible Transformer do not come at the expense of accuracy.
Reversible layers in machine translation. also evaluate reversible layers context of an mecoder-decoder Transformer model for machine translation from English to German. We start by making both the encoder and the decoder fully reversible in the Transformer-buse architecture, and
https: http::/gggtttttt//eteee httpp////ttttatoooooo/////////es
aLBU+Case 1c+ le+lang.eeee www.laag..eeeeeeeeee detnumrefs. 1+mmooth. htt.teesst.ttt//uulleeeee int.e
sr.eeeoaee xed+lang. an-deenumre 1+mooohh. www.ttsssettttttttt..o... wmt 1//fo1 Teook. +
8

---

Published as a conference paper at ICLR 2020
LSH Attention on Imagenet64
attention
hashes
hashes
hashes
16 hashes
4.0
bpd
3.9
3.7
20%
80K
100K
steps
Figure 4: LSH attention performance as function of hashing rounds on imagenet64.
Performance Scaling with enwik
Attertion Dependence Sequence Synthetic
attendion
lyeers
55-b 1ush
12 layers
t5 2hhoh
23-
16 ktyers
-5+0 hash
layirs
Biiseh
2.0
1.0
14.
12
56 3X
10K
12E
13C
17C
20C
220
23K
102432
2042/18
4020-8
$192
363542
32768/
steps
sequenice length/ botch
Figure 5: Left: LSH attention performance as a function of number of layers on enwik8. Right: Speed of attention evaluation as a function of input length for full- and LSH- attention.
see that the resulting model performs comparably to Vaswani et al. (2017) when trained for 100K steps. We also evaluate training for greater number of steps and with larger model. Reformer models are very memory-efficient, so for the latter two experiments we do not need to save mem- ory by sharing embedding and output projection weight matrices throughout the model. Results are shown in Table 4 We do not apply LSH attention in this setting because examples are single sentences, and sentences tend to be relatively short. Our typical LSH attention configuration uses chunks of 128 tokens after hashing and sorting, whereas the examples in the WMT14 4 test set are all shorter than 128 tokens.
LSH attention Transformer. LSH attention is an approximation for full attention that, as evi- denced in Figure 4 becomes more accurate as the number hashes increases. At Brounds already almost matches full attention. The computational cost of a model grows with the number of hashes, this hyperparameter can be adjusted depending on the available compute budget. Ad- ditionally, as in Table 2, the number of hashes can be increased at evaluation time to produce more accurate results. On the right half of Figure 3 we plot the speed of different attention types vs. the sequence length, while holding the total number of tokens fixed. We see that while regular attention becomes slower at longer sequence length, LSH attention speed remains flat.
Large Reformer models. To verify that the Reformer can indeed fit large models single core and train fast on long sequences, we train up to 20-layer big Reformers on enwik8 and imagenet64. As can be seen in Figure 5 these models fit into memory and train. We were not able train Trans- former baselines in this case as they are too slow and memory memory-hungry, but we see clear improvement with the number of layers. A 12-layer model on enwik8 trained for 20K steps with dropout rate of 0.1 achieves 1.19 bits/dim on the test set. We also trained a 12-layer Reformer model for longer with further tuning and improvements and reached 1.05 bits/dim the enwiki8 test set.

---

Published as a conference paper at ICLR 2020
6 CONCLUSION
Reformer combines the modeling capacity of'a Transformer with an architecture that can be executed efficiently on long sequences and with small memory use even for models with a large number of layers. We believe that this will help large, wehly-parameterizzed Transformer models become more widespread and accessible. Also, the ability to handle long sequences opens the way for the use of the Reformer on many generative tasks. In addition to generating very long coherent text, the Reformer can bring the power of Transformer models other domains like time-series forecasting, music, image and video generation.
REFERENCES
Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. language modeling with deeper self-attention. CoRR, 0a/18888.4444, 2018. / //axxii.rr/aab///a.0 04444
Character-level
URL http:
Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, Ilya P. Razenshteyn, and Ludwig Schmidt. Practical and optimal LSH for angular distance. CoRR, 0os/1500.008880 2015. URL http://arxiv.
org/ abs/ 1509. 02897
Jimmy Lei Ba. Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 2016. http:tpp//arxx...rr wtt.o////g////..........
abs 1607. 06450
Antoine Bordes, Nicolas Usunier, Sumit Chopra, and Jason Weston. Large-scale simple question answering with memory networks. CoRR, abs/ 1506002075 2015. uRtptpp://rx.i
abs/1506.0200 02075
Sarath Chandar, Sungjin Ahn, Hugo Larochelle, Pascal Vincent, Gerald Tesauro, and Yoshua Ben- gio. Hierarchical memory networks. arXiv preprint arXiv:1605.07427, 2016.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. URL
2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, fos//880.04805, 2018. URL
http://arxiv. wwwpxxi.org/abbs/100
Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger Grosse The reversible residual net- work: Backpropagation without storing activations. Advances in neural information processing systems, pp. 22 2214-2224, 2017.
Felix Hill, Antoine Bordes, Sumit Chopra, and Jason Weston. The goldilocks principle: Reading children's books with explicit memory representations. CoRR, abs/1551 2015. URL
hhtt:://arxiv: wwwpxii.orgg/ab/////30 .23001
Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Curtis Hawthorne, An- drew M Dai, Matthew D Hoffman, and Douglas Eck. Music transformer: Generating music with long-term structure. arXiv preprini 1rXiv:1809.0444 180.:04888 2018.
Guillaume Lample, Alexandre Sablayrolles, Marc' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. CoRR, abs/1907.05244 055242, 2019. URL http:
Peter J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. CoRR, abs/18011.10198, 2018.
URL http:://pxiiiii org/ abs/ abs/1801 100//181110 10198
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine trans- lation. In Proceedings of the Third Conference Machine Translation: Research Papers, pp. 1-9, Brussels, Belgium, October 2018. Association for Computational Linguistics. doi: 10. N0.1865/v1//11-6650
http:t www. acl www.://www.oor//tooooooooooo /w1 8- 6301
10

---

Published as a conference paper at ICLR 2020
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, and Alexander Ku. Image transformer. CoRR, abs/1802.05751, 2018. URL htttp://axxiv.orgg
05751
Matt Post. A call for clarity in reporting BLEU scores. In Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 186-191, Belgium, Brussels, October 2018, As-
sociation for Computational Linguistics. URL wwtpps::/ww..ttt. /tt.ttttss//ww.............eaa.a www.w//wwww....ttttoooo.........
018-6319
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Jack W Rae, Jonathan Hunt, Tim Harley, Ivo Danihelka, Andrew Senior, Greg Wayne, Alex Graves, and Timothy Lillierap. Scaling semmory-augmented neural networks with sparse reads and writes. In Advances in Neural Information Processing Systems, (NIPS), 2016.
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone aetf-attention vision models. CoRR, htts/1906.5900% 2019. URL http:
7/ 7/arxiv. wttpxx..wggaaa////111 09919 05909
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy P. Lillicrap. One- shot learning with memory-augmentedd neural networks. CoRR, abs/ 66660, 2016. URL
http:///rxiv> org/ 16065
Noam Shazeer and Mitchell Stem. Adafactor: Adaptive learning rates with sublinear memory cost. CoRR, abs/1804.042333 2018.
org/ abs /1044. 04235
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake Hecht- man. Mesh-tensorflow: Deep learning for supercomputers. CoRR, abs/181 .02084, 2018. URI.
http::/arxiv..og// /aaas///1l... 1821.
Nimit Sharad Sohoni, Christopher Richard Aberger, Megan Leszezynski, Jian Zhang, and Christo- pher Ré. Low-memory neural network training: technical report. CoRR, abs/ 190.19944411111 2019.
URL http:httpp//aaii...oorroo.e
9904.994400 10631
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive atten- tion span transformers. CoRR, abs/1905.07799, 2019a. URL http::/aaxiv..oet
1905 07799
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Hervé Jégou, and Armand Joulin. Aug- menting self-attention with persistent memory. CoRR, 0bb//1907.0140,, 2019b. URL http:
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, 2017. URL http:
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks, CoRR, abs/1410.3916, 2014.
URL http:://axxii....mm/// wwtp:/iiiorrg//oooo wttp///g////araa///e 3916
11

---

Published as a conference paper at ICLR 2020
A MULTI-ROUND LSH ATTENTION
In this section we describe in more detail the multi-hash version of our LSH attention mechanism. We first repeat Equation from the main text, which describes general formulation of attention with sparsity:
foo
P otherwise
Σ JEP.
exp (q: h2 mmj, P;) where m(j, P,)
0j
0
In the multi-round case, query position i can attend to key positions Pa as defined in 6 which also repeat here:
Mroando
P.
p(
where P!! - ()) (a))
6
For batching purposes, attention is performed chunks of sorted queries/keys:
(r)
p(:)
j
8;
8
(11)
m
m
Combining 3 and (6 gives:
0
Σ exp (q: hi míj, P:) - ((),
(12)
jEP:
Rrounds
Σ exp
((i,p(!))=((1
- 2(1, Σ
Nis exp (a: m(j p! =(1, p!:)))
JEP(
(13)
Arounda
Σ exp (x() z(i, =(i,Pi) Pi)) 6.")
(14)
o
(r)
Σ exp 4: 26.k5 kj m
(r) - 2(2, p(:))) 4s
(15)
(
x 10s log Ni.j
if $ p(()) if jj otherwise
P(
where No. [ [ and m
(16)
(r) that can be computed independently from other rounds, except for the inclusion of term Ni,, to avoid double-counting elements when constructing
Each round of L.SH attention produces vector o
the union of ((r) sets. our implementation we fold the N.,; factor into the masking term mn)r)
We also modify m introduce a special case for = j. This case is added because causal masking in standard Transformer allows position i to attend to itself, which is not desirable in shared-QK -QK formulation. We set the mask to a large but finite value to disallow attention-in-plas except in the situation where token has other valid attention targets. For example, the first token in sequence attends only to itself, because no prior context is available.
12


Reviews:
[{"review_id": "r1lXpbp3YS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents an attempt to reduce the memory complexity of Transformers. The authors call their model the Reformer. It presents a LSH based self-attention mechanism, along with reversible adaptation of Transformers. The Locality sensitive hashing scheme reduces complexity from L^2 to L which is pretty neat. \n\nTackling the quadratic complexity of self-attention is indeed an important and nice direction. I think the LSH based attention quite novel and is a natural solution to reducing the complexity of the self-attention module.  However, I think the technical description could be improved as the current form is quite confusing and difficult to parse.\n\nThe experiments are a little on the weaker side. Authors presented results on imagenet, enwiki and a synthetic task. I am mainly concerned if the Reformer works on tasks such as machine translation or other NLP tasks. The paper does not present much evidence that the effectiveness of LSH is broad and versatile.\n\nMy current vote is a weak accept, based on some preliminary understanding and the general novelty of the idea. \n\nI do have some questions/issues/comments:\n\n1) Given that there is some form of QK sorting, how is it possible to mask the future? Is this because tokens are sorted within buckets?\n2) Can the authors clarify what \"Causal masking on the Transformer is typically implemented to allow a position i to attend to itself.\" mean?\n3) I'm a little confused about how the sorting is being done. Can this be done in an end-to-end differentiable manner?\n4) Can the authors present some results on other tasks? While neat, I think other tasks (e.g., MT or QA) can be investigated to further ascertain that the LSH attention works well. Current experimental results are not too convincing.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "Syg6AbTp5H", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. ", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "H1e_gXhRYS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a method to make Transformer models more efficient in time and memory. The proposed approach consists mainly of three main operations: \n- Using reversible layers (inspired from RevNets) in order to prevent the need of storing the activations of all layers to be reused for back propagation; \n- Using locality sensitive hashing to approximate the costly softmax(QK^T) computation in the full dot-product attention;\n- Chunking the feed-forward layers computations to reduce their cost.\nThis approach is first applied to a toy dataset to analyze its complexity, then tested on enwik8 language modelling task and imagenet-64 image generation task for ablation study and performance assessment. \n\nThe problem approached by the paper is interesting and the proposed approach is novel to the best of my knowledge. The paper is well structured and clearly written a part from some small typos (see minor comments below).   \n\nWhile the analysis of complexity is sound and convincing, and the fact of being able to train larger Reformers is very interesting, I have some questions and concerns about the approach and experiments. \n- Effect of reversible layers: It is clear for the experiment of Imagenet64 that the effect is negligible, but the experiment on enwik8 in the paper seems unfinished. Did the authors manage to finish the training, and does it confirm the observation? \n- Sharing QK: I am a bit confused about the effect and usefulness of this operation. Can the authors comment on why it is needed for LSH attention? It seems to me that the same operations can be achieved with different Q and K. Indeed, doing so, the authors slightly reduce the capacity of the model. The observed non-significantly decreased performance can be an effect of using only 3-layers. This may explain why the results reported for larger models in figure 5 show higher bpc than similar size state of the art models.\n- Time per iterations: Can the authors report the time per iteration for the larger hash rounds (8 and 16) that are closer to full attention? For the highest reported number (4), from a quick and not precise look at figure 4, it seems that the performance achieved by the proposed method after 140k iterations is achieved by the full attention after ~40k iterations. The gain in time per iteration for this particular number of hash rounds can be lost by the loss in performance. \n- Can the authors detail how they chose the hyperparameters of their approach? e.g. the size of hash buckets, the distribution used to generate the random matrix R .. \n- The reported results can be made stronger by reporting average/error bars across several trial to show consistency. \n\nMinor: typos: \nDimension of matrix R [d_k, d_b/2] -> [d_k, b/2]\nLast paragraph of page 6: state of these art -> state of the art\n\n——————————————— \nAfter rebuttal:\nI have read the authors answer, and found they addressed my concerns. I'm therefore increasing my score.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Transformer models have proven to be quite successful when applied to a variety of ML tasks such as NLP.  However, the computational and memory requirements can at times be prohibitive, such as when dealing with long sequences.  This paper proposes locality-sensitive hashing to reduce the sequence-length complexity, as well as reversible residual layers to reduce storage requirements.  Experimental results confirm that the performance of Transformer models can be preserved even with these new efficiencies in place, and hence, this paper will likely have significant impact within the community.  

Some relatively minor points notwithstanding, all reviewers voted for acceptance which is my recommendation as well.  Note that this paper was also vetted by several detailed external commenters.  In all cases the authors provided reasonable feedback, and the final revision of the work will surely be even stronger.
Decision: Accept (Talk)