PaperlayerView exact revision on arXiv

Source-linked AI summary

Reformer: The Efficient Transformer

Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya

arXiv:2001.04451v2cs.LGcs.CLstat.ML

TL;DR

Large Transformer models achieve state-of-the-art results but are costly to train, especially on long sequences. Reformer combines locality-sensitive hashing attention and reversible layers, matching full Transformers while running faster and using substantially less memory.

  • Problem

    Large Transformer models achieve state-of-the-art results, but their attention and layer activations make training increasingly costly for long sequences.

  • Method

    Reformer combines locality-sensitive hashing attention with reversible layers to reduce attention and activation-memory costs.

  • Results

    Reformer matches full Transformer results on synthetic, text, and image-generation tasks while running much faster and achieving orders-of-magnitude better memory efficiency.

  • Takeaways & Limitations

    Reformer enables efficient Transformer execution on long sequences and may extend Transformer-based modeling to long-text and other generative tasks.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Large Transformer models achieve strong results but become costly to train as model size and sequence length grow. Reformer addresses these constraints with reversible layers, chunked feed-forward computation, and locality-sensitive hashing attention, while preserving near-standard training behavior and strong results.

  • Motivation: Large Transformers strain memory and computation because they combine many layers, wide feed-forward components, and quadratic attention over sequence length.Activations scale with the number of layers, feed-forward memory scales with d_ff, and attention costs O(L^2).
  • Approach: Reformer uses reversible layers so the N-fold activation-storage cost disappears during training.Reversible layers store only a single copy of activations across the whole model.
  • Approach: Chunking feed-forward computation removes the d_ff factor from the largest intermediate activations.The method processes feed-forward activations in chunks to reduce memory use inside those layers.
  • Approach: Locality-sensitive hashing replaces the O(L^2) attention factor with O(L log L), enabling more efficient processing of long sequences.The approximation focuses each query on nearby keys rather than materializing all query-key interactions.
  • Results: The proposed techniques have negligible training impact overall, although hashing attention can affect training dynamics depending on the number of concurrent hashes.The authors identify a hash setting that is efficient and produces results close to full attention.
  • Results: Reformer matches full Transformer results on 64K-token text and 12K-token image-generation sequences while running faster and using orders-of-magnitude less memory.The reported gains are especially pronounced on the text task.

2 LOCALITY-SENSITIVE HASHING ATTENTION

LSH attention approximates full attention by hashing nearby query-key vectors into shared buckets, then sorting and chunking them to restrict attention. The approach targets the quadratic cost and supports duplication-task accuracy that improves with more hashes.

  • Standard attention: Attention computes query-key interactions across a length-by-length matrix, creating O(L^2) computational and memory complexity.For 64K tokens, the QK^T matrix alone would require 16GB in 32-bit floats.
  • LSH attention: LSH attention focuses each query on nearby keys because softmax is dominated by the largest query-key similarities.This reduces the nearest-neighbor search problem to one addressed with locality-sensitive hashing.
  • LSH attention: Locality-sensitive hashing assigns nearby vectors to the same bucket with high probability while discouraging distant vectors from sharing buckets.The paper uses random projections with h(x) = arg max([xR; −xR]) to generate hashes.
  • Causal masking: The method preserves causal attention by associating vectors with position indices, applying the same sorting permutation, and masking future positions.Shared-QK attention additionally forbids self-attention except when no other valid target exists.
  • LSH attention: LSH attention restricts each query to targets within a single hash bucket, approximating full attention after queries and keys are sorted by bucket.The sorted structure clusters same-bucket pairs near the diagonal, enabling chunked computation.
  • Synthetic-task analysis: On the duplication task, four-hash training achieved almost perfect accuracy, while evaluation with eight hashes made accuracy perfect.Models trained with fewer hashes performed worse, but even one-hash training performed almost perfectly when evaluated with eight hashes.

3 REVERSIBLE TRANSFORMER

The Reversible Transformer reduces activation memory by combining reversible residual blocks with chunked feed-forward computation. This removes the layer-count memory term while preserving standard Transformer performance under matched parameter counts.

  • 3 REVERSIBLE TRANSFORMER: Reversible layers and chunking address the nl and dff contributions to Transformer memory use, respectively.The paper summarizes their effects on memory and time complexity in Table 3.
  • 3 REVERSIBLE TRANSFORMER: The Reversible Transformer combines attention and feed-forward layers inside a reversible residual block.The attention layer serves as F, the feed-forward layer as G, and layer normalization moves inside the residual blocks.
  • 3 REVERSIBLE TRANSFORMER: Reversible layers reduce whole-network activation storage by eliminating the number-of-layers term.Activations are recovered during back-propagation from later-layer activations and model parameters.
  • 3 REVERSIBLE TRANSFORMER: The reversible Transformer performs the same as the normal Transformer when using the same number of parameters.Both reversible streams are given size dmodel to preserve the parameter count.
  • 3 REVERSIBLE TRANSFORMER: Chunking reduces memory for wide feed-forward layers because their computations are independent across sequence positions.The computation is split into c chunks, with reverse computation and backward passes also chunked.
  • 3 REVERSIBLE TRANSFORMER: Layer parameters can be swapped with CPU memory when inactive, with larger batch-size-by-length compute amortizing transfer costs.This addresses parameter memory, which still grows with the number of layers.

4 RELATED WORK

Related work has reduced Transformer resource requirements through checkpointing, precision reduction, sparse attention, and alternative attention mechanisms. The paper positions locality-sensitive hashing as a distinct application to Transformer attention.

  • 4 RELATED WORK: Transformers have been extended beyond natural language to music, images, and self-supervised large-language-model training.The cited applications illustrate the model class’s use across diverse data and training settings.
  • 4 RELATED WORK: Prior methods reduce Transformer memory and computation through precision reduction, gradient checkpointing, and more efficient self-attention.These approaches respond to the computational requirements of state-of-the-art sequence models.
  • 4 RELATED WORK: Sparse Transformer exploits a factorized sparse representation of attention.Product-key attention has also reduced feed-forward memory requirements without loss in performance.
  • 4 RELATED WORK: Locality-sensitive hashing had not previously been directly applied to Transformer attention layers, to the authors’ knowledge.Earlier external-memory neural networks handled large memories but required the memory to be fixed before training.

5 EXPERIMENTS

Experiments assess reversible layers, shared query-key spaces, LSH attention, and full Reformer models across synthetic, text, image, and translation tasks. The techniques preserve or match Transformer performance while enabling long-sequence training with improved efficiency.

  • Experimental setup: Experiments use imagenet64, enwik8-64K, and WMT 2014 English-German tasks, with 3-layer models for ablations against regular Transformers.The ablations use dmodel = 1024, dff = 4096, 8 attention heads, and total batch size 8.
  • Shared query-key space: Shared query-key attention does not reduce accuracy relative to regular attention and appears to train slightly faster on enwik8.The method prevents tokens from attending to themselves except when no other context is available.
  • Reversibility: Reversible layers have nearly identical learning curves and preserve accuracy while providing memory savings compared with standard residual layers.The reversible and regular Transformer models have identical parameter counts.
  • LSH attention: At nrounds = 8, LSH attention almost matches full attention, while its computational cost can be increased through additional hashing rounds.Evaluation can use more hashes for greater accuracy, depending on the available compute budget.
  • LSH attention: LSH attention speed remains flat as sequence length increases, whereas regular attention becomes slower.The comparison holds the total number of tokens fixed.
  • Large Reformer models: A 12-layer Reformer reaches 1.19 bits/dim on enwik8 after 20K steps and 1.05 bits/dim after longer training and further tuning.The authors could train models up to 20 layers but could not train Transformer baselines because they were too slow and memory-hungry.

6 CONCLUSION

Reformer combines Transformer modeling capacity with an architecture designed for efficient execution on long sequences and low memory use, including for deep models. The authors connect this capability to broader use of large Transformers and long-sequence generative tasks.

  • Conclusion: Reformer combines Transformer modeling capacity with efficient execution on long sequences and small memory use for models with many layers.The authors state that this could help large, richly parameterized Transformer models become more widespread and accessible.
  • Conclusion: Handling long sequences opens potential applications in long-text, time-series, music, image, and video generation.These applications are presented as directions enabled by Reformer’s long-sequence capability.

A MULTI-ROUND LSH ATTENTION

The multi-round LSH mechanism computes attention over position sets formed across hashing rounds and combines the resulting outputs while correcting for duplicate elements. Its implementation uses sorted chunks and masking to support batching and shared-query-key behavior.

  • Multi-round attention: Multi-round LSH attention lets each query position attend to key positions defined by the union of its round-specific position sets.The position sets are specified by the paper’s multi-round formulation.
  • Implementation: Attention is performed on chunks of sorted queries and keys to support batching.Sorting follows the hashing-based organization of positions.
  • Multi-round attention: Combining the sparse attention formulation with multi-round position sets yields the multi-round attention computation.Each hashing round produces an independently computable output vector.
  • Duplicate handling: A correction term prevents double-counting elements when the union of round-specific attention sets is constructed.The implementation folds this factor into the masking term.
  • Masking: The masking rule disallows attention to the current position in shared-query-key causal attention except when no other valid target exists.The first token therefore attends only to itself because no prior context is available.
Loading 2001.04451v2…