Attention Mechanisms · January 2020

Reformer — LSH Attention

advanced

Introduced by Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya (Google Research) in Reformer: The Efficient Transformer (arXiv:2001.04451)

efficiency

Replace dense O(N²) attention with locality-sensitive hashing — bucket similar Q, K together and only attend within buckets. O(N · log N) attention, learnable rather than fixed sparsity pattern.

§ 1 · Premise

Learn the sparsity pattern instead of fixing it

A dense attention layer at sequence length $N = 65{,}536$ with head dimension $d_h = 64$ materializes a $65{,}536 \times 65{,}536$ logit matrix per head — 17 GB at fp32, dwarfing even the activations of the rest of the model. Quadratic compute is the obvious bottleneck; quadratic memory is the actual one in 2020, when single-GPU memory was 16–32 GB.

The 2019–2020 efficient-attention literature converged on two responses. Sparse Transformer (entry, Child et al. 2019) hand-picks a sparsity pattern (strided + fixed) and reaches $\mathcal{O}(N \sqrt{N})$ cost. Linformer (entry, Wang et al. 2020) projects keys/values down to a fixed rank $k$ and reaches $\mathcal{O}(N k)$ . Both fix the sparsity pattern at training time and accept whatever inductive bias that imposes.

Kitaev et al.’s observation: in a trained attention head, only a small fraction of (query, key) pairs have meaningful similarity. The identity of that fraction is content-dependent — it varies by sample, by layer, by head. A hand-designed pattern necessarily wastes compute on low-similarity pairs while missing some high-similarity ones. The right primitive is to discover the similarity structure on the fly.

Locality-sensitive hashing (LSH; Indyk & Motwani 1998, “Approximate Nearest Neighbors”, STOC) provides that primitive. An LSH family is a distribution over hash functions $h: \mathbb{R}^{d_h} \to \mathbb{Z}$ with the property that for any pair $\mathbf{x}, \mathbf{y}$ ,

\Pr_h\!\bigl[h(\mathbf{x}) = h(\mathbf{y})\bigr] \;=\; f\bigl(\mathrm{sim}(\mathbf{x}, \mathbf{y})\bigr),

where $\mathrm{sim}$ is some similarity measure and $f$ is monotone-increasing. For cosine similarity, Charikar’s random-hyperplane hashing (Charikar 2002, STOC) gives an LSH family $h(\mathbf{x}) = \mathrm{sign}(\boldsymbol{\omega}^\top \mathbf{x})$ with $\Pr[h(\mathbf{x}) = h(\mathbf{y})] = 1 - \theta(\mathbf{x}, \mathbf{y})/\pi$ .

Reformer’s bet: use cosine LSH on $Q$ and $K$ to bucket similar tokens together, then attend within each bucket. Compute drops to $\mathcal{O}(N \log N)$ (sort) plus $\mathcal{O}(N \cdot B)$ (bucket-local attention, $B$ = bucket size). With $B \approx \sqrt{N}$ this is $\mathcal{O}(N \sqrt{N})$ — the same as Sparse Transformer, but the sparsity is learned via the projection that defines the LSH.

Reformer’s second contribution is orthogonal to the first: reversible residual layers (Gomez et al. 2017, “The Reversible Residual Network”, arXiv 1707.04585) that eliminate the $\Theta(L)$ activation memory of a depth- $L$ transformer.

§ 2 · Derivation

Shared-QK projections, angular LSH, chunked attention

Step 1: tie $W_Q$ and $W_K$ . A small but load-bearing simplification. Standard MHA uses separate $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_h}$ projections, so $\mathbf{q}_t$ and $\mathbf{k}_t$ for the same token live in different directions. Reformer sets $W_Q = W_K$ (Kitaev et al. § 2.2). This means each token has a single hashable vector $\mathbf{x}_t = W_Q \mathbf{h}_t$ and the LSH bucketing is well-defined: queries and keys either land in the same bucket or they don’t. Kitaev et al. (Table 1) show that shared-QK costs $\le 0.1$ BPC on enwik8 — a small price.

Step 2: angular LSH via random rotation. Given the shared-QK vectors $\mathbf{x}_t$ , define a random rotation matrix $R \in \mathbb{R}^{d_h \times b}$ where $b$ is the number of buckets (Kitaev et al. use $b = 64$ ). The hash is

h(\mathbf{x}) \;=\; \arg\max_i \bigl[\,R^\top \mathbf{x}\,\bigr]_i \quad \text{or equivalently} \quad \arg\max_i \bigl[\,R^\top \mathbf{x};\, -R^\top \mathbf{x}\,\bigr]_i,

selecting the index of the largest signed projection. (The signed version doubles the effective bucket count and reduces tie-breaking ambiguity.) The result is an integer in $\{1, \ldots, b\}$ . Cosine-similar vectors land in the same bucket with high probability — formally, $\Pr[h(\mathbf{x}) = h(\mathbf{y})]$ is monotone in $\cos(\angle(\mathbf{x}, \mathbf{y}))$ (Andoni & Indyk 2008, “Near-Optimal Hashing”, CACM).

Step 3: sort by hash, chunk into fixed-size groups. Sort token positions by $h$ :

\pi \;=\; \mathrm{argsort}\bigl(h(\mathbf{x}_1), h(\mathbf{x}_2), \ldots, h(\mathbf{x}_N)\bigr).

The sort costs $\Theta(N \log N)$ comparisons. Reformer then partitions the sorted sequence into chunks of fixed size $C$ (Kitaev et al. use $C = N / b = 128$ at $N = 8{,}192$ ). Within each chunk, attention is computed densely. Adjacent chunks are also attended to (causally) to handle the case where a hash bucket straddles a chunk boundary.

The chunking is essential: hash buckets in the wild have uneven sizes — some buckets contain many tokens, others contain few. A dense attention over the largest bucket would re-introduce quadratic cost in the worst case. Fixed-size chunks bound the per-step compute at $\Theta(C d_h)$ per token regardless of bucket variability.

Step 4: causal masking and the bucket-spans-chunk-boundary case. After sorting, the original temporal order is lost. The causal mask must be reconstructed in the sorted index space: token $\pi(i)$ may attend to token $\pi(j)$ only if $\pi(i) > \pi(j)$ in the original positions. This is implemented by computing the dense attention within a chunk, then masking out post-sort positions that correspond to pre-sort future tokens.

Within-chunk attention plus one adjacent backward chunk gives effective receptive field $2C$ in the sorted axis, but at low chunk-boundary recall when a hash bucket spans the boundary. Kitaev et al. mitigate this with multi-round hashing (their § 2.4):

\mathrm{out}_t \;=\; \sum_{r=1}^{R} \mathrm{Attn}_r(\mathbf{q}_t, K_{B_r(t)}, V_{B_r(t)}),

where each round $r$ uses an independently sampled rotation matrix $R_r$ and produces its own bucketing $B_r$ . The final output is the average (or attention-weighted combination) across rounds. Multi-round hashing trades compute for recall: $R = 8$ rounds at $b = 64$ gives near-perfect nearest-neighbor recovery (Kitaev et al. § 4.2).

Step 5: total complexity. With $N$ tokens, $b$ buckets, $C = N/b$ chunk size, $R$ hash rounds:

\mathrm{FLOPs}_{\text{LSH}} \;=\; \underbrace{\Theta(R \cdot N \cdot d_h \cdot b)}_{\text{hash projection}} \;+\; \underbrace{\Theta(R \cdot N \log N)}_{\text{sort}} \;+\; \underbrace{\Theta(R \cdot N \cdot C \cdot d_h)}_{\text{within-chunk attention}}.

Choosing $C = \Theta(\sqrt{N})$ gives the canonical $\Theta(R \cdot N \sqrt{N} \cdot d_h)$ cost. For $C = \log N$ (more buckets, smaller chunks) the cost is $\Theta(R \cdot N \log N \cdot d_h)$ . Memory is $\Theta(N d_h)$ — the $N \times N$ logit matrix is never materialized.

Step 6: the second contribution — reversible residuals. Independent of LSH, Reformer adopts RevNets’ reversible residual structure to remove activation memory. A reversible layer takes $(\mathbf{x}_1, \mathbf{x}_2)$ and produces

\mathbf{y}_1 \;=\; \mathbf{x}_1 + F(\mathbf{x}_2), \qquad \mathbf{y}_2 \;=\; \mathbf{x}_2 + G(\mathbf{y}_1),

where $F$ is the attention block and $G$ is the feed-forward block. The inverse is

\mathbf{x}_2 \;=\; \mathbf{y}_2 - G(\mathbf{y}_1), \qquad \mathbf{x}_1 \;=\; \mathbf{y}_1 - F(\mathbf{x}_2).

This means the input activations need not be stored during the forward pass — they can be recomputed during backprop from the output. Activation memory drops from $\Theta(L N d_{\text{model}})$ to $\Theta(N d_{\text{model}})$ — independent of depth. The trade is an extra forward pass per layer during backprop, ~ $33\%$ more FLOPs.

Combining LSH attention ( $\mathcal{O}(N \sqrt{N})$ compute, $\mathcal{O}(N)$ attention memory) with reversible layers ( $\mathcal{O}(N)$ activation memory regardless of $L$ ) lets Reformer train transformers at $N = 65{,}536$ on a single 16 GB GPU — vs. ~ $N = 2{,}048$ for the dense baseline at the same memory budget (Kitaev et al. § 4.1).

§ 3 · Reference implementation

Sketch

def lsh_attention(x, R_rounds, num_buckets, chunk_size):
    # x: [B, N, d_h]  (shared-QK vectors, so q == k == x)
    # R_rounds: list of random rotation matrices, each [d_h, num_buckets/2]
    outputs = []
    for R in R_rounds:
        # Step 1: hash each token to a bucket via signed argmax projection
        proj   = x @ R                                # [B, N, num_buckets/2]
        signed = cat([proj, -proj], dim=-1)            # [B, N, num_buckets]
        hashes = signed.argmax(dim=-1)                # [B, N]
        # Step 2: sort tokens by hash so same-bucket tokens are adjacent
        order  = hashes.argsort(dim=-1)
        x_srt  = gather(x, dim=1, index=order)        # [B, N, d_h]
        # Step 3: split sorted sequence into fixed-size chunks
        chunks = x_srt.view(B, N // chunk_size, chunk_size, d_h)
        # Step 4: dense causal attention within each chunk + one prior chunk
        attn   = chunked_causal_attention(chunks, look_back=1)
        # Step 5: unsort back to original positions
        outputs.append(scatter(attn.view(B, N, d_h), dim=1, index=order))
    # Multi-round: average across rounds
    return stack(outputs).mean(dim=0)

The load-bearing differences vs. dense attention are the sort (which makes same-bucket tokens adjacent), the fixed-size chunking (which bounds per-step compute regardless of bucket imbalance), and the optional multi-round averaging (which improves recall at compute cost).

§ 4 · Empirical evidence

Results

enwik8 character-level language modeling (Kitaev et al. 2020, Table 2). A 3-layer Reformer at $N = 65{,}536$ reaches 1.05 BPC on enwik8, vs. 1.06 for the dense Transformer-XL baseline at $N = 3{,}072$ . With $R = 8$ hash rounds the gap closes to within 0.01 BPC. Compute scales as expected: $\Theta(N \sqrt{N})$ vs. the dense $\Theta(N^2)$ .

Image generation on ImageNet-64 (Kitaev et al. Table 3). A 12-layer Reformer reaches 3.65 bits/dim on ImageNet-64 with $N = 12{,}288$ (the full pixel sequence) — beating the 12-layer dense Image Transformer at 3.77 bits/dim. The model fits in 11 GB of GPU memory due to reversible layers; the dense baseline would have required >40 GB at the same depth.

Quality vs. number of hash rounds (Kitaev et al. Figure 5). With $R = 1$ hash round, Reformer underperforms dense attention by ~0.05 BPC on enwik8. The gap shrinks monotonically with $R$ , reaching parity at $R = 8$ . Beyond $R = 16$ there is no further improvement and the compute exceeds dense attention. The recommended default is $R = 8$ .

Long Range Arena (Tay et al. 2020, arXiv 2011.04006, Table 1). Independent reproduction scores Reformer at average 50.67 across the five LRA tasks (ListOps 37.3, Text 56.1, Retrieval 53.4, Image 38.1, Pathfinder 68.5) vs. 54.39 for the dense softmax baseline — a ~3.7-point gap, the worst among the efficient variants the LRA authors test except for Synthesizer. Reformer underperforms Linformer, Performer, BigBird, and Longformer on this benchmark. Tay et al. attribute Reformer’s weakness to its sensitivity to hash collisions on tasks where attention is sharp.

Bucket imbalance and load balancing. Roy et al. (2021, “Efficient Content-Based Sparse Attention with Routing Transformers”, arXiv 2003.05997, § 4.1) analyze Reformer’s bucket distribution and find that, on natural-language data, $30\%$ of hash buckets contain $> 2 C$ tokens (oversubscribed) and $40\%$ contain $< 0.5 C$ tokens (undersubscribed). The fixed chunking truncates oversubscribed buckets, dropping ~10% of relevant attention. Routing Transformer replaces LSH with k-means clustering to address this and improves on Reformer’s LRA score by 1.5 points.

FlashAttention dominance after 2022. Dao et al. 2022 (arXiv 2205.14135) showed that exact softmax attention with I/O-aware tiling achieves the same memory complexity as Reformer ( $\Theta(N)$ logits in HBM, $\Theta(\sqrt{N})$ in SRAM) without approximation. On A100/H100 at $N \le 64K$ , FlashAttention is faster than any Reformer variant and produces exact softmax. Once FlashAttention shipped in PyTorch (May 2022), demand for approximate sub-quadratic attention at moderate context collapsed.

Reversible residuals survived. The LSH-attention component did not transfer to frontier LLMs, but reversible residuals are still cited as a memory-saving technique in long-context training research (e.g., Memorizing Transformers, Wu et al. 2022, arXiv 2203.08913, adopts the reversible layer structure independently of LSH attention).

No public production-LLM adoption. I do not know of any frontier decoder-only LLM that ships LSH attention. The technique persists as a baseline in efficient-attention research and as a historical reference; the broader idea — content-based sparse attention via clustering or hashing — survives in routing transformers and mixture-of-experts attention proposals.

Lineage

Predecessors: Sparse TransformerSparse Transformer

Cite

BibTeX entry for the original paper

@article{arxiv2001_04451,
  title  = {Reformer: The Efficient Transformer},
  author = {Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya (Google Research)},
  year   = {2020},
  eprint = {2001.04451},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2001.04451}
}

Or cite the paper directly: arXiv:2001.04451.

Export

BibTeX

@article{arxiv_2001_04451,
  title         = {Reformer: The Efficient Transformer},
  author        = {Nikita Kitaev and Łukasz Kaiser and Anselm Levskaya (Google Research)},
  year          = {2020},
  eprint        = {2001.04451},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2001.04451}
}

CSL JSON

{
  "id": "arxiv_2001_04451",
  "type": "article-journal",
  "title": "Reformer: The Efficient Transformer",
  "author": [
    {
      "literal": "Nikita Kitaev"
    },
    {
      "literal": "Łukasz Kaiser"
    },
    {
      "literal": "Anselm Levskaya (Google Research)"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2020
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/2001.04451",
  "number": "2001.04451",
  "source": "arXiv"
}

RIS

TY  - JOUR
TI  - Reformer: The Efficient Transformer
AU  - Nikita Kitaev
AU  - Łukasz Kaiser
AU  - Anselm Levskaya (Google Research)
PY  - 2020
JO  - arXiv
AN  - arXiv:2001.04451
UR  - https://arxiv.org/abs/2001.04451
ER  -