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), 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. 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 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 , 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 research 1 . Many large Transformer models can only realistically be trained in large industrial research laboratories and such models trained with model parallelism cannot 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 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 a single-layer model due to the fact that activations need to be stored for back-propagation. • Since the depth d f f of intermediate feed-forward layers is often much larger than the depth d model of attention activations, it accounts for a large fraction of memory use. • Attention on sequences of length L is O(L 2 ) in both computational and memory complexity, so even for a single sequence of 64K tokens can exhaust accelerator memory. We introduce the Reformer model which solves these problems using the following techniques: 1 https://hackingsemantics.xyz/2019/leaderboards/ • 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 d f f factor and saves memory inside feed-forward layers. • Approximate attention computation based on locality-sensitive hashing replaces the O(L 2 ) factor in attention layers with O(L) 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. 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, richly-parameterized 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 to other domains like time-series forecasting, music, image and video generation.
Efficient Transformer with locality-sensitive hashing and reversible layers
{"src": "scitldr-AIC_train", "paper_id": "rkgNKkHtvB"}