Performer: Linearizing Attention Complexity with Random Projections

A Beautiful Random Mapping

The primary contribution of Performer is finding a very elegant mapping scheme: \begin{aligned} e^{\boldsymbol{q}\cdot \boldsymbol{k}}&=\mathbb{E}_{\boldsymbol{\omega}\sim \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d)}\left[e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\right]\\[6pt] &\approx\underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{q}}} \cdot \underbrace{\frac{1}{\sqrt{m}}\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2} \end{pmatrix}}_{\tilde{\boldsymbol{k}}} \end{aligned} \label{eq:core} Let’s analyze what this equation says. The first equality means the two sides are identical. This implies that if we sample infinitely many \boldsymbol{\omega} from the standard normal distribution \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d) and calculate the average of e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}, the result equals e^{\boldsymbol{q}\cdot \boldsymbol{k}}. Written in integral form: \begin{aligned} &\frac{1}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}\Vert^2 / 2}\times e^{\boldsymbol{\omega}\cdot \boldsymbol{q}-\Vert \boldsymbol{q}\Vert^2 / 2} \times e^{\boldsymbol{\omega}\cdot \boldsymbol{k}-\Vert \boldsymbol{k}\Vert^2 / 2}d\boldsymbol{\omega} \\ =&\frac{1}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}-\boldsymbol{q}-\boldsymbol{k}\Vert^2 / 2 + \boldsymbol{q}\cdot \boldsymbol{k}}d\boldsymbol{\omega}\\ =&\, e^{\boldsymbol{q}\cdot \boldsymbol{k}} \end{aligned} In practice, we can only sample a finite number m, leading to the second approximation. This can be expressed as the dot product of two m-dimensional vectors, which is exactly the e^{\boldsymbol{q}\cdot \boldsymbol{k}}\approx \tilde{\boldsymbol{q}}\cdot\tilde{\boldsymbol{k}} we need! Thus, with this approximation, we can transform the exponent of the dot product of two d-dimensional vectors into the dot product of two m-dimensional vectors. Theoretically, we can convert standard Attention with a head_size of d into Linear Attention with a head_size of m. This is the core idea of the paper.

Discussion on the Derivation

Some readers might be interested in the origin of Equation [eq:core]. While the validity of Equation [eq:core] can be verified by direct integration, can a similar linear approximation be found for any arbitrarily defined \text{sim}(\boldsymbol{q}, \boldsymbol{k})? We will show that a similar linearization scheme can be found through a generalized process, although the resulting form might be far less elegant or effective than Equation [eq:core].

Specifically, for any \text{sim}(\boldsymbol{q}, \boldsymbol{k}), we can rewrite it as: \text{sim}(\boldsymbol{q}, \boldsymbol{k}) = \frac{\beta(\boldsymbol{q})\gamma(\boldsymbol{k})\text{sim}(\boldsymbol{q}, \boldsymbol{k})}{\beta(\boldsymbol{q})\gamma(\boldsymbol{k})} Then we can perform a Fourier transform on \beta(\boldsymbol{q})\gamma(\boldsymbol{k})\text{sim}(\boldsymbol{q}, \boldsymbol{k}): \mathcal{F}(\boldsymbol{\omega}_q, \boldsymbol{\omega}_k)=\frac{1}{(2\pi)^{d/2}}\int \beta(\boldsymbol{q})\gamma(\boldsymbol{k})\text{sim}(\boldsymbol{q}, \boldsymbol{k})e^{-i\boldsymbol{\omega}_q\cdot \boldsymbol{q}-i\boldsymbol{\omega}_k\cdot \boldsymbol{k}}d\boldsymbol{q}d\boldsymbol{k} The reason for multiplying by \beta(\boldsymbol{q})\gamma(\boldsymbol{k}) first is that the direct Fourier transform of \text{sim}(\boldsymbol{q}, \boldsymbol{k}) might not be well-behaved or might not exist. Multiplying by a function can fix this; for example, letting \beta(\boldsymbol{x})=\gamma(\boldsymbol{x})=e^{-\lambda\Vert x\Vert^2}. If \lambda is large enough, many \text{sim}(\boldsymbol{q}, \boldsymbol{k}) can undergo a Fourier transform.

Next, we perform the inverse transform and substitute it back: \text{sim}(\boldsymbol{q}, \boldsymbol{k})=\frac{1}{(2\pi)^{d/2}}\int \mathcal{F}(\boldsymbol{\omega}_q, \boldsymbol{\omega}_k)\frac{e^{i\boldsymbol{\omega}_q\cdot \boldsymbol{q}}}{\beta(\boldsymbol{q})} \frac{e^{i\boldsymbol{\omega}_k\cdot \boldsymbol{k}}}{\gamma(\boldsymbol{k})}d\boldsymbol{\omega}_q d\boldsymbol{\omega}_k If we can calculate \mathcal{F}(\boldsymbol{\omega}_q, \boldsymbol{\omega}_k) and normalize it, it becomes a distribution from which we can sample random vectors \boldsymbol{\omega}_q, \boldsymbol{\omega}_k. This allows an approximation as the dot product of vector sequences composed of \frac{e^{i\boldsymbol{\omega}_q\cdot \boldsymbol{q}}}{\beta(\boldsymbol{q})} and \frac{e^{i\boldsymbol{\omega}_k\cdot \boldsymbol{k}}}{\gamma(\boldsymbol{k})}. While this involves imaginary numbers, we usually deal with real numbers. This is not a major issue; we can expand using Euler’s formula e^{i \theta}=\cos\theta + i\sin\theta and keep only the real part. Theoretically, this process works, but compared to Equation [eq:core], it has two problems: 1. It requires sampling two sets of random variables \boldsymbol{\omega}_q, \boldsymbol{\omega}_k, which increases the estimation variance; 2. After keeping the real part, the result is a combination of \sin and \cos, which does not guarantee non-negativity and requires extra clipping.

The special thing about Equation [eq:core] is that e^{\boldsymbol{q}\cdot \boldsymbol{k}} can be rewritten as: e^{\boldsymbol{q}\cdot \boldsymbol{k}} = e^{\Vert \boldsymbol{q}\Vert^2 / 2 + \Vert \boldsymbol{k}\Vert^2 / 2 - \Vert\boldsymbol{q}-\boldsymbol{k}\Vert^2 / 2} So it only needs to be transformed into a problem of a single variable \boldsymbol{q}-\boldsymbol{k}. The Fourier transform of e^{-\Vert\boldsymbol{q}-\boldsymbol{k}\Vert^2 / 2} is exactly e^{-\Vert\boldsymbol{\omega}\Vert^2 / 2}. Performing the inverse transform, we have: e^{\boldsymbol{q}\cdot \boldsymbol{k}}=\frac{e^{\Vert \boldsymbol{q}\Vert^2 / 2 + \Vert \boldsymbol{k}\Vert^2 / 2}}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}\Vert^2 / 2 + i \boldsymbol{\omega}\cdot (\boldsymbol{q} - \boldsymbol{k})} d\boldsymbol{\omega} At this point, if we directly take the real part, we get a combination of \sin and \cos, which is the \text{trig} projection scheme mentioned in the original paper. However, a more ingenious property changes everything! Note that the above is an identity, so we can substitute \boldsymbol{q}\to -i\boldsymbol{q} and \boldsymbol{k}\to i\boldsymbol{k} on both sides. The result is: e^{\boldsymbol{q}\cdot \boldsymbol{k}}=\frac{e^{-\Vert \boldsymbol{q}\Vert^2 / 2 - \Vert \boldsymbol{k}\Vert^2 / 2}}{(2\pi)^{d/2}}\int e^{-\Vert\boldsymbol{\omega}\Vert^2 / 2 + \boldsymbol{\omega}\cdot (\boldsymbol{q} + \boldsymbol{k})} d\boldsymbol{\omega} This is Equation [eq:core]. The substitution \boldsymbol{q}\to -i\boldsymbol{q}, \boldsymbol{k}\to i\boldsymbol{k} keeps the left side unchanged while the right side completely avoids imaginary numbers and maintains non-negativity. It is a remarkable convergence of coincidences—truly "one of a kind."

Orthogonalization to Reduce Variance

In addition to proposing Equation [eq:core] for linearizing standard Attention, the original paper includes a further enhancement. In Equation [eq:core], \boldsymbol{\omega}_1, \boldsymbol{\omega}_2, \dots, \boldsymbol{\omega}_m are sampled independently and identically from \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d). The paper points out that if these \boldsymbol{\omega}_i are orthogonalized, the estimation variance can be effectively reduced, improving the average accuracy of a single estimation.

Note that orthogonalization here refers to keeping the norm of \boldsymbol{\omega}_i constant and only performing Gram-Schmidt orthogonalization on their directions. This operation was first proposed in another Google paper, "The Unreasonable Effectiveness of Structured Random Orthogonal Embeddings". Performer’s authors spent six pages in the appendix proving this. Rather than reproducing the proof, how should we understand this strategy?

The fundamental reason this strategy works is the isotropy of the sampling distribution \mathcal{N}(\boldsymbol{\omega};0,\boldsymbol{1}_d); its probability density function (2\pi)^{-d/2}e^{-\Vert\boldsymbol{\omega}\Vert^2 / 2} depends only on the norm \Vert\boldsymbol{\omega}\Vert, making it uniform in direction. To reduce estimation variance, we should reduce sampling randomness and make the samples more uniform. Orthogonalizing the vectors is one way to achieve directional uniformity. In other words, orthogonalizing \boldsymbol{\omega}_i promotes uniformity in the sampling results, thereby reducing variance. Additionally, orthogonalization is generally only effective for m \leq d. If m > d, the paper suggests orthogonalizing in groups of d vectors.

We can imagine that while orthogonalization makes the sampling directions more uniform, we could also make the norms more uniform to be even more thorough. Specifically, transforming the standard normal distribution to d-dimensional spherical coordinates gives the probability element: \frac{1}{(2\pi)^{d/2}} r^{d-1} e^{-r^2 /2} dr dS where dS = \sin^{d-2}\varphi_{1} \sin^{d-3}\varphi_{2} \cdots \sin\varphi_{d-2}\,d\varphi_{1}\,d\varphi_{2}\cdots d\varphi_{d-1} represents the integration element on the d-dimensional sphere. This shows that the standard normal distribution is uniform in direction, and the probability density function of the norm is proportional to r^{d-1} e^{-r^2 /2}. We define its cumulative distribution function: P_d(r\leq R) = \frac{\int_0^R r^{d-1} e^{-r^2 /2} dr}{\int_0^{\infty} r^{d-1} e^{-r^2 /2} dr} If we want to sample m points, we can set P_d(r\leq R_i) = \frac{i}{m+1}, \, i=1,2,\dots,m, and solve for m values of R_i to use as the norms.

Evaluation in the Original Paper

The original paper’s evaluation is quite extensive but somewhat disorganized, which might be confusing for readers primarily interested in NLP.

First is the speed evaluation. Unsurprisingly, Performer shows a significant advantage over standard Transformer as the sequence length increases:

Next is the comparison of approximation accuracy, demonstrating the effectiveness of orthogonal sampling and the precision of Equation [eq:core] compared to the older \sin, \cos-based trigonometric forms:

Can it achieve the initial goal of not needing to retrain an already trained model? Unfortunately, no. The paper conducted two experiments showing that Performer cannot reproduce existing results by directly loading Transformer weights, but it can recover quickly after fine-tuning. The paper does not provide a detailed analysis of why it fails to reproduce results initially.

Finally, the paper includes experiments on protein sequences and images, proving that Performer is effective for long sequences—specifically, more effective than Reformer. The experiments are numerous but can feel a bit disorienting.

Evaluations from Other Papers

Perhaps even their colleagues felt the need for clarity, as Google later released two more papers: "Efficient Transformers: A Survey" and "Long Range Arena: A Benchmark for Efficient Transformers". These systematically evaluate and compare existing methods for improving Transformer efficiency, including Performer. In comparison, the results in these two papers are much more intuitive, using simple charts to clearly show the position of each model.

For more detailed evaluation information, readers are encouraged to check those two papers.