Efficient Transformers

Sub-quadratic and IO-aware attention mechanisms that conquer the quadratic memory and computation bottleneck. They scale transformers to massive context lengths.

Efficient transformer architectures reduce quadratic complexity through low-rank projection, hashing, kernel linearization, and IO-aware SRAM tiling.

Why Does This Exist?

Standard self-attention in vanilla transformers computes compatibility across all pairs of sequence tokens:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V

Given an input sequence of length $N$ and hidden dimension $d$ , the product $Q K^\top$ produces an explicit $N \times N$ attention matrix. Both the computational time to calculate this matrix and the memory storage to materialize it scale quadratically as $\mathcal{O}(N^2)$ .

For short sequences ( $N = 512$ ), this quadratic overhead is negligible. However, as language and multimodal models scale to long documents, codebases, and audio streams ( $N = 32{,}768$ to $N = 1{,}000{,}000$ tokens), $\mathcal{O}(N^2)$ becomes completely intractable: an $N = 64\text{k}$ sequence would require materializing a 4-billion-element matrix per attention head at every layer, triggering catastrophic GPU Out-Of-Memory (OOM) failures.

Efficient transformer architectures tackle this bottleneck through four distinct mathematical and systems paradigms: low-rank projection (Linformer), sparse hashing (Reformer), kernel feature linearization (Performer), and hardware IO-aware memory tiling (FlashAttention).

Think of It Like This

A town hall debate compared across four crowd-management strategies

Imagine a town hall meeting with 4,000 citizens where every citizen needs to hear and respond to everyone else's ideas:

In the vanilla setup, every citizen conducts a private one-on-one consultation with every other person in the room. This requires $4{,}000 \times 4{,}000 = 16{,}000{,}000$ separate conversations. The building fills with noise, gridlock ensues, and nobody finishes before midnight.

To make the assembly viable, four distinct committee designs emerge:

First, Linformer (The Representative Delegation): rather than debating 4,000 citizens individually, the crowd elects 128 regional delegates. Every citizen speaks only to the 128 delegates, who aggregate the viewpoints. Total discussions drop from sixteen million to roughly half a million.

Second, Reformer (Interest Circles): citizens are sorted into topical breakout rooms using a quick personality questionnaire (Locality-Sensitive Hashing). Citizens debate only within their specific interest room, ignoring unrelated groups.

Third, Performer (The Summary Scoreboard): instead of pairwise chats, each speaker writes their views as numerical attributes on an open community scoreboard. Listeners read the aggregated column totals directly, bypassing pairwise interaction entirely through associative addition.

Fourth, FlashAttention (The High-Speed Micro-Conference): rather than changing the debate rules or taking approximations, the organizers realize the bottleneck was citizens walking back and forth to an off-site archive. Organizers divide attendees into small room batches (SRAM tiles), let them compute exact consensus in small fast sprints, and immediately record final tallies without ever cluttering the main hallway.

How It Actually Works

Four Paradigms for Breaking Quadratic Scaling

The literature on efficient transformers divides into mathematical approximation methods and exact hardware-aware kernel optimizations.

1. Linformer: Low-Rank Sequence Projection

Wang et al. (2020) demonstrated theoretically and empirically that the $N \times N$ attention matrix $P = \text{softmax}(Q K^\top / \sqrt{d})$ has low intrinsic rank: its singular values decay exponentially. Linformer exploits this by projecting the sequence length of Keys and Values from $N$ down to a fixed constant dimension $k \ll N$ using learned linear projection matrices $E, F \in \mathbb{R}^{k \times N}$ :

\tilde{K} = E K \in \mathbb{R}^{k \times d}, \quad \tilde{V} = F V \in \mathbb{R}^{k \times d}

The reduced attention operation evaluates to:

\text{LinformerAttention}(Q, K, V) = \text{softmax}\left(\frac{Q \tilde{K}^\top}{\sqrt{d}}\right) \tilde{V}

The attention matrix $Q \tilde{K}^\top$ shrinks from $N \times N$ to $N \times k$ . This achieves $\mathcal{O}(N \cdot k)$ linear computational time and memory footprint with respect to sequence length $N$ .

2. Reformer: Locality-Sensitive Hashing (LSH)

Kitaev et al. (2020) observed that in the softmax calculation, tokens with large positive dot products dominate the output, while distant orthogonal tokens receive near-zero weights. Reformer uses Locality-Sensitive Hashing (LSH) via random spherical projections to assign query and key vectors into discrete hash buckets:

h(x) = \arg\max \left[ [x R \,;\, -x R] \right]

where $R \in \mathbb{R}^{d \times b/2}$ is a random Gaussian projection matrix. Only queries and keys falling into the identical hash bucket are permitted to attend to one another. Sorting by bucket reduces complexity to $\mathcal{O}(N \log N)$ . In addition, Reformer replaces standard residual layers with reversible RevNet layers, eliminating the need to cache intermediate activations for backpropagation.

3. Performer: Fast Attention via Positive Orthogonal Random Features (FAVOR+)

Choromanski et al. (2020) approximate the softmax attention kernel using positive random feature maps $\phi: \mathbb{R}^d \to \mathbb{R}^m$ :

\text{softmax}\left(\frac{q^\top k}{\sqrt{d}}\right) \approx \phi(q)^\top \phi(k)

Because the kernel decomposes into an inner product, matrix multiplication associativity allows re-ordering the computation:

\text{Attention}(Q, K, V) = (Q' K'^\top) V \equiv Q' (K'^\top V)

where $Q' = \phi(Q) \in \mathbb{R}^{N \times m}$ and $K' = \phi(K) \in \mathbb{R}^{N \times m}$ . By computing $K'^\top V \in \mathbb{R}^{m \times d}$ first ( $\mathcal{O}(N m d)$ operations), followed by multiplication with $Q'$ , the full $N \times N$ matrix is never formed, delivering strict $\mathcal{O}(N)$ linear time.

4. FlashAttention: IO-Aware Tiling

Dao et al. (2022, 2023) showed that on modern GPUs, attention runtime is bounded by memory bandwidth (IO) rather than arithmetic compute throughput (FLOPS). Standard attention repeatedly reads and writes the intermediate $N \times N$ matrix between slow GPU High-Bandwidth Memory (HBM, 80 GB at ~2 TB/s) and fast on-chip SRAM cache (approx. 20 MB per Streaming Multiprocessor at ~19 TB/s).

FlashAttention computes the exact, unapproximated softmax attention by tiling $Q, K, V$ into blocks that fit entirely within SRAM. It uses the online softmax trick to incrementally update running maximums and normalizers across tiles:

m_{\text{new}} = \max(m_{\text{old}}, x_i), \quad d_{\text{new}} = d_{\text{old}} e^{m_{\text{old}} - m_{\text{new}}} + e^{x_i - m_{\text{new}}}

By fusing softmax scaling and value multiplication into a single GPU CUDA kernel, FlashAttention avoids ever writing the $N \times N$ matrix to HBM, cutting memory accesses from $\mathcal{O}(N^2)$ to $\mathcal{O}(N)$ and delivering 2-4x wall-clock speedups.

Worked Example

Compare the computational floating-point operations (FLOPs) of Vanilla Attention versus a Linear Associative Kernel (Performer) on a sequence of length $N = 4096$ , feature dimension $d = 64$ , and feature projection dimension $m = 64$ .

1. Vanilla Attention $\mathcal{O}(N^2 d)$

Step A: Compute $Q K^\top$ : Multiply $Q \in \mathbb{R}^{4096 \times 64}$ by $K^\top \in \mathbb{R}^{64 \times 4096}$ :
$\text{FLOPs}_1 = 2 \times N \times N \times d = 2 \times 4096 \times 4096 \times 64 = 2{,}147{,}483{,}648 \approx 2.15 \times 10^9 \text{ FLOPs}$
Step B: Multiply $A V$ : Multiply Attention matrix $A \in \mathbb{R}^{4096 \times 4096}$ by $V \in \mathbb{R}^{4096 \times 64}$ :
$\text{FLOPs}_2 = 2 \times N \times N \times d = 2{,}147{,}483{,}648 \approx 2.15 \times 10^9 \text{ FLOPs}$
Total Vanilla FLOPs:
$\text{Total}_{\text{vanilla}} \approx 4.29 \times 10^9 \text{ FLOPs}$

2. Linear Kernel Associative Attention $\mathcal{O}(N m d)$

With $Q' \in \mathbb{R}^{4096 \times 64}$ , $K' \in \mathbb{R}^{4096 \times 64}$ , and $V \in \mathbb{R}^{4096 \times 64}$ :

Step A: Compute $M = K'^\top V$ : Multiply $K'^\top \in \mathbb{R}^{64 \times 4096}$ by $V \in \mathbb{R}^{4096 \times 64}$ :
$\text{FLOPs}_1 = 2 \times m \times N \times d = 2 \times 64 \times 4096 \times 64 = 33{,}554{,}432 \approx 3.36 \times 10^7 \text{ FLOPs}$
Step B: Multiply $Q' M$ : Multiply $Q' \in \mathbb{R}^{4096 \times 64}$ by $M \in \mathbb{R}^{64 \times 64}$ :
$\text{FLOPs}_2 = 2 \times N \times m \times d = 2 \times 4096 \times 64 \times 64 = 33{,}554{,}432 \approx 3.36 \times 10^7 \text{ FLOPs}$
Total Linear Kernel FLOPs:
$\text{Total}_{\text{linear}} = 3.36 \times 10^7 + 3.36 \times 10^7 = 6.71 \times 10^7 \text{ FLOPs}$

Speedup Ratio:

\frac{\text{Total}_{\text{vanilla}}}{\text{Total}_{\text{linear}}} = \frac{4.29 \times 10^9}{6.71 \times 10^7} \approx 64.0

By swapping evaluation order via kernel associativity, the operation requires 64 times fewer computations on an $N = 4096$ sequence.

Code

1import numpy as np2
3def linformer_attention(4    Q: np.ndarray,5    K: np.ndarray,6    V: np.ndarray,7    E_proj: np.ndarray,8    F_proj: np.ndarray9) -> np.ndarray:10    """Compute Linformer attention using low-rank projection matrices."""11    # Q: (N, d), K: (N, d), V: (N, d)12    # E_proj, F_proj: (k, N) where k << N13    N, d = Q.shape14    k = E_proj.shape[0]15
16    # Project Keys and Values from N -> k17    K_proj = np.dot(E_proj, K)  # (k, d)18    V_proj = np.dot(F_proj, V)  # (k, d)19
20    # Compute reduced (N x k) attention scores21    scale = 1.0 / np.sqrt(d)22    scores = np.dot(Q, K_proj.T) * scale  # (N, k)23
24    # Softmax along projected key dimension k25    scores_shifted = scores - np.max(scores, axis=-1, keepdims=True)26    exp_scores = np.exp(scores_shifted)27    attn_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)28
29    # Output: (N, k) x (k, d) -> (N, d)30    output = np.dot(attn_weights, V_proj)31    return np.round(output, 4)32
33def linear_kernel_attention(34    Q_prime: np.ndarray,35    K_prime: np.ndarray,36    V: np.ndarray37) -> np.ndarray:38    """Compute associative linear attention: Q' * (K'^T * V)."""39    # Q_prime: (N, m), K_prime: (N, m), V: (N, d)40    # Associative contraction: (m, N) x (N, d) -> (m, d)41    context_kv = np.dot(K_prime.T, V)42
43    # Final projection: (N, m) x (m, d) -> (N, d)44    output = np.dot(Q_prime, context_kv)45
46    # Normalization by row sums47    normalizer = np.dot(Q_prime, np.sum(K_prime, axis=0, keepdims=True).T)48    output = output / (normalizer + 1e-6)49    return np.round(output, 4)50
51# Demonstration: N = 4, d = 2, k = 252Q_demo = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5]])53K_demo = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5]])54V_demo = np.array([[2.0, 1.0], [1.0, 2.0], [3.0, 3.0], [1.5, 1.5]])55
56# Down-projection matrix from N=4 to k=257E_demo = np.array([[0.5, 0.5, 0.0, 0.0], [0.0, 0.0, 0.5, 0.5]])58F_demo = E_demo.copy()59
60out_linformer = linformer_attention(Q_demo, K_demo, V_demo, E_demo, F_demo)61print(f"Linformer output shape: {out_linformer.shape}")62# -> Linformer output shape: (4, 2)63
64print(f"First row output: {out_linformer[0].tolist()}")65# -> First row output: [1.8385, 1.8385]

Watch Out For

Assuming kernel or low-rank approximations preserve exact needle-in-a-haystack retrieval

A common failure mode is deploying low-rank (Linformer) or kernel-based (Performer) models for tasks requiring precise token retrieval across long documents (such as searching for a specific variable definition in a 50,000-token repository).

Because low-rank projections and random kernel maps smooth out high-frequency dot-product peaks, they blur distinct sharp retrieval spikes. Consequently, while they excel at language modeling perplexity and document classification, they perform poorly on fine-grained retrieval benchmarks. For applications requiring exact associative recall, exact IO-aware methods like FlashAttention maintain zero mathematical loss while delivering high throughput.

The Quick Version

Vanilla self-attention scales quadratically ( $\mathcal{O}(N^2)$ ), causing compute and memory bottlenecks on long sequences.
Linformer and Performer achieve linear complexity through low-rank sequence projection and associative random feature kernels.
FlashAttention delivers exact mathematical attention with 2-4x wall-clock speedups by tiling operations inside fast GPU SRAM.