Linear attention

Overview

Softmax attention scores every query against every key, so training costs grow with the square of the length and decoding keeps a cache that grows with it 2103.02143 §2.1. Linear attention writes the score as a product of a feature map on the query and a feature map on the key. Because of that product, the sum over keys can be formed once and reused for every query, so the cost is linear in length 2103.02143 §3.1, eq. 5. In causal (left-to-right) use, the sums become a running state that is updated per token 2103.02143 eq. 6. The state has a fixed size, which is what makes generation cheap 2405.21060 §1.

Not every linear-cost approximation uses a feature map. The landmark (Nyström) approximation in TransMIL replaces the score matrix with a product through a small set of landmark tokens, and the paper states that this lowers the cost from quadratic to linear 2106.00908 §3.3. XCiT takes a different route: its attention matrix is over feature channels rather than tokens, so the matrix is d × d and its cost grows linearly with the number of tokens 2106.09681 §3.1.

The papers differ in four choices: the kernel (a feature map, a landmark approximation, a channel-wise cross-covariance or a Gumbel-noised feature map), whether the running state forgets old tokens, whether the state is a matrix with gates (mLSTM), and whether attention is exact inside blocks and linear only between them (FLASH). This page rests on eight cards, so it covers those variants and not the full literature. LoFTR, TransMIL, XCiT and NodeFormer use their attention inside a larger model, so their headline results do not isolate the attention itself (see the notes under each variant). FLASH's headline results compare whole models with Transformer baselines too.

Variants

Random Fourier features (RFA)

The feature map is a set of random Fourier features. Random projections are sampled from a Gaussian, so the inner product of two features estimates a Gaussian kernel, which stands in for the softmax exponential 2103.02143 §2.2, eq. 2–4. The main derivation assumes unit-norm queries and keys. An unnormalised form exists, at the cost of an extra per-token scalar 2103.02143 §3.1, app. A.4. An optional sigmoid gate decays the old state and scales the new term, which gives a recency bias 2103.02143 §3.2, eq. 7. Card: Random Feature Attention.

Author-reported results, small models unless stated:

• WikiText-103 test perplexity: softmax base 34.5, RFA-Gaussian 35.7, gated RFA-Gaussian 32.7 2103.02143 tab. 1. The stateful gated variant, which carries state across blocks, reaches 30.5 2103.02143 tab. 1 §4.1.
• Machine translation: gated RFA-Gaussian matches the base on WMT14 EN-DE (28.1 vs 28.1 BLEU) and decodes 1.8x faster 2103.02143 tab. 2.
• Decoding at 2,048 tokens in the authors' conditional setting: about 12x faster with under 10% of the memory 2103.02143 §5, fig. 2.

Feature size matters. Cross attention with 32 or 64 random features failed to converge, and causal attention with 32 failed 2103.02143 tab. 9. The paper's own practical starting points are 128 for cross attention and 64 for causal attention in its MT setup 2103.02143 §4.2.

Elu-plus-one feature map

The elu-plus-one feature map from Katharopoulos et al. (2020) is the baseline in the RFA paper 2103.02143 §4.1. It is the cheapest option in that paper, because it needs no random projections. It was clearly weaker on quality in the authors' runs: 40.1 test perplexity against 34.5 for the softmax base 2103.02143 tab. 1, and 21.3 BLEU on EN-DE against 28.1 2103.02143 tab. 2. Its decoding speed was 2.0x in the same table. The original paper has no card here, so nothing beyond these numbers is used.

LoFTR uses the same map, φ(x) = elu(x) + 1, in the self and cross attention of its coarse module 2104.00680 §3.2. Its official code, which is Apache-2.0, adds a small epsilon to the normaliser and scales the values before the reduction, to avoid fp16 overflow. The paper's formula has neither 2104.00680 §3.2. The LoFTR card has no ablation that replaces this attention with full softmax attention, so the effect of the approximation in that matcher is not measured 2104.00680 §4.4. Card: LoFTR.

Decaying linear attention with a scalar state (SSD)

Write a state-space model with one scalar decay per head and the output becomes masked linear attention. The plain causal form is the special case with decay 1 2405.21060 §5.2. The paper calls this link state space duality (SSD). It restricts the decay to one scalar per head, so Mamba's per-channel diagonal decay is not covered 2405.21060 §5, §7.2. Card: Transformers are SSMs.

The chunked SSD algorithm uses quadratic blocks inside each chunk and a short recurrence between chunks. Most of its work is matrix products 2405.21060 §6. Author-reported results:

• Speed: 2–8x faster than Mamba's fused scan at state size 64, and faster than FlashAttention-2 from sequence length 2K 2405.21060 fig. 10 §9.3.
• Quality: zero-shot average 60.2 for Mamba-2-2.7B against 58.3 for Pythia-6.9B 2405.21060 tab. 1 tab. 9.
• Hybrid: six attention layers in a 2.7B model reach Pile perplexity 5.95, against 6.13 for Transformer++ 2405.21060 tab. 3. The paper reports about 10% attention layers as best in its 350M sweep 2405.21060 §9.2.3 tab. 2.

The feature map has a small effect in this form. At 130M, the kernels fall between 11.57 and 11.97 perplexity, except Performer at 12.21. The no-kernel baseline is 11.58 2405.21060 tab. 6. The default is Swish 2405.21060 §7.3, and the authors conclude that the kernels do not beat a plain pointwise activation 2405.21060 §9.4.3.

Landmark (Nyström) approximation

TransMIL approximates softmax self-attention with m landmark tokens picked from the queries and keys. The paper states that this reduces the cost from O(n²) to O(n) in the number of tokens 2106.00908 §3.3. The landmark count, the pseudo-inverse iterations and the head count are not stated 2106.00908 eq. 9, app. B. Card: TransMIL.

The attention sits in a two-layer Transformer over all patch features of a whole-slide image. A bag holds about 8,800 patches on average in CAMELYON16 2106.00908 §1 §4. The model reaches AUC 0.9309 on CAMELYON16, the highest of the compared methods 2106.00908 tab. 1. The paper does not compare the Nyström layers with softmax attention, so these results do not isolate the approximation.

Gumbel-Softmax over random features (NodeFormer)

NodeFormer lets every node attend to every other node in each layer, and approximates the softmax kernel with positive random features, so the cost is linear in the number of nodes 2306.08385 §3.1, eq. 3–5. The sums over nodes are formed once per layer and shared by all queries, which avoids the node-by-node similarity matrix 2306.08385 §3.1, eq. 5. To learn a sparse latent graph, the paper adds Gumbel noise to the logits with a temperature τ. The kernelized form multiplies each key feature by e^(g/τ) before the sums are taken, and the output is averaged over K draws 2306.08385 §3.1, eq. 6–7. The approximation error is bounded in terms of the number of random features m and τ, not the node count 2306.08385 §3.2. The paper notes that a large τ makes every edge weight close to 1/N, which turns the layer into mean pooling, and that a small τ needs a large m 2306.08385 §3.2. Optional terms use an observed graph as a relational bias and as an edge likelihood loss, at a cost of O(N + E) 2306.08385 §3.3, §3.4. Card: NodeFormer.

Author-reported results, five runs on random 50/25/25 splits 2306.08385 §4:

• OGB-Proteins ROC-AUC: 77.45 for NodeFormer against 74.63 for GraphSAINT-GAT 2306.08385 tab. 2.
• Amazon2M accuracy: 87.85 for NodeFormer against 85.17 for GraphSAINT-GAT 2306.08385 tab. 3.
• Replacing the Gumbel operator with an ordinary softmax, at τ = 1.0 or τ = 0.25, gives 75.50–76.18 on OGB-Proteins and 87.02–87.55 on Amazon2M 2306.08385 tab. 2 tab. 3. The gap to the full model is 1.3 to 1.9 points on OGB-Proteins and 0.3 points on Amazon2M, so the margin over softmax is small on the second task.
• Cost: the reported savings, up to 93.1% in training time, are against LDS and IDGL, which are other structure learners 2306.08385 §4.4 fig. 3. The paper does not report a direct cost comparison with full softmax attention on these tasks, so the table does not isolate the linear kernel's speed benefit.

The released code differs from the paper in two places that matter when reimplementing. It scales queries and keys by d^(−1/4), so the effective exponent is q·k/(τ√d) rather than q·k/τ nodeformer.py. It also draws Gumbel noise only in training, and predicts at test time with the deterministic kernel softmax at the same τ nodeformer.py. The code applies e^(g/τ) without a maximum subtraction, which can overflow in float32 for large draws nodeformer.py. The repository has no LICENSE file, so this page cites the code by path and line and quotes nothing from it.

Cross-covariance attention over channels (XCA)

XCiT builds the attention matrix from the queries and keys across all tokens, but its size is set by the channels: the d × d product of transposed keys and queries replaces the N × N token-to-token matrix 2106.09681 §3.1. Each channel of queries and keys is scaled to unit length over the tokens, which keeps every entry in [−1, 1]. A learned temperature then scales the matrix before a softmax over key channels 2106.09681 §3.2, eq. 1. The channels are split into heads, which divides the aggregation cost by the head count 2106.09681 §3.2. Total cost is O(N d² / h) time and O(d² / h + N d) memory 2106.09681 §3.2.

Because the matrix sums over every token, XCA is not a running state and has no causal form. XCiT also adds a depth-wise 3 × 3 convolution block so that neighbouring patches can exchange information directly, since XCA alone does not mix them 2106.09681 §3.3. Card: XCiT.

Author-reported results, supervised ImageNet-1k with distillation tab. 2: XCiT-S12/16 at 224 px reaches 83.3 top-1 with 4.8B FLOPs, against 81.2 with 4.6B FLOPs for DeiT-S (distilled). The paper's larger XCiT models are reported at higher resolutions, which raises FLOPs 2106.09681 tab. 2. These numbers come from a whole image classifier, so they do not isolate the attention.

Matrix memory with exponential gates (mLSTM)

The mLSTM cell keeps a d × d matrix memory C and a normaliser vector n. Keys and values are written with an outer-product update, and the query reads the matrix back 2405.04517 §2.3:

• C_t = f_t C_{t−1} + i_t v_t k_tᵀ, and n_t = f_t n_{t−1} + i_t k_t 2405.04517 §2.3.
• The read-out is divided by max(|nᵀq|, 1) 2405.04517 §2.3.

The paper's version has no link from the previous output into the gates, so training can use a parallel form over the whole sequence [arxiv:2405.04517 app. A.3]. The input gate is exponential, and the forget gate may be too, with a stabiliser state that keeps the exponentials from overflowing 2405.04517 §2.2, app. A.2. This is a causal recurrence with input-dependent decay, so it sits in the running-state family above, but it writes a matrix rather than a feature map.

Author-reported ablation, 15B SlimPajama tokens 2405.04517 tab. 2: exponential gating with no matrix memory gives 17.70 validation perplexity, and adding the matrix memory gives 13.48. The full comparison at 15B tokens puts xLSTM[1:0] at 13.43, against 13.70 for Mamba and 14.25 for Llama 2405.04517 tab. 1. At 350M and 300B tokens the same table reports 11.31 for xLSTM[1:0], against 11.64 for Mamba and 12.19 for Llama 2405.04517 tab. 3. The paper's notation xLSTM[a:b] puts mLSTM and sLSTM blocks in a ratio, so [1:0] uses only mLSTM blocks 2405.04517 §2.4. Card: xLSTM.

Chunked attention: exact inside blocks, linear across them (FLASH)

FLASH cuts the sequence into non-overlapping chunks. Inside each chunk it runs exact squared-ReLU attention, which is quadratic in the chunk length. Across chunks it adds a linear term built from one summary per chunk 2202.10447 §3.2, eq. 7–8. With a fixed chunk size the cost is linear in length. In causal use the prefix sum runs over chunks rather than tokens, so the sequential loop is shorter by the chunk size 2202.10447 §3.2.1. The attention sits inside a gated attention unit: one softmax-free head with a learned gate, and a shared width of 128 2202.10447 §2 eq. 3–5. The authors say decoding needs constant work per step, because the linear cache is one matrix and the local attention covers one chunk 2202.10447 §3.2.1.

Author-reported results, against Transformer baselines of about 110M parameters for C4 and Wiki-40B 2202.10447 §4:

• Wiki-40B language model at context 8192, on 64 TPU-v4 cores: perplexity 15.109 for FLASH against 15.254 for Transformer++, with step latency 62.5 ms against 261.0 ms 2202.10447 tab. 13.
• Training-cost speedup to reach Transformer++ quality on Wiki-40B, contexts 512 to 8192: 1.2x to 4.9x 2202.10447 §4.2 fig. 1. On C4 masked language modelling the range is 1.0x to 4.8x 2202.10447 §4.1.
• PG-19 at context 8192 with about 500M parameters: 12.12x speedup against Transformer+ 2202.10447 tab. 3.
• TriviaQA: the variant with a first-to-all token and state size 512 reaches pretraining perplexity 3.24 against 3.48, but fine-tuning F1 is 73.9 against 74.2 for Transformer+ 2202.10447 tab. 4. The pretraining gain did not carry over to F1 in that setting.

The paper's ablations keep only one of the two parts, local or global 2202.10447 §4.4. The authors also argue that the usual Transformer baselines are weak, and that adding rotary positions and a GLU makes them much stronger 2202.10447 §4. Compare against the stronger version.

The paper's equations and its appendix listings disagree on details that change the output. The listings divide the scores by the length and the chunk summaries by the chunk size. Equations 5, 7 and 8 use neither factor 2202.10447 eq. 5, eq. 7–8, app. D. The paper does not say which version produced its numbers, so test both. Card: Transformer Quality in Linear Time.

When to use which

Implementation sketch

The core is a running sum over the prefix. The feature map and the optional decay are the two places the variants differ.

# Causal linear attention with a pluggable feature map  [arxiv:2103.02143 eq. 6]
# Input:  queries, keys (T, d); values (T, d_v); feature map phi
# Output: outputs (T, d_v)
# Decay per token is optional  [arxiv:2405.21060 §5.2]

def causal_linear_attention(queries, keys, values, phi, decay=None):
    # Running state: a feature-by-value matrix and a normaliser  [eq. 6]
    state = zeros(D, d_v)                    # D = feature size
    norm = zeros(D)
    outputs = []
    for t in range(T):
        # Map the current key and query into feature space  [eq. 2]
        k_feat = phi(keys[t])                # (D,)
        q_feat = phi(queries[t])             # (D,)
        # Optional scalar decay forgets old tokens; 1 gives plain sums  [§5.2]
        a = 1.0 if decay is None else decay[t]
        # Add this key-value term to the running state  [eq. 6]
        state = a * state + outer(k_feat, values[t])
        norm = a * norm + k_feat
        # Normalised read-out for step t  [eq. 6]
        # The denominator has no epsilon in the paper  [eq. 6]
        outputs.append((q_feat @ state) / (q_feat @ norm))
    return stack(outputs)

The RFA gate scales the new term by one minus the gate instead of using the decay directly 2103.02143 eq. 7. The chunked SSD form 2405.21060 §6 computes the same sums in blocks and is needed for training speed on GPUs. The mLSTM step fits this loop with the identity as its feature map and input-dependent gates, plus a normaliser read-out 2405.04517 §2.3. The landmark variant does not follow this loop: it forms the score matrix through m landmarks, so it is not a running sum 2106.00908 eq. 9. XCA is not a loop over tokens at all: it forms the d × d matrix from every token before the read-out 2106.09681 §3.1. NodeFormer's step fits the loop with a random-feature map and one extra multiplier per key, e^(g/τ), drawn fresh for each sample 2306.08385 eq. 7.

FLASH fits the loop at chunk granularity. The exact part runs inside each chunk, and the linear state is updated once per chunk rather than once per token. In the causal version the linear term reads only earlier chunks 2202.10447 §3.2 eq. 8, §3.2.1.

Contradictions

None between the numbers in the cards, because the papers compare against different baselines. RFA measures against softmax at matched sizes. SSD measures against Mamba's scan, FlashAttention-2 and Pythia 2103.02143 §4 2405.21060 §9.2 §9.3. Their speed numbers are not comparable. The xLSTM table ranks two papers' linear Transformers, RetNet at 16.23 and GLA at 16.15, behind Llama at 14.25 on 15B tokens 2405.04517 tab. 1. The RFA and SSD cards report no such comparison, so this ranking does not contradict them, but it does show that linear variants can trail softmax under one author's setup. LoFTR, TransMIL and XCiT compare whole models with other matchers, MIL methods and classifiers 2104.00680 §4.1 §4.2 2106.00908 §4 2106.09681 tab. 2, and NodeFormer compares against GNN and sampling baselines 2306.08385 tab. 2 tab. 3. Their headline results do not test the attention alone.

Within the RFA paper, the abstract and conclusion say "twice as fast", while Table 2 gives 1.8–1.9x 2103.02143 §7 tab. 2.

Within the FLASH paper, the introduction attributes the 12.1x PG-19 speedup to Transformer++, while Table 3 gives it against Transformer+ [arxiv:2202.10447 §1 tab. 3]. The table is the baseline to use. The scaling and segment-mask conventions also differ between its equations and appendix listings, as noted above 2202.10447 eq. 7–8, app. D.

Papers

• Random Feature Attention: random Fourier features as the feature map, and an optional gate for recency.
• LoFTR: elu-plus-one linear attention in the coarse module of a detector-free matcher.
• TransMIL: Nyström landmark approximation of self-attention over the patches of a slide.
• NodeFormer: Gumbel-Softmax over random features for all-pair message passing on graphs, linear in nodes.
• XCiT: cross-covariance attention over channels, linear in the number of tokens.
• xLSTM: mLSTM matrix memory with exponential gates, next to the sLSTM scalar-memory cell.
• Transformers are SSMs: scalar-decay linear attention as SSD, and the chunked algorithm behind Mamba-2.
• Transformer Quality in Linear Time: FLASH, a gated attention unit with exact attention in chunks and linear attention across them.