Reformer
Reformer is a memory-efficient transformer that introduces two key innovations: Locality-Sensitive Hashing (LSH) attention (reducing complexity from O(n²) to O(n log n)) and reversible residual layers (reducing memory from O(n_layers × n) to O(n)) — targeting extremely long sequences (64K+ tokens) where both compute and memory are prohibitive, by replacing exact full attention with an efficient approximation that attends only to similar tokens.
What Is Reformer?
- Definition: A transformer architecture (Kitaev et al., 2020, Google Research) that addresses two memory bottlenecks: (1) the O(n²) attention matrix is replaced by LSH attention that groups similar tokens into buckets and computes attention only within buckets, and (2) the O(L × n) activation storage for backpropagation is eliminated by reversible residual layers that recompute activations during the backward pass.
- The Two Memory Problems: For a sequence of 64K tokens with 12 layers: (1) Attention matrix = 64K² × 12 × 2 bytes ≈ 100 GB (impossible). (2) Stored activations = 64K × hidden_dim × 12 layers × 2 bytes ≈ 6 GB (significant). Reformer attacks both simultaneously.
- The Approximation: Unlike FlashAttention (which computes exact attention efficiently), LSH attention is an approximation — it assumes that tokens with high attention weights tend to have similar Q and K vectors, and groups them via hashing.
Innovation 1: LSH Attention
| Concept | Description |
|---|---|
| Core Idea | Tokens with similar Q/K vectors will have high attention weights. Hash Q and K into buckets; only attend within same bucket. |
| LSH Hash | Random projection-based hash function that maps similar vectors to the same bucket with high probability |
| Bucket Size | Sequence divided into ~n/bucket_size buckets; attention computed within each bucket |
| Multi-Round | Multiple hash rounds (typically 4-8) for coverage — reduces chance of missing important attention pairs |
| Complexity | O(n log n) vs O(n²) for full attention |
How LSH Attention Works
| Step | Action | Complexity |
|---|---|---|
| 1. Hash | Apply LSH to Q and K vectors → bucket assignments | O(n × rounds) |
| 2. Sort | Sort tokens by bucket assignment | O(n log n) |
| 3. Chunk | Divide sorted sequence into chunks | O(n) |
| 4. Attend within chunks | Full attention within each chunk (small, ~128-256 tokens) | O(n × chunk_size) |
| 5. Multi-round | Repeat with different hash functions, average results | O(n × rounds × chunk_size) |
Innovation 2: Reversible Residual Layers
| Standard Transformer | Reformer (Reversible) |
|---|---|
| Store activations at every layer for backpropagation | Only store final layer activations |
| Memory: O(L × n × d) where L = layers | Memory: O(n × d) regardless of depth |
| Forward: y = x + F(x) | Forward: y₁ = x₁ + F(x₂), y₂ = x₂ + G(y₁) |
| Backward: need stored activations | Backward: recompute x₂ = y₂ - G(y₁), x₁ = y₁ - F(x₂) |
Reformer vs Other Efficient Attention
| Method | Complexity | Exact? | Memory | Best For |
|---|---|---|---|---|
| Full Attention | O(n²) | Yes | O(n²) | Short sequences (<2K) |
| FlashAttention | O(n²) FLOPs, O(n) memory | Yes | O(n) | Standard training (exact, fast) |
| Reformer (LSH) | O(n log n) | No (approximate) | O(n) | Very long sequences (64K+) |
| Longformer | O(n × w) | Exact (sparse) | O(n × w) | Long documents (4K-16K) |
| Performer | O(n) | No (approximate) | O(n) | When linear complexity critical |
Reformer is the pioneering memory-efficient transformer for very long sequences — combining LSH attention (O(n log n) approximate attention that groups similar tokens via hashing) with reversible residual layers (O(n) activation memory regardless of depth), demonstrating that both the compute and memory barriers of standard transformers can be dramatically reduced for processing sequences of 64K+ tokens, trading exact attention for efficient approximation.