Readers who have seen my previous articles "Exploration of Linear Attention: Does Attention Need a Softmax?" and "Performer: Linearizing Attention Complexity with Random Projections" might find the title of this article a bit unnatural. Since Linear Attention preceded the Performer, and their relationship is that "Performer is an implementation of Linear Attention that maintains an approximation of standard Attention while ensuring linear complexity," it would normally be "From Linear Attention to Performer."
However, this article does not intend to trace the history of Linear Attention. Instead, it aims to think in reverse about the insights that Performer brings to Linear Attention. Thus, the title is "From Performer to Linear Attention."
Activation Functions
The common form of Linear Attention is: 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)} where \phi(\cdot) and \varphi(\cdot) are activation functions with non-negative ranges. How should we choose these activation functions? Performer tells us that we should choose the exponential function: \phi(x)=\varphi(x)=e^x
First, let’s see how this differs from existing results. In "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention", the choice provided is: \phi(x)=\varphi(x)=1 + \text{elu}(x) = \left\{\begin{aligned}1 + x,\, x \geq 0\\ e^x,\, x < 0\end{aligned}\right. We know that 1+x is exactly the first-order Taylor expansion of e^x at x=0. Therefore, the choice of 1+\text{elu}(x) is actually quite close to e^x.
Furthermore, the scheme \phi(x)=\varphi(x)=e^x is very similar to the design of Linear Attention using dual softmax introduced in "Efficient Attention: Attention with Linear Complexities". In that design, \phi(\boldsymbol{q})=\text{softmax}(\boldsymbol{q}) and \varphi(\boldsymbol{k})=e^{\boldsymbol{k}}. Compared to directly using \phi(x)=\varphi(x)=e^x, the only difference is the position of normalization.
Simple Derivation
Why do we say that Performer tells us the optimal choice for the activation function is e^x? Let’s look at the mapping found by Performer to linearize standard Attention: \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} Simply put, Performer found a mapping that maps d-dimensional vectors \boldsymbol{q}, \boldsymbol{k} to m-dimensional vectors \tilde{\boldsymbol{q}}, \tilde{\boldsymbol{k}}, satisfying the approximation e^{\boldsymbol{q}\cdot \boldsymbol{k}}\approx \tilde{\boldsymbol{q}}\cdot\tilde{\boldsymbol{k}}. At this point: 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} The last equality shows that multiplying \tilde{\boldsymbol{q}} by a constant (even if this constant depends on \tilde{\boldsymbol{q}}) does not change the Performer’s result at all. This means that if we change the mapping to: \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} The result of the Performer will not change. Of course, the \Vert \boldsymbol{k}\Vert^2 term cannot be removed yet. However, if we assume that \Vert \boldsymbol{k}\Vert^2 does not fluctuate too much and is not the primary factor in Attention, then this term also acts like a constant. Thus, the final mapping is (approximately) equivalent to: \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} How should we understand this much-simplified mapping? In fact, the m random vectors \boldsymbol{\omega}_1, \boldsymbol{\omega}_2, \dots, \boldsymbol{\omega}_m form a d \times m matrix that maps the d-dimensional \boldsymbol{q}, \boldsymbol{k} to m-dimensional vectors, which are then passed through the activation function e^x to obtain \tilde{\boldsymbol{q}}, \tilde{\boldsymbol{k}}. We know that \boldsymbol{q}, \boldsymbol{k} in Attention already undergo a fully connected layer transformation. If we integrate this d \times m mapping matrix into the fully connected layer, what remains is simply the activation function e^x!
So, this is the origin of e^x as the optimal activation function. As long as we change the output dimension of \boldsymbol{q}, \boldsymbol{k} from d to m and use the e^x activation function, it theoretically possesses the fitting capability of Performer, or even stronger. This is because Performer’s d \times m matrix is a fixed random matrix, whereas here we treat the matrix as trainable and remove the low-rank constraint, providing a larger space than Performer.
The Low-Rank Problem
Whether it is the Performer discussed here or the Nyströmformer introduced previously in "Nyströmformer: A Linearized Attention Scheme Based on Matrix Decomposition", the idea is to "find a linear Attention that can approximate standard Attention." A natural question arises: what is so good about standard Attention? Why is everyone trying to align with it?
From the perspective of information loss, the "rank" of the standard Attention matrix may be larger, meaning it is closer to being an invertible matrix, which allows it to retain more effective information. Specifically, the Attention matrix is an n \times n matrix derived from \boldsymbol{Q}, \boldsymbol{K} \in \mathbb{R}^{n \times d} via \text{softmax}(\boldsymbol{Q}\boldsymbol{K}^{\top}). Note that d is the key_size of the Attention (e.g., for BERT-base, it is only 64), while n is often quite large. This indicates that the rank of \boldsymbol{Q}\boldsymbol{K}^{\top} does not exceed d, and since d \ll n, it is far from full rank. However, the key operation of \text{softmax} is e^{\boldsymbol{Q}\boldsymbol{K}^{\top}}. If every element of a matrix is exponentiated, the rank of the new matrix can potentially increase! Therefore, the standard Attention matrix has the potential for "rank elevation," implying it possesses a more effective capacity for processing information.
In contrast, the Linear Attention matrix takes the form \tilde{\boldsymbol{Q}}\tilde{\boldsymbol{K}}^{\top}, so its rank definitely does not exceed m. To compensate for the loss of rank, m is generally set such that m > d. In Performer’s experiments, m = 4d was chosen, meaning the key_size was expanded fourfold, highlighting the importance of rank. Of course, a direct consequence of expanding the key_size is that Linear Attention may be slower than standard Attention when processing short sequences—this is an inherent bottleneck of Linear Attention.
Regarding the theoretical analysis of the rank of the Attention matrix, there are several papers for reference. For instance, "Low-Rank Bottleneck in Multi-head Attention Models" points out that even in standard Attention, low-rankness is a serious bottleneck, and increasing the key_size can improve performance. A paper from last month, "Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth", points out that without residuals and FFNs, standard Attention runs a high risk of degenerating into a simple transformation with a rank of 1. If even standard Attention, which has "rank elevation potential," suffers from low-rank issues, then Linear Attention, which has an inherent upper bound on its rank, is even more susceptible.
In short: 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 is an "Attention mechanism," it must "concentrate attention." If it is too dispersed, it might be equivalent to average pooling. "Concentrating attention" means that each token should only be significantly associated with a few other tokens. In mathematical terms, this means the Attention matrix should be sparse, or at least have the potential to become sparse.
For standard Attention, it normalizes via softmax: a_{i,j} = \frac{e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}}{\sum\limits_j e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}} The exponential function e^x plays an amplifying role. As long as the individual \boldsymbol{q}_i\cdot \boldsymbol{k}_j values have some gap, e^{\boldsymbol{q}_i\cdot \boldsymbol{k}_j} will further amplify this gap. The result is that after normalization, except for the positions with the largest values, the remaining probabilities are very close to 0. This shows that standard Attention has the potential to "concentrate attention." For Linear Attention, it is the direct result of an inner product without the further amplification of e^x, so its attention is relatively dense. When the sequence length is large, it often approaches average pooling. To alleviate this, one still needs to increase the key_size to amplify the differences. Intuitively, n vectors are too "crowded" in a low-dimensional space; moving to a higher-dimensional space makes it "looser."
How can we verify the importance of sparsity? I have tried calculating the Linear Attention matrix first and then forcibly truncating it (i.e., each token only attends to a few tokens before and after it, turning it into a local Attention). The results showed that this truncated Linear Attention performed significantly better than the full-matrix Linear Attention. This confirms the importance of sparsity. Of course, calculating the Attention matrix first and then truncating it means the complexity is no longer linear, so it has no practical value and was only used for theoretical verification.
Another experimental phenomenon that supports the importance of sparsity is that when Linear Attention is used for language models or decoders, its performance is not much different from standard Attention. In this case, Linear Attention becomes a unidirectional RNN (refer to "Transformers are RNNs"), which is equivalent to the Attention matrix becoming a lower triangular matrix, which is also sparser. In contrast, if non-sparse bidirectional Linear Attention is used directly for an MLM model, the performance drop is quite significant.
More importantly, sparsity and the rank mentioned in the previous section are closely related; one could even say they are "two sides of the same coin": appropriate sparsification methods can increase the rank of a matrix! For example, in the lower triangular Attention matrix used in language models, as long as the diagonal elements are non-zero (which is usually the case), the matrix is a full-rank invertible matrix! My experiments with local Attention truncation also increased the rank of the matrix. For instance, in the extreme case where each token only attends to itself, the Attention matrix becomes a full-rank identity matrix!
Summary
This article reflects on some issues of Linear Attention starting from the Performer, including the choice of activation functions and the bottlenecks of Linear Attention (low-rankness and sparsity). The general conclusion is that the optimal activation function for Linear Attention should be the exponential function, and an effective Attention mechanism should possess higher rank and greater sparsity.
Reprinted from: https://kexue.fm/archives/8338
For more details on reprinting, please refer to: "Scientific Space FAQ"