The main motivation of the reformer paper is that it reduces the memory footprint as compared to a normal transformer with little compromise in performance. In the previous sections we have seen that the reformer can achieve similar performance as a normal transformer on many tasks, and that is scales well to deep architecture (number of layers experiment) and long sequences (validation speed experiment). But how is memory actually allocated during training?

Claim: The reformer reduces the memory footprint compared to a normal transformer. The claim is summarized in table 5 of the reformer paper.

We write $dmodel$ and $d_{ff}$ for model depth and assume $d_{ff}$ ≥ $dmodel$. $b$ stands for batch size, $l$ for length, $n_l$ for the number of layers. We assume $n_c=l/32$ so $4l/nc=128$ and we write $c=128^2$

To investigate the claim we log the memory allocation of various reformer combinations during 0.1 % of an epoch of training the tiny Shakespeare dataset. The sequence length for all experiments was set to 4096, with a batch size of 1. We investigate the following combinations:

1. Comparing the Transformer LM, LSH LM, Reversible LM and full Reformer LM
2. Comparing Reformer LMs with different number of hashes
3. Comparing Reformer LMs with different number of layers

1. Comparing Transformer LM, LSH LM, Reversible LM and the full Reformer LM

The figure below shows the peak memory usage for the Transformer, LSH LM, Reversible LM and the full Reformer. We see that the transformer stores activations for each forward pass during training, and that these are gradually released as the backward pass is completed. In this case we have 6 distinct jumps in memory which corresponds to the number of layers 6.

With the LSH LM, with 8 hashing rounds, memory allocation is mostly parallel to the transformer, but since the LSH Attention computation is cheaper, the actual memory consumption is lower. Here we also observe that activations are stored for each layer.

For the Reversible LM memory doesn't accumulate over the forward passes, since it recalculates gradients for the backward pass. But since it needs to store 2 intermediate layers, the actual memory allocation per layer (i.e. the step size) is approximately twice that of the transformer. Note that we observe 4 peaks in the chart for the Reversible LM. Each peak corresponds to a forward and backward pass. The timing isn't directly comparable to that of the Transformer LM, so the plot for the Revesible LM starts in the middle of a forward and backward pass.

The full Reformer includes both LHS-attention and Reversible Residual Layers. Memory allocation therefore doesn't increase like the Transformer, and since LHS-attention is cheaper than standard attention, the peak memory usage is smaller than for the Reversible LM.

2. Reformer Memory vs Number of Hashes

In previous sections we have seen that increasing the number of hashes lead to increased performance as the attention approximation of LSH-attention approaches that of classic dot product attention. But since we have to store the result of each hashing round we would expect memory to grow linearly with the number of hashes. Note that it's only during the actual LSH-attention calculation that the number of hashes matter. I.e. the intermediate shape of LSH-attention is [batch_size, n_chunks, chunk_size, chunk_size*2], where n_chunks is the product of the number of hash buckets per round and the number of hash rounds, n_chunks = n_buckets * n_hashrounds.

The figure below, from our experiments, confirms that peak memory scales linearly with number of hashes. The memory peak happens during the forward pass during calculation of LSH-attention. Note that the output of LSH-attention is [bs, seq_len, d_model] and is independent of the number of hashes. This explains the drop in memory when LSH-attention calculation is finished. A 6-layer ReformerLM with a sequence length of 4096 was used for this experiment.

3. Reformer with changing number of layers

In this experiment compare Reformer LMs with n_layers = 2,4,6,8,10,12. 8 hashes are used in all cases. We see the same memory peak during LSH-attention calculation as in the figure above, and also that each added layer increases the memory footprint equally during the entire process. The figure show two forward and backward passes.

Since the reformer isn't accumulating forward and backward passes through all it's layers, but only stores intermediate activations, the peak memory allocation scales linearly to each extra layer.

Summary

Our experiments have verified that the reformer has a much smaller peak memory allocation as compared to the transformer. We have also shown that the reformer can scale to much deeper models than the transformer within a fixed budget. The main bottleneck of the reformer w.r.t. memory is the number of hash rounds used. For practical purposes this means that a user will have to strike a balance between performance and budget to fit their particular need.