Sparse Attention Patterns

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Fixed Sparse Patterns
5. Strided Attention (Sparse Transformer)
6. BigBird: Random + Window + Global
7. Hash-Based Attention (Reformer)
8. Routing Transformer
9. Learned Sparsity (Adaptively Sparse)
10. Theoretical Guarantees
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Classify sparse attention methods as fixed-pattern vs content-based.
2. Derive the $O(T\sqrt{T})$ complexity of strided attention.
3. Prove that BigBird's pattern can simulate any Turing machine (universal approximation).
4. Explain LSH attention and its collision probability guarantee.
5. Compare the expressiveness-efficiency tradeoffs of each method.

Notation

• $\mathcal{N}(i) \subseteq \{1, \ldots, T\}$ — the set of positions token $i$ attends to
• $s$ — stride length
• $g$ — number of global tokens
• $r$ — number of random connections
• $b$ — number of hash buckets

Core Intuition

Full attention attends to all $T$ positions — but most attention weights are near-zero (softmax concentrates on few keys). Sparse attention pre-selects which positions to attend to, reducing from $T^2$ to $T \cdot |\mathcal{N}|$ computations. The challenge: choose $\mathcal{N}(i)$ to preserve the important connections while being efficient to compute.

Fixed Sparse Patterns

Definition: The attention connectivity $\mathcal{N}(i)$ is determined by position alone (independent of content).

Advantages: Simple, hardware-efficient, predictable memory. Disadvantages: May miss content-dependent important connections.

Strided Attention (Sparse Transformer)

OpenAI's Sparse Transformer (2019): Factorize attention into two patterns:

Pattern 1 (local): Attend to the previous $\sqrt{T}$ positions:

Pattern 2 (strided): Attend to every $\sqrt{T}$-th position:

Alternating layers: Odd layers use pattern 1, even layers use pattern 2.

Complexity: Each pattern has $|\mathcal{N}| = \sqrt{T}$ connections per token. Total: $O(T\sqrt{T})$ per layer.

Effective full connectivity: After 2 layers (one local + one strided), information can flow between any pair of positions through an intermediate "stride" position.

BigBird: Random + Window + Global

BigBird (2020) combines three patterns:

1. Window attention (local): Each token attends to $w$ neighbors.
2. Global tokens: $g$ designated tokens attend to/from all positions.
3. Random connections: Each token randomly attends to $r$ other positions.

Complexity: $O(T(w + g + r))$ — linear in $T$.

Theoretical result (Zaheer et al.): BigBird with $g \geq 1$ global token is a universal approximator of sequence-to-sequence functions, and can simulate any Turing machine.

Proof sketch: The global tokens act as a "memory bus" — any position can write to and read from them, enabling arbitrary communication patterns through multiple layers.

Hash-Based Attention (Reformer)

Key idea: Similar queries and keys should attend to each other. Use Locality-Sensitive Hashing (LSH) to group similar vectors into the same bucket, then attend within buckets.

LSH with random projections: Hash function $h(\mathbf{x}) = \arg\max[\mathbf{Rx}; -\mathbf{Rx}]$ where $\mathbf{R} \in \mathbb{R}^{b/2 \times d}$ is random.

Collision probability: $\Pr[h(\mathbf{q}) = h(\mathbf{k})] \propto 1 - \frac{\angle(\mathbf{q}, \mathbf{k})}{\pi}$ — similar vectors hash together with high probability.

Algorithm:

1. Hash all queries and keys into $b$ buckets.
2. Sort by bucket (tokens in same bucket are adjacent).
3. Apply attention within each bucket (size $\approx T/b$).

Complexity: $O(T \cdot T/b) = O(T^2/b)$. With $b = \sqrt{T}$: $O(T^{3/2})$. With multiple rounds of hashing: further improvement.

Issue: Approximate — may miss important key-query pairs that hash to different buckets. Multiple hashing rounds reduce miss probability.

Routing Transformer

Content-based sparse attention: Learn which tokens should attend to each other.

Algorithm:

1. Cluster tokens using k-means on their representations: $C_1, \ldots, C_k$.
2. Each token attends only to tokens in the same cluster.

Clustering: Online k-means updated during training. Cluster assignments based on current token representations.

Complexity: $O(T \cdot T/k)$ per layer. With $k = \sqrt{T}$: $O(T\sqrt{T})$.

Advantage over LSH: Learned clustering adapts to data distribution; hash functions are data-independent.

Learned Sparsity (Adaptively Sparse)

Top-k attention: Compute full attention scores, keep only top-$k$:

Problem: Still requires $O(T^2)$ to compute all scores (then discard most). Useful only for reducing memory, not compute.

Entmax / $\alpha$-entmax: Replace softmax with a mapping that produces exact zeros:

where $H_\alpha$ is the Tsallis $\alpha$-entropy. For $\alpha > 1$, the solution has exact zeros — sparse attention without explicit top-k.

Theoretical Guarantees

Theorem (BigBird). Any function computable by a full-attention transformer with $L$ layers can be approximated by a BigBird sparse transformer with $O(L)$ layers, window $w = O(\sqrt{T})$, and $g = O(1)$ global tokens.

Theorem (Reformer, probabilistic). With $m$ rounds of LSH hashing, the probability of missing an important key (one with attention weight $> \epsilon$) is at most $(1 - p_\epsilon)^m$, where $p_\epsilon$ depends on the cosine similarity threshold.

Information-theoretic lower bound: Any sparse attention pattern with $|\mathcal{N}| = k$ per token cannot distinguish more than $\binom{T}{k}$ different attention configurations — suggesting a minimum sparsity level for given tasks.

Common Pitfalls

Pitfall 1. Using sparse attention for tasks requiring precise long-range retrieval. If the task requires finding a specific needle in a haystack, and the sparse pattern doesn't include that position, performance drops catastrophically.

Pitfall 2. Ignoring hardware efficiency. Strided access patterns are less hardware-friendly than dense or block-sparse patterns. Block-sparse attention (attending to contiguous blocks) is much faster in practice despite similar theoretical complexity.

Pitfall 3. Comparing sparse methods only on perplexity. Sparse attention often matches full attention on perplexity but fails on downstream tasks requiring long-range reasoning.

Summary

• Strided: $O(T\sqrt{T})$, simple, effective for image/audio.
• BigBird: $O(T)$ with universal approximation guarantee (window + global + random).
• Reformer (LSH): $O(T\sqrt{T})$, content-aware, approximate.
• Routing: Learned clustering, content-adaptive, $O(T\sqrt{T})$.
• Entmax: Exact zeros in attention weights, but still $O(T^2)$ to compute scores.
• Modern trend: FlashAttention makes full $O(T^2)$ attention fast enough for most lengths; sparse methods mainly for $T > 64K$.

Exercises

Exercise 1. For a sequence of length $T=16384$ with strided attention ($s=\sqrt{T}=128$): compute the total attention FLOPs and compare to full attention.

Exercise 2. Prove that BigBird with at least one global token can propagate information between any two positions in 2 layers.

Exercise 3. For Reformer with $b=64$ buckets and $T=8192$: compute the expected bucket size and attention complexity per token.

Exercise 4. Explain why block-sparse attention is faster on GPUs than element-wise sparse attention (same number of non-zero entries).

Exercise 5. Design a sparse attention pattern for a model that must process a 200K-token document but answer questions about specific sections.