Reformer

The Problem

Standard Transformers scale poorly with sequence length. The self-attention mechanism computes a dot product between every query and every key, resulting in an $O(N^2)$ time and memory complexity for a sequence of length $N$. Additionally, standard residual networks store activations for every layer to compute gradients during the backward pass. For a deep Transformer with $L$ layers, memory usage scales linearly with depth, making it impossible to process very long sequences (e.g., 64K tokens) on a single GPU.

The Idea

The authors introduced two core architectural changes to break these bottlenecks. First, instead of computing all possible query-key pairs, the model uses Locality-Sensitive Hashing (LSH) to quickly find queries and keys that are similar and only computes attention within those localized buckets. Second, they replaced standard residual blocks with Reversible Layers, which allow the network to recalculate activations during the backward pass instead of storing them in memory.

How It Works

Shared Queries and Keys In standard attention, queries and keys are projected separately. The Reformer uses a single projection for both, meaning queries and keys are identical. This allows the model to group them together.

Locality-Sensitive Hashing (LSH) LSH is an algorithmic technique where similar vectors are assigned the same hash with high probability. The Reformer applies random LSH to the query/key vectors. Vectors that fall into the same hash bucket are highly likely to have a high attention weight.

Chunked Attention Tokens are sorted by their hash bucket and then chunked. Attention is computed only within the same chunk and the immediate neighboring chunk (to account for bucket boundaries). Because most keys outside the bucket would have near-zero attention weights anyway, ignoring them avoids $O(N^2)$ computations without severely degrading the model's output quality, dropping the complexity to $O(N \log N)$.

Reversible Residual Layers In a normal residual network, $y = x + F(x)$, we must store $x$ to compute the gradient of $F$. A reversible layer processes the input as two halves: $y_1 = x_1 + F(x_2)$ and $y_2 = x_2 + G(y_1)$. During backpropagation, $x_1$ and $x_2$ can be perfectly reconstructed from $y_1$ and $y_2$. This means only the final layer's activations need to be stored, decoupling memory usage from the network's depth.

Why It Mattered

Reformer proved that the absolute memory limits of Transformers were architectural choices, not strict mathematical necessities. By solving both the attention bottleneck and the activation bottleneck, it enabled training on contexts of up to 64,000 tokens on a single GPU—an order of magnitude larger than the standard BERT or GPT-2 models of the time.

What Came After

Reformer kicked off a wave of "Efficient Transformers" (like Linformer, Longformer, and BigBird) seeking to lower the $O(N^2)$ attention cost. While LSH attention was eventually superseded by hardware-aware exact attention methods like FlashAttention (which optimized memory IO rather than changing the math), the Reversible Layers concept remains highly influential and was adopted by later architectures like Routing Transformers and heavily scaled variants.