Reformer: Pushing the Limits of Language Modeling with Memory-Efficient Transformers

Overview

The Reformer model addresses the memory bottleneck of standard Transformers when processing long sequences. By redesigning self‑attention, feed‑forward layers, residual connections, and positional encodings, it allows training on up to 500 000 tokens while staying under 8 GB of RAM, far beyond the 512‑token limit of models like BERT.

Reformer Self-Attention Layer

The Reformer replaces global self‑attention with two memory‑efficient alternatives: local self‑attention and LSH self‑attention.

Local Self-Attention

Local self‑attention splits the input into chunks of length config.local_chunk_length and applies global self‑attention inside each chunk. By adding overlap (config.local_num_chunks_before and config.local_num_chunks_after) each token can attend to a limited context of neighboring chunks, reducing the quadratic memory cost to O(n × chunk_length). This alone is insufficient for tasks needing long‑range dependencies.

LSH Self-Attention

LSH self‑attention approximates global attention by hashing query (and key) vectors into buckets using locality‑sensitive hashing. Vectors that hash to the same bucket are treated as similar, so self‑attention is computed only within each bucket. The permuted input is then chunked (with overlap) and attended to, allowing information from distant tokens to mix while keeping memory usage O(n × hashes × chunk_length). Multiple hash rounds (config.num_hashes) can be combined to improve accuracy.

Benchmark

Benchmarking the google/reformer-enwik8 model shows the memory savings:

  • With global self‑attention (chunk lengths set to 8192) the model exceeds GPU memory at ~16 K tokens.
  • Using the default local + LSH self‑attention the same model runs up to ~16 K tokens on a 11 GB GPU before out‑of‑memory, demonstrating the reduced growth of memory consumption with sequence length.

Chunked Feed Forward Layers

Large feed‑forward intermediate matrices dominate memory in wide Transformers. The Reformer introduces chunked feed‑forward layers, where the linear layers are processed in small chunks of size config.chunk_size_feed_forward. This avoids storing the full intermediate tensor, trading extra computation for lower memory usage.

Benchmark

When the feed‑forward size is inflated (e.g., to 16384) and attention heads are reduced, enabling chunking (chunk_size_feed_forward=1) drops peak memory from ~9 GB to ~6 GB for a sequence length of 4096 with batch size 8, confirming the benefit for models where the feed‑forward layer is the bottleneck.

Reversible Residual Layers

Training a standard Transformer stores all intermediate activations, causing memory to grow linearly with depth. The Reformer uses reversible residual layers, which allow activations to be recomputed during the backward pass instead of being saved. Only the final layer’s outputs need to be retained, reducing per‑layer memory overhead from several hundred megabytes to under 100 MB.

Benchmark

Comparing BERT and Reformer with increasing layers (4, 8, 12) at sequence length 512 and batch size 8 shows:

  • BERT‑12‑Layers: ~7.4 GB
  • Reformer‑12‑Layers: ~5.4 GB Thus Reformer adds far less memory per layer, enabling deeper models within the same hardware limits.

Axial Positional Encodings

Standard positional embeddings grow linearly with maximum sequence length, becoming prohibitive for very long inputs (e.g., 0.5 M tokens would need ~2 GB). Axial positional encodings factorize the position space into two smaller dimensions (config.axial_pos_shape) and split the hidden size (config.axial_pos_embds_dim). The resulting embedding table size is shape[0]×dim[0] + shape[1]×dim[1] instead of max_len×hidden_size.

Benchmark

For the google/reformer-crime-and-punishment model (capable of 0.5 M tokens, hidden size 256):

  • Default positional embeddings: 524 288 × 256 parameters (~136 M).
  • Axial positional embeddings with shape (512, 1024) and dim (64, 192): ~2.6 M parameters. Memory at inference drops from ~959 MB to ~447 MB for batch size 8 and sequence length 512, illustrating the substantial reduction.

Implications

By combining these four mechanisms, the Reformer makes it feasible to train language models on very long documents—such as full books or codebases—without requiring specialized hardware. This expands the range of NLP tasks that can benefit from full‑context understanding, including long‑range summarization, question answering over large corpora, and modeling of sequential data where dependencies span hundreds of thousands of tokens.

Sources