Transformer Upgrade Path: 3. From Performer to Linear Attention

Transformer升级之路：3、从Performer到线性Attention 苏剑林 · 2021-04-22 · Machine Learning language model attention

Readers who have read my earlier posts Exploring Linear Attention: Does Attention Necessarily Need a Softmax? and Performer: Linearizing the Complexity of Attention via Random Projections might find the title of this post a bit odd, since linear attention actually came before Performer, and their relationship is that "Performer is one particular implementation of linear attention, one that preserves linear complexity while approximating standard attention." So logically it ought to be "from linear attention to Performer."

However, this post isn't meant to trace the developmental history of linear attention. Instead, it's meant to think in reverse about what insights Performer offers for linear attention — hence "from Performer to linear attention."

Activation Function

The common form of linear attention is

\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i = \frac{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\boldsymbol{v}_j}{\sum\limits_{j=1}^n \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)} = \frac{\sum\limits_{j=1}^n \phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)\boldsymbol{v}_j}{\sum\limits_{j=1}^n \phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)}\end{equation}more

where $\phi(\cdot)$ and $\varphi(\cdot)$ are activation functions with non-negative range. So how should we choose this activation function? Performer tells us that we should choose the exponential function

\begin{equation}\phi(x)=\varphi(x)=e^x\end{equation}

First, let's look at how this differs from existing results. The choice given in Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention is:

\begin{equation}\phi(x)=\varphi(x)=1 + \text{elu}(x) = \left\{\begin{aligned}1 + x,\, x \geq 0\\ e^x,\, x < 0\end{aligned}\right.\end{equation}

We know that $1+x$ is exactly the first-order Taylor expansion of $e^x$ at $x=0$, so the choice $1+\text{elu}(x)$ is actually already quite close to $e^x$.

Moreover, the scheme $\phi(x)=\varphi(x)=e^x$ is also quite similar to the double-softmax construction of linear attention introduced in Efficient Attention: Attention with Linear Complexities, where we have $\phi(\boldsymbol{q})=softmax(\boldsymbol{q}),\varphi(\boldsymbol{k})=e^{\boldsymbol{k}}$ — compared to directly using $\phi(x)=\varphi(x)=e^x$, the only difference is where the normalization is applied.

A Simple Derivation

Why do I say that Performer tells us the best choice of activation function is $e^x$? Let's look at the mapping that Performer found for linearizing standard attention:

\begin{equation}\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}\end{equation}

In short, Performer found a mapping such that the $d$-dimensional vector $\boldsymbol{q},\boldsymbol{k}$ is mapped to the $m$-dimensional vector $\tilde{\boldsymbol{q}},\tilde{\boldsymbol{k}}$, satisfying the approximate relation $e^{\boldsymbol{q}\cdot \boldsymbol{k}}\approx \tilde{\boldsymbol{q}}\cdot\tilde{\boldsymbol{k}}$. In this case,

\begin{equation}a_{i,j} = \frac{e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}{\sum\limits_j e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}\approx \frac{\tilde{\boldsymbol{q}}_i\cdot\tilde{\boldsymbol{k}}_j}{\sum\limits_j \tilde{\boldsymbol{q}}_i\cdot\tilde{\boldsymbol{k}}_j} = \frac{(\lambda(\tilde{\boldsymbol{q}}_i)\tilde{\boldsymbol{q}}_i)\cdot\tilde{\boldsymbol{k}}_j}{\sum\limits_j (\lambda(\tilde{\boldsymbol{q}}_i)\tilde{\boldsymbol{q}}_i)\cdot\tilde{\boldsymbol{k}}_j}\end{equation}

The last equality shows that multiplying $\tilde{\boldsymbol{q}}$ by a constant (even one that depends on $\tilde{\boldsymbol{q}}$) doesn't change Performer's result at all, which means that changing the mapping to

\begin{equation} \tilde{\boldsymbol{q}} = \begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}} \end{pmatrix},\qquad \tilde{\boldsymbol{k}}=\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} \end{equation}

leaves Performer's result completely unchanged. Of course, the term $\Vert \boldsymbol{k}\Vert^2$ here can't be dropped just yet, but if we assume that $\Vert \boldsymbol{k}\Vert^2$ doesn't fluctuate too much and isn't a major factor in attention, then this term also behaves like a constant, so the final mapping is (approximately) equivalent to

\begin{equation} \tilde{\boldsymbol{q}} = \begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{q}} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{q}}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{q}} \end{pmatrix},\qquad \tilde{\boldsymbol{k}}=\begin{pmatrix}e^{\boldsymbol{\omega}_1\cdot \boldsymbol{k}} \\ e^{\boldsymbol{\omega}_2\cdot \boldsymbol{k}}\\ \vdots\\ e^{\boldsymbol{\omega}_m\cdot \boldsymbol{k}} \end{pmatrix} \end{equation}

How should we understand this now much-simplified mapping? In effect, the $m$ random vectors $\boldsymbol{\omega}_1,\boldsymbol{\omega}_2,\cdots,\boldsymbol{\omega}_m$ are stacked into a $d\times m$ matrix, which maps the $d$-dimensional $\boldsymbol{q},\boldsymbol{k}$ into a $m$-dimensional vector, to which the activation function $e^x$ is then applied to obtain $\tilde{\boldsymbol{q}},\tilde{\boldsymbol{k}}$. We know that in attention, $\boldsymbol{q},\boldsymbol{k}$ each already go through a fully-connected layer transformation. If we fold this $d\times m$ mapping matrix into that fully-connected layer, all that's left is an activation function $e^x$!

So this is where the optimal activation function $e^x$ comes from: as long as we change the output dimension of $\boldsymbol{q},\boldsymbol{k}$ from $d$ to $m$, and pair it with the activation function $e^x$, then in theory it should have Performer's fitting capacity, or even stronger — because Performer's $d\times m$ matrix is a fixed random matrix, whereas here we're effectively making that matrix trainable too, and removing the low-rank constraint, giving us a larger space than Performer.

The Low-Rank Problem

Whether it's Performer, the star of this post, or Nyströmformer, which I introduced earlier in Nyströmformer: A Matrix-Decomposition-Based Linearized Attention Scheme, both share the same idea: "find a linear attention that approximates standard attention." A natural question then arises: what's so good about standard attention? What exactly is worth aligning with?

From the perspective of information loss, the "rank" of the standard attention matrix can be larger — that is, closer to being an invertible matrix — which means it can retain more effective information. Specifically, the attention matrix is a $n\times n$ matrix, obtained from $\boldsymbol{Q},\boldsymbol{K}\in\mathbb{R}^{n\times d}$ via $softmax(\boldsymbol{Q}\boldsymbol{K}^{\top})$. Note that here $d$ is attention's key_size — for BERT base, for instance, it's just 64 — while $n$ tends to be fairly large. This means the rank of $\boldsymbol{Q}\boldsymbol{K}^{\top}$ doesn't exceed $d$, and moreover $d\ll n$, i.e., it's far from full rank. However, the key operation in $softmax$ is $e^{\boldsymbol{Q}\boldsymbol{K}^{\top}}$ — and taking the exponential of every element of a matrix can potentially increase the rank of the resulting matrix! So the standard attention matrix has the potential for rank elevation, meaning it embodies a more effective capacity for processing information.

By contrast, the linear attention matrix takes the form $\tilde{\boldsymbol{Q}}\tilde{\boldsymbol{K}}^{\top}$, so its rank is guaranteed not to exceed $m$. To compensate for this loss of rank, one generally has to set $m > d$; in Performer's experiments the choice was $m = 4d$, meaning key_size is expanded fourfold — which shows just how much rank matters. Of course, an immediate consequence of enlarging key_size is that, for short sequences, linear attention ends up slower than standard attention — an inherent bottleneck of linear attention.

There are also some papers worth referencing for theoretical analysis of the rank of the attention matrix. For example, Low-Rank Bottleneck in Multi-head Attention Models points out that even in standard attention, low rank is a serious bottleneck, and increasing key_size can improve performance. Last month's Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth pointed out that, without residual connections and FFNs, standard attention runs a great risk of degenerating into a simple rank-1 transformation. Even standard attention — a model with "rank-elevation potential" — has low-rank issues, let alone linear attention, whose rank is inherently capped.

So, in a nutshell: using linear attention requires a larger key_size to maintain the rank of the matrix.

Concentrating Attention

We can also understand the benefits of standard attention from the perspective of sparsity. Intuitively, since it's called an "attention mechanism," it surely needs to "concentrate attention" — if it's too spread out, it might as well be average pooling. "Concentrating attention" means that each token should only be able to significantly relate to a handful of other tokens. Mathematically, this means the attention matrix should be sparse, or at least have the potential to become sparse.

For standard attention, normalization is done via softmax:

\begin{equation}a_{i,j} = \frac{e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}{\sum\limits_j e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}\end{equation}

where the exponential function $e^x$ plays an amplifying role: as long as the individual $\boldsymbol{q}_i\cdot \boldsymbol{k}_j$ values already have some spread among them, $e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}$ will further amplify that spread. The result is that, after normalization, all probabilities except for those at the few largest positions end up very close to 0 — which shows that standard attention has the potential to "concentrate attention." Linear attention, on the other hand, is simply the result of a direct inner product, without the further amplification from $e^x$, so its attention tends to be fairly dense — for long sequences, it often ends up very close to average pooling. To mitigate this, one again needs to increase key_size to amplify the gaps; intuitively, it's like the $n$ vectors being too "cramped" in a low-dimensional space, and moving to a higher-dimensional space gives them more "room to breathe."

How can we verify the importance of sparsity? I once tried the following: first compute the linear attention matrix in full, then forcibly truncate it (i.e., have each token only attend to a handful of neighboring tokens, turning it into a local form of attention) to make it sparse. The result was that this truncated linear attention performed noticeably better than the full-matrix version. This confirms the importance of sparsity. Of course, computing the full attention matrix first and then truncating it means linear attention is no longer linear in complexity, so this approach has no practical value — it's only useful for theoretical verification.

There's another experimental observation that helps support the importance of sparsity: when linear attention is used for language modeling or as a decoder, its performance is on par with standard attention. In that setting, linear attention becomes a unidirectional RNN (see Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention), equivalent to the attention matrix becoming lower-triangular — again, sparser. By contrast, if we use non-sparse, bidirectional linear attention directly for an MLM model, the performance drop is quite noticeable.

More importantly, sparsity is closely related to the rank discussed in the previous section — one could even say they're "two sides of the same coin": an appropriate sparsification method can increase the rank of a matrix! For instance, the lower-triangular attention matrix used in language modeling — as long as the diagonal elements are nonzero (which is usually the case) — is directly a full-rank invertible matrix! Likewise, the local attention truncation I experimented with can also increase the matrix's rank; in the extreme case, if each token only attends to itself, the attention matrix becomes the identity matrix, which is full rank!

Summary

This post started from Performer to think through some issues with linear attention, including the choice of activation function for linear attention and where its bottlenecks lie (low rank, sparsity). The overall conclusion is that the best activation function for linear attention should be the exponential function, and an effective attention mechanism should have both higher rank and greater sparsity.