Random feature approximation

Overview

Many models need the value of a kernel or an exponential for every pair of inputs. Computing that for all pairs is expensive. Random feature approximation draws a fixed set of random projections once and maps each input to a short vector, so the kernel value becomes a plain inner product of two vectors. With random Fourier features, that inner product is an unbiased estimate of a Gaussian kernel 2103.02143 §2.2, eq. 3. Because the map is a finite product, the sums over pairs can be regrouped: sum the keys once, then read that summary for each query. In attention this makes the cost linear in length 2103.02143 §3.1, eq. 5.

The same idea appears outside sequence attention. A Gaussian process surrogate with a shift-invariant kernel can use random Fourier features, so that drawing a function sample becomes a Gaussian draw over a short weight vector 2205.14090 §2.1, §3.1, eq. 5. A graph transformer uses positive random features in the same way, so that every node can attend to every other node in a layer at cost linear in the number of nodes 2306.08385 §3.1, eq. 5. Three cards use it so far: two for attention-style models and one for Bayesian optimization. The page does not cover the wider kernel-approximation literature.

Variants

Gaussian random Fourier features

Sample D random directions from a Gaussian. The feature of an input is the sine and cosine of its projections onto those directions, scaled by the square root of 1/D, which gives a vector of length 2D [arxiv:2103.02143 §2.2, eq. 2]. This is the main variant of the paper. Card: Random Feature Attention.

Its reported behaviour, small models unless stated:

• WikiText-103 test perplexity: softmax base 34.5, Gaussian features 35.7, gated Gaussian features 32.7 2103.02143 tab. 1.
• Machine translation: the gated Gaussian variant matches the base on WMT14 EN-DE (28.1 vs 28.1 BLEU) and decodes about 1.8x faster [arxiv:2103.02143 tab. 2].
• Long Range Arena: the Gaussian variant is close to the Transformer in accuracy while its speed and memory scale better with length 2103.02143 tab. 3 §4.3.

Learned per-dimension scale

Instead of one fixed scale, the paper learns a scale per dimension. The directions stay fixed unit Gaussians, and the projection is the learned scale times each direction 2103.02143 eq. 8 §3.3. The authors note this departs from the isotropic case in the theorem behind the method 2103.02143 §3.3 fn.. Initial values of the scale are not stated.

Arc-cosine features

The paper also tests a ReLU-based feature map. In language modeling it trails the Gaussian map, with 36.0 vs 33.6 small dev perplexity [arxiv:2103.02143 tab. 1]. The authors report it is cheaper per sample [arxiv:2103.02143 §5 app. C.1].

Unit-norm and unnormalised forms

The main derivation assumes queries and keys have unit norm [arxiv:2103.02143 §3.1]. Without that step, an extra per-token scalar appears in the numerator and the denominator 2103.02143 app. A.4. The paper says normalisation has little effect once the scale is set or learned 2103.02143 §3.3.

Causal running state with a gate

For left-to-right use, the sums become a running state, updated per token 2103.02143 eq. 6. An optional sigmoid gate decays the old state and scales the new term, which gives a recency bias [arxiv:2103.02143 §3.2, eq. 7]. In language modeling, the gated Gaussian variant beats the base at both sizes, and the stateful version (which carries state across blocks) gets the lowest perplexity in the table 2103.02143 tab. 1 §4.1.

Positive features with Gumbel sampling for graphs

Here the kernel is the softmax exponential, written as an inner product of positive random features with m random directions drawn from a standard Gaussian 2306.08385 §3.1, eq. 3, eq. 4. Applied to a graph, every node attends to every other node in each layer. The sums over all nodes are computed once per layer, which gives linear cost in the node count and avoids the N by N similarity matrix 2306.08385 §3.1, eq. 5.

To make each layer learn a sparse latent neighbourhood, the paper adds Gumbel noise to the logits, scaled by a temperature τ 2306.08385 §3.1, eq. 6. In the kernelised form, each key feature is multiplied by a noise factor before the sums, and the output is averaged over K draws 2306.08385 §3.1, eq. 7. At τ = 0.25, the paper's ablation with an ordinary softmax at the same τ scores lower on both large graphs, 76.18 vs 77.45 ROC-AUC on OGB-Proteins and 87.55 vs 87.85 accuracy on Amazon2M 2306.08385 tab. 2 tab. 3. These are author-reported figures. Card: NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification.

Posterior function samples for Thompson sampling

Here the features are not used for attention. A Gaussian process surrogate with a shift-invariant kernel is written as a linear model over D random frequencies, drawn from the kernel's normalised spectral density. The feature map has 1/sqrt(D) scaling and length 2D, and the weights get a Gaussian prior 2205.14090 §2.1, §3.1, eq. 5. Because the model is linear and Gaussian, the posterior and its updates have closed forms [arxiv:2205.14090 §3.1, §8]. To choose the next query, the method draws one weight vector from the posterior and maximises the resulting function 2205.14090 eq. 12. The paper's experiments use a mixture of several such Gaussian processes, with D = 50 per component 2205.14090 §3, §5. The authors report lower simple regret than single-GP Thompson sampling and GP-EI on their synthetic and robot-pushing tasks 2205.14090 §5, fig. 1–3. Card: Surrogate modeling for Bayesian optimization beyond a single Gaussian process.

The paper's regret bounds cover the exact Gaussian process, not this approximation 2205.14090 §4, §7. The paper gives the sampling rule only through the spectral density. It does not say how to sample frequencies for Matérn kernels 2205.14090 §2.1. Check the card's Paper vs code section before copying scale choices from the released code.

When to use which

Feature size is the main accuracy knob. Cross attention with 32 or 64 features failed to converge, and causal attention with 32 failed 2103.02143 tab. 9. The paper's MT setup starts at 128 for cross attention and 64 for causal attention 2103.02143 §4.2. In the Bayesian optimization paper, the experiments use D = 50 per component 2205.14090 §5. In the graph paper, the error bound depends on m and τ, and a small τ needs a large m for an accurate approximation [arxiv:2306.08385 §3.2]. Its main text does not give a value for m 2306.08385 §4.

Implementation sketch

The common core is a fixed projection, a feature map, and a sum over keys that can be computed once. Language-agnostic pseudocode, with notes where the variants differ.

# Gaussian random features for one input  [arxiv:2103.02143 §2.2 eq. 2]
# Input:  x (d,); projections W (D, d), drawn once from a Gaussian
# Output: features (2D,)

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

# Approximate attention with the feature inner product  [arxiv:2103.02143 eq. 5]
# Keys and values are summarised once, then read per query

def approx_attention(q, keys, values, W):
    # Sum over keys first: a matrix and a vector  [eq. 5]
    S = sum(outer(random_features(k, W), v)
            for k, v in zip(keys, values))
    z = sum(random_features(k, W) for k in keys)
    # Read the summary for one query and normalise  [eq. 5]
    phi_q = random_features(q, W)
    return (phi_q @ S) / (phi_q @ z)

For the causal form, replace the two sums with per-token updates 2103.02143 eq. 6. For the gated form, decay the old state by a learned gate before adding the new term 2103.02143 eq. 7.

For the graph variant, the same sums are reused, but each key feature is multiplied by a per-node Gumbel factor before the sums, and the output is averaged over K draws 2306.08385 eq. 7, alg. 1. The card has a full pseudocode version of this layer.

For posterior sampling, the same map is reused with a different use. Draw the weights from the Gaussian posterior and maximise the resulting function over the candidate set 2205.14090 eq. 12. The paper's notation is different, so check the symbols before copying.

Contradictions

• Speed claim. The abstract and conclusion say "twice as fast" 2103.02143 §7. The machine-translation table reports 1.8x to 1.9x decoding speed 2103.02143 tab. 2.
• Feature-size unit. The map produces 2D entries 2103.02143 §3.1, but the complexity analysis is written in D 2103.02143 §3.4. The language-modeling text says "64-dimensional" features [arxiv:2103.02143 §4.1], and its table says "Random Feature Map Size 64" [arxiv:2103.02143 tab. 7]. The feature-size table uses the size of the map itself 2103.02143 tab. 9. Whether the language-modeling run used a map of size 64 or 128 is not stated.
• State shape. The algorithms give the running state as D by d 2103.02143 alg. 1 alg. 2, while §3.1 gives 2D by d 2103.02143 §3.1. The sketch above uses 2D, following §3.1.
• Feature form across papers. The Gaussian map and the posterior-sampling map both use sine and cosine with length 2D 2103.02143 §2.2 2205.14090 §3.1. The graph paper uses positive features with length m 2306.08385 §3.1 eq. 4. Check which form a source uses before reusing its scale or its feature count.
• Paper and released code (graph paper). The released code differs from the paper's kernel scaling and redraws its random directions on every forward pass, among other differences. The card's Paper vs code section lists them. Check it before copying scale choices from the code Card.

Papers

• Random Feature Attention: replaces softmax attention with random Fourier features, with a causal form and an optional gate.
• Surrogate modeling for Bayesian optimization beyond a single Gaussian process: uses random Fourier features to draw Gaussian process posterior samples for Thompson sampling over a mixture of kernels.
• NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification: uses positive random features and Gumbel-Softmax sampling so that each graph layer can attend across all nodes in linear time.