Reproducibility Challenge: Reformer
Anonymous Author(s)
AfﬁliationAddress email
Reproducibility Summary
We attempt to reproduce the central claims of ICLR 2020 Paper "Reformer: The Efﬁcient Transformer" (Kitaevet al.
[2020]); that the techniques introduced enable performance on par with a traditional Transformer model whilebeing much more memory-efﬁcient and much faster on long sequences.
This community effort reproduced claimsaround speed for long sequences and observed a reduction in memory usage.
We could not match the performanceof a traditional Transformer with Reformer.
Finally, substantial coding effort was required, a lack of implementationdocumentation compounded this.
We provide all code documentation in our GitHub 1 .
Scope of reproducibility The scope of this work is to verify the claims of memory efﬁciency and speed on longersequences of the Reformer.
We replicated only the NLP experiments due to limited computational resources.
Methodology We ﬁrst reimplemented the original Transformer model (Vaswani et al.
[2017]) and which we thenmodiﬁed.
We referred to the authors’ code for the model and data pipeline.
We used the fastai library (Howard and
Gugger [2020]) for training, Weights and Biases (Biewald [2020]) for experiment tracking and nbdev (Howard [2019])for development.
All experiments were done in a single GPU setting.
Results Claims around speed on longer sequences and reduced memory footprint were validated; as sequence lengthincreased,
Locality Sensitive Hashing ("LSH")
Attention became faster and increasing the number of hashes improvedperformance.
We could not achieve the performance of a traditional Transformer with Reformer.
Some experimentswere not run for as long as in the paper due to a lack of computational resources.
Potentially the under-performanceof our Reformer may be due to under-training, implementation differences or nuances in JAX vs Pytorch.
Also,exploding gradients were encountered with mixed precision training and several model settings were found to beunstable depending on the random seed or learning rate.
What was easy Obtaining the data was straightforward as they are commonly used benchmarks.
There were noissues reproducing the data pipeline or Chunked Feed Forward layers and code for the Axial Positional Encodings wasimported.
What was difﬁcult Substantial effort was made to ensure a correct reimplementation.
It was challenging due to manyengineering design decisions or hyperparameters not being fully documented.
Signiﬁcant hyperparameter tuning was also needed.
Communication email correspondence
Introduction
One of the main characteristics of the Transformer (Vaswani et al.
[2017]) is a large number of parameters, and costlytraining (see e.g. Brown et al.
Finding ways of matching the performance of the Transformer, but reducingthe cost of training is, therefore, a democratizing effort.
The Reformer (Kitaev et al.
[2020], referred to later as the
Reformer paper) attempts to achieve this by reducing the memory footprint of the Transformer and thus enablingmodelling of deeper models and longer sequences within a ﬁxed resource budget.
Scope of Reproducibility
The Reformer introduces several techniques for reducing the memory footprint of the Transformer while retainingacceptable performance.
The Reformer paper claims that the Reformer “performs on par with Transformer modelswhile being much more memory-efﬁcient and much faster on long sequences”.
We break down this claim into thefollowing sub-claims, based on the experiments in the paper:
Proof of concept: A Transformer with LSH attention can solve a sequence copy task that requires non-localattention lookups.2.
Attention approximation: The performance of the Reformer increases with the number of hashing rounds andis close to full attention performance at 8 hashes.3.
Scaling: Deep Reformer models: a) can be trained on very long sequences using a single accelerator (GPU orTPU core) and b) adding layers improves performance.4.
The performance of Reversible Transformers are similar to the Transformer on language modelling tasks.
Shared queries and keys in the attention module has minimal effect on performance in a Transformer.6.
Evaluation speed: The evaluation speed of LSH attention grows with the number of hashes but not withincreased sequence length if the number of tokens is kept constant.
The Reformer reduces the memory footprint compared to the transformer.
We provide all code in an anonymous github repository.
We have built an additional documentation web page ofhow experiments were performed.
The documentation will be made available after the review period has ended.
Thenotebooks which the documentation is based on is available in the /nbs folder of the repo.
Finally, a pypi package withthe reformer implementation will also be released.
Methodology3.1 Reformer Techniques
The Reformer paper introduces ﬁve techniques which can be implemented separately or combined together into a full
Reformer model:
LSH attention is an approximation of full attention that is more memory and time efﬁcient for longer sequences.
Shared keys and queries in the attention module are introduced to make LSH work better by ensuring thatthe number of queries is equal to the number of keys within each bucket.3.
Reversible Residual Layers were ﬁrst presented in Gomez et al.
Adding Reversible Residual Layersenables us to perform backpropagation without storing the activations in memory.
Chunked feed-forward layers.
Processing large feed forward layers in smaller chunks.5.
Axial positional encodings are a more memory efﬁcient way of encoding positions than standard learnablepositional encodings.
For a more in-depth explanation of these techniques refer to our GitHub repository.
3.2 Datasets used
Synthetic dataset : A sequence of integers of the form 0w0w, where w is a sequence of integers from 1 to 128of some length.
enwik8 is a dump of English language Wikipedia from 2006 (Mahoney [2006]).
WMT 2014 English-German is a dataset of 4.5 million English and German sentence pairs (Bojar et al.[2014]).
Downloaded from Huggingface.
Tiny Shakespeare is a subselection of Shakespeare’s works, downloaded from Huggingface.
For all datasets, our code and documentation details the data pipeline.
3.3 Experimental setup and code
Our code and documentation with supplementary material such as full experiment notebooks are available on Github.
All experiments are logged to Weights & Biases.
All experimentation was carried out with single GPUs, unlike the
Reformer paper which used 8 x V100s.
The total number of GPU hours was approximately 1 600 hours.
Results
We attempt to verify each of the claims stated in section 2.
4.1 Language modelling with enwik8 - experiment settings
For all enwik8 language modelling experiments we train a 3-layer Transformer language model ("TransformerLM")with a model dimension of 1024 and axial positional encoding, modiﬁed with the relevant Reformer technique, on theenwik8 dataset with a sequence length of 4096, unless noted otherwise.
The Adafactor optimizer and a batch size of 8was used, via Gradient Accumulation.
Experiments were run on 15GB and 12GB GPUs and training was carried out infull-precision.
Experiment conﬁguration is detailed in our documentation.
4.2 LSH attention analysis on synthetic task
The synthetic task is a sequence copy task. The input sequences are of the form 0w0w, where w is n random integersbetween 1 and 128.
We use a language model to predict the next token in such sequences. The learning objective is tocopy the ﬁnal half of a sequence from the ﬁrst half.
See our documentation for details on the task.
Claim: A full attention Transformer can solve this task perfectly.
An LSH-attention Transformer can also solve it, butwith decreasing performance as the number of hashing rounds decrease.
A model trained with one type of attention canalso be evaluated with a different type of attention.
Experiments For this task we train a shared-QK TransformerLM and three LSH Language Models ("LSH-LM") with1,2 and 4 hashing rounds respectively.
We trained for 150 000 steps with a batch size of 64, a training size of 12 800,and a validation size of 1 280.
We used hyperparameters as described in the paper, and other defaults suggested in the
Trax Github repository.
4 See our SyntheticConﬁg class for experiment defaults.
We evaluated the models with fullattention and LSH attention with 8,4,2 and 1 hashing rounds respectively.
Results Results from a single run are summarized in table 1.
The ﬁrst column shows how the model was trained, thesubsequent ones how it was evaluated.
See our documentation for analysis setup.
[Table]
We can partly support the claim, but there are three clear differences when we compare our results to Table 2 of the
Reformer paper.
First of all, our results are worse than the Reformer paper.
However appendix A also shows thatmodels using shared-QK attention have large variation in training loss.
This suggests results are somewhat random,and depend on the speciﬁc training setup.
Secondly, models trained with LSH attention and validated with standardattention do much better in our experiments.
And thirdly, the model trained with standard attention and validated with
LSH attention does much worse in our experiment.
We have been unable to explain these differences.
4.3 Effect of sharing QK
Shared Query-Key ("shared-QK")
Attention was used in the Reformer paper as it is needed for LSH. It also helps tofurther reduce the memory footprint of the model.
Claim: A shared query-key space does not perform worse than regular attention
Experiments To investigate this we train a standard TransformerLM and a TransformerLM with shared-QK attentionon the enwik8 dataset.
The Bits-Per-Character (BPC) metric was used.
They were trained for 10 epochs and the meanvalidation BPC of 3 training rounds was taken, results are shown in ﬁgure 1.
Results After experimentation, we cannot validate the claim that shared-QK attention does not perform worse thanstandard attention for this experimental setting.
Nor did we see the effect of shared-QK attention training slightly fasteras noted in the Reformer paper.
Potentially with additional training shared-QK attention performance will convergewith standard attention, we leave this to future work.
4.4 Effect of reversible layers
[Figure]
Claim: Reversible Residual Layers in a Transformer enable more memory efﬁcient training and do not come at theexpense of model accuracy
To validate this claim we train a TransformerLM and full sequence-to-sequence Transformer ("ReversibleTransformer")using Reversible Residual Layers and compared them to their standard Transformer equivalents.
Reversible Language Model on enwik8 experiment To investigate this claim we trained on sequences of length4096, since training this model ("ReversibleLM") on full 64k sequences from enwik8 as per the Reformer paper wasnot feasible given our computational budget.
Figure 2 shows the mean of 3 runs for each model-type.
Results We could not validate the claim that Reversible Residual Layers do not have a signiﬁcant impact on languagemodel performance due to the sizeable difference of 0.11 mean validation BPC between the ReversibleLM and thebaseline TransformerLM.
Potentially this difference can be reduced with further training, we will leave this to furtherwork.
4.5 Reversible Transformer on translation task experiment
We train a full Reversible Transformer, a 12-layer sequence-to-sequence Transformer model with Reversible Residuallayers on the WMT-14 en-de translation dataset.
Given the short sequences in this dataset, a sequence length of 256 wasused with a batch size of 64 and the AdamW optimizer.
Gradient Accumulation was used when training with a 12GB
GPU.
Training was carried out for 2 epochs for the Reversible Transformer and 1 epoch for the standard Transformer.
Full-precision was used in both cases.
The one-cycle learning rate schedule from fast.ai was used, with a maximumlearning rate of 1e-4, initial div of 5, and a percent start of 12.5%
Results From our experiments, see ﬁgure 3, we could not not validate the claim that Reversible Residual Layers donot have a signiﬁcant impact on Transformer model performance.
Our validation CorpusBLEU score for a standard
Transformer trained for 1 epoch outperformed a ReversibleTransformer that was trained for 2 epochs.
4.6 Effect of number of hashing rounds on the performance
When using LSH for chunking there is a small probability that similar tokens will end up in different buckets, andtherefore attention will not be computed for these tokens.
To reduce this probability we can do multiple rounds ofhashing.
Claim: Performance of Transformer with LSH attention increases with the number of hashing rounds and is close tofull attention performance at n_hashes = 8 .
Experiment To validate this claim we train and compare models with different numbers of hashing rounds on theenwik8 dataset.
We train a LSH-LM and a stanard TransformerLM for 10 epochs.
Results Figure 4 shows training and validation loss for LSH-LM runs with a number of hashing rounds from 2 to16 compared to the TransformerLM.
Losses for one run of each setting are reported.
Our results verify the claim andillustrate that training and validation losses tend to improve with an increasing number of hashing rounds.
To be notedthat both time and memory requirements also increase with n_hashes .
Also while training losses are very close for larger n_hashes and full attention, models using LSH seem to have slightlyhigher generalization error.
4.7 LSH attention evaluation speed
Claim: The evaluation time for LSH attention increases with the number of hashes but not with increased sequencelength when the total numbers of tokens are kept constant.
Experiment To verify the claim we test evaluation time of the LSH-LM on the synthetic task with hyperparametersas indicated in the right part of Figure 5 of the Reformer paper.
We were unable to complete the longest sequencelengths for full attention due to memory limitations on a single GPU.
Results Our results support the claim and are summarized in ﬁgure 5.
Claims: 1) Deep Reformer models can be trained on very long sequences using single accelerator (GPU or TPU core);2) Adding layers to Reformer models improves performance.
Experiment To test these claims we train full ReformerLMs of increasing depth on sequences of length 16 384.
Weuse the enwik8 dataset.
Our experiments are designed to run on a single 12GB GPU.
We trained models with 3, 6 and12 layers for 4 epochs, training deeper models was not possible with our computational budget.
Results We conﬁrm claim (1), deeper models can be trained on a single accelerator as expected, given the factthat Reversible layers have O (1) memory complexity.
For claim (2), our results in ﬁgure 6 are indicative of betterperformance for deeper models, but longer training runs are needed to conﬁrm this. [Figure]
A central motivation of the Reformer paper is that it has a reduced memory footprint as compared to a standard
Transformer.
Claim: The Reformer reduces the memory footprint compared to the Transformer.
The claim is summarized in Table 5of the Reformer paper.
Experiments To investigate the claim we log the memory allocation of various reformer combinations for one trainingstep with the tiny Shakespeare dataset.
The sequence length for all experiments was set to 4096, with a batch size of1.
We investigate the following combinations: a) Comparing the TransformerLM, LSH-LM, ReversibleLM and full
ReformerLM; b) Comparing ReformerLMs with different number of hashes.
Comparing TransformerLM, LSH-LM, ReversibleLM and the full ReformerLM Figure 7 (left) shows thememory usage for the models compared.
We see that the Transformer stores activations for the forward pass duringtraining, and that memory is gradually released as the backward pass is completed.
In this case we have six distinctjumps in memory which correspond to its 6 layers.
With the LSH-LM, memory allocation is mostly parallel to thetransformer, but since the LSH attention computation is cheaper, the actual memory consumption is lower.
For the
ReversibleLM memory doesn’t accumulate over the forward passes, since it recalculates intermediate activations forthe backward pass.
Note that we observe 6 peaks for the forward pass, each peak corresponds to full attention matrixcomputation.
Another 6 higher peaks correspond to backward pass.
The full ReformerLM includes both LSH attentionand Reversible Residual Layers.
Memory allocation therefore doesn’t increase with number of layers, and since LSHattention is cheaper than standard attention, the peak memory usage is smaller than for the ReversibleLM.
Reformer Memory vs Number of Hashes
A 2-layer ReformerLM with a sequence length of 4096 was used for thisexperiment.
In section 4.6 we see that increasing the number of hashes leads to better performance.
But since we storethe result of each hashing round we expect memory to grow linearly with the number of hashes.
Note that it’s onlyduring the LSH attention calculation that the number of hashes matter.
I.e. the intermediate shape of LSH attentionis [batch_size, n_chunks, chunk_size, chunk_size*2] , where n_chunks is the product of the number ofhash buckets per round and the number of hash rounds, n_chunks = n_buckets * n_hashes .
The output of LSHattention is [bs, seq_len, d_model] and is independent of the number of hashes.
Figure 7 (right) conﬁrms thatpeak memory scales linearly with number of hashes.
The memory peak happens at the backward pass during calculationof LSH attention.
Appendix B shows a similar experiment that investigates how the number of layers affect memoryconsumption.
Results Our experiments have veriﬁed that the Reformer has a much smaller peak memory allocation as compared tothe Transformer.
We have also shown that the reformer can scale to much deeper models than the Transformer within aﬁxed budget.
The main bottleneck of the reformer w.r.t.
memory is the number of hash rounds used.
Discussion and Conclusion
We have validated a number of claims such as superior evaluation speed for longer sequence lengths, how LSH attentionis close to approximating full attention, that scaling deeper Reformer models works on a single GPU and that Reformerhas a reduced memory footprint compared to the Transformer.
However we cannot validate the claim that performance of a Transformer model with either Reversible Residual Layersor Shared-QK attention performs on par with a traditional Transformer.
There may be a number of explanations, ourlack of compute compared to the original authors being the most likely factor.
This reproducibility work involved a non-trivial amount of engineering to ensure our reimplementation was as faithfulto the original paper as we could make it.
The Reformer paper introduced ﬁve new techniques meaning that substantialeffort had to be made to ensure a correct reimplementation of each.
LSH attention is, for example, quite complexcompared to normal dot-product attention.
While there were Reformer code resources available, many design decisionsor hyperparameters were not fully documented in the paper.
Several of the experiments are also quite computationally intensive for a single GPU setup, taking over 30 hours perepoch.
The enwik8 64K sequence length and imagenet-64 experiments could not be tested for this reason.
Neverthelessthe use of shorter sequence lengths enabled us to thoroughly investigate the claims made.
In our experimentation we observe that Reformer enables training of deep models with long sequences on a single GPU,supporting the democratization of machine learning.
Practitioners interested in using Reformer should note howeverthat we also observed a non-negligible increase in time complexity due to the trade-off of memory for compute in the
Reversible Residual Layers and the trade-off of time for increased accuracy as the number of hashes used increases.
This behaviour is also observed elsewhere, e.g. Katharopoulos et al.
[2020] and Tay et al.
As this was a community effort with many contributors we were forced to improve our development processes duringthe project.
We started using nbdev (Howard [2019]), an open-source dev-ops system with documentation, testing andfull CI/CD integrated into GitHub.
This resulted in a much smoother development process for the team.
All of our code and documentation is available on GitHub 1 .
Of particular interest to future practitioners might be ourintegration of the ReformerLM and Transformer sequence-to-sequence training into the fastai training loop, as well asthe memory proﬁling scripts used in our memory consumption experiment.
Acknowledgments
There are a number of contributions that we would like to acknowledge; Weights & Biases for supporting this communitywork by providing a team account, DAGsHub (Pleban and Smoilovsky [2019]) for providing computational resourcesfor experimentation and Phil Wang (@lucidrains on GitHub 6 ) for his Pytorch Reformer implementation as well as anumber of other transformer implementations which made this complex implementation much smoother.
Finally, wewould like to acknowledge that this work originates in the fastai community and to thank Jeremy Howard, Sylvain
Gugger and all of the contributors to the library and the forums for continuing to make deep learning more accessible,more friendly and more performant.
A Training dynamics
Supplementary Shared-QK Experiment During the experiment phase of the synthetic task (section 4.2) , we useda TransformerLM with Shared Query-Keys.
We expected this model to solve the relatively simple synthetic task toperfection, but we found it difﬁcult to train the model to convergence even when training for 750 epochs.
When weinstead used a standard TransformerLM (i.e. with separate Queries and keys), the model consistently converges in 7-8epochs.
Figure 8 illustrates several runs on identical setup, changing only the model seed.
Out of ﬁve runs with shared query-keyonly one model converged after about 200 epochs.
For our standard TransformerLM the three runs converge within 7-8epochs.
B Reformer vs Number of Layers
In this experiment we compare ReformerLMs with n_layers =
All models use 8 hashes.
Theresults in ﬁgure 9 show a slight increase in memory as the number of layers grow.
The Reformer isn’t accumulatingmemory during forward and backward passes, but additional memory is used because more layers increases the numberof model parameters.