Sparse Attention Mechanisms

When neighbours aren't enough

Consider a legal-tech startup in Bengaluru building a research assistant over Supreme Court of India judgments. A single constitutional-bench ruling routinely runs past 40,000 words. The bench states the question of law in the opening paragraphs, spends the middle third summarising rival counsel's arguments and the precedents each side relied on, and only in the final quarter lays out the ratio decidendi, the actual binding reasoning. A lawyer asks the assistant: "Does the Court's reasoning in paragraph 312 rely on the precedent cited in paragraph 9?" Those two paragraphs sit tens of thousands of tokens apart. A model that lets each token look only at nearby neighbours, however wide that neighbourhood is dialled, will never let paragraph 9 and paragraph 312 exchange information unless the window is widened to cover the whole document, at which point the cost is back to full quadratic attention anyway. What the assistant actually needs is a way for a token to find whichever other tokens are relevant to it, wherever they sit in the sequence, without comparing itself to all of them. That is the problem this chapter solves: sparsity chosen not by position, but by content.

Two ways to be sparse: position rules versus content matching

Every sparse-attention scheme starts from the same observation: in a full self-attention layer, each of the n tokens computes a score against all n tokens, so the score matrix has n² entries and computing it costs O(n²) time and, if materialised, O(n²) memory. Almost every technique for cutting that cost falls into one of two families, and the difference between them is the single most important idea in this chapter.

The first family fixes the sparsity pattern in advance, purely as a function of token position, before the model has looked at a single embedding. A token is told at design time exactly which offsets it is allowed to attend to: its immediate neighbours, tokens a fixed stride away, a handful of globally visible tokens, or some combination of these. The pattern is baked into the architecture and does not change from one input to the next.

The second family, the subject of this chapter, decides sparsity from the content of the tokens themselves, at inference time, for that specific input. The model does not hard-code which positions may talk to which; instead it computes, from the actual query and key vectors, which tokens are likely to have high attention scores, and restricts the expensive score computation to that predicted subset. Paragraph 9 and paragraph 312 in the judgment can end up compared to each other precisely because their content is similar, even though no fixed rule anticipated that they would need to be. This is strictly more expressive than position-based sparsity for exactly the retrieval-shaped queries a legal or research assistant faces, and it is the family this chapter builds from first principles.

Locality-sensitive hashing as an attention router

The central trick, introduced as LSH attention in the Reformer architecture (Kitaev, Kaiser & Levskaya, 2020), reframes attention as an approximate nearest-neighbour search. Two vectors that point in similar directions have a large dot product and therefore a large (pre-softmax) attention score; two vectors pointing in very different directions have a small one. If a hash function could be found that places similar vectors in the same bucket with high probability and dissimilar vectors in different buckets with high probability, then a token would only need to compute exact attention scores against the other tokens sharing its bucket, and against a great many fewer tokens elsewhere in the sequence.

Reformer uses angular locality-sensitive hashing, built from a random rotation. Fix a projection matrix R of shape (d, b/2), where d is the vector dimension and b is the number of buckets. For a vector x, compute the rotated projection xR, concatenate it with its negation to get a length-b vector [xR, -xR], and take the index of the largest entry as the bucket:

h(x) = argmax([xR, -xR])

Two vectors that are close in angle will, after the same rotation R, still be close in angle, and are likely to have their largest rotated coordinate land in the same slot; two vectors pointing in opposite directions are pushed toward opposite ends of the concatenated vector and almost never collide. Because Reformer ties queries and keys, computing Q = K from the same projection, every token's own vector is trivially the best match for itself, which creates a subtlety worth flagging now and returning to at the end of the chapter: the paper has to explicitly mask a token's attention to itself whenever other, more informative tokens share its bucket, otherwise the softmax collapses onto the guaranteed-identical self-comparison and starves out everything else.

Because any single hash draws a hard boundary and a truly similar pair can, by bad luck, land on opposite sides of it, Reformer repeats the hash n_rounds times with independently drawn rotation matrices and unions the bucket-mates found across all rounds before attending. This trades a roughly linear increase in compute for a large increase in the probability that genuinely similar tokens end up compared.

Worked example: hashing and sorting eight tokens by hand

To make the mechanism completely concrete, trace it through a toy sequence of eight token vectors in two dimensions, standing in for eight paragraphs of the judgment at positions p1 through p8 in document order. Fix R to the 2×2 identity matrix (a real model draws R at random; the identity is used here purely so every step can be checked by hand). With R = I, the hash reduces to h(x) = argmax([x1, x2, -x1, -x2]), so a vector hashes to bucket 0 if its first coordinate is its largest positive value, bucket 1 if its second coordinate is, bucket 2 if the negative of its first coordinate is largest (meaning x1 is strongly negative), and bucket 3 if the negative of its second coordinate is largest.

import numpy as np

def lsh_bucket(x, R):
    """Angular LSH hash (Reformer, Kitaev et al. 2020).
    x: token vector, shape (d,)
    R: projection matrix, shape (d, b/2)
    returns: bucket index in [0, b)
    """
    proj = x @ R
    candidates = np.concatenate([proj, -proj])
    return int(np.argmax(candidates))

R = np.eye(2)  # identity rotation, used here only so the hash can be traced by hand
tokens = {
    "p1": np.array([0.90, 0.10]),
    "p2": np.array([-0.20, 0.90]),
    "p3": np.array([-0.85, -0.20]),
    "p4": np.array([0.15, -0.85]),
    "p5": np.array([0.80, 0.30]),
    "p6": np.array([-0.10, 0.95]),
    "p7": np.array([-0.70, -0.40]),
    "p8": np.array([0.05, -0.90]),
}

buckets = {name: lsh_bucket(vec, R) for name, vec in tokens.items()}
print(buckets)
# {'p1': 0, 'p2': 1, 'p3': 2, 'p4': 3, 'p5': 0, 'p6': 1, 'p7': 2, 'p8': 3}

Check p1 = (0.90, 0.10) by hand: the candidate vector is [0.90, 0.10, -0.90, -0.10], and its largest entry, 0.90, sits at index 0, so p1 hashes to bucket 0. Check p3 = (-0.85, -0.20): the candidate vector is [-0.85, -0.20, 0.85, 0.20], whose largest entry, 0.85, sits at index 2, giving bucket 2. Every other row in the printed dictionary can be verified the same way, and notice the pattern: p1 and p5 both point mostly in the +x direction and land in bucket 0 even though they are four positions apart in the document; p3 and p7 both point mostly in the -x direction and land in bucket 2 despite being separated by three other paragraphs. This is exactly the non-local grouping the legal-document example needed: relevance, not proximity, decides who gets compared.

Reformer's next step is to sort the sequence by bucket index so that same-bucket tokens become contiguous, then run ordinary local attention over fixed-size chunks of the sorted sequence:

sorted_positions = sorted(buckets, key=lambda name: buckets[name])
print(sorted_positions)
# ['p1', 'p5', 'p2', 'p6', 'p3', 'p7', 'p4', 'p8']

With a chunk size m = 2, the sorted sequence splits into four chunks, [p1,p5], [p2,p6], [p3,p7], [p4,p8], each of which attends only within itself. (The real Reformer implementation also lets each chunk peek at the previous chunk, as insurance against a bucket boundary that would otherwise be split awkwardly across two chunks; the count below omits that extra look-back for clarity, so it understates Reformer's true operation count by a small constant factor without changing the asymptotic comparison.) Counting query-key pairs actually scored:

n, b = 8, 4
full_pairs = n ** 2
bucket_size = n // b
lsh_pairs = b * bucket_size ** 2
print(full_pairs, lsh_pairs, full_pairs / lsh_pairs)
# 64 16 4.0

Full attention scores all 8 × 8 = 64 pairs. LSH attention, with four balanced buckets of size two, scores 4 × 2² = 16 pairs, a 4× reduction. Notice that the reduction factor equals b, the bucket count, exactly when buckets are balanced: pairs = b · (n/b)² = n²/b, so full/LSH = b.

Why the sort makes it O(n log n)

The toy example fixes n = 8 and b = 4, but the interesting question is what happens as the document grows. In practice the chunk size m is held roughly constant (a hardware-friendly block size, independent of sequence length), which forces the number of buckets b to grow linearly with n, as b ≈ n/m. Substituting into pairs = n²/b gives pairs = n² / (n/m) = n·m, which is linear in n once m is fixed. The attention computation itself is therefore O(n). What is left is the cost of sorting n tokens into their buckets, which is O(n log n) with any comparison sort, and that term dominates the linear attention cost for large n. That is the origin of Reformer's headline O(n log n) complexity, in contrast to full attention's O(n²): it is the sort, not the attention arithmetic, that sets the asymptotic bound.

Scaling the eight-token example up makes the payoff concrete. Take n = 1,000,000 tokens, roughly the length of a large multi-volume case file, with a chunk size of m = 128. Full attention scores n² = 10¹² pairs. LSH attention scores n·m = 1,000,000 × 128 = 1.28 × 10⁸ pairs. The ratio is 10¹² / 1.28 × 10⁸ ≈ 7,813× fewer score computations, and unlike a fixed sliding window, none of those retained pairs were chosen by how close two tokens sit in the document; every one of them was chosen because the hash predicted the two vectors were likely to be similar.

Misconception: sparse does not mean random

A common misreading, especially for a student who has already met random-edge sparse-attention schemes that mix a handful of literally randomly chosen key positions into an otherwise fixed pattern, is to assume that all sparse attention works by picking connections at random or by some other schedule fixed at design time, and that the model simply has to learn to work around whatever fixed or random subset it was given. LSH attention is neither. The bucket assignment is a deterministic function of that input's own query and key vectors, recomputed fresh on every forward pass; feed the model a different document and the same token position will very likely land in a different bucket, next to different partners, because the hash responds to what the vector actually contains, not to where it sits or to a coin flip. The rotation matrix R is the only thing fixed ahead of time (per layer, drawn once), and R only defines how similarity is measured, not which tokens end up similar for a given input. So the sparsity pattern is content-adaptive even though the hash function that produces it is fixed, which is the crucial distinction a student should walk away with: fixed hash function, input-dependent pattern.

The wider family: learned routing and hardware-aligned sparsity

LSH is one way to convert "find similar tokens" into "compute a cheap key and look it up," but it is not the only one. The Routing Transformer (Roy, Saffar, Vaswani & Grangier, 2021) replaces the fixed random hash with online k-means clustering: a set of cluster centroids is maintained and updated as training proceeds, and each query and key is assigned to its nearest centroid, so tokens attend within their cluster. Unlike Reformer's hash, which has no trainable parameters and only reacts to whatever vectors it is given, Routing Transformer's centroids are themselves learned, so the very notion of "which tokens count as similar" adapts over the course of training rather than being fixed by a one-time random draw.

Both of these ideas predate the current generation of frontier models, but the underlying principle, letting the model choose its own sparse attention pattern from content rather than from a hand-designed schedule, reappears in production systems. DeepSeek-AI's 2025 paper on Native Sparse Attention (NSA) pushes content-based routing further by making the selection mechanism itself part of the trained network and by designing the block-sparse memory access pattern to match GPU hardware, so that the theoretical compute savings of sparsity actually translate into measured wall-clock speedups rather than being lost to irregular memory access. It combines a coarse, compressed summary of far-away context with a learned mechanism that selects a small number of fine-grained blocks worth attending to in full detail, which is a direct descendant of the hash-and-chunk idea in this chapter: cheaply narrow down the candidates first, then spend full attention only on the survivors. (Reformer also pairs its LSH attention with reversible residual connections to cut activation memory during backpropagation; that is a memory-engineering technique orthogonal to attention sparsity itself and is mentioned here only so it is not confused with the hashing mechanism this chapter is about.)

Active recall

Attempt each question before reading its answer.

1. Using R = I and h(x) = argmax([x1, x2, -x1, -x2]), which bucket does x = (-0.6, 0.7) hash to?
2. In the worked example, suppose the bucket-labelling convention were flipped so bucket indices are computed as argmax([-x1, -x2, x1, x2]) instead. Does the actual attention connectivity (who attends to whom) change?
3. Ripple effect: still with n = 8 tokens, suppose the number of buckets is increased from b = 4 to b = 8, so that in the best case every token gets its own bucket. Recompute the number of score-pairs, and explain what breaks about the model's ability to answer the paragraph-9-to-paragraph-312 style question from the opening example.
4. For n = 1,000,000 tokens and chunk size m = 128, compute the LSH pair count and the full-attention pair count, and state the speedup factor.
5. True or false, with justification: "Because Reformer ties queries and keys, a token always attends to itself, so self-attention needs no special handling."
6. In one sentence each, contrast how LSH attention (Reformer) and Routing Transformer decide which tokens belong together, and whether that grouping mechanism is itself learned.

Answers.

1. The candidate vector is [-0.6, 0.7, 0.6, -0.7]. The largest entry is 0.7 at index 1, so x hashes to bucket 1.

2. No. Flipping the labelling convention only relabels which integer names each group (bucket 0 might become what was previously called bucket 2, and so on), but it does not change which vectors end up grouped together, since the underlying partition is still determined by which of {x1, x2, -x1, -x2} is largest. The attention pattern depends only on the partition, not on the names assigned to its parts.

3. With b = 8 and n = 8, the best case gives one token per bucket, so bucket_size = 1 and pairs = b · 1² = 8, meaning every token only ever attends to itself. Increasing the bucket count without bound does shrink the pair count further, but it destroys the entire purpose of attention: a paragraph can no longer gather information from any other paragraph, including a genuinely similar one elsewhere in the document. This is why chunk size and bucket count are tuned as a real hyperparameter trade-off rather than pushed to the sparsest possible extreme: too coarse a hash wastes compute on dissimilar pairs, too fine a hash isolates every token from the context it needs.

4. LSH pairs = n·m = 1,000,000 × 128 = 128,000,000. Full pairs = n² = 1,000,000,000,000. Speedup = 10¹² / 1.28×10⁸ ≈ 7,813×.

5. True that a token's own vector is always in its own bucket (a vector trivially has maximal similarity with itself), but false that this needs no special handling. Because the self-comparison score is essentially guaranteed to be the largest one available, the softmax would otherwise concentrate almost all of its weight on the token attending to itself and starve out attention to every other token in the bucket. Reformer explicitly masks a token's attention to itself whenever the bucket contains other tokens, so the model is forced to route attention to the genuinely informative neighbours the hash found, not to a comparison that carries no new information.

6. LSH attention groups tokens using a fixed random rotation applied fresh to each input's own vectors, with no trainable parameters in the hash itself, so the grouping mechanism is fixed while its output is input-dependent; Routing Transformer groups tokens by assigning them to the nearest of a set of cluster centroids that are updated during training, so the grouping mechanism itself is learned and changes as training progresses.

Think About It

Think about this: How would you explain sparse attention mechanisms to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.