Sparse & Linear Attention (2D): Taming the O(n²) Transformer Bottleneck

Sequence Length

Tokens N 24

Attention Pattern

Window radius w 3

Animate query sweep

Drag the heatmap to pan, scroll to zoom, click a row to inspect its attention profile.

Live Cost

Materialized pairs576

Est. ops vs full O(n²)100%

Speedup ×1.0×

Score memory2.3 KB

How it works

Standard self-attention scores every query token against every key token:

Attention(Q,K,V) = softmax(QKᵀ/√d) V
cost(full)    = O(n² · d)      — n×n score matrix
cost(window)  = O(n·w · d)     — each token sees only 2w+1 neighbours
cost(strided) = O(n·(n/s) · d) — every s-th key, dilated
cost(linear)  = O(n·r · d)     — Performer-style random features φ(x),
                                  never materializes the n×n matrix:
                                  φ(Q)(φ(K)ᵀV) computed in O(n·r·d)

Full — every one of the n² query-key pairs is computed; this is what BERT/GPT-2-era transformers do, and it's why context length used to be capped near a few thousand tokens.
Window — a token only attends to its 2w+1 nearest neighbours (Longformer-style local attention). Cost grows linearly in n.
Strided — a token attends to every s-th token across the whole sequence (dilated/strided attention), trading density for long-range reach at lower cost than full.
Linear — a Performer/Linformer-style kernel approximation. Queries and keys are projected through r random features φ(·); the n×n heatmap shown here is only a reconstruction for comparison — the real algorithm computes φ(K)ᵀV first (r×d) and never forms the full grid, which is why its cost line is flat in n.

Each heatmap cell (i, j) is the softmax attention weight query token i places on key token j (brighter = higher weight); cells the pattern skips are flattened to near-zero and dimmed. The left mini-panel plots the currently swept row's full weight profile; the right mini-panel compares materialized pairs across all four patterns at the current N (log scale).

Full Attention

drag · scroll

High attention weight Skipped / masked pair

Row i attention profile

Pairs by pattern (log)