Transformer Quality in Linear Time

In plain words

Standard Transformers compare every word with every other word, so the work grows with the square of the text length. That is why most models read only a few thousand words at a time, and why efficient variants that cut this cost often lose quality.

The paper makes two changes. First, it replaces the usual pair of attention and feed-forward blocks with one simpler block that gates the output of a single, weak attention head. The authors find that the gate lets the attention be much weaker without hurting quality. Second, it makes that block cheap on long text. The text is cut into fixed-size chunks. Inside each chunk the model keeps exact attention. Between chunks it passes a compact summary, which keeps the cost linear in length. The design is written so that accelerator chips run it efficiently.

The authors report that the resulting model, FLASH, matches strong Transformer baselines on language and masked-language modeling for lengths from 512 to 8K tokens, and that it trains much faster on long text. Most of the evidence is in training curves and cost comparisons.

TL;DR

FLASH replaces each Transformer block with gated attention units (GAUs). A GAU uses one squared-ReLU attention head with a learned gate §2, eq. 3, eq. 5. For long inputs, the attention is split into exact quadratic attention inside non-overlapping chunks and linear attention between chunks, with the chunk-level sums taken in causal order for language modeling §3.2, eq. 7-8. The reported gains come from training-cost comparisons on TPU-v4 against Transformer baselines, with results in tables and curves tab. 3, tab. 12, tab. 13. Model sizes: a 24-layer, 768-wide model (about 112M parameters) for C4 and Wiki-40B tab. 8, tab. 9, and roughly 500M parameters for PG-19 tab. 10.

Problem

Full softmax attention costs quadratic time and memory in the sequence length, which limits context §1. Many linear-cost attention variants exist, but the authors report three problems §1: their quality is well below augmented Transformers; their memory re-arrangements make them slower than their theory suggests on accelerators such as TPUs and GPUs; and linear attention is slow to train autoregressively, because its prefix sums run as a sequential loop over tokens §3.1. The authors also argue that the common Transformer baselines are weak, and that adding rotary position embeddings and a gated linear unit (GLU) makes them much stronger §4.

Method

Gated attention unit (GAU). A GLU computes U = φ_u(XW_u) and V = φ_v(XW_v), then outputs (U ⊙ V)W_o §2, eq. 1-2. GAU replaces V with a token-mixed value, V̂ = AV, and outputs O = (U ⊙ V̂)W_o §2, eq. 3. A is a T×T matrix of attention weights computed from a shared low-width representation Z = φ_z(XW_z) with width s = 128 §2, eq. 4; §2 footnote. Queries and keys come from Z by a per-dimension scale and offset, and A = relu²(QKᵀ + b), where b is a relative position bias §2, eq. 5. The attention is single-head and softmax-free §2.

A = \mathrm{relu}^2\left(\mathcal{Q}(Z)\,\mathcal{K}(Z)^\top + b\right) eq. 5

The attention adds only W_z (d·s parameters) on top of a GLU §2. With e = 2d, two GAUs replace one MLP-plus-attention block at similar size and speed §2. FLASH-Quad, the quadratic variant, is a stack of GAUs §4.

Mixed chunk attention (FLASH). The sequence is cut into G non-overlapping chunks of length C §3.2, "Preparation". Each chunk gets U_g, V_g and Z_g as in GAU. Four heads are derived from Z: Q and K for the quadratic part, and Q and K for the linear part §3.2. The only extra parameters over GAU are the scale and offset for the linear Q and K, 4s in total §3.2.

1. Local part: exact squared-ReLU attention inside each chunk, V̂_quad,g = relu²(Q_quad,g K_quad,gᵀ + b) V_g. Its cost is O(G·C²·d) = O(TCd), linear in T when C is fixed §3.2.
2. Global part, non-causal: V̂_lin,g = Q_lin,g (Σ_h K_lin,hᵀ V_h), summed over all chunks §3.2, eq. 7.
3. Global part, causal: the same sum over earlier chunks only, h = 1 to g−1 §3.2, eq. 8.

\hat V_g^\text{lin} = Q_g^\text{lin}\Big(\sum_{h=1}^{g-1} {K_h^\text{lin}}^\top V_h\Big) eq. 8

4. Output: O_g = [U_g ⊙ (V̂_quad,g + V̂_lin,g)] W_o §3.2, display after eq. 8. Global sums are taken per chunk, so the causal prefix sum has G steps instead of T, shortening the sequential loop by a factor of C §3.2, §3.2.1.

The authors state that a chunk size in {128, 256, 512} gives large training speedups §3.2.1. Decoding needs constant work per step, O(C·d²), because the linear cache is one matrix and the local attention is limited to one chunk §3.2.1.

Optional variant: first-to-all. In TriviaQA fine-tuning, the first token of each chunk can also attend to the whole sequence with one softmax head §4.3; tab. 4 caption.

Ablation variants. FLASH (LocalOnly) keeps only the local part; FLASH (GlobalOnly) keeps only the global part and uses chunk size 64 §4.4. MC-TFM++ applies the same chunk scheme to a multi-head Transformer++ §4.4.

# Gated attention unit, one head (FLASH-Quad)  [§2 eq. 1-5] [app. D, Code 6]
# Input:  x (B, N, d)     Output: (B, N, d)
# e = 2d and s = 128  [§2] [app. D]

def gau_block(x, causal):
    # Pre-norm; the skip connection is added back at the end  [app. D, Code 6]
    shortcut = x
    x = norm(x)                                  # (B, N, d)
    # One projection, then SiLU: gives U, V and the base Z  [§2] [app. D]
    u, v, base = split(silu(dense(x, 2 * e + s)), [e, e, s])
    # Scale and offset give queries and keys, then RoPE  [app. D]
    queries, keys = rope(scale_offset(base))     # (B, N, s) each
    # Squared-ReLU scores; the listing divides QK^T by N  [§2 eq. 5] [app. D]
    scores = queries @ transpose(keys) / N + rel_pos_bias(N)  # (B, N, N)
    kernel = relu(scores) ** 2                   # (B, N, N)
    # Causal models zero out future positions  [app. D, Code 6]
    if causal:
        kernel = kernel * lower_triangular_mask(N)
    # Gate the attended values with U  [§2 eq. 3]
    gated = u * (kernel @ v)                     # (B, N, e)
    # Project back to width d and add the skip  [§2] [app. D]
    return dense(gated, d) + shortcut

# FLASH block: exact attention inside chunks, linear attention across them
# [§3.2 eq. 7-8] [app. D, Code 8]
# Input:  x (B, G, C, d), G chunks of C tokens   Output: (B, G, C, d)
# e = 2d and s = 128  [app. D]; segment masking for packed text is omitted

def flash_block(x, causal):
    # Pre-norm; the skip connection is added back at the end  [app. D]
    shortcut = x
    x = norm(x)                                  # (B, G, C, d)
    u, v, base = split(silu(dense(x, 2 * e + s)), [e, e, s])
    # Four heads: quadratic query and key, linear query and key  [§3.2]
    # RoPE runs over the whole sequence, not restarted per chunk  [app. D]
    quad_q, quad_k, lin_q, lin_k = rope(scale_offset(base))
    # Local part: squared-ReLU attention inside each chunk  [§3.2] [app. D]
    scores = quad_q @ transpose(quad_k) / C + rel_pos_bias(C)  # (B, G, C, C)
    kernel = relu(scores) ** 2                   # (B, G, C, C)
    if causal:
        kernel = kernel * lower_triangular_mask(C)
    quadratic = kernel @ v                       # (B, G, C, e)
    # One summary per chunk: lin_k^T V, divided by C in the listing  [app. D]
    # The equations use a plain sum here instead  [§3.2 eq. 7]
    summaries = transpose(lin_k) @ v / C         # (B, G, s, e)
    # Global part: earlier chunks only when causal, all chunks otherwise
    # [§3.2 eq. 7-8]
    if causal:
        past = exclusive_cumsum(summaries, axis=G)   # (B, G, s, e)
    else:
        past = sum(summaries, axis=G)            # (B, 1, s, e), broadcast
    linear = lin_q @ past                        # (B, G, C, e)
    # Gate the sum of both parts, project back, add the skip  [§3.2]
    out = u * (quadratic + linear)
    return dense(out, d) + shortcut

Hyperparameters

Values for the pretraining runs. "Not stated" means the paper does not give it.

Results

Author-reported. Setup: all models trained with the same code and tokenizer; baselines are Transformer (GELU), Transformer+ (RoPE), Transformer++ (RoPE and GLU), Performer (ReLU kernel) and Combiner (axial), at about 110M parameters for C4 and Wiki-40B §4, tab. 8 tab. 9. Perplexity is token-level for Wiki-40B and word-level for PG-19 §4.2. Speedups are wall-clock time to reach a reference perplexity; per-step latency is a separate number §4.2 tab. 3.

Other results: the GAU ablation in tab. 1 shows that replacing the squared-ReLU kernel with softmax, or using multi-head attention, each worsens perplexity; the TriviaQA gap between FLASH-Quad and FLASH (s=512) is 74.8 versus 73.3 F1 tab. 4. The per-step GPU latency on one V100 at 4096 × 1 is 452.9 ms for FLASH, while Transformer++ runs out of memory there tab. 11.

Paper vs code

No official code.

Searches: the arXiv abstract, the paper text and the LaTeX source contain no code link, and no project page is linked. GitHub search with the method name returned nothing under the google-research organisation. A search for "transformer quality linear time" returned one repository, lucidrains/FLASH-pytorch (MIT), a third-party reimplementation, which is not used. Searches for "gated attention unit transformer quality" and "mixed chunk attention linear" returned nothing. A user search for the first author returned unrelated accounts.

The paper's appendix listings app. D, Code 2-8 and the main-text listing §3.2, Code 1 are the only code in the source. They do not match the equations or each other:

• Scaling: Code 6 and Code 8 divide QKᵀ by the sequence or chunk length. Code 1 and the listing in fig. 2 do not. Equations 5, 7 and 8 do not either.
• Chunk summaries: Code 8 divides the chunk summary by C. Equation 7 uses a plain sum.
• Segment masks: Code 7 normalises each row of the chunk mask to sum to 1 (divide_no_nan). Code 8 applies this mask to the summaries, so the global term becomes an average over visible chunks, not the sum in equations 7 and 8. Code 8 uses segment ids only in the linear part; the local quadratic part has no segment mask.
• Activation: the listing in fig. 2 applies no activation to U and V. Code 6 and Code 8 apply SiLU to U, V and Z, as tab. 8 to tab. 10 state.
• Position: fig. 2 and equation 5 show no RoPE. Code 6 and Code 8 apply RoPE to queries and keys, and the tables list RoPE.
• Chunk sizes: the caption of fig. 4 says C is at least 128 in every experiment. §4.4 uses C = 64 for the GlobalOnly ablation.
• Typos: Code 8 has a missing closing parenthesis in its non-causal quadratic kernel line. Code 5 returns scale * x, but the variable is named scalar. Code 4 uses an undefined dtype in the branch for length 512 and above.

Pitfalls

• Scaling factor. Equations 5, 7 and 8 have no 1/N or 1/C factor, but the appendix listings use one. The paper does not say which version produced the reported numbers. Implementers should choose and test both. The pseudocode above follows the listings.
• Segment masks. Row-normalising the chunk mask (Code 7) changes the linear term from a sum to an average, as described under Paper vs code. Pick one and test it on packed text.
• Relative bias. The paper calls b a "relative position bias" but does not define it. The listing builds it in two ways, switched on the length argument: a learned Toeplitz matrix with 2n−1 weights when n < 512, and a Toeplitz matrix from RoPE-rotated 128-dimensional vectors a and b otherwise app. D, Code 4. In FLASH the argument is the chunk length, so C = 512 uses the RoPE branch.
• Position indices. The listing applies RoPE across chunk index and offset together, so positions are global app. D, Code 8. The local bias is indexed within the chunk. Which mix the paper used is not stated.
• Causal global term. The sum runs over strictly earlier chunks §3.2 eq. 8; app. D, Code 8. Including the current chunk leaks future tokens inside it.
• Init and scale. The listings start the query and key scales at stddev 0.02, not at one app. D, Code 6. My inference from this: scores start near zero, so the squared-ReLU weights start small. The paper does not discuss the start of training.
• Local gradient clipping. The paper gives 0.1 but not the norm, the scope (global, per-tensor or per-layer) or whether it applies to the scales. The paper says it is disabled for FLASH fine-tuning §4.3.
• Chunk size. The paper uses 256 for C4 and Wiki-40B, 512 for PG-19 and 128 for TriviaQA fine-tuning. The text says larger chunks help above length 1024 §4.4, but the four chunk sizes in the appendix sweep are only in fig. 8, not in the text.
• Speedup baseline. The PG-19 speedups (1.66x to 12.12x) are relative to Transformer+ tab. 3. The introduction says the 12.1x is relative to Transformer++ §1. Use the table's baseline.
• Speedup vs latency. The per-step latencies in tab. 12 and tab. 13 are not training-cost speedups. The speedups in fig. 1, fig. 5, fig. 6 and tab. 3 combine step latency with steps to reach a reference quality.
• Metric naming. The captions of fig. 5 and fig. 6 say "negative log perplexity"; the tables report perplexity. tab. 1 and tab. 2 give no context length, step count or seed.
• Cost reporting. §4.1 says total cost is reported in TPU-v4 core-days, but no core-day figures appear in the paper.
• Small differences. §4.3 calls a 0.3 F1 difference within variance, but no variance or number of seeds is given. The fine-tuning batch size, step count and schedule are not stated.
• Overlapping local attention. §3.2 says overlapping chunks improve quality but slow training. The preliminary numbers are not reported, so the trade-off cannot be checked.
• Decode memory. The claim of constant per-step decoding cost, O(C·d²) §3.2.1, assumes a cache of the current chunk's keys and values. How the paper implements this cache is not stated.

Lineage

• Builds on: the GLU of Shazeer (2020) §2; the gating argument that a weaker attention suffices, citing Liu et al. (2021, "Pay attention to MLPs") §2; the squared-ReLU kernel §2 footnote; RoPE §4; ScaleNorm app. B.2; chunked attention in the style of Combiner §3.2, app. A.
• Compared against: Transformer (Vaswani et al., 2017); Transformer+ and Transformer++ (the paper's augmented baselines) §4; Performer (Choromanski et al., 2020) with the ReLU kernel; Combiner (Ren et al., 2021) with the axial variant §4.
• Related: Longformer and BigBird, as overlapping local attention §3.2.
• Superseded by: not stated in the source.