Akshath Tiwari

This is an explainer of the 2020 Reformer by Kitaev, Kaiser, and Levskaya. Its central idea is that softmax attention is sharply peaked so nearly all the attention weight lands on each query’s nearest keys and the rest round to zero, which means attention is really an approximate nearest-neighbor search; exact nearest-neighbor is itself quadratic, but locality-sensitive hashing does it approximately in order n log n by sending similar vectors to the same bucket so each query only compares against keys in its own bucket. It explains locality-sensitive hashing as the opposite of an ordinary hash: the probability two vectors share a bucket rises as they become more similar, and for angular similarity the standard construction is random-projection or angular LSH, assigning each vector to a bucket by which angular sector it falls into after a random rotation, with collision probability equal to one minus theta over pi where theta is the angle between the vectors. Reformer ties queries and keys, shared-QK, so a token’s query and key always hash together, then sorts tokens by bucket and processes them in equal chunks so each attends to its bucket-mates. The complexity is order n log n because balanced buckets have order log n occupancy so attention costs n times log n, and the dominant sort by bucket is also n log n. Because LSH is probabilistic and a single hash can split a query from a relevant key across a sector boundary, Reformer runs multiple rounds of hashing with independent rotations and unions the buckets to raise recall at the cost of more compute. The other half of fitting long sequences is memory: Reformer uses reversible residual layers from RevNets so a layer’s input can be reconstructed from its output, letting activations be recomputed in the backward pass instead of stored, cutting activation memory from order number-of-layers to order one, plus a chunked feed-forward; together with LSH attention this let Reformer train on sequences up to roughly 64 thousand tokens on a single accelerator. Caveats: LSH only approximates full attention so a missed near-neighbor is a real error and high recall needs several hashing rounds; shared-QK constrains expressivity and the sort-and-chunk machinery adds complexity and irregular memory access; it only helps when attention is genuinely sparse; and once FlashAttention made exact attention memory-cheap and fast the motivation to approximate at moderate lengths largely evaporated so LSH attention saw little production adoption, its lasting contribution being the conceptual view of attention as retrieval. In the lineage Reformer is in the attend-to-less family alongside the Sparse Transformer but with content-based rather than fixed hand-designed sparsity, the hash discovering per input which tokens to compare, an idea continued in the Routing Transformer, while the approximate low-rank and kernel branch that pushes cost to order n is the subject of the Linformer and Performer posts. An interactive angular-LSH ring lets the reader rotate the hash, change the number of buckets, add hashing rounds, and click any vector to make it the query and see which same-bucket keys it attends to.