Reformer

The Reformer is a memory-efficient transformer that Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya of UC Berkeley and Google Research introduced in Reformer: The Efficient Transformer at ICLR 2020. It attacks the two costs that limit sequence length. Locality-sensitive hashing (LSH) attention replaces the full n × n attention matrix and costs about L log L for a sequence of length L, and reversible residual layers store the activations needed for backpropagation once for the whole network instead of once per layer. Google reported that the Reformer handles context windows of up to 1 million words on a single accelerator with 16 GB of memory, and it processed the entire novel Crime and Punishment as a single training example.

LSH Attention#

After the softmax, each query's attention is dominated by the few keys most similar to it, so the Reformer looks only for those. It uses one shared projection for queries and keys and hashes every vector with angular locality-sensitive hashing: multiply by a random matrix R and take the bucket as the position of the largest value in [xR, −xR], so vectors that point in similar directions usually land in the same bucket. The tokens are then sorted by bucket, the sorted sequence is cut into equal chunks, and each token attends only within its own chunk and the chunk before it. Several hash rounds run in parallel to catch similar pairs that a single round splits apart.

LSH is a classic method for approximate nearest-neighbor vector search: it finds most of a query's k-nearest neighbors at a fraction of the cost of an exact search. The bucketing step is short enough to run on its own:

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

The near vector, with a cosine similarity of 0.96 to the query, shares its bucket in 72% of the 1,000 rounds, and the unrelated vector, at 0.04, in 9%. Attention restricted to a bucket therefore keeps most of the pairs that matter.

Reversible Layers and Chunking#

A standard neural network keeps every layer's activations for the backward pass, so memory grows with depth as well as length. The Reformer adapts reversible residual networks (RevNets): it splits the activations into two halves and updates each half with a function of the other.

During the backward pass, each layer recovers its inputs from its outputs, x2 = y2 − FeedForward(y1) and then x1 = y1 − Attention(x2), so only the final activations are stored and each layer is recomputed once. Gradient checkpointing makes a similar trade by storing activations at a few checkpoints and recomputing the layers between them. The Reformer also runs its feed-forward layers in chunks of positions, because their hidden layer is several times wider than the model. Activation memory is what caps the batch size when training Ultralytics YOLO26: with batch=-1, AutoBatch measures GPU memory at several batch sizes, fits a straight line, and picks the batch predicted to fill 60% of the free memory.

Results and Implementations#

The Reformer matched standard transformers on character-level text from enwik8 and on pixel-by-pixel image generation on ImageNet-64 while using far less memory and running faster on long sequences. Hugging Face Transformers implements it as ReformerModel, with released checkpoints trained on enwik8 and on Crime and Punishment.

Reformer vs. Related Models#

• Longformer: Fixes its attention pattern in advance: a sliding window around each token plus a few global tokens.
• Sparse attention: The family of methods that score only some token pairs. LSH attention is a content-based member of it.
• Linear attention: Replaces the softmax with a kernel feature map so attention can be computed in linear time.
• Flash attention: Computes exact attention in GPU on-chip memory tiles without storing the score matrix, so memory grows linearly while the number of scores stays quadratic.