Random Feature Attention

In plain words

Standard attention lets every word in a text look at every other word. That is what makes Transformers strong, but the work grows with the square of the text length, so long inputs and long generated outputs get expensive quickly.

This paper swaps the attention step for an approximation. Each query and key is passed through a fixed set of random "projections" that turn it into a short vector. Once keys are turned into vectors, the model can add them up into a running total and reuse that total for every query. The cost then grows in a straight line with the length of the text instead of with its square, and generating one word at a time does not have to re-read the whole past each step.

The authors also add an optional gate. It is a learned dial that lets the running total forget older words a little at each step, so the model prefers recent context, similar to a recurrent network. The paper reports that the approximation matches standard Transformers on language modeling and translation, with the biggest speed and memory gains on long sequences. Everything here is the authors' own reported result.

TL;DR

• RFA is a drop-in replacement for softmax attention, built on random Fourier features §3.1.
• Cost is linear in sequence length in time and space §3.4 tab. 4.
• A sigmoid gate on the causal variant gives recency bias and improves language modeling §3.2 tab. 1.
• Author-reported: on MT it matches the Transformer baseline and decodes about 1.8-1.9x faster tab. 2. The gains grow with length fig. 2.
• The approximation is not a faithful copy of softmax: swapping weights between the two attentions does not transfer app. C.3.

Problem

Softmax attention computes, for each query, a weighted sum over all M keys. That takes O(M) time and space per query, and O(MN) space for N queries. When the queries cannot run in parallel, as in autoregressive decoding, time is quadratic in sequence length §2.1. Efficient variants often save asymptotic cost but add overhead that cancels the gain at moderate lengths §1. The paper targets both long and moderate inputs §1.

The method assumes the softmax temperature and feature size are the only knobs, and that the feature map approximates the exponential of a dot product with a kernel §2.2 eq. 4.

Method

Feature map (random Fourier features). Let phi map a d-dimensional vector to a 2D-dimensional feature vector built from D random directions w_i. The authors sample each w_i from a Gaussian, and the inner product of two features is an unbiased estimate of a Gaussian kernel §2.2 eq. 2 eq. 3. The features are scaled by sqrt(1/D).

Approximating softmax. The exponential of a dot product is split into two norm terms times a Gaussian kernel eq. 4. Replacing the kernel with the feature inner product gives an attention that has no interaction between query and key features. Summing over keys first gives a linear-time form eq. 5 §3.1. This form assumes queries and keys have unit norm §3.1. The paper notes this constraint is optional and an unnormalised version exists §3.3 app. A.4.

Causal (autoregressive) form. With a prefix-only attention, the sums become recurrent: a matrix state S_t and a vector z_t are updated per token, and the output is phi(q_t) times S_t, divided by phi(q_t) dotted with z_t eq. 6 alg. 1.

Gate (RFA-Gate). A learned scalar gate g_t = sigmoid(w_g . x_t + b_g) mixes the previous state with the new key-value term eq. 7 §3.2. The gate decays history exponentially §3.2.

Learned variance. Instead of a fixed scalar sigma, the authors learn a per-dimension vector sigma. The random directions are sampled once as standard Gaussian vectors that are not learned, and w_i = sigma elementwise-times w_tilde_i eq. 8 §3.3. This departs from the isotropic assumption of the theorem §3.3 fn..

Feature maps tested. Gaussian random features (the main variant) and arc-cosine features (ReLU) §3.3. The elu-plus-one feature map of Katharopoulos et al. serves as a baseline §4.1.

Multi-head. Like standard multi-head attention, RFA runs per head. During training each head samples its random projections from a pool of matrices created offline (typically 200), drawing one per step; at test time each head uses one fixed projection app. B.

Algorithm 1 gives the causal procedure and Algorithm 2 the cross-attention procedure alg. 1 alg. 2 app. A.1. Cross attention first summarises all source keys and values, then reads them out for each query alg. 2.

# Gaussian random feature map  [§2.2] [eq. 2] [eq. 8]
# Input:  x (d,); projections (D, d) = sigma * w_tilde  [eq. 8]
# Output: features (2D,)

def phi(x, projections):
    # Project the input onto each random direction  [eq. 2]
    angles = projections @ x                      # (D,)
    # Sine and cosine parts, scaled by sqrt(1/D)  [eq. 2]
    scale = sqrt(1 / D)
    return scale * concat(sin(angles), cos(angles))

# Causal random feature attention  [alg. 1] [eq. 6]
# Input:  queries, keys (N, d); values (N, d_v)
# Output: outputs (N, d_v)
# state is (2D, d_v); norm_sum is (2D,)

def causal_rfa(queries, keys, values, projections):
    # Running sums over the prefix start at zero  [alg. 1]
    state = zeros(2 * D, d_v)
    norm_sum = zeros(2 * D)
    outputs = []
    for t in range(N):
        # Random feature maps of the current key and query  [eq. 2]
        key_feat = phi(keys[t], projections)
        query_feat = phi(queries[t], projections)
        # Add this key-value term to the running sums  [eq. 6]
        state = state + outer(key_feat, values[t])
        norm_sum = norm_sum + key_feat
        # Normalised read-out for step t  [eq. 6]
        denom = dot(query_feat, norm_sum)
        outputs.append((query_feat @ state) / denom)
    return stack(outputs)

# Gated update of the running sums  [§3.2] [eq. 7]
# Replaces the two running-sum updates in the loop above

def gated_update(state, norm_sum, key_feat, value, w_g, b_g, x_t):
    # Learned scalar gate in (0, 1)  [eq. 7]
    gate = sigmoid(dot(w_g, x_t) + b_g)
    # Decay the history and add the new term, both scaled  [eq. 7]
    state = gate * state + (1 - gate) * outer(key_feat, value)
    norm_sum = gate * norm_sum + (1 - gate) * key_feat
    return state, norm_sum

Key equations, written out:

\mathrm{attn}(q_t,\{k_i\},\{v_i\}) = \sum_i \frac{\exp(q_t\cdot k_i/\tau)}{\sum_j \exp(q_t\cdot k_j/\tau)}\, v_i^\top eq. 1

\phi(x) = \sqrt{1/D}\,[\sin(w_1\cdot x),\dots,\sin(w_D\cdot x),\cos(w_1\cdot x),\dots,\cos(w_D\cdot x)]^\top eq. 2 §2.2

\mathbb{E}_{w_i}[\phi(x)\cdot\phi(y)] = \exp(-\lVert x-y\rVert^2/2\sigma^2),\quad w_i \sim \mathcal{N}(0,\sigma^2 I_d) eq. 3 §2.2

\exp(x\cdot y/\sigma^2) = \exp\!\big(\lVert x\rVert^2/2\sigma^2+\lVert y\rVert^2/2\sigma^2\big)\exp\!\big(-\lVert x-y\rVert^2/2\sigma^2\big) \approx \exp\!\big(\lVert x\rVert^2/2\sigma^2+\lVert y\rVert^2/2\sigma^2\big)\,\phi(x)\cdot\phi(y) eq. 4

\mathrm{RFA}(q_t,\{k_i\},\{v_i\}) = \frac{\phi(q_t)^\top\sum_i \phi(k_i)\otimes v_i}{\phi(q_t)\cdot\sum_j \phi(k_j)} eq. 5

S_t = S_{t-1} + \phi(k_t)\otimes v_t,\qquad z_t = z_{t-1} + \phi(k_t),\qquad \mathrm{out}_t = \frac{\phi(q_t)^\top S_t}{\phi(q_t)\cdot z_t} eq. 6

g_t = \sigma(w_g\cdot x_t + b_g),\quad S_t = g_t S_{t-1} + (1-g_t)\,\phi(k_t)\otimes v_t,\quad z_t = g_t z_{t-1} + (1-g_t)\,\phi(k_t) eq. 7

\tilde w_i \sim \mathcal{N}(0, I_d),\qquad w_i = \boldsymbol{\sigma}\circ\tilde w_i eq. 8

Equation numbers follow the compiled paper's order. Appendix equations are cited by section.

Hyperparameters

Language modeling §4.1 app. B.1 tab. 7:

Machine translation §4.2 app. B.2 tab. 8:

Long Range Arena §4.3:

Dataset statistics tab. 5: WikiText-103 has 103M training tokens, 218K dev and 246K test tokens, and a 268K vocabulary. WMT14 EN-DE and EN-FR each have 4.5M training pairs, 3K dev, 3K test, and a 32K vocabulary. IWSLT14 DE-EN has 160K training pairs, 7K dev, 7K test, and 9K/7K BPE types tab. 5 §4.2.

Results

Author-reported throughout. All numbers are from the paper's own experiments. Nothing here has been reproduced.

Language modeling, WikiText-103 perplexity (lower is better) tab. 1 Setup: Transformer-based LM, 512-token blocks, two sizes §4.1.

The table's bold marks the models that beat Base. The paper states that the gated Gaussian variant beats Base by at least 1.2 test perplexity at both sizes §4.1. The stateful variant carries the last hidden state into the next block during training and evaluation §4.1. The paper reports no speedup for these LM runs and a training slowdown of about 15% §4.1 fn..

Machine translation, test BLEU tab. 2 Setup: base Transformer; RFA in decoder cross and causal attention only §4.2. Speed is relative to Base, batch size 32, one TPU v2 tab. 2.

Long Range Arena, accuracy, speed and memory tab. 3 Setup: text classification on ListOps (LO), IMDb and AAN; speed and memory measured on IMDb at 1K to 4K length. Speed is relative (higher is better) and memory is relative (lower is better) to the Transformer. Rows other than RFA-Gaussian are the baselines' numbers, taken from Tay et al. (2020) tab. 3.

The table's underlining marks models above the Transformer in accuracy, and its bold marks the best per column tab. 3. The paper states RFA is the only model close to the baseline in both accuracy and efficiency among those compared §4.3.

Decoding time and memory versus length fig. 2 fig. 3. Figure 2 shows conditional decoding (sequence-to-sequence, same input and output length, greedy, batch 16, one TPU v2). RFA's speed and memory stay roughly constant as length grows, while the softmax baseline slows and its memory grows. At 2,048 tokens the paper reports about 12x faster decoding with under 10% of the memory §5. Figure 3 (unconditional decoding) shows similar trends app. C.1.

Softmax-to-RFA transfer app. C.3 app. C.4 fig. 4. On IWSLT14, models trained with one attention and evaluated with the other reach 2.3 BLEU (RFA trained, softmax evaluated) and 1.1 BLEU (softmax trained, RFA evaluated) without finetuning app. C.3. Both pretrained models exceed 35.2 dev BLEU. Starting from a softmax checkpoint, RFA takes about 1,500 steps to reach similar training loss, but resetting the attention projections gives a curve of similar shape app. C.4 fig. 4.

Paper vs code

No official code.

The paper names no code repository, and the arXiv metadata has no comment or link to one. The searches found nothing that clearly belongs to the authors. The paper's links point only to JAX, SacreBLEU and the Long Range Arena repository app. B §4.3. GitHub searches with WebFetch covered "rfa" and "random feature attention" (plain and with an organisation filter for google-research), plus an author-name search. None returned a repository from the authors. The implementation is described as JAX-based app. B, so the official code was not located and nothing is cited from it.

Pitfalls

• Unit-norm assumption. The main derivation eq. 5 needs unit-norm queries and keys. The paper's normalisation is a choice, not a requirement. Without it, an extra per-token scalar C(x) = exp(||x||^2 / 2 sigma^2) appears in numerator and denominator app. A.4. The paper says normalisation has little effect once sigma is set or learned §3.3. Implement it as the paper does, or follow app. A.4 exactly.
• Feature-map size unit is inconsistent. §3.1 says phi maps to 2D dimensions, and the complexity analysis uses D §3.4. The LM text says "64-dimensional random feature maps" §4.1, which could mean D or 2D, and the LM table says "Random Feature Map Size 64" tab. 7. The MT text says "D = 128/64" §4.2, while its table lists "Feature Map 128/64" tab. 8. Table 9 uses the size of phi (2D) tab. 9. Whether the LM run uses phi of size 64 or 128 is not stated.
• Algorithm dimensions. Algorithms 1 and 2 say the state is D x d alg. 1 alg. 2, while §3.1 says S_t is 2D x d §3.1. The pseudocode above uses 2D, matching §3.1.
• Randomness schedule. Each head draws from a pool of about 200 offline matrices during training, and one fixed matrix is used at test time app. B. This is an unstated implementation choice in many setups and its effect is reported only on one MT dev run tab. 6.
• Denominator stability. The denominator phi(q) . z is a sum of random-feature dot products. The paper states no epsilon or clamp, and none is described in the algorithms. Behaviour at near-zero denominators is not discussed.
• Feature-map sensitivity. Small feature maps do not converge: for cross attention, sizes 32 and 64 failed, and for causal attention size 32 failed tab. 9 app. C.2. Practical sizes start around 128 for cross and 64 for causal in the MT setup. The paper says D should be at least d for a reasonable approximation §3.4.
• Gaussian versus arc-cosine. Gaussian features were more stable to train, which the authors attribute to unit-norm outputs (conjecture, not tested) §4.1 fn.. The arc-cosine variant is cheaper per sample app. C.1 §5.
• Speed claims differ. The abstract and conclusion say "twice as fast" §7. Table 2 gives 1.8-1.9x tab. 2. The benchmark batch size (16 in §5, 32 in tab. 2) and the length matter, and the paper says batch size was held fixed and that larger batches would favour RFA §5.
• Speed is not free at short lengths. At 512 tokens the LM has no speedup, and the LM runs trained about 15% slower than Base §3.4 §4.1 fn.. Gains appear for long inputs, roughly above 1,000 tokens §3.4.
• Gating is causal-only in practice. The gate is defined for the causal recurrence §3.2. The gated softmax counterpart is quadratic in time, so it cannot be used as a cheap baseline app. A.5. Gate initialisation (w_g, b_g) is not stated.
• Sigma is learned per dimension. The paper learns a vector sigma and samples fixed unit Gaussian directions, which departs from the isotropic Bochner theorem §3.3 fn.. Initial values of sigma are not stated.
• Stateful training. The stateful LM carries (S_t, z_t) across blocks during training and evaluation §4.1. The non-stateful Base cannot do this §4.1 fn.. The state must be detached or handled consistently in implementation: not stated.
• Not a softmax approximation in practice. Trained RFA models do not transfer to softmax attention, so RFA should not be treated as a drop-in fine-tuning replacement for a pretrained softmax model without finetuning app. C.3 app. C.4.
• Learning-rate and dropout selection. The paper gives grids and says models were tuned, but not which setting each reported run used tab. 7.
• Long Range Arena baselines are not rerun. The baseline numbers are taken from Tay et al. (2020) tab. 3. RFA's LRA numbers come from the authors' own runs §4.3. Speed and memory were measured only on IMDb §4.3.

Lineage

• Builds on: random Fourier features, Rahimi and Recht §2.2 eq. 2. It also builds on the kernel view of softmax §1.
• Inspired by: the causal linear-attention form of Katharopoulos et al. (2020), which uses an elu-based feature map. The paper reports that the elu-based variant underperforms RFA in its controlled experiments §3.3 §4.1.
• Related: Choromanski et al. (2020) also approximate attention with random features, using positive random features. RFA normalises queries and keys before random projection instead §3.3.
• Compared against: Transformer baselines, Katharopoulos et al. (2020), and the Long Range Arena baselines listed in tab. 3.
• Superseded by: not stated in the paper, and no card in this wiki covers a successor.