Linformer — Low-Rank Attention Projection

§ 1 · Premise

Attention is approximately low-rank — bake that in

Standard self-attention over a sequence of length $N$ with head dimension $d_h$ computes an $N \times N$ probability matrix $P = \mathrm{softmax}(QK^\top/\sqrt{d_h})$ and then applies it to $V \in \mathbb{R}^{N \times d_h}$. Time and memory are $\Theta(N^2 d_h)$, and the dominant cost at $N \gtrsim d_h$ is the matrix $P$ itself, which carries $N^2$ logits regardless of content. At $N = 4{,}096$, a single attention head materializes 67 MB of fp32 logits per batch element before softmax — and the model has $H$ of these per layer.

Wang et al. (§ 2 and Figure 2 of arXiv 2006.04768) ran a singular-value analysis of $P$ in pretrained RoBERTa and observed that its singular spectrum decays rapidly: ~90% of the spectral energy is captured by the top-128 singular vectors at $N = 512$. They proved (their Theorem 1, via Johnson–Lindenstrauss) that for any softmax attention matrix and any $\varepsilon > 0$, there exists a rank-$k$ approximation with $k = \Theta(d_h \log(d_h) / \varepsilon^2)$ that achieves $\varepsilon$ error in the resulting output — and crucially, $k$ does not grow with $N$.

The architectural question: can the low rank be enforced at training time rather than discovered post hoc on a trained model? Linformer’s bet: yes, by projecting $K$ and $V$ along the sequence axis to a fixed dimension $k$ via two learned matrices. The resulting attention has $\Theta(N k d_h)$ cost — linear in $N$ — and trains end-to-end.

The predecessor lineage is the broader “fast transformer” line of 2019–2020 — Sparse Transformer (entry) commits to a structured sparsity pattern, Reformer (entry) hashes to learned buckets. Linformer differs by attacking the low-rank structure of the dense probability matrix itself, not the sparsity of its support.

§ 2 · Derivation

Project the sequence axis, not the feature axis

Let $Q \in \mathbb{R}^{N \times d_h}$, $K \in \mathbb{R}^{N \times d_h}$, $V \in \mathbb{R}^{N \times d_h}$ be the queries, keys, and values for one head ($H$ heads in parallel). Standard scaled dot-product attention is

Step 1: insert a sequence-axis projection. Introduce two learnable matrices $E \in \mathbb{R}^{k \times N}$ and $F \in \mathbb{R}^{k \times N}$, where $k$ is a chosen rank (Wang et al. use $k = 128$ or $256$ regardless of $N$). Apply them to $K$ and $V$ along the sequence axis:

Note that $E$ and $F$ act on the sequence dimension, not the feature dimension. The features $d_h$ are untouched. This is the load-bearing design choice: a feature-axis projection (reducing $d_h$ to $d' < d_h$) loses the model’s representational dimensionality; a sequence-axis projection only loses the number of distinct “memory slots” the queries can attend to, which the low-rank result says is safe.

Step 2: attend against the projected sequence. Replace $K, V$ with $\tilde{K}, \tilde{V}$:

The logit matrix $Q \tilde{K}^\top$ is now $N \times k$, not $N \times N$. The softmax is over $k$ entries per row. The final matmul against $\tilde{V}$ is $N \times k \cdot k \times d_h$.

Step 3: count the FLOPs. Each query computes $k$ inner products against the $k$ projected keys, then weights $k$ projected values:

Linear in $N$ at fixed $k$. The $EK$ and $FV$ projections themselves are $k \times N \cdot N \times d_h = \Theta(N k d_h)$ matmuls, so they cost the same as the attention step rather than dominating.

Step 4: why is the approximation tight? Wang et al. (Theorem 1 of arXiv 2006.04768, restated):

For any $Q, K \in \mathbb{R}^{N \times d_h}$ and any $\varepsilon, \delta \in (0, 1)$, there exists a matrix $E \in \mathbb{R}^{k \times N}$ with $k = \Theta(d_h \log d_h / \varepsilon^2)$ such that, with probability $\ge 1 - \delta$, $\bigl\lVert \mathrm{softmax}(QK^\top) v - \mathrm{softmax}(Q (EK)^\top) (E v) \bigr\rVert \;\le\; \varepsilon \lVert \mathrm{softmax}(QK^\top) v \rVert$ simultaneously for every column $v$ of $V$.

The proof reduces softmax-attention error to Johnson–Lindenstrauss inner-product preservation on the unnormalized logits, then exploits the Lipschitz constant of softmax. The key takeaway is that the required $k$ depends on $d_h$ (head dimension) and the target error, not on $N$ — the projection rank can be a fixed hyperparameter and quality does not degrade as $N$ grows. Empirically (Figure 4 of the paper), $k = 128$ is sufficient up to $N = 4{,}096$.

Step 5: parameter sharing. A naive instantiation has $H$ pairs of $(E_h, F_h)$ matrices per layer, $2 H k N$ parameters total. Wang et al. sweep three sharing modes (§ 4.2):

1. Headwise: separate $(E_h, F_h)$ per head — most parameters, best quality.
2. Key-Value shared: $E_h = F_h$ within a head — half the params, ~0.1 GLUE points worse.
3. Layerwise shared: one $(E, F)$ pair shared across all heads in a layer — fewest params, best Pareto point.

The recommended default is layerwise sharing, which adds $2 k N$ parameters per layer — roughly $0.5\%$ of the FFN’s parameter count at typical $k, N$ scales.

Why this works only for encoders. $E$ has shape $k \times N$, hard-coded at training time to the maximum sequence length $N_{\text{max}}$. Two consequences:

• At inference time, the model cannot handle sequences longer than $N_{\text{max}}$ — the projection has no defined behavior on out-of-range positions.
• Causal masking is incompatible with $E K$. The projection mixes keys across all positions, so a query at position $t$ ends up attending to $\tilde{K}_j = \sum_{s=1}^N E_{js} \mathbf{k}_s$ — a key that has already been contaminated by future positions $s > t$.

Wang et al. address only the bidirectional (encoder) case. Subsequent work — Linear Attention (Katharopoulos et al. 2020), Performer (entry) — solves the decoder case via a feature-map kernel rather than a sequence-axis projection.

Parameter count and complexity summary. Per layer at head dim $d_h$, $H$ heads, sequence length $N$, rank $k$, with layerwise-shared projection:

The FLOP saving vs. standard attention is the factor $k / N$.

§ 3 · Reference implementation

Sketch

def linformer_attention(Q, K, V, E, F):
    # Q, K, V: [B, T, d_h]            T <= T_max, fixed at training time
    # E, F:    [k, T_max]              shared across heads in this layer
    K_proj = einsum("kt,btd->bkd", E, K)            # [B, k, d_h]
    V_proj = einsum("kt,btd->bkd", F, V)            # [B, k, d_h]
    logits = einsum("btd,bkd->btk", Q, K_proj)      # [B, T, k]
    logits = logits / d_h**0.5
    attn   = logits.softmax(-1)                     # softmax over k, not over T
    return einsum("btk,bkd->btd", attn, V_proj)      # [B, T, d_h]

The load-bearing change vs. MHA is the two einsum projections that compress the sequence axis from $T$ to $k$ before any $Q$-dependent computation. The remaining attention is identical in shape to standard MHA with a “key sequence length” of $k$. The implementation fits in a single Hugging Face transformers patch (see linformer-pytorch by Lucidrains for a production-style version: github.com/lucidrains/linformer).

§ 4 · Empirical evidence

Results

Pretraining quality at fixed $N$ (Wang et al. 2020, Table 3). A 12-layer Linformer with $k = 128$ trained on the same RoBERTa pretraining recipe at $N = 512$ matches the dense RoBERTa baseline within 0.5 GLUE points on average and matches on SQuAD-1.1. At $k = 256$ the gap closes to within 0.1 GLUE points. The paper reports a 1.5× wall-clock speedup at $N = 512$ that grows to 5.5× at $N = 4{,}096$ (their Figure 3).

Quality at long sequences (Table 3 again). At $N = 1024$ and $N = 2048$, Linformer at $k = 256$ remains within ~0.3 GLUE points of the dense baseline that was not trained at these lengths — Linformer is the first transformer in their suite to support these sequence lengths on the same hardware. The paper does not have a dense-attention comparison at $N \ge 4{,}096$ because it does not fit on the V100 hardware they used.

Long Range Arena (Tay et al. 2020, arXiv 2011.04006, Table 1). Independent reproduction by the LRA benchmark suite scores Linformer at average 51.36 across the five LRA tasks (ListOps 35.7, Text 53.9, Retrieval 52.3, Image 38.6, Pathfinder 76.3), vs. 54.39 for the dense softmax-attention baseline — a ~3-point gap. Linformer is mid-pack among the LRA variants: better than Reformer (50.67) and Performer (51.41), worse than BigBird (55.0) and Longformer (53.5). LRA uses a fixed $N = 4{,}096$ across tasks, which is exactly the regime Linformer was designed for.

Sensitivity to $k$ (Wang et al. Figure 4). Quality is monotone-increasing in $k$ with a clear knee around $k = 128$ for $N \le 1{,}024$ and around $k = 256$ for $N \le 4{,}096$. Past $k = 512$ there is no measurable improvement and the speedup over dense attention shrinks toward 1×.

Sharing scheme (Wang et al. Table 5). Layerwise sharing loses ~0.1–0.2 GLUE points vs. headwise sharing but uses $1/H$ as many projection parameters. Key-value sharing within a head ($E = F$) is essentially free at $\le 0.1$ point cost. Their recommended default ships layerwise-shared projections.

Why no decoder LLM ships it. Beyond the causality and variable-length issues laid out in § 2, Wang et al. did not run a decoder-side ablation. Subsequent surveys (Tay et al. 2022, “Efficient Transformers: A Survey”, arXiv 2009.06732, §3.4) classify Linformer as a “fixed-pattern + low-rank” hybrid suitable only for the encoder setting. The lineage descendant that does work autoregressively is Linear Attention (Katharopoulos et al. 2020), which uses a feature-map kernel $\phi(\mathbf{q})^\top \phi(\mathbf{k})$ instead of a sequence-axis projection.

Independent reproduction. The HuggingFace nielsr/linformer implementation reproduces the GLUE numbers within 0.3 points using $k = 128$; the FAIR fairseq reference shipped a Linformer variant in 2021 that reproduces the speedup curves on TPU v3 hardware.

No public production-LLM adoption. I do not know of any frontier decoder-only LLM that ships Linformer. The technique persists in long-document encoder applications and as a baseline in efficient-attention research, but not in production generative models.

@article{arxiv2006_04768,
  title  = {Linformer: Self-Attention with Linear Complexity},
  author = {Sinong Wang and others (Facebook AI)},
  year   = {2020},
  eprint = {2006.04768},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2006.04768}
}

@article{arxiv_2006_04768,
  title         = {Linformer: Self-Attention with Linear Complexity},
  author        = {Sinong Wang et al. (Facebook AI)},
  year          = {2020},
  eprint        = {2006.04768},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2006.04768}
}

{
  "id": "arxiv_2006_04768",
  "type": "article-journal",
  "title": "Linformer: Self-Attention with Linear Complexity",
  "author": [
    {
      "literal": "Sinong Wang et al. (Facebook AI)"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2020
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/2006.04768",
  "number": "2006.04768",
  "source": "arXiv"
}

TY  - JOUR
TI  - Linformer: Self-Attention with Linear Complexity
AU  - Sinong Wang et al. (Facebook AI)
PY  - 2020
JO  - arXiv
AN  - arXiv:2006.04768
UR  - https://arxiv.org/abs/2006.04768
ER  -