ArXiv: 2001.04451
π― Pitch
The Transformerβs quadratic attention bottleneck vanishes when you replace exhaustive pairwise comparisons with locality-sensitive hashing, slashing complexity from O(LΒ²) to O(L log L). Combined with reversible layers that drop memory from N copies of activations to just one, the Reformer trains on 64K-token sequences on a single accelerator with no accuracy penaltyβmaking large-scale long-context models accessible outside industrial labs.
1. Executive Summary
This paper introduces the Reformer, a Transformer variant that replaces standard dot-product attention with locality-sensitive hashing (LSH) attention (restricting each query to attend only to keys within the same hash bucket, sorted and chunked along the diagonal) and replaces standard residual layers with reversible residual layers (allowing activations from any layer to be recovered from the subsequent layer during backpropagation, eliminating the need to store activations per layer). Evaluated on enwik8-64K (sequences of length 64K), imagenet64 (sequences of length 12K), and the WMT 2014 English-to-German translation task, the Reformer reduces attention complexity from O(LΒ²) to O(L log L) and memory consumption from being proportional to the number of layers to being independent of layer count β achieving performance on par with full Transformers while running "much faster" and with "orders of magnitude better memory efficiency," establishing that large-scale long-sequence Transformer models can be trained on single accelerators with negligible accuracy loss when LSH attention uses a sufficient number of hashing rounds (n_rounds = 8 nearly matching full attention).
2. Context and Motivation
The Core Problem: Transformers Scale Poorly on Long Sequences
The fundamental problem this paper addresses is straightforward to state but profoundly constraining in practice: standard Transformer architectures have computational and memory requirements that grow quadratically with sequence length, making them prohibitively expensive β or outright infeasible β to train on long sequences, even on modern accelerator hardware.
This is not merely an implementation inconvenience. It represents a genuine architectural bottleneck that limits what researchers can do with Transformers. The quadratic scaling appears in the attention mechanism's core operation: computing , which produces a matrix of shape [batch_size, length, length]. For each position in the sequence, standard dot-product attention computes a compatibility score with every other position β including positions that are entirely irrelevant to the current token's meaning. This dense, all-to-all attention pattern is what gives Transformers their celebrated ability to model long-range dependencies, but it also means that doubling the sequence length quadruples the attention cost.
The authors concretize this with a stark calculation in Section 1. Consider a single sequence of 64K tokens (a realistic length for document-level NLP, music generation, or image processing where each pixel or patch is a token). Even with a batch size of 1, the matrix for that sequence would be entries. Stored in 32-bit floating point, this single matrix requires 16 GB of GPU memory β for just one attention head of one layer of one sequence. A typical Transformer has 8 or more heads per layer and 6β24 layers. The memory cost is simply untenable on any single accelerator available when this paper was written.
But the quadratic attention cost is only one part of the picture. The paper identifies a second, equally critical bottleneck: the requirement to store intermediate activations for backpropagation. In a standard Transformer with layers, every layer's activations must be retained in memory during the forward pass so they can be accessed during the backward pass for gradient computation. This means the memory footprint scales linearly with the number of layers, not just sequence length. The authors illustrate this with another concrete calculation: a model with , , and would require an additional 16 GB of memory just for feed-forward layer activations, separate from the attention cost.
These two bottlenecks β quadratic attention scaling and layer-proportional activation storage β combine to create what the paper describes as a situation where "many large Transformer models can only realistically be trained in large industrial research laboratories and such models trained with model parallelism cannot even be fine-tuned on a single GPU as their memory requirements demand a multi-accelerator hardware setup even for a single training step." This is the practical consequence: the resource requirements gatekeep Transformer research, limiting who can experiment with large models on long sequences.
Why This Problem Matters: The Trend Toward Longer Sequences and Larger Models
The problem's importance is amplified by clear trends in the field at the time of writing (late 2019). The authors document three converging developments that make efficient long-sequence Transformers urgently necessary:
Models are getting deeper and wider. Shazeer et al. (2018) reported Transformer configurations exceeding 0.5 billion parameters per layer, while Al-Rfou et al. (2018) pushed depth to 64 layers. This is not arbitrary scaling β researchers were finding that larger models consistently improved performance, creating a powerful incentive to keep growing. But each additional layer compounds the activation storage problem: memory use is -times larger in an -layer model than in a single-layer model due to backpropagation's storage requirements.
Sequences are getting longer. Liu et al. (2018) processed up to 11,000 tokens of text in a single example for document-level tasks. Going beyond text, Huang et al. (2018) applied Transformers to music generation, where capturing long-range musical structure demands handling very long sequences. Parmar et al. (2018) and Ramachandran et al. (2019) applied Transformers to images, where an image can be represented as a sequence of patches or pixels β easily producing sequences far longer than those in NLP. The trend was clear: Transformers were being applied to modalities where long sequences were the norm, not the exception, and the quadratic attention bottleneck was becoming the limiting factor.
The gap between capability and accessibility is widening. The authors frame this as a concern for the health of the research community, noting that "some argue that this trend is breaking NLP research" (footnote referencing a community discussion about how resource requirements concentrate research power in a few well-funded industrial labs). The practical implication is stark: if only large companies with multi-accelerator clusters can train state-of-the-art models, the research ecosystem becomes less competitive, less diverse, and less innovative.
Beyond the accessibility argument, there is a deeper theoretical question the paper implicitly raises: does the Transformer architecture fundamentally require quadratic complexity, or is that complexity an artifact of a particular implementation choice? The fact that attention matrices are empirically sparse (most tokens only attend meaningfully to a small subset of other tokens) suggests the latter. If the architecture is inherently inefficient rather than inherently expensive, then there is substantial room for improvement without sacrificing the modeling capacity that makes Transformers effective. This is the optimistic hypothesis that motivates the entire paper.
Where Existing Approaches Fall Short
The paper does not operate in a vacuum. By late 2019, several lines of work had attempted to address Transformer efficiency, each with specific limitations that the Reformer aims to overcome. The authors organize these prior approaches and identify where they leave gaps.
Memory-efficient attention as a partial solution. The paper itself credits an existing technique for computing attention for each query independently rather than materializing the full matrix (Section 2). This "memory-efficient" implementation β which computes for one query at a time, then recomputes it during backpropagation β reduces the attention memory cost to being proportional to sequence length rather than its square. However, it does nothing for the computational cost: you still perform dot products, just one row at a time. It also doesn't address the layer-proportional activation storage problem. The authors use this memory-efficient implementation for their full-attention baselines, but it is clearly insufficient as a complete solution β it makes long-sequence attention possible but not fast.
Sparse attention patterns. The most direct antecedent to LSH attention is the Sparse Transformer (Child et al., 2019), which replaces dense attention with a fixed sparse pattern: each position attends to a pre-specified subset of other positions based on a factorized scheme (e.g., strided patterns that alternate between attending to nearby positions and attending to positions at regular intervals). This reduces complexity from to or similar. The key insight is that attention matrices in practice are sparse β most positions only need to attend to a small fraction of other positions β so a fixed sparse pattern can approximate full attention without computing it.
The limitation, which the Reformer paper makes explicit through its synthetic duplication task (Section 2.1), is that fixed sparse patterns can miss important non-local dependencies. The duplication task requires the model to copy a sequence of symbols from the first half of the input to the second half. This demands attention across an arbitrary distance β a token at position 512 in the second half needs to attend to position 0 in the first half, which a fixed local-window pattern would miss. The authors frame this limitation clearly: "it requires non-local attention lookups, so it cannot be solved by any model relying on sparse attention with a limited span." In other words, fixed-pattern sparsity imposes a structural prior about which long-range dependencies matter, and that prior may be wrong for some tasks. The Reformer's ambition is to achieve sparsity without imposing a rigid structural prior β using data-dependent hashing to decide which positions attend to which.
Other attention efficiency approaches. The paper references several other lines of work that attempted to address Transformer efficiency but stops short of deep engagement because they target different bottlenecks or use fundamentally different mechanisms:
- Adaptive attention span (Sukhbaatar et al., 2019a) learns per-head, per-layer attention spans so that each head attends only to a context window of learned size. This reduces computation but still enforces a contiguous local window β it cannot capture long-range dependencies that skip over irrelevant intermediate tokens.
- Product-key attention (Lample et al., 2019) reduces memory requirements in feed-forward layers (not attention layers) by using a product-key mechanism to increase the key space. This addresses the memory bottleneck but leaves the quadratic attention problem untouched.
- External memory with nearest neighbors (Rae et al., 2016) used locality-sensitive hashing and random kd-trees for memory lookups in memory-augmented neural networks. This is the closest prior work in terms of mechanism, but it applied LSH to external, fixed-size memory modules rather than to the self-attention mechanism itself. The Reformer paper can be seen as bringing this idea inside the Transformer's core computation.
Gradient checkpointing as a memory-saving technique. The paper acknowledges standard methods for reducing memory, including "gradient checkpointing" (referenced via Sohoni et al., 2019). Gradient checkpointing trades computation for memory: instead of storing all intermediate activations, you store a subset and recompute the rest during backpropagation. This can reduce the memory factor substantially, but it comes at the cost of additional forward passes β typically increasing computation by about 33% for a well-chosen checkpointing schedule. Moreover, it doesn't address the attention bottleneck at all; it only reduces activation storage. The Reformer's reversible layers (Section 3) can be seen as a more elegant solution to the same problem: rather than storing or recomputing activations, make the layers mathematically invertible so that activations can be reconstructed exactly from the next layer's activations at no extra computational cost.
The fundamental gap: no unified solution. Reading across these prior approaches, the paper's implicit critique is that each addresses one piece of the puzzle in isolation. Sparsity methods reduce attention compute but don't address memory. Reversible methods (from the RevNet literature) reduce memory but haven't been applied to Transformers. Memory-efficient attention reduces memory but not compute. What's missing is a unified architecture that addresses both the quadratic attention complexity and the layer-proportional activation storage simultaneously, without imposing rigid structural assumptions about which long-range dependencies matter. This is the gap the Reformer aims to fill.
How the Reformer Positions Itself
The paper's positioning can be understood along three axes: technical novelty, practical ambition, and philosophical stance.
Technical novelty: combining two ideas for the first time. The paper is refreshingly candid that neither of its core techniques is individually novel. Reversible layers were introduced by Gomez et al. (2017) for ResNets and applied to image classification. Locality-sensitive hashing for approximate nearest neighbors has a long history in computer science, and its application to neural network memory was explored by Rae et al. (2016). What is novel is the combination: adapting both techniques to the Transformer architecture, showing that they can coexist without interfering with each other, and demonstrating that the combination yields multiplicative efficiency gains β memory savings from reversibility compound with computational savings from LSH attention.
The paper also introduces several architecture-specific adaptations that are genuinely novel in their details: the shared-QK formulation that makes LSH attention applicable (since queries and keys must be identical for hashing to make sense), the sorting-and-chunking scheme that enables efficient batching of bucket-based attention, the multi-round LSH procedure with the double-counting correction factor (derived in Appendix A), and the specific causal masking adaptation that prevents self-attention in the shared-QK setting. These are not just implementation details β they are necessary engineering contributions that transform "LSH could work for attention" into "LSH does work for attention in a trainable system."
Practical ambition: single-accelerator training of large long-sequence models. The paper's stated goal is explicitly practical and measurable. The authors want to bring large Transformer training "within the reach of a single academic researcher with a single GPU." This is a concrete accessibility goal, not a vague efficiency claim. The paper repeatedly anchors its contributions in memory calculations (the 16 GB matrix, the 16 GB feed-forward activation storage) to make clear that the target is not an asymptotic complexity improvement on paper but a realizable reduction in hardware requirements for models people actually want to train.
This practical orientation also explains the paper's choice of experiments. Enwik8-64K and imagenet64 are not just convenient benchmarks β they represent the kind of long-sequence tasks that were becoming important (document-level language modeling, image generation as sequence modeling) but that standard Transformers struggled to handle. The WMT translation experiment serves a different purpose: showing that the reversible Transformer component works even when LSH attention is not needed (because translation sequences are typically short). This is a deliberate strategy to demonstrate that the components are independently useful and don't require the full long-sequence setting to provide value.
Philosophical stance: the architecture is inefficient, not the task. Perhaps the most important element of the paper's positioning is its implicit argument that the Transformer's resource demands reflect architectural inefficiency rather than a fundamental computational requirement of sequence modeling. The authors open Section 1 with this rhetorical framing: "Do large Transformer models fundamentally require such huge resources or are they simply inefficient?" The rest of the paper is structured to argue the latter.
The synthetic duplication task (Section 2.1) is the clearest articulation of this stance. The task is deliberately trivial β a single-layer Transformer can solve it perfectly β but it requires non-local attention that fixed sparsity patterns would miss. By showing that LSH attention with 4 hashing rounds achieves 99.9% accuracy on this task (Table 2), the paper demonstrates that data-dependent sparsity can preserve the essential capability that dense attention provides (arbitrary long-range dependency modeling) while avoiding the quadratic cost. The intuition is that LSH attention is "sparse but adaptive" β it doesn't impose a prior about which positions should attend to each other; it lets the data determine that through the hashing mechanism.
The paper also positions itself relative to the "more compute" solution that was becoming standard in the field. Rather than scaling hardware (model parallelism, TPU pods, multi-GPU clusters) to accommodate inefficient algorithms, the Reformer scales down the algorithm's resource requirements to fit available hardware. This is a conscious alternative to the dominant paradigm, and the paper's detailed FLOP and memory calculations serve to make this case quantitatively rather than rhetorically.
3. Technical Approach
3.1 Reader Orientation
This is an architectural redesign paper β the authors take the standard Transformer and surgically replace two of its most resource-intensive components (dot-product attention and standard residual layers) with more efficient alternatives that preserve modeling capacity while dramatically reducing computational and memory costs. The Reformer is not a new model family or training procedure; it is a drop-in replacement for the Transformer's internal mechanisms that makes training on sequences of length 64K feasible on a single accelerator, solving the quadratic attention bottleneck through locality-sensitive hashing and the layer-proportional memory bottleneck through mathematical reversibility of residual blocks.
3.2 Big-Picture Architecture (Diagram in Words)
The Reformer has two independently developed but complementary modifications to the standard Transformer, each targeting a different resource bottleneck. You can think of the architecture as having these major components:
-
LSH Attention (replaces standard dot-product attention): Instead of computing dot products between every query and every key (), the system hashes query and key vectors into buckets using random projections, sorts them by bucket assignment, chunks the sorted sequence into blocks, and only computes attention within each chunk plus one adjacent chunk. This reduces attention to . The hashing uses multiple rounds in parallel to reduce the probability that similar vectors end up in different buckets.
-
Shared-QK Formulation (enables LSH attention): To make hashing work, queries and keys must be identical vectors (otherwise a query and its corresponding key could hash to different buckets). The authors use a single linear projection to produce both Q and K from the input activations, then normalize K to unit length. This has negligible impact on model quality.
-
Reversible Residual Layers (replaces standard residual blocks): Instead of the standard residual, each reversible block operates on pairs of activation streams , combining them through attention and feed-forward layers in a way that makes the transformation mathematically invertible. During backpropagation, activations are reconstructed on-the-fly from the next layer rather than stored, eliminating the memory factor.
-
Chunked Feed-Forward Computation (reduces memory inside reversible blocks): The feed-forward sublayer, which typically has dimension much larger than the model dimension , is computed in smaller chunks across the sequence dimension. Since feed-forward operations are position-independent, this is numerically equivalent to batched computation but uses less peak memory.
-
Multi-Round LSH Aggregation (improves approximation quality): To compensate for hash collisions that separate similar vectors, LSH attention is run times with independent hash functions, and the results are combined with a correction factor that prevents over-counting items that appear in multiple rounds.
Information flows through the architecture as follows: input tokens are embedded β activations are split into two streams for the reversible block β passes through LSH attention (shared-QK, hashed, sorted, chunked) β the attention output is added to to form β passes through chunked feed-forward β the feed-forward output is added to to form β the pair becomes the input to the next reversible block. During backpropagation, this process runs in reverse: given and the stored model parameters, the system reconstructs by subtracting the residuals in reverse order.
3.3 Roadmap for the Deep Dive
- First, the shared-QK formulation β why queries and keys must be identical, how this is achieved, and the self-attention masking adjustment it requires. This is the prerequisite for making LSH attention possible.
- Second, the LSH mechanism itself β the random projection hash function, why it is locality-sensitive, and the step-by-step transformation from standard attention (Equation 1) through hashed attention (Equation 4) to the batched chunked version (Equation 5). This is the core computational innovation.
- Third, multi-round LSH attention β how multiple independent hash functions are combined, the correction factor that prevents double-counting, and the parallelization strategy. This is where the approximation quality is controlled.
- Fourth, the reversible Transformer block β the mathematical formulation in Equations 7β9, how attention and feed-forward layers are assigned to the F and G roles, how layer normalization is relocated, and why this eliminates activation storage.
- Fifth, chunking β how feed-forward layers are split across the sequence dimension, why this reduces peak memory, and how it interacts with backward computation.
- Sixth, the full Reformer integration β how LSH attention and reversible layers compose, the memory analysis in Table 3, and the parameter swapping strategy that enables large models on single accelerators.
This order follows the logical dependency chain: shared-QK enables LSH attention; LSH attention reduces the computational bottleneck; reversible layers with chunking reduce the memory bottleneck; the combination yields the full Reformer. Each component can be understood and evaluated independently, which is why the experimental section studies them in isolation before combining them.
3.4 Detailed, Sentence-Based Technical Breakdown
Shared-QK: Making Queries and Keys Identical
The standard Transformer produces queries (), keys (), and values () from the input activations (shape [batch_size, length, d_model]) using three separate learned linear projections:
where , , and are learned weight matrices of shape [d_model, d_k] (for Q and K) and [d_model, d_v] (for V). In multi-head attention, this projection is done times in parallel with different weight matrices, producing sets of queries, keys, and values.
Why this must change for LSH. Locality-sensitive hashing assigns each vector to a bucket based on its direction. If and come from different linear projections, then even for the same position , the query vector and the key vector will generally be different vectors pointing in different directions β meaning they could hash to different buckets. Since the whole premise of LSH attention is that a query should attend to keys in the same hash bucket (which are likely to be similar to the query), having queries and keys from different spaces breaks this logic. The fix is conceptually simple: use the same projection for both.
The shared-QK formulation. The authors modify the projection step so that:
That is, . There is now a single linear layer that produces both queries and keys. The value projection remains separate. The authors also normalize the key vectors to unit length: . This normalization is important for the LSH scheme (described in the next subsection) because the hash function operates on angular distance β normalizing to the unit sphere makes the hashing depend only on direction, not magnitude, which is the intended behavior for the random-projection LSH family.
Empirical justification: shared-QK doesn't hurt performance. The authors validate this design choice experimentally (Figure 3, left panel, discussed in Section 5). They train a standard Transformer and a shared-QK Transformer on enwik8 and imagenet64 and find that "a shared query-key space does not perform worse than regular attention; in fact, for enwik8 it appears to train slightly faster." This is a critical result: it means the architectural constraint imposed by LSH (forcing Q and K to be the same) comes at no accuracy cost. The slight speedup on enwik8 is likely because sharing weights reduces the total parameter count slightly or because the normalization improves training dynamics β but the key takeaway is that there is no downside.
Causal masking modification. Standard Transformer decoders use causal masking to prevent position from attending to positions (future positions). The mask is implemented by adding (or a very large negative value) to the attention logits for forbidden positions before the softmax. In a shared-QK formulation, a new problem arises: self-attention dominates. The dot product of a query vector with itself, , will almost always be larger than for any , because a vector is maximally correlated with itself. If position is allowed to attend to itself, the softmax will concentrate almost all probability mass on the self-position, drowning out attention to other positions that may be more informative.
The fix: disallow self-attention by modifying the causal mask to also exclude position from attending to itself. Formally, the mask term in Equation 3 takes value (or a large finite constant) not only when (future positions) but also when (self-position). There is one exception: if a token has no other valid attention targets β for instance, the very first token in a sequence, which has no preceding context β then self-attention is permitted because the alternative is to attend to nothing, which would produce undefined behavior. This is handled as a special case in the masking logic.
This modification is subtle but essential. Without it, LSH attention would degenerate into each token attending almost exclusively to itself, providing no meaningful contextual information. With it, attention is forced to look at other positions in the same hash bucket, which is the intended behavior.
Locality-Sensitive Hashing: The Core Mechanism
What LSH solves. The key observation driving LSH attention is stated in Section 2: "Since softmax is dominated by the largest elements, for each query we only need to focus on the keys in that are closest to ." In dense attention, we compute for all , but the vast majority of these dot products are small β tokens in one paragraph have almost nothing to do with tokens in a distant paragraph. The softmax exponentially amplifies the differences, so the attention weights are concentrated on a small number of keys. If we could efficiently find, for each query, the subset of keys with the largest dot products, we could compute attention only over that subset and ignore the rest. LSH gives us a way to approximately find these nearest neighbors without computing all pairwise dot products.
The hash function: angular LSH with random projections. The specific LSH family used is based on random projections onto the unit sphere (Andoni et al., 2015). The construction is:
-
Generate random projection vectors. Fix a random matrix of shape
[d_k, b/2], where is the desired number of hash buckets and is the key/query dimension. Each column of is a random direction in . -
Compute signed projections. For an input vector , compute (shape
[b/2]) and (also shape[b/2]). Concatenate these into a vector of length . -
Assign hash bucket. The hash value is the index of the maximum element: .
Why this is locality-sensitive. The hash function effectively partitions the unit sphere into regions using random hyperplanes through the origin. Each hyperplane is defined by a column of β its normal vector. The projection measures how aligned is with direction . By taking both and (which equals ), we are checking alignment with both directions along each axis. The argmax picks the direction (among all signed axes) with which is most aligned.
Two vectors and that are close in angular distance (small angle between them) will have similar projections onto the random directions β they will tend to be most aligned with the same few axes. Therefore, they will tend to receive the same hash value with high probability. Conversely, vectors that are far apart (pointing in very different directions) will tend to align with different axes and receive different hash values. Figure 1 in the paper illustrates this: in a simplified 2D depiction, two points and that are close on the sphere share the same hash bucket for all three depicted hashes, while distant points do not.
Why normalized keys. The hash function depends only on direction (which axis a vector is most aligned with), not magnitude. If keys were not normalized, two vectors pointing in the same direction but with very different lengths would get the same hash value β which is fine for angular LSH β but their dot products would differ substantially, making bucket co-membership a poor proxy for attention relevance. By normalizing keys to unit length (), the dot product reduces to the cosine similarity between the query and key directions. This makes the LSH bucket assignment (based on angular proximity) directly correspond to attention score magnitude (also based on angular proximity, since and is constant for a given query).
From standard attention to hashed attention. The authors formalize this transformation starting from standard attention for a single query position :
where is the set of positions that query can attend to (for causal masking, this is all positions ), is the log-normalizer (logarithm of the sum of exponentiated dot products over , ensuring the weights sum to 1), and is the value vector at position . The scaling by is omitted for clarity in this notation but is present in the actual implementation.
In standard attention, is large β it includes all preceding positions, up to the entire sequence length. The key idea of LSH attention is to restrict to only those positions whose keys hash to the same bucket as query :
where is the LSH function described above and we have enforced (shared-QK), so . This is Equation 4. Now, instead of attending over all , position attends only over positions whose keys fall in the same hash bucket. For a well-chosen hash function, this bucket will contain exactly the keys that are most similar to β the ones that would have received the largest attention weights anyway.
The sorting-and-chunking trick for efficient batching. The formulation is conceptually clean but computationally awkward. Hash buckets vary in size β some may be nearly empty, others may be very large (if many vectors happen to align with the same random direction). This makes it difficult to batch computation across buckets efficiently on GPU/TPU hardware, which prefers regular, fixed-shape tensor operations.
The paper's solution is elegant and practical. After computing hash values for all queries/keys:
-
Sort by hash bucket, then by sequence position within each bucket. This produces a permutation of the sequence indices. Let be the new position of original index after sorting. In the sorted order, all queries/keys that share a hash bucket become contiguous.
-
Chunk the sorted sequence into blocks of size . The authors set , where is the sequence length and is the number of hash buckets ( in the random projection scheme). The average bucket size is , so is twice the expected bucket size. The assumption is that the probability of any bucket exceeding its expected size is sufficiently low.
-
Define attention scope. In the sorted order, a query at sorted position attends to keys in the same chunk and the immediately preceding chunk. Formally, the attention set becomes:
This is Equation 5. In words: query attends to all positions whose chunk index (floor of sorted position divided by chunk size) is either equal to the query's chunk index or one less.
Why one chunk back. The sorted order places queries/keys from the same hash bucket contiguously. Most of the bucket will fall within a single chunk (since is twice the average bucket size), but some bucket members may spill into adjacent chunks. By allowing attention to the previous chunk as well as the current one, the scheme covers the full bucket with high probability while keeping the attention computation to rather than . Figure 2 in the paper visualizes this: the sorted attention matrix (panel c) shows hash-bucket clusters near the diagonal, and the chunked version (panel d) restricts attention to a band around the diagonal of width .
Memory and time complexity. The LSH attention mechanism has complexity where is the number of chunks (approximately , so ) and is the number of hashing rounds. This is , or when is chosen proportional to . Table 1 in the paper summarizes the complexity comparison with standard and memory-efficient attention.
The causal masking implementation under sorting. The sorting operation permutes the sequence. To apply causal masking correctly after sorting, the original position indices must be preserved and re-ordered according to the same permutation. After sorting, for each query-key pair at sorted positions and , the mask is applied by comparing their original indices: if (query comes before key in the original sequence), the attention weight is masked. Similarly, the self-attention prohibition () uses the original indices. The paper notes that this is implemented as a "comparison operation" on the re-ordered position indices.
Multi-Round LSH Attention: Reducing Hash Miss Probability
The problem: similar items in different buckets. Any single hash function has a non-zero probability of assigning two similar vectors to different buckets β this is called a "hash miss." The angular LSH scheme partitions the sphere into regions using random hyperplanes. Two vectors that are close in angle could fall on opposite sides of one of these hyperplanes, especially if they happen to straddle a decision boundary. The probability of a miss depends on the angle between the vectors and the number of hash buckets, but it is never zero for a single hash function.
The solution: multiple independent hash functions. The paper runs the LSH attention mechanism times in parallel, each with a different, independently generated random projection matrix . For each round , a separate hash function is used, producing a separate attention pattern . The final attention set for query is the union across all rounds:
This is Equation 6. If a key is missed by one hash function, it may be caught by another. With enough rounds, the probability that a truly relevant key is missed by all hash functions becomes negligibly small.
The double-counting correction. A key that is similar to query may appear in the hash bucket for multiple rounds β it is in the union, but it appears multiple times. If we simply summed the attention outputs from each round, such keys would be over-counted. The correction, derived in Appendix A, introduces a factor which counts how many rounds the key appears in for query :
For a key that appears in exactly one round, , and its contribution is counted once. For a key that appears in rounds, , and its contribution must be scaled down by so that the total contribution across all rounds sums to the correct value. The full multi-round attention formula (Equation 13) is:
where:
- is the log-normalizer for round (the log-sum-exp of dot products over )
- is the log-normalizer for the union over all rounds
- The outer exponential re-weights each round's contribution according to its relative normalizer, so that the combined softmax is properly normalized across the full union
- prevents double-counting of keys that appear in multiple rounds
- is the masking term that excludes positions not in (set to or a large constant for excluded positions)
What this equation computes operationally. For each hashing round :
- Perform LSH attention (sort, chunk, compute dot products within chunks) to produce a per-round output vector that is a weighted sum of value vectors for keys in the round- attention set, normalized by the round- partition function.
- Scale each round's output by the ratio of its partition function to the overall partition function β this ensures that rounds with larger attention sets (more keys in the union) don't dominate.
- Sum the scaled outputs across rounds.
The double-counting correction is folded into the masking term in the actual implementation (Equation 16), where it appears as an additive in the exponent, equivalent to dividing by in linear space.
Why this form (rather than simpler alternatives). A naive alternative would be to simply concatenate the attention sets from all rounds and compute softmax over the concatenated set. But this would require materializing the full attention matrix across all rounds, losing the chunking efficiency. The multi-round formulation in Equation 13 allows each round to be computed independently in its own sorted-and-chunked representation, with only the per-round normalizers and the counts needing to be shared. This preserves the per-round complexity and enables parallel computation of rounds.
Practical choices for . The experimental section explores . Figure 4 shows that performance improves with more rounds, with "almost matching full attention." Table 2 on the synthetic duplication task shows that a model trained with achieves 99.9% evaluation accuracy, while drops to 77.9%. The number of rounds can also be increased at evaluation time even if fewer were used during training β the model trained with 1 round achieves 99.9% accuracy when evaluated with 8 rounds, showing that the multi-round mechanism is robust to distribution shift in the attention pattern.
Reversible Transformer: Eliminating Per-Layer Activation Storage
The problem: memory grows with depth. In a standard Transformer with layers, the forward pass must store the input activations to every layer (both the attention sublayer and the feed-forward sublayer) because the backward pass needs them to compute gradients. The memory required for activations scales as for the attention outputs and for the feed-forward intermediate activations. Since is typically 4Γ larger than (e.g., , ), the feed-forward layers dominate. For , , , and 32-bit floats, this is approximately 16 GB just for feed-forward activations β before accounting for attention, parameters, or optimizer state.
The RevNet idea applied to Transformers. Reversible residual networks (Gomez et al., 2017) solve this by making each layer's transformation mathematically invertible. In a standard residual block, , you cannot recover from without knowing β which requires storing or recomputing from scratch. In a reversible block operating on pairs of activations, the forward computation is:
where and are arbitrary sublayers (in the Reformer, is the attention sublayer and is the feed-forward sublayer). This is Equation 7.
Why this is invertible. Given and the functions and (which depend only on the layer's parameters, which are stored regardless), you can recover by working backwards:
This is Equation 8. The key property: can be computed from alone (no need for ), so can be recovered first. Then can be computed from the recovered , allowing to be recovered. At no point do you need to have stored or β they are reconstructed on-the-fly during the backward pass by running the layer in reverse.
How the Reformer assigns F and G. In the Reformer's reversible block (Equation 9):
The input activations are split into two halves, and , each of size . The attention sublayer takes as input (producing queries, keys, and values from ) and its output is added to . The feed-forward sublayer takes (which already incorporates the attention output) as input and its output is added to . Both and have the full model dimension , so the total parameter count is the same as a standard Transformer β the split is in the activation representation, not in the model size.
Where layer normalization goes. In the standard Transformer, layer normalization (Ba et al., 2016) is applied before each sublayer (pre-norm) or after each sublayer (post-norm, as in the original Vaswani et al. 2017 formulation). In the Reformer, layer normalization is "moved inside the residual blocks" β applied to the inputs of the attention and feed-forward sublayers within and , rather than as separate operations outside the reversible block. This is necessary because the reversible formulation requires and to be deterministic functions of their inputs (and the model parameters) so that the inverse computation is exact. If normalization had access to running statistics or stochastic elements that differ between forward and reverse passes, the inversion would not be exact.
What this buys in practice. During training, the forward pass proceeds normally through all reversible blocks, producing the final output at the top layer. Only this final pair of activation streams needs to be stored for backpropagation β not the intermediate activations from every layer. When the backward pass reaches layer , it uses the stored and the layer's parameters to reconstruct via Equation 8, computes gradients with respect to the reconstructed inputs and the parameters, and passes the gradients to the previous layer. This process repeats layer by layer, with each layer's inputs being reconstructed from the outputs of the layer above.
The memory savings are dramatic: instead of storing (per-layer activation size), you store only (per-layer activation size) β the outputs of the final reversible block. This eliminates the factor entirely from activation memory. The computational cost of the reconstruction is exactly one additional forward pass through each sublayer: computing and during the reverse pass re-uses exactly the same operations as the forward pass, just in reverse order and with subtraction instead of addition. There is no extra computational overhead compared to storing activations β you are computing and during backpropagation instead of retrieving their outputs from memory.
Connection to gradient checkpointing. Standard gradient checkpointing also avoids storing all activations, but it does so by recomputing them from stored checkpoints at a subset of layers. This typically requires one extra forward pass for each checkpointed segment, adding ~33% computational overhead. The reversible approach achieves the same memory reduction with zero additional computation β the reconstruction pass IS the computation needed for backpropagation anyway (you need to compute and for the gradients), just organized differently.
Chunking: Reducing Feed-Forward Memory Without Changing Computation
The problem: is large. Even with reversible layers eliminating the factor, the feed-forward sublayer still requires memory proportional to for its intermediate activations. In a typical configuration, while , so the feed-forward layer uses 4Γ more memory for its internal computations than the attention layer uses for its activations.
The solution: position-wise chunking. The key observation is that the feed-forward layer operates independently on each position in the sequence. The standard feed-forward computation is:
where has shape [d_model, d_ff] and has shape [d_ff, d_model]. This is applied to every position independently β there is no interaction between positions. This means the computation can be split into chunks along the sequence dimension without any change to the mathematical result.
In the Reformer, the feed-forward computation on is chunked into pieces:
This is Equation 10. Each chunk has shape [batch_size, l/c, d_model], and the feed-forward computation on this chunk requires intermediate memory proportional to [batch_size, l/c, d_ff] rather than [batch_size, l, d_ff]. By processing chunks sequentially β computing one chunk, releasing its intermediate memory, then moving to the next β the peak memory is reduced by a factor of .
Why this is numerically identical. Unlike the LSH attention approximation, chunking the feed-forward layer is an exact transformation. The output is bit-for-bit identical to processing all positions at once, because the computation for position does not depend on position . The only change is the order of operations and the peak memory usage. This is a pure engineering optimization with no accuracy tradeoff.
Chunking during the reverse pass. The reversible block's inverse computation (Equation 8) requires computing , which is the feed-forward sublayer. During backpropagation, when reconstructing activations layer by layer, the feed-forward computation is again chunked β the system reconstructs from the layer above, then computes in chunks to recover , then computes to recover . The chunking applies in both forward and reverse directions.
Additional chunking: log-probabilities at the output. For language modeling tasks with large vocabularies ( word types), the final projection from to vocabulary size can also be memory-intensive. The paper mentions chunking "the log-probabilities at the output and calculate the loss for sections of the sequence at a time." This applies the same principle β compute softmax and cross-entropy loss for subsets of positions rather than all positions simultaneously β to avoid materializing a [batch_size, l, V] tensor.
Full Reformer Integration and Parameter Swapping
Composability of the components. The Reformer is the combination of all the components described above: LSH attention (with shared-QK and multi-round hashing) plus reversible layers (with chunked feed-forward). The components are designed to be independent β LSH attention reduces the factor in attention complexity regardless of whether reversible layers are used, and reversible layers eliminate per-layer activation storage regardless of whether LSH attention is used. When combined, the memory and computational savings multiply.
Memory complexity analysis (Table 3). The paper provides a systematic breakdown:
- Standard Transformer: Memory complexity is β the factor multiplies everything because all activations are stored for backpropagation.
- Reversible Transformer: Memory complexity drops to β the factor is eliminated because only the top layer's outputs are stored.
- Chunked Reversible Transformer: Memory drops further to β the factor is replaced by because feed-forward is chunked, reducing peak memory to the attention activations (which are -dimensional).
- LSH Transformer: Memory is β the factor in attention is replaced by , where is the number of chunks. With the paper's settings (, so ), the attention memory is rather than .
- Full Reformer: Memory is β combining both the reversible elimination and the LSH elimination yields memory that is independent of both the number of layers and the quadratic sequence length term.
Parameter swapping for very large models. Even with activations under control, model parameters themselves consume memory β and the number of parameters grows with the number of layers. The paper describes a parameter swapping strategy: "we can swap layer parameters to and from CPU memory when this layer is not computing." In a standard Transformer with small batch sizes, this would be inefficient because the time spent transferring parameters to CPU would be comparable to the time spent computing with them. But in the Reformer, the batch size multiplied by sequence length is "much larger" β because long sequences are now feasible β meaning the amount of computation done with a layer's parameters before they need to be evicted is substantial enough to amortize the transfer cost.
The layer execution schedule. During training, only the currently active layer's parameters need to be on the accelerator. As the forward pass proceeds layer by layer, parameters for completed layers can be offloaded to CPU memory and parameters for upcoming layers can be prefetched. The same applies during the reverse pass. This enables training models with more parameters than can fit in accelerator memory at once, similar in spirit to model parallelism but without requiring multiple accelerators β it uses the CPU as a parameter cache.
Causal masking in the full Reformer. The masking mechanism described for LSH attention (preventing self-attention and maintaining causal ordering through sorted position indices) integrates directly with the reversible block structure. The attention sublayer within each reversible block receives as input, applies the full LSH attention pipeline (hashing, sorting, chunking, multi-round aggregation with masking), and outputs a tensor of the same shape that is added to . The causal masking operates identically in every layer β there is no interaction between the masking logic and the reversible architecture.
Training and inference differences. The reversible formulation primarily benefits training by reducing activation storage. At inference time, there is no backward pass, so reversible layers do not provide memory savings (the forward pass in a reversible block is identical in cost to the forward pass in a standard block). However, LSH attention provides speed benefits at inference as well as training, because the reduced computational complexity applies to the forward pass. Figure 5 (right panel) shows that LSH attention evaluation speed remains flat as sequence length increases, while full attention becomes linearly slower β this applies to both training and inference forward passes.
The computational trade-off. The Reformer does not reduce the total number of floating-point operations required for a forward-backward pass compared to a standard Transformer with the same hyperparameters β it actually performs slightly more due to the multi-round hashing and the sorting/chunking overhead. The gain is in memory efficiency (enabling larger models and longer sequences to fit on a single accelerator) and wall-clock time (because the reduced memory pressure allows larger batch sizes and eliminates costly memory transfers between accelerator and host). The paper's experiments demonstrate these gains empirically: on enwik8-64K, the Reformer can train with sequences of length 64K at batch size 8 on 8 GPUs, while the standard Transformer baseline cannot fit in memory at all for the larger model configurations (20 layers).
4. Key Insights and Innovations
Innovation 1: Attention Sparsity Can Be Data-Dependent Rather Than Structurally Predefined
The dominant approach to efficient attention before the Reformer was to impose a fixed sparse pattern β each position attends to a predetermined subset of other positions based on structural assumptions about which dependencies matter. The Sparse Transformer (Child et al., 2019) exemplifies this strategy: it uses strided patterns that alternate between local windows and regularly-spaced long-range connections, reducing complexity from to . The implicit assumption is that important dependencies follow these structural patterns β nearby tokens matter, and tokens at regular intervals matter, and everything in between can be ignored.
The Reformer makes a fundamentally different bet: let the data determine which positions attend to which, via a randomized hashing mechanism that is sensitive to the content of the vectors, not their positions. Locality-sensitive hashing does not encode any assumption about which positions are likely to be relevant to each other. It encodes only the assumption that vectors pointing in similar directions (similar content) are likely to be relevant β which is precisely the assumption that dot-product attention itself makes. In other words, LSH attention approximates full attention by exploiting the same signal that full attention uses, rather than substituting a different structural prior.
This is a fundamental shift in how sparsity is conceptualized in attention mechanisms. Fixed-pattern sparsity asks: "which positions should this position attend to, based on their relative locations in the sequence?" LSH attention asks: "which positions are similar enough to this position that their dot products would be large anyway?" The first question imposes a topology on the attention graph that is independent of the input; the second lets the attention graph adapt to the specific content of each input. The synthetic duplication task (Table 2) makes this distinction concrete: a fixed local-window pattern fails because it cannot connect position 512 in the second half to position 0 in the first half β the input-dependent pattern of LSH succeeds because the query vector at position 512 (encoding "I need to copy the symbol from earlier") is content-similar to the key vector at position 0 (encoding "this is the symbol to be copied"), and hashing places them in the same bucket.
The significance extends beyond the specific LSH mechanism. This paper, along with contemporaneous work on content-based sparsity, opened a design space that is now standard in efficient Transformers: rather than pre-specifying which positions attend to which, use some fast approximate similarity mechanism (hashing, clustering, routing) to dynamically determine the attention pattern per input. The Reformer's LSH attention is one of the earliest and cleanest instantiations of this principle, and the paper's exposition β particularly the visual argument in Figure 2 showing how the sorted attention matrix concentrates similarity along the diagonal β provided a compelling intuitive framework for why data-dependent sparsity can work.
The paper's own evidence anchors this claim. Table 2 shows that LSH attention with four hashing rounds achieves 99.9% accuracy on the duplication task (which requires non-local attention), demonstrating that data-dependent sparsity preserves the essential capability that fixed-pattern sparsity sacrifices. The difficulty-bin analysis in Figure 4 β showing that performance improves monotonically with the number of hashing rounds β confirms that the approximation is controllable and can approach full attention arbitrarily closely.
Innovation 2: The Two Bottlenecks of Transformer Scaling Are Independent and Can Be Addressed Orthogonally
Before the Reformer, efficiency improvements to Transformers typically targeted a single bottleneck in isolation. The Sparse Transformer addressed attention complexity but did nothing for activation storage. Gradient checkpointing (Sohoni et al., 2019) reduced memory but added computational overhead and left attention complexity untouched. The field lacked a unified analysis showing that the computational bottleneck ( attention) and the memory bottleneck ( activation storage) are architecturally independent problems that admit independent solutions, and that combining those solutions yields multiplicative rather than merely additive efficiency gains.
The Reformer paper makes this insight explicit through its architecture and, more importantly, through its systematic complexity analysis in Tables 1 and 3. Table 1 breaks down attention variants by both memory and time complexity, showing that LSH replaces the term with independently of other architectural choices. Table 3 shows how the factor in memory complexity is eliminated by reversible layers, independent of whether attention is LSH or full. The full Reformer shows both terms reduced simultaneously: the factor becomes and the factor disappears entirely, yielding a memory complexity of that is linear in sequence length and independent of depth.
This is more than just combining two existing ideas. It is a diagnostic framework that decomposes the Transformer's resource requirements into independent components, each with its own scaling behavior, and shows how to address each component with the appropriate tool. The paper's structure β introducing LSH attention (Section 2) and reversible layers (Section 3) as separate sections with separate analyses, then combining them β reflects this diagnostic thinking. It teaches the reader to see Transformer efficiency not as one monolithic problem but as a set of separable sub-problems.
The distinction from incremental refinement is important here. Neither component is individually novel (LSH for nearest neighbors dates to the 1990s; reversible networks date to Gomez et al., 2017). What is novel is the systems-level insight that these two independently developed ideas address complementary bottlenecks in the Transformer architecture, and that combining them does not create interference or trade-offs β each component's benefits are preserved in the combined system. This is not obvious a priori. One might worry that the sorting-and-chunking of LSH attention would break the reversibility of the residual block, or that the chunking of feed-forward layers would interact badly with the reversible formulation. The paper shows empirically and analytically that these concerns do not materialize, establishing a composability principle that subsequent work on efficient Transformers has built upon.
The evidence for this claim is distributed across the paper's experiments. Figure 3 (right panel) shows that reversible layers alone have "negligible effect on training" compared to standard Transformers at the same parameter count. Figure 5 (left) shows that LSH attention continues to improve with more layers, confirming that the reversible architecture scales to depths that would be memory-prohibitive for standard Transformers. And the enwik8 result β where a 12-layer Reformer achieves 1.05 bits/dim, competitive with much larger standard models β implicitly validates the composability claim, since achieving that depth on 64K-length sequences would be impossible without both components working together.
Innovation 3: Reversibility as an Exact, Computation-Neutral Alternative to Gradient Checkpointing
The standard approach to reducing activation memory in deep networks at the time was gradient checkpointing: store activations at a subset of layers, and recompute the missing activations during backpropagation from the nearest stored checkpoint. This trades memory for computation β typically requiring one extra forward pass per checkpointed segment, adding roughly 33% computational overhead depending on the checkpointing schedule. The implicit assumption was that activation memory could only be reduced by spending computation to reconstruct what was not stored.
The Reformer's adoption of reversible residual layers challenges this assumption at a fundamental level. Reversibility achieves the same memory reduction as checkpointing β eliminating the factor from activation storage β but with zero additional computation. The reconstruction of from during backpropagation (Equation 8) requires computing and , which are precisely the same operations whose gradients are needed for the backward pass anyway. The computation that checkpointing would perform redundantly (re-running forward passes) is, in a reversible architecture, the computation that must be done regardless to obtain the gradients.
This is not merely an implementation optimization. It is a conceptual reframing: the memory required for backpropagation is not an inherent requirement of training deep networks, but an artifact of using non-invertible layer transformations. If each layer's transformation is designed to be invertible, the storage-vs-computation trade-off that gradient checkpointing navigates simply disappears β you get both low memory and no extra computation by construction.
The paper's application of this idea to the Transformer is particularly elegant because the Transformer's architecture naturally decomposes into paired sublayers (attention and feed-forward) that map cleanly onto the and roles in the reversible block (Equation 9: , ). The standard Transformer already applies these two sublayers sequentially; the only change is splitting the activations into two streams so that each sublayer's input and output are paired with a different stream. This is a minimal architectural modification that preserves the Transformer's computational structure exactly.
The distinction from incremental improvement is subtle but significant. Gradient checkpointing was already known and widely used. The Reformer's contribution is not "reversible layers are new" (they aren't) but "reversible layers solve the Transformer's memory problem better than the standard solution (checkpointing) because they eliminate the recomputation overhead, and the Transformer's paired-sublayer structure makes them a natural fit." This is a case where adapting an idea from a different domain (RevNets for image classification) to a new architecture (Transformers) required recognizing a structural affinity that was not previously obvious.
The evidence is indirect but compelling. The paper does not provide a head-to-head comparison of reversible layers versus gradient checkpointing. Instead, it demonstrates (Figure 3, right panel) that the reversible Transformer matches the standard Transformer in learning curves, while Table 3 shows the asymptotic memory reduction. The implicit argument is: you get the memory benefit of checkpointing (eliminating the factor) without the ~33% slowdown, and without any accuracy penalty. The WMT translation results (Table 4) further validate that this holds for encoder-decoder architectures, not just decoder-only language models, with a reversible Transformer-base matching the standard Transformer-base at 27.6 vs. 27.3 BLEU and a reversible Transformer-big reaching 29.1 BLEU.
Innovation 4: The Transformer's Attention Can Be Understood Through the Lens of Approximate Nearest Neighbor Search
This is the most intellectually generative contribution of the paper: framing the attention mechanism as an approximate nearest neighbor (ANN) search problem and showing that established ANN techniques (locality-sensitive hashing) can be dropped into the attention computation with minimal modification. Before the Reformer, attention was typically understood as a matrix multiplication problem: compute , apply softmax, multiply by . Efficiency improvements focused on making this matrix multiplication faster or sparser, but the conceptual framework remained rooted in linear algebra.
The Reformer reframes attention as: for each query , find the keys with the largest dot products , then compute the attention-weighted sum over only those keys. This is exactly the nearest neighbor search problem: given a query point, retrieve the most similar points from a database (the set of keys). The innovation is recognizing that this framing is not just an analogy β it is mathematically precise, because the softmax function is dominated by the largest dot products. Keys with small dot products contribute negligibly to the output, so ignoring them introduces negligible error.
This reframing opens a door that was previously closed. The ANN literature has developed decades of techniques for efficient similarity search in high dimensions β LSH, kd-trees, product quantization, graph-based indices β that had never been applied to the Transformer's self-attention mechanism. By casting attention as ANN, the paper makes this entire body of work available as a design space for future attention mechanisms. The specific LSH scheme the paper uses (angular LSH with random projections) is just one instantiation; the conceptual contribution is the framing itself.
The significance goes beyond the Reformer. This paper, along with contemporaneous work on routing-based attention and memory-augmented Transformers, established a new way of thinking about attention efficiency that has become dominant in the field. Subsequent work on clustering-based attention, product-key attention, and learned routing all operate within this framework: attention is retrieval, and making attention efficient means making retrieval efficient. The Reformer's particular contribution is showing that the retrieval can be exact enough to match full attention performance β the worry that ANN approximation would degrade model quality is empirically refuted by the results in Figure 4.
The paper's evidence for this framing is distributed. The step-by-step transformation from Equation 1 (full attention) through Equation 4 (hashed attention) to Equation 5 (chunked LSH attention) is presented not just as a mathematical derivation but as a conceptual argument: each step approximates the true attention distribution more coarsely, and the question is how much approximation the model can tolerate. Table 2 shows that even 1 hashing round suffices for 77.9% accuracy on the duplication task, and 4 rounds achieves 99.9%, directly quantifying the approximation quality. The language throughout Section 2 β "since softmax is dominated by the largest elements," "we only need to focus on the keys in K that are closest to qi" β consistently frames attention as retrieval, not as matrix multiplication. This is a pedagogical choice that reshapes how the reader thinks about the problem, not just how the code is implemented.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The paper uses three main tasks. The primary long-sequence benchmarks are enwik8-64K (a variant of the enwik8 text compression dataset chunked into subsequences of K tokens; the full enwik8 contains 100M characters of English Wikipedia text) and imagenet64 (a 64Γ64-pixel variant of ImageNet used for image generation, with sequences of length K tokens). Additionally, the paper evaluates on the WMT 2014 English-to-German translation task (newstest2014) for the reversible Transformer component, since translation sequences are short and do not benefit from LSH attention. Each serves a different role: enwik8-64K for text modeling on very long sequences, imagenet64 for demonstrating applicability beyond text, and WMT for isolating the effects of reversible layers.
-
Base model(s). All experiments use standard Transformer architectures from the Vaswani et al. (2017) family as baselines. The base model configuration for ablations uses , , , and , with a total batch size of 8 sequences. The 3-layer depth is chosen deliberately: "so as to make it tractable to compare with the regular Transformer, which has high memory usage and performs full attention." For the large-scale Reformer experiments, models are scaled up to 20 layers. The WMT translation experiments use the Transformer-base (Vaswani et al., 2017) and Transformer-big configurations, with both the encoder and decoder made fully reversible. The base model is not a specific pre-trained checkpoint β all models are trained from scratch for each experiment.
-
Metrics. The paper reports bits per dimension (bits/dim) for the enwik8 and imagenet64 tasks, which is standard for evaluating language modeling and density estimation β lower is better. For the WMT translation task, the primary metric is BLEU score on the newstest2014 test set, with both standard BLEU and detokenized sacreBLEU scores reported (Post, 2018) to enable fair comparison with prior work. For the synthetic duplication task (Section 2.1), the metric is accuracy (%) on the second half of the input sequence (the positions that can actually be predicted from prior context). All metrics are computed on held-out test data.
-
Baselines. The paper compares against several baselines:
- Standard Transformer (Vaswani et al., 2017): the full dot-product attention Transformer with standard residual layers, used for all ablation comparisons. This is the primary baseline.
- Memory-efficient attention Transformer: the same standard Transformer but using the memory-efficient attention implementation described in Section 2 (computing attention per query without materializing the full matrix). This is noted as the implementation used for full-attention baselines in the long-sequence experiments.
- Sparse Transformer (Child et al., 2019): referenced conceptually in the motivation but not directly compared in experiments, since the synthetic duplication task serves as a proxy argument that fixed-pattern sparsity would fail on non-local attention dependencies.
- For the WMT translation task, the baselines are the published results from Vaswani et al. (2017) (base: 27.3 BLEU, big: 28.4 BLEU) and Ott et al. (2018) (big: 29.3 BLEU).
-
Generation budget / compute accounting. The paper measures compute in several complementary ways: (1) Wall-clock speed of attention evaluation as a function of sequence length (Figure 5, right panel), holding the total number of tokens fixed. (2) Memory complexity analysis (Tables 1 and 3) measuring how memory requirements scale with sequence length , batch size , number of layers , and attention type. (3) Training throughput indirectly via the feasibility of experiments β the paper notes that standard Transformer baselines "cannot fit in memory" for configurations with 12+ layers and 64K-length sequences, while the Reformer can. No FLOP counting is provided; the focus is on memory and wall-clock time rather than arithmetic operations. For LSH attention, the number of hashing rounds serves as a compute-quality tradeoff parameter, with sweeps over .
-
Cross-validation / statistical protocol. The paper does not employ cross-validation in a formal sense. The synthetic duplication task uses separate train and evaluation settings (Table 2) where models trained with a specific number of LSH rounds are evaluated with varying numbers of rounds, serving as a generalization check on the attention mechanism. For enwik8 and imagenet64, models are trained for a fixed number of steps and evaluated on held-out data, with learning curves plotted to compare convergence. The WMT experiments report test-set BLEU at specific training step milestones (100K, 300K, 500K steps). No multiple random seeds or confidence intervals are reported, which is a notable methodological limitation.
Main Quantitative Results
Shared-QK and Reversible Layer Ablations: Establishing That the Components Do No Harm
The experimental strategy is to first validate that each architectural modification β shared query-key projection and reversible residual layers β does not degrade performance compared to the standard Transformer before evaluating LSH attention, which is the more significant approximation.
Shared-QK attention (Figure 3, left panel). On both enwik8 and imagenet64, the shared-QK Transformer trains "slightly faster" on enwik8 and matches the standard Transformer on imagenet64. The bits/dim learning curves (plotted against training steps) are essentially overlapping, with the shared-QK variant showing a small advantage on enwik8. The paper states: "A shared query-key space does not perform worse than regular attention; in fact, for enwik8 it appears to train slightly faster." This validates the prerequisite for LSH attention β forcing queries and keys to be identical does not sacrifice model capacity.
Reversible layers (Figure 3, right panel). The standard Transformer and the reversible Transformer (with identical parameter counts) produce "nearly the same" learning curves on both enwik8 and imagenet64. The bits/dim trajectories overlap almost perfectly, as shown in the two right-hand plots of Figure 3. This is a critical finding: the memory savings from reversibility come with "negligible effect on training" β the reversible blocks are a drop-in replacement that maintains accuracy while eliminating the factor in activation storage.
Reversible layers in machine translation (Table 4). In the encoder-decoder setting for WMT English-to-German:
- A reversible Transformer-base trained for 100K steps achieves 27.6 BLEU (sacreBLEU uncased: 27.4, cased: 26.9), compared to 27.3 BLEU from Vaswani et al. (2017) base model.
- Extended training for 500K steps with no weight sharing between embedding and output projection (possible because of memory efficiency) reaches 28.0 BLEU (sacreBLEU: 27.9 uncased, 27.4 cased).
- A reversible Transformer-big trained for 300K steps achieves 29.1 BLEU (sacreBLEU: 28.9 uncased, 28.4 cased), which is competitive with Ott et al. (2018) big model (29.3 BLEU) and surpasses Vaswani et al. (2017) big model (28.4 BLEU).
These translation results demonstrate that reversible layers work for both encoder and decoder components, and that the architectural change does not interfere with the attention mechanism in a cross-attention setting. The paper notes that LSH attention is not applied for translation because "examples are single sentences, and sentences tend to be relatively short" (the typical LSH chunk size of 128 exceeds the test set sequence lengths).
LSH Attention Performance: The Approximation-Controlled Scaling
Dependence on hashing rounds (Figure 4). On imagenet64, LSH attention performance improves monotonically with the number of hashing rounds. At , LSH attention "almost matches full attention." With fewer rounds, performance degrades gracefully β this is not a brittle all-or-nothing approximation. The paper also shows that the number of hashing rounds can be increased at evaluation time to improve results: a model trained with 4 hashes achieves "almost perfect accuracy" on the duplication task, which "becomes perfect when evaluated with 8 hashes" (Table 2). Conversely, a model "trained with just 1 hash performs almost perfectly when evaluated with 8 hashes," showing that the attention mechanism generalizes well across different sparsity levels.
Synthetic duplication task (Table 2). This task provides the cleanest analysis of LSH attention's approximation quality because the correct attention pattern is known β the model must attend non-locally from the second copy to the first. Key results:
- Full attention achieves 100% accuracy both at training and evaluation.
- LSH with achieves 100% training accuracy and 99.9% evaluation accuracy.
- LSH with achieves 77.9% evaluation accuracy when trained with 1 round β substantial degradation, showing that a single hash function has non-negligible miss probability.
- Cross-evaluation reveals an interesting asymmetry: a model trained with full attention and evaluated with LSH- drops to 94.8% (indicating the approximation does lose some information), but models trained with LSH are remarkably robust to increased hashing at test time. For instance, LSH- trained model achieves 99.9% when evaluated with LSH- β the extra hash rounds compensate for the coarser training-time approximation.
Speed vs. sequence length (Figure 5, right panel). The paper measures evaluation speed (not training throughput) while holding the total number of tokens fixed. The result is striking: "while regular attention becomes slower at longer sequence length, LSH attention speed remains flat." This is because LSH attention's complexity is rather than , so per-token cost decreases as sequence length increases relative to the fixed token budget, keeping total time approximately constant. The rightmost point (longest sequences) shows the largest speed gap in favor of LSH attention.
Large Reformer Models: Scaling Depth and Long Sequences
Training feasibility (Figure 5, left panel). The paper trains "up to 20-layer big Reformers on enwik8 and imagenet64" and states: "As can be seen in Figure 5, these models fit into memory and train. We were not able to train Transformer baselines in this case as they are too slow and memory-hungry, but we see clear improvement with the number of layers." This is the central practical demonstration β the Reformer enables training configurations that are simply infeasible for standard Transformers, and performance continues to improve with depth, showing the architecture does not sacrifice scaling capability.
The learning curves in Figure 5 (left) plot bits/dim against training steps for a 3-layer reversible Transformer (with and without LSH attention) and for deeper models. The 3-layer LSH Transformer achieves competitive performance with the 3-layer full-attention Transformer, while deeper Reformers (the exact layer counts are not specified in the text for this figure, but the paper mentions training up to 20 layers) show progressively lower bits/dim, confirming that depth continues to improve performance even with LSH attention.
Final results on enwik8. The paper reports two key numbers:
- A 12-layer Reformer trained for 20K steps with dropout 0.1 achieves 1.19 bits/dim on the test set.
- With further tuning and training ("trained for longer with further tuning and improvements"), a 12-layer model reaches 1.05 bits/dim on the enwik8 test set.
These results are presented without direct comparison to published baselines at the same sequence length (enwik8-64K is a non-standard chunked variant), but the paper frames them as demonstrating that Reformers can achieve competitive performance with standard Transformers while being far more memory-efficient. The 1.05 bits/dim figure represents the best result reported for the Reformer in this paper.
Ablation Studies and Robustness Checks
-
Shared vs. separate query-key projection (Figure 3, left): The shared-QK Transformer trains equivalently to or slightly better than the standard Transformer on enwik8 and imagenet64. This is a critical ablation because shared-QK is a necessary constraint for LSH attention β if sharing QK had degraded performance, LSH attention would be fundamentally limited. The result shows the constraint is benign.
-
Reversible vs. standard residual layers (Figure 3, right; Table 4): Across three tasks (enwik8, imagenet64, WMT translation) and two model sizes (base and big), reversible layers produce learning curves and final performance indistinguishable from standard residual layers at the same parameter count. This validates that the memory reduction does not come at a capacity cost.
-
Number of LSH hashing rounds (Figure 4; Table 2): This is the primary control knob for the attention approximation. Performance improves monotonically with from 1 to 8, with nearly matching full attention. The computational cost scales linearly with , making this an adjustable parameter for trading off accuracy and speed. Table 2 demonstrates that the number of rounds can differ between training and evaluation β a model trained with fewer rounds benefits from additional rounds at test time, which is practically useful for deployment scenarios where training must be efficient but evaluation can be more thorough.
-
Causal masking modification for shared-QK (implicit ablation): The paper describes the self-attention prohibition in Section 2 ("We therefore modify the masking to forbid a token from attending to itself, except in situations where a token has no other valid attention targets") but does not provide an explicit ablation. The justification is mathematical rather than empirical β the dot product of a vector with itself dominates, so allowing self-attention would drown out cross-position attention. The fact that models train successfully with this modification serves as implicit validation.
-
Chunking in feed-forward layers (Table 3, implicit): The chunking transformation is described as "numerically identical to the layers used in the Transformer" β it is an implementation detail that changes operation order but not the computed values. The paper does not provide a separate ablation for chunking because there is nothing to ablate: the computation is mathematically equivalent.
-
Negative result: standard Transformer baselines cannot be trained at scale (Section 5): The paper explicitly notes that "we were not able to train Transformer baselines" for the large-scale experiments (20-layer models on 64K-length sequences) because they are "too slow and memory-hungry." This is a negative result that serves as the motivation for the Reformer β it is not just faster but enables experiments that are otherwise infeasible.
Critical Assessment
Does the paper demonstrate that the Reformer performs on par with full Transformers?
The claim that the Reformer "performs on par with Transformer models" (abstract) is supported with specific, credible evidence, but the scope of the comparison matters. For the individual components, the evidence is strong:
- Shared-QK matches or slightly exceeds standard attention (Figure 3, left) on exactly matched hyperparameters.
- Reversible layers match standard residuals (Figure 3, right; Table 4) with identical parameter counts.
- LSH attention with nearly matches full attention on imagenet64 (Figure 4) and achieves near-perfect accuracy on the synthetic duplication task (Table 2).
However, the full Reformer (combining all components) is never directly compared to a full Transformer at the same scale, because the standard Transformer cannot be trained at those scales on the available hardware. The paper is candid about this: "We were not able to train Transformer baselines in this case as they are too slow and memory-hungry." This means the claim of parity rests on the component-level ablations (which show each piece individually matches the standard Transformer) plus the observation that deeper Reformers continue to improve with depth (Figure 5, left). The inference that a 12-layer standard Transformer on enwik8-64K would achieve similar bits/dim to the Reformer's 1.05 is plausible but untestable within the paper's experimental setup.
The WMT translation results (Table 4) provide the strongest direct evidence of parity because translation sequences are short and enable full Transformer baselines. Here, the reversible Transformer matches standard Transformer performance at base size and approaches published big-model results, establishing that at least the reversible component performs on par in a setting where direct comparison is possible.
The enwik8 final result of 1.05 bits/dim is reported without comparison to a same-scale standard Transformer baseline, which limits how strongly we can conclude parity. To strengthen this claim, the paper would need either (a) a way to train the standard Transformer baseline (perhaps with model parallelism) for direct comparison, or (b) evidence that existing published results for standard Transformers on enwik8-64K are comparable to 1.05 bits/dim when trained at similar scale and compute budget. Neither is provided.
Does the paper demonstrate that the Reformer is "much more memory-efficient"?
This is the strongest claim in the paper, and the evidence is overwhelming β but with an important caveat about what is actually measured. The memory complexity analysis (Tables 1 and 3) is precise and mathematically sound, showing asymptotic reductions from to independently of . The paper also provides concrete memory calculations (the 16 GB matrix, the 16 GB feed-forward activation storage) that ground these asymptotic claims in real hardware numbers.
The practical evidence is in what the Reformer can train that the standard Transformer cannot. The 20-layer, 64K-sequence-length experiments are feasible on 8 GPUs or TPU v3 cores with the Reformer but infeasible with the standard Transformer. This is a genuine demonstration of orders-of-magnitude improvement in memory efficiency.
The caveat is that no quantitative memory measurements are reported β no peak GPU memory usage graphs, no memory-vs-sequence-length scaling plots for the full Reformer vs. baselines, no profiling data showing where memory is actually consumed in practice. The paper relies on analytical complexity (Tables 1 and 3) and the existence proof ("we were able to train, they were not"). For a paper whose primary contribution is memory efficiency, the absence of actual memory measurements is a notable gap. Table 3 provides complexity formulas, not empirical memory usage.
Does the paper demonstrate that the Reformer is "much faster" on long sequences?
The speed evidence is mixed. The attention evaluation speed plot (Figure 5, right) convincingly shows that LSH attention speed stays flat with increasing sequence length while full attention becomes slower, when holding total token count fixed. This is a clean demonstration of the asymptotic advantage.
However, this is an evaluation-only measurement, not a training throughput comparison. The paper does not report training steps per second, wall-clock time to reach a given bits/dim, or end-to-end training speed comparisons for any model configuration. The complexity analysis in Table 3 shows that LSH attention has the same factor as full attention for time complexity β the memory savings do not directly reduce FLOP counts for a forward-backward pass. In fact, the full Reformer has time complexity , which includes the factor (multiple hash rounds) that full attention does not have.
The "much faster" claim in the abstract likely refers to the combination of (a) reduced attention cost on long sequences and (b) the ability to use larger batch sizes due to memory savings, which improves hardware utilization and thus wall-clock training speed. But this second effect is neither quantified nor discussed in the paper's timing analysis. The speed advantage for attention at evaluation time (Figure 5, right) is clear; the speed advantage for end-to-end training is asserted but not measured.
How robust are the results across tasks and configurations?
The evaluation is spread across qualitatively different domains β text compression (enwik8), image generation (imagenet64), and machine translation (WMT) β which demonstrates some generality. The synthetic duplication task adds a controlled setting for understanding attention patterns. However, each domain is represented by a single dataset:
- enwik8 is a byte-level text compression task. Results on word-level or subword-level NLP tasks (translation, summarization, question answering) are not reported beyond the WMT setting (which only tests reversible layers, not LSH attention).
- imagenet64 is a 64Γ64 image generation task. Results on standard image classification, object detection, or higher-resolution generation are not reported.
- WMT translation tests reversible layers but explicitly avoids LSH attention because sequences are short. This means LSH attention is never tested on a standard NLP benchmark β text generation quality (BLEU, ROUGE, perplexity on standard validation sets) with LSH attention is unexplored.
The single model configuration (, , for most experiments) is reasonable but narrow. The paper does not explore how LSH attention's approximation quality varies with model dimension, number of heads, or the ratio of to . The WMT translation experiments use the Vaswani et al. base and big configurations (which have different dimensions: and 1024), but only for reversible layers.
What experiments are missing?
Several experiments would have strengthened the paper's claims:
- Memory profiling: peak GPU memory usage for the Reformer vs. standard Transformer across sequence lengths, showing the actual memory reduction rather than just the asymptotic analysis.
- Training throughput: steps per second or wall-clock time to reach a target bits/dim for Reformer vs. standard Transformer at the same layer count and sequence length, where both can be trained (i.e., at shorter sequences).
- Full attention vs. LSH attention on the same model size: a scaled-down experiment where both fit in memory (e.g., 6 layers, length 16K) to directly compare final performance and training dynamics.
- LSH attention on standard NLP tasks: applying LSH attention in the WMT translation setting (perhaps with smaller chunk sizes to handle short sequences) to test whether the approximation degrades BLEU.
- Ablation on bucket/chunk size: the paper sets chunk size to but does not explore sensitivity to this choice. A larger chunk size would improve hash-miss tolerance at computational cost; a smaller chunk would increase speed at the risk of missing relevant keys.
- Ablation on the LSH hash function: the angular LSH scheme with the argmax over signed projections is one specific LSH family. Comparing against alternatives (e.g., standard random hyperplane LSH, or learned hash functions) would clarify whether this particular choice matters.
What are the genuine weaknesses of the experimental design?
The 3-layer ablation models are unnaturally shallow. The paper uses 3-layer models for its component ablations (shared-QK, reversible layers) "so as to make it tractable to compare with the regular Transformer." This is pragmatically necessary β full attention on 64K sequences is extremely expensive β but it raises a question: do the component-level findings at 3 layers generalize to deeper models? The interaction between LSH attention's approximation quality and network depth is not explored. It is possible that approximation errors compound with depth, or that deeper models require more hashing rounds to maintain parity with full attention.
The enwik8-64K is a non-standard benchmark split. Chunking enwik8 into 64K-token sub-sequences creates an artificial long-sequence task that may not reflect the actual dependency structure of natural language β the model cannot attend across chunk boundaries, which may break coherent multi-document context that a truly long-context model would capture. The paper uses this because it enables controlled long-sequence experiments, but results on this task may not transfer to settings where the long sequences are semantically coherent (e.g., entire books, long conversations).
No confidence intervals, no multiple seeds. The paper reports point estimates (bits/dim, BLEU) without error bars, standard deviations, or multi-seed averages. For the WMT translation results (Table 4), the BLEU differences between configurations are small (27.6 vs. 27.3, 29.1 vs. 28.4), and without variance estimates it is impossible to determine whether these differences are statistically meaningful or within the noise of training randomness.
Training steps are not held constant across comparisons. The WMT results report different training durations (100K steps for the base reversible model, 500K for the no-weight-sharing variant, 300K for big). The learning curves (Figure 3) are plotted against steps but the final numbers are taken at different points, making it unclear whether differences reflect architectural quality or simply training duration.
The 1.05 bits/dim result lacks methodological detail. The paper states this result was achieved after "further tuning and improvements" without specifying what tuning was performed (learning rate? dropout? number of hashing rounds? model size?). This makes the result non-reproducible and potentially cherry-picked from a hyperparameter sweep that is not described.
6. Limitations and Trade-offs
LSH Attention Requires Shared Queries and Keys β A Structural Constraint with Unknown Long-Range Implications
The assumption or constraint. LSH attention fundamentally requires queries and keys to be identical so that a token's query and its own key hash to the same bucket, enabling the sorting-and-chunking scheme to work. The paper enforces this through the shared-QK formulation (Section 2): a single linear projection produces both Q and K, key vectors are normalized to unit length, and tokens are forbidden from attending to themselves (except the first token, which has no other valid targets). The paper validates that shared-QK "does not perform worse than regular attention" on two tasks (Figure 3, left panel), but this is an empirical observation on specific benchmarks, not a guarantee.
The consequence. The shared-QK formulation removes a degree of freedom from the attention mechanism. In standard multi-head attention, each head learns separate projections for queries and keys β this allows a head to query for one type of information while offering a different type of information as a key to other positions. For example, a head might use its query to search for syntactic dependencies while using its key to advertise semantic content. Forcing Q = K means a token's query (what it's looking for) must equal its normalized key (what it offers to others), imposing a symmetry that may limit the expressiveness of individual attention heads. The paper's ablation (Figure 3, left) only tests this at 3 layers on enwik8 and imagenet64 β it is unknown whether deeper models, more diverse tasks, or tasks requiring asymmetric query-key relationships (e.g., cross-modal attention, retrieval-augmented generation) would suffer from this constraint. The paper offers no theoretical argument for why shared-QK should be capacity-neutral; the evidence is purely empirical on a narrow task set.
What evidence exists in the paper. Only the learning curves in Figure 3 (left), which show shared-QK training equivalently to or slightly better than standard attention on enwik8 and imagenet64 at 3 layers with . No analysis is provided of what individual attention heads learn under the shared-QK constraint versus standard attention, no probing tasks test whether specific attention patterns are disrupted, and no results exist for deeper models or standard NLP benchmarks (translation, summarization, QA) under shared-QK.
Mitigation status. Not addressed. The paper treats the empirical finding that shared-QK doesn't hurt performance as sufficient, but does not explore why it doesn't hurt or under what conditions it might. The normalization of keys to unit length () partially decouples query and key magnitudes while preserving direction, which may mitigate the symmetry constraint somewhat (the query can still have arbitrary magnitude, affecting the softmax temperature implicitly), but this is not discussed as a mitigation. The self-attention prohibition is a necessary patch for the shared-QK formulation to work at all, acknowledged as a modification to "typical implementations of the Transformer," but no ablation compares performance with and without this prohibition in a setting where it's not required (i.e., standard attention with shared-QK but allowing self-attention).
The Difficulty Estimation Cost Is Not Accounted for in the Resource Budget β Controls Approximation Quality but Scales Cost Linearly
The assumption or constraint. The quality of the LSH attention approximation is controlled by the number of hashing rounds . Figure 4 shows that "almost matches full attention," while produces substantially degraded performance (77.9% vs. 99.9% on the duplication task, Table 2). The paper presents as a tuneable parameter β "this hyperparameter can be adjusted depending on the available compute budget" β implying that practitioners can trade off accuracy for speed. The time complexity analysis (Table 1) confirms that LSH attention scales linearly with : the term includes as a multiplicative factor. At , LSH attention performs more hash computations, sorting operations, and chunked attention evaluations than at .
The consequence. The headline complexity reduction from to hides a constant factor that grows with . For (the setting needed to match full attention quality), the actual computational cost includes 8 parallel LSH attention passes. This means the constant-factor overhead of hashing, sorting, and chunking is multiplied by 8. The paper's speed comparison (Figure 5, right panel) β showing LSH attention speed remaining flat while full attention slows with sequence length β is an evaluation-only measurement whose value is not explicitly stated for that plot. A practitioner reading "O(L log L)" might expect an order-of-magnitude speedup at long sequences; the actual wall-clock improvement depends heavily on the sequence length and the needed for acceptable accuracy on their task, and the paper provides no systematic speed-vs-accuracy trade-off curves that would enable this estimation.
Moreover, Table 2 reveals an important subtlety: the number of hashing rounds at training and evaluation can differ, with models trained at lower benefiting from higher at test time. This is presented as a feature, but it means the training-time speed advantage (using fewer rounds) depends on the model generalizing to a denser attention pattern at inference β a capability that is demonstrated only on the synthetic duplication task and not validated on enwik8 or imagenet64.
What evidence exists in the paper. Figure 4 shows performance vs. on imagenet64, establishing the quality-sensitivity relationship. Table 2 provides the training-evaluation mismatch results on the synthetic duplication task. Table 1 provides the asymptotic complexity analysis showing linear scaling with . Figure 5 (right panel) shows evaluation speed vs. sequence length, but without specifying and without corresponding accuracy-at-that-speed data. No experiment systematically varies and measures both speed and accuracy to produce a Pareto frontier.
Mitigation status. The paper frames as a tuneable hyperparameter ("this hyperparameter can be adjusted depending on the available compute budget") but does not provide guidance on how to choose it for a new task. There is no attempt to reduce the required through improved hash function design, learned hash functions, or adaptive round allocation (e.g., using more rounds for tokens where the attention distribution has high entropy and fewer where it is concentrated). The observation that evaluation-time can exceed training-time is noted but not systematically exploited or validated beyond the synthetic task.
Reversible Layers Address Training Memory but Provide No Benefit at Inference β Most Deployments See No Gain from Half the Architecture
The assumption or constraint. The reversible residual layers (Section 3) reduce memory during training by enabling activation reconstruction during backpropagation. At inference time, there is no backward pass, so the reversible formulation provides no memory savings over standard residual layers β the forward pass through a reversible block (Equation 9) computes exactly the same operations as a standard block with paired residual connections. The paper is implicitly aware of this but never states it explicitly; the entire reversible layers motivation is framed in terms of training memory (Section 3: "back-propagation proceeds from the output of the network to its input"). The WMT translation experiments train reversible models but do not report inference speed or memory.
The consequence. For any deployment scenario β serving a trained model for inference, fine-tuning on a downstream task with a pre-trained checkpoint, or running inference on edge devices β the reversible layers contribute zero efficiency gain. Only the LSH attention component provides inference benefits (reduced attention complexity on long sequences). This means the Reformer's efficiency claims are asymmetric: the full architecture provides dramatic training benefits (memory independent of layer count, enabling deeper models on limited hardware), but a deployed Reformer model is only faster than a standard Transformer if sequences are long enough for LSH attention's advantage to outweigh its constant-factor overhead. For short-sequence tasks (the dominant use case for most deployed Transformer models β chatbots, translation, text classification), a trained Reformer provides no advantage over a standard Transformer at inference.
This asymmetry matters for the paper's stated goal of accessibility. The introduction argues that "these large-scale long-sequence models... cannot even be fine-tuned on a single GPU." Reversible layers address the fine-tuning memory problem, which is valuable. But the broader claim that the Reformer will "help large, richly-parameterized Transformer models become more widespread and accessible" must be qualified: it makes training and fine-tuning more accessible, but the resulting model is not inherently more efficient to deploy unless long sequences are involved. A practitioner who trains a Reformer on long documents and then wants to serve it for short-query responses gains inference benefit only from LSH attention, not from reversibility.
What evidence exists in the paper. None directly. The paper does not report inference-time memory or speed measurements. The WMT translation experiments (Table 4) use reversible layers but only report training BLEU β no inference latency comparisons. The speed plot (Figure 5, right panel) is labeled "Speed of attention evaluation" and compares LSH vs. full attention mechanisms, not reversible vs. standard layer architectures. The paper provides no ablation or discussion of inference-time implications for reversibility.
Mitigation status. Not addressed. The paper never acknowledges this asymmetry. The reversible layers are presented as a strict improvement over standard residuals, without qualification that the benefit is training-only. A complete efficiency analysis would separate training-time and inference-time gains, showing what fraction of the total resource savings comes from each component in each phase. This is particularly relevant because for many production systems, inference cost dominates total cost (models are trained once and served millions of times).
All Long-Sequence Evaluations Use a Single Model Family on Non-Standard Benchmark Splits β Generality to Standard NLP Tasks and Other Architectures Is Unestablished
The assumption or constraint. Every experiment in the paper uses a single architecture family β the Vaswani et al. (2017) Transformer with and β with the exception of the WMT translation experiments (which test only reversible layers, not LSH attention). The long-sequence tasks (enwik8-64K, imagenet64) are non-standard variants: enwik8 is artificially chunked into 64K-token segments, and imagenet64 at 12K tokens is a specific image generation formulation. The behavior of LSH attention on standard NLP benchmarks with naturally long sequences (document-level translation, summarization of long texts, multi-turn dialogue) is not tested. The paper states that LSH attention is not applied for WMT translation because "examples are single sentences, and sentences tend to be relatively short" β but this means the component that provides the complexity reduction is never evaluated on a standard NLP task where translation quality, summarization accuracy, or question-answering performance can be compared against published baselines.
The consequence. A practitioner considering the Reformer for a standard NLP pipeline β say, training a document-level translation model where sequences are long enough to benefit from LSH attention β cannot estimate from this paper what accuracy degradation to expect. The enwik8 bits/dim metric measures compression quality, which correlates imperfectly with downstream task performance. The imagenet64 task is a generative modeling task, not the classification or detection tasks that dominate vision benchmarks. The synthetic duplication task is deliberately simple. There is a significant gap between "LSH attention works on enwik8-64K" and "LSH attention works on document-level machine translation" β the latter requires the attention mechanism to capture cross-lingual alignment patterns that may have different sparsity characteristics than monolingual next-byte prediction.
Furthermore, the model is always a standard Transformer with the specific hyperparameter choices , , . The paper does not explore how LSH attention's approximation quality scales with model dimension, number of heads, or the ratio. A wider model with more heads might tolerate coarser hashing (because each head captures a narrower slice of the representation space, making hash buckets more discriminative), or it might require more rounds because each head has fewer dimensions to distinguish relevant from irrelevant keys. Without ablations on these hyperparameters, a practitioner cannot adapt the Reformer to their own model architecture with confidence.
What evidence exists in the paper. Table 2 (synthetic duplication) demonstrates that LSH attention can capture non-local dependencies in a controlled setting. Figure 4 (imagenet64) shows scaling on an image generation task. Figure 5 (left, enwik8) shows that deeper Reformers continue to improve on text compression. Table 4 (WMT) shows reversible layers working on translation but explicitly excludes LSH attention. The paper acknowledges the single-model limitation implicitly by using only one model configuration throughout, but does not discuss the generalizability concern.
Mitigation status. Not addressed. The paper frames enwik8 and imagenet64 as representative long-sequence tasks without arguing why results on these tasks should transfer to NLP or vision benchmarks more broadly. The claim that Reformer "performs on par with Transformer models" (abstract) implies generality, but the evidence covers two non-standard task formulations and one component-only evaluation on a standard task. No zero-shot or transfer experiments are conducted. The absence of LSH attention results on any standard sequence-to-sequence or sequence classification benchmark is a significant gap in the evidence for generality.
The Full Reformer Is Never Directly Compared to a Standard Transformer at Scale β Parity Is Inferred from Shallow Ablations and Untestable Claims
The assumption or constraint. The paper's central claim β that the full Reformer (LSH attention + reversible layers + chunking) "performs on par with Transformer models" β rests on a chain of inference rather than direct comparison. The reasoning is: (1) shared-QK matches standard attention at 3 layers (Figure 3, left); (2) reversible layers match standard residuals at 3 layers (Figure 3, right) and on WMT translation (Table 4); (3) LSH attention with matches full attention on imagenet64 at 3 layers (Figure 4); therefore (4) the combination at 12β20 layers should match a hypothetical same-scale standard Transformer. But step 4 is untested: "We were not able to train Transformer baselines in this case as they are too slow and memory-hungry" (Section 5). The very memory bottleneck that the Reformer solves prevents the controlled experiment that would validate the full architecture.
The consequence. There is no direct evidence that a 12-layer LSH Reformer achieves the same bits/dim as a 12-layer standard Transformer trained identically on enwik8-64K. It is possible that the approximation errors from LSH attention compound across layers β an attention miss in layer 1 changes the key vectors in layer 2, which changes hash bucket assignments, which causes further misses, etc. Or it is possible that the reversible block's paired-stream formulation (, ) interacts with LSH attention's sorting-and-chunking in ways that subtly alter optimization dynamics at scale, even though both components work well individually at 3 layers. Without a scaled comparison, we cannot distinguish between "the Reformer is slightly worse than an equivalently-sized standard Transformer but makes up for it by enabling deeper models" and "the Reformer is genuinely performance-equivalent."
The best result reported (1.05 bits/dim on enwik8, 12-layer Reformer) lacks any standard Transformer baseline at comparable compute. The paper does not cite or compare against published enwik8 results for standard Transformers at similar model sizes, nor does it train a standard Transformer at a smaller scale (e.g., 6 layers) where both models fit in memory and then extrapolate the scaling trend. The claim of parity is therefore a reasonable hypothesis that the component-level ablations support but that the paper does not β and given the hardware constraints, cannot β definitively prove.
What evidence exists in the paper. The component ablations at 3 layers (Figure 3, Figure 4) provide the only controlled comparisons. Table 4 provides a controlled comparison of reversible vs. standard Transformers on WMT translation, but without LSH attention. Figure 5 (left) shows that deeper Reformers improve over shallower ones β "we see clear improvement with the number of layers" β but does not plot any standard Transformer curve for reference, because none was trainable. The 1.05 bits/dim result is reported without methodological detail about the "further tuning and improvements" applied.
Mitigation status. The paper is honest about the missing baseline ("we were not able to train Transformer baselines in this case as they are too slow and memory-hungry") but does not attempt alternative validation strategies. Possibilities not pursued include: training a smaller standard Transformer (fewer layers, shorter sequences) and the equivalent Reformer to establish scaling trends; using model parallelism to train a standard Transformer baseline despite the expense, for a single data point; or comparing against published enwik8 results from prior work at similar model scales. The WMT results partially mitigate the concern for reversible layers specifically, but not for the full LSH + reversible combination.
The Parameter Swapping Strategy for Memory Reduction Is Described but Not Evaluated β Practical Feasibility Depends on Hardware Transfer Bandwidths That Are Not Measured
The assumption or constraint. Section 3 describes a strategy for handling models whose parameter count exceeds accelerator memory: "we can swap layer parameters to and from CPU memory when this layer is not computing." The argument for why this works is that in the Reformer, "the batch size multiplied by length... is much larger and therefore the amount of compute done with the parameters amortizes the cost of their transfer." The paper assumes that the time to transfer a layer's parameters over the PCIe bus (or equivalent interconnect) is small compared to the time to compute that layer's forward and backward passes on a long-sequence batch, so the GPU/TPU is never idle waiting for parameters.
The consequence. This strategy's feasibility is entirely hardware-dependent and is never verified. A typical PCIe 3.0 x16 connection provides ~16 GB/s bandwidth. A single Transformer layer with and contains approximately million parameters (accounting for QKV projections, output projection, and two feed-forward weight matrices), which at 4 bytes per float32 parameter is ~67 MB. At 16 GB/s, transferring this layer takes ~4 ms. The forward+backward computation for this layer on a batch of 8 sequences of length 64K with involves roughly FLOPs for the feed-forward sublayer alone. On a GPU capable of ~10 TFLOPS, this computation takes ~430 ms. So the 4 ms transfer is indeed amortized β but this calculation assumes perfect overlap of transfer and computation (double-buffering), no contention on the PCIe bus from other operations, and that the computation time dominates by nearly two orders of magnitude. For shorter sequences, smaller batches, or slower interconnects (common in consumer GPU setups where PCIe lanes may be shared or limited to 8x or 4x), the amortization ratio shrinks rapidly.
The paper provides no measurements of actual parameter transfer times, no profiling of CPU-GPU communication overhead in the Reformer training loop, and no ablation showing that training throughput with parameter swapping is comparable to training without it (when both fit in memory). The strategy is presented as a conceptual possibility, not as a validated technique.
What evidence exists in the paper. None. The parameter swapping strategy is described qualitatively in Section 3 ("we can swap layer parameters to and from CPU memory") and mentioned in Table 3's complexity analysis, but there are no timing measurements, no throughput comparisons, and no experiments that isolate the effect of parameter swapping on training speed. The paper reports training on "8 GPUs or 8 TPU v3 cores" without specifying whether parameter swapping was used for any of the reported experiments or whether the models fit entirely in accelerator memory.
Mitigation status. Not addressed. The paper treats parameter swapping as an optional extension that further improves memory efficiency beyond what reversible layers and chunking provide, but does not validate it. This is a significant gap because the paper's accessibility argument β enabling large model training on "a single accelerator" β depends partly on parameter swapping when the model's total parameter count exceeds single-accelerator memory. Without validation, a practitioner cannot rely on this technique working at acceptable throughput in their hardware environment. The paper also does not discuss the engineering complexity of implementing parameter swapping (which requires managing asynchronous CPU-GPU transfers, double-buffering layer parameters, and ensuring that gradient computations correctly accumulate across swapped layers), all of which represent practical barriers to adoption.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper's most enduring contribution is not any single technique but rather the conceptual reframing of Transformer attention as an approximate nearest neighbor (ANN) search problem. Before the Reformer, attention was understood as a matrix multiplication bottleneck β the question was how to compute more efficiently. After the Reformer, attention could be understood as a retrieval problem: for each query, find the most similar keys and ignore the rest. This shift is more than a notational convenience; it opens a door that had been closed, making decades of ANN research β locality-sensitive hashing, product quantization, graph-based indices, learned hash functions β available as a design space for attention mechanisms. The Reformer is the paper that explicitly made this connection and demonstrated its viability at Transformer scale.
The magnitude of this shift is best characterized as a reframing that unlocked a new subfield rather than a paradigm shift that displaced existing approaches. Full attention remained the gold standard; the Reformer showed that an ANN-based approximation could approach full attention quality when the approximation was sufficiently careful ( nearly matching full attention on imagenet64, as shown in Figure 4). But the Reformer did not render full attention obsolete β it established that data-dependent sparsity through hashing is a viable alternative whose trade-offs (approximation quality vs. speed) are controllable. The many subsequent papers on routing-based attention, clustering attention, and learned sparsity patterns (including the Performer, Linformer, and Routing Transformer, all of which appeared within a year of this work) operate in the intellectual space that the Reformer opened.
The paper also provides a concrete resolution to a tension that was implicit in prior work but rarely articulated: fixed-pattern sparsity can handle local dependencies but fails on non-local relationships that cross arbitrary distances. The Sparse Transformer (Child et al., 2019) had demonstrated that fixed strided patterns achieve strong results on many tasks, suggesting that most important dependencies are local or regularly-spaced. But the Reformer's synthetic duplication task (Table 2) shows a clear counterexample: a task that is trivial for full attention and solvable by LSH attention (99.9% with 4 hashing rounds) but impossible for any fixed-pattern sparse attention with limited span. This is not just a theoretical point β it demonstrates that data-dependent sparsity preserves a capability (arbitrary long-range non-local attention) that fixed-pattern sparsity sacrifices. The practical implication: if you don't know in advance which long-range dependencies your task requires, data-dependent sparsity is safer than structural sparsity.
A second, equally important reframing is the paper's diagnostic decomposition of Transformer resource bottlenecks. By tabulating memory and time complexity for attention variants (Table 1) and full Transformer variants (Table 3) as separable terms β the factor from attention, the factor from activation storage, the factor from feed-forward layers β the paper teaches the reader to see Transformer efficiency not as one monolithic problem but as a set of independent sub-problems, each with its own scaling behavior and its own solution. This diagnostic framework has proven more influential than any specific technique: subsequent work on efficient Transformers routinely analyzes complexity by separating attention cost from feed-forward cost from activation storage, treating each as an independent axis for improvement. The Reformer's particular combination β LSH for attention, reversibility for activation storage, chunking for feed-forward memory β is one instantiation of this framework, but the framework itself is the deeper contribution.
The paper also changes the economics of Transformer research in a specific, measurable way. The introduction's complaint β that large Transformer models "can only realistically be trained in large industrial research laboratories" and "cannot even be fine-tuned on a single GPU" β was not rhetorical flourish; it described a genuine barrier to entry that concentrated research power in a few well-funded labs. The Reformer's memory reductions (eliminating the factor and reducing the factor to ) directly address this barrier. A 12-layer Transformer on 64K-length sequences that previously required multi-accelerator setups becomes trainable on 8 GPUs or TPU v3 cores β still not "a single GPU" for the largest configurations, but dramatically more accessible. The concrete number from the paper: the matrix for a 64K-length sequence would require 16 GB at batch size 1 in 32-bit precision, and the Reformer eliminates this matrix entirely. This is a specific, quantitative improvement in the hardware threshold for long-sequence Transformer research.
Finally, the paper establishes that reversibility is a computation-neutral alternative to gradient checkpointing for Transformers. Before the Reformer, the standard approach to reducing activation memory was checkpointing β trading computation for memory with a ~33% overhead. The Reformer's reversible layers achieve the same memory reduction (eliminating the factor) with zero additional computation, because the reconstruction of activations during backpropagation uses exactly the operations whose gradients are required anyway. This is not just an incremental improvement over checkpointing; it eliminates the storage-vs-computation trade-off entirely, making it strictly preferable for architectures (like Transformers) whose paired-sublayer structure maps naturally onto the reversible block formulation. The evidence that reversible layers match standard residual layers in learning curves (Figure 3, right) and in final translation quality (Table 4, where reversible Transformer-big reaches 29.1 BLEU, competitive with non-reversible baselines) makes this a practical choice, not just a theoretical possibility.
Follow-Up Research This Work Enables
Cheap and adaptive difficulty estimation for LSH attention. The paper shows that controls the approximation quality of LSH attention β 8 rounds nearly match full attention, 1 round degrades substantially (Figure 4, Table 2) β but treats as a fixed global hyperparameter. A natural extension is per-token or per-head adaptive round allocation: use a lightweight signal (e.g., the entropy of the attention distribution from a single hash round, or the variance of attention weights across rounds) to decide whether a particular query needs additional hashing rounds to resolve its attention pattern accurately. Queries whose attention is naturally concentrated on a small number of keys (low entropy) might achieve accurate attention with 1β2 rounds, while queries with diffuse attention (high entropy) might need 4β8 rounds. A strong follow-up would train a Reformer on enwik8-64K or a standard long-document NLP task with an adaptive-round policy, measure the average used per token, and show that the adaptive policy achieves the same bits/dim as fixed while using, say, 3.2 average rounds β a ~2.5Γ speedup in the attention computation. The paper's demonstration that evaluation-time can differ from training-time (Table 2) suggests that dynamically varying at the token level is feasible without degrading training stability.
LSH attention applied to standard long-document NLP tasks with direct baseline comparison. The paper evaluates LSH attention on enwik8-64K (a byte-level compression task) and imagenet64 (image generation), but explicitly avoids applying it to the WMT translation task because sequences are short (Section 5: "examples are single sentences, and sentences tend to be relatively short"). The most important missing experiment is LSH attention on a standard NLP benchmark where sequences are naturally long and where full-attention Transformer baselines are well-characterized. Concrete candidates: document-level machine translation (e.g., the WMT document-level tasks, or IWSLT TED talk translation where entire talks form sequences of thousands of tokens), long-form summarization (e.g., PubMed, arXiv, or BookSum datasets where input documents exceed 4Kβ16K tokens), or long-range text classification (e.g., the Long Range Arena benchmark introduced subsequently by Tay et al., 2020). A strong follow-up would train Reformer and standard Transformer models at matched parameter counts on these tasks, report both task-specific metrics (BLEU, ROUGE, accuracy) and wall-clock training/inference time, and determine whether the asymptotic advantage translates to practical speedups at the ~4Kβ16K sequence lengths typical of long-document NLP, where the constant-factor overhead of hashing and sorting may eat into the theoretical gains.
Combining LSH attention with other efficient attention mechanisms to study complementarity. The paper's framing of attention as ANN search suggests that LSH is one of many possible retrieval mechanisms. A natural experiment is hybrid attention heads: some heads in each layer use LSH attention (capturing content-based long-range dependencies), while others use local sliding-window attention (capturing positional short-range dependencies efficiently and exactly). This would test whether the information captured by LSH attention and local attention is complementary β i.e., whether a hybrid model outperforms either pure LSH or pure local attention at the same computational budget. A strong follow-up would train hybrid models on enwik8-64K or a long-document NLP task, sweep the ratio of LSH to local heads, and compare against (a) pure LSH with matched total FLOPs, (b) pure local attention with matched total FLOPs, and (c) the Sparse Transformer's fixed strided pattern. The Reformer's synthetic duplication task (Table 2) demonstrates that LSH handles a specific failure mode of local-only attention (non-local dependencies); a hybrid experiment would quantify how often that failure mode matters in realistic tasks and whether a small number of LSH heads suffices to capture it.
Learned hash functions versus random projections for LSH attention. The paper uses a fixed random projection matrix to define the LSH function: . This hash function is data-independent β the random directions are fixed before training and do not adapt to the distribution of query/key vectors. An obvious extension is learned hash functions: replace the random with a learned projection (or a small neural network) trained end-to-end with the rest of the model, so that the hash buckets reflect the actual similarity structure of the learned representations rather than arbitrary random directions. This could reduce the number of hashing rounds needed to achieve a given approximation quality (since learned hash functions can be more discriminative) or enable more even bucket sizes (addressing the uneven-bucket problem that forced the chunking approximation in the Reformer). A strong follow-up would compare random vs. learned hashing on enwik8-64K at fixed , measuring both bits/dim and bucket size variance (the paper notes that uneven bucket sizes motivated the sorted-chunked batching; learned hashing might produce more balanced buckets and enable simpler batching). A negative result β learned hashing provides no improvement over random projections β would be equally informative, suggesting that random projections already capture the relevant similarity structure for Transformer representations.
Memory profiling and throughput benchmarking of the full Reformer against standard Transformers at equal scale. The paper's central claim is that the Reformer is "much more memory-efficient" and "much faster on long sequences," but the evidence is largely analytical (Tables 1 and 3) and asymptotic (Figure 5, right, shows evaluation speed for attention only). No experiment measures actual GPU memory usage or training throughput for the full Reformer versus a standard Transformer at a scale where both can be trained. A follow-up study that is purely empirical β profiling peak memory, time per training step, and total wall-clock time to reach a target bits/dim for a 6-layer Reformer vs. a 6-layer standard Transformer on enwik8-64K at batch size 8, with the standard Transformer using memory-efficient attention and gradient checkpointing to fit in the same memory budget β would provide the quantitative comparison the paper lacks. This study should report both the memory savings (validating the elimination from reversibility and the reduction from LSH) and the wall-clock implications (quantifying the constant-factor overhead of hashing and sorting vs. the asymptotic advantage, and separating training-time from inference-time measurements). This is not glamorous research, but it is essential for practitioners deciding whether to adopt the Reformer architecture, and the paper's omission of these measurements is its most significant empirical gap.
Practical Applications and Downstream Use Cases
Training long-document NLP models on academic hardware budgets. The most direct application enabled by this work is training Transformer models on document-level tasks (long-form text generation, document summarization, document-level translation) using hardware available to academic researchers. Before the Reformer, training a 12-layer Transformer on 64K-token sequences required multi-GPU setups that were out of reach for most academic labs; the Reformer's memory reductions β eliminating the storage factor and reducing attention from to β bring this within the range of 8 consumer GPUs or cloud instances that a typical research grant can fund. The enwik8-64K result (1.05 bits/dim with a 12-layer Reformer on 8 GPUs) serves as a proof of concept: a model that would have been infeasible to train with standard Transformers on this hardware can be trained with the Reformer. For a researcher wanting to experiment with long-document summarization on PubMed articles (median length ~3K words, or ~5Kβ8K BPE tokens), the Reformer architecture enables models with effective context windows of 8Kβ16K tokens without requiring model parallelism or industrial-scale compute clusters.
Fine-tuning large pre-trained Transformers on long sequences with a single GPU. The paper's observation that large Transformer models "cannot even be fine-tuned on a single GPU" is a practical pain point for many practitioners. The reversible layers component of the Reformer directly addresses this: by eliminating the factor from activation storage, a 24-layer Transformer that previously required multi-GPU fine-tuning (because activations for all 24 layers had to fit in GPU memory) can now be fine-tuned with activation storage equivalent to a single layer. The WMT translation experiments (Table 4) demonstrate that reversible layers work for both encoder and decoder in sequence-to-sequence tasks. A practitioner with a pre-trained BART or T5 model (12β24 layers) who wants to fine-tune on long sequences (e.g., scientific document summarization) can convert the model to reversible blocks β a structural change that preserves parameter counts and, per Figure 3 (right), should maintain performance β and reduce memory by roughly a factor of , enabling fine-tuning on a single consumer GPU (e.g., RTX 3090 with 24 GB) that would otherwise require a multi-GPU setup. The chunked feed-forward component (Equation 10) further reduces peak memory for the largest activations, contributing to this feasibility.
Image and video generation as sequence modeling with long-range dependencies. The paper's imagenet64 experiments (sequences of length 12K) demonstrate that the Reformer handles modalities beyond text where long sequences arise naturally from treating pixels or patches as tokens. This opens a practical path for high-resolution image generation and video generation using autoregressive Transformers. At the time of writing, Image Transformer (Parmar et al., 2018) and similar models were limited by the quadratic attention bottleneck to relatively small images or aggressive downsampling. The Reformer's attention makes it practical to train autoregressive Transformers on 256Γ256 or 512Γ512 images (where treating each pixel as a token gives 65K or 262K tokens, respectively), or on video frames where sequence lengths quickly reach hundreds of thousands of tokens. The practical value is not just research-scale training but also deployment: LSH attention provides inference speed benefits (Figure 5, right panel, showing flat evaluation speed vs. sequence length), which matters for interactive image generation applications where generation latency must be manageable.