Reformer and LSH Attention

Full self-attention makes every token compute a similarity score with every other token, which is wasteful — for most queries, only a handful of keys actually matter, and the rest contribute near-zero attention weight after the softmax. Reformer's idea is to skip computing the scores that will end up near zero, by first grouping similar queries and keys together and only attending within each group.

It does this with locality-sensitive hashing (LSH): a random hash function is chosen so that vectors pointing in similar directions land in the same hash bucket with high probability, and dissimilar vectors usually land in different buckets. Tokens are hashed into buckets, sorted by bucket, and attention is computed only between tokens sharing a bucket (plus a neighboring chunk, to soften boundary effects) — never across the full sequence. Since bucket sizes are kept roughly constant regardless of how long the input is, total attention cost grows close to linearly with sequence length rather than quadratically.

Reformer pairs this with reversible residual layers, which recompute activations during the backward pass instead of storing them, cutting memory further — the two tricks together are what let it handle sequences of tens of thousands of tokens on ordinary hardware.