Exploring Linear Attention: Does Attention Require a Softmax?

Removing the Softmax

Readers might not expect that the key factor restricting Attention performance is actually the Softmax in its definition! In fact, a simple derivation leads to this conclusion. The step \boldsymbol{Q}\boldsymbol{K}^{\top} yields an n\times n matrix, and it is this step that determines the \mathcal{O}(n^2) complexity of Attention. If there were no Softmax, it would be a continuous multiplication of three matrices \boldsymbol{Q}\boldsymbol{K}^{\top}\boldsymbol{V}. Since matrix multiplication satisfies the associative law, we could first calculate \boldsymbol{K}^{\top}\boldsymbol{V} to get a d\times d matrix, and then left-multiply it by \boldsymbol{Q}. Since d \ll n, the approximate complexity of this calculation is only \mathcal{O}(n) (dominated by the left-multiplication of \boldsymbol{Q}).

In other words, the complexity of Attention without Softmax can be reduced to the ideal linear level \mathcal{O}(n)! This is clearly our ultimate pursuit: Linear Attention, an Attention mechanism with linear complexity. Therefore, the theme of this article is to explore Linear Attention after removing the Softmax.

General Definition

The question is, can it still be considered Attention if the Softmax is directly removed? Can it still achieve the effect of standard Attention? To answer this, let us first rewrite the definition of Scaled-Dot Attention [eq:std-att] equivalently as (vectors in this article are column vectors): Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i = \frac{\sum\limits_{j=1}^n e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}\boldsymbol{v}_j}{\sum\limits_{j=1}^n e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}}\label{eq:std-att-2} Thus, Scaled-Dot Attention is essentially a weighted average of \boldsymbol{v}_j using e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j} as weights. Therefore, we can propose a generalized definition of Attention: 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)}\label{eq:gen-att} This replaces e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j} with a general function \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) of \boldsymbol{q}_i and \boldsymbol{k}_j. To preserve distribution characteristics similar to Attention, we require \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0 to hold always. That is, if we want to define a new type of Attention, we should maintain the form of equation [eq:gen-att] and satisfy \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0.

This general form of Attention is also known as a Non-Local Network in Computer Vision (CV), originating from the paper "Non-local Neural Networks".

Kernel Function Form

A natural idea is: if every element of \boldsymbol{q}_i and \boldsymbol{k}_j is non-negative, then the dot product will naturally be non-negative. To achieve this, we can apply activation functions \phi and \varphi to \boldsymbol{q}_i and \boldsymbol{k}_j respectively: \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)\label{eq:gen-att-2} where \phi(\cdot) and \varphi(\cdot) are activation functions with non-negative ranges. The paper mentioned at the beginning, "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention", chooses \phi(x)=\varphi(x)=\text{elu}(x)+1.

To frame it conceptually, equation [eq:gen-att-2] can be linked to "kernel methods". Especially when \phi=\varphi, \phi corresponds to a kernel function, and \langle \phi(\boldsymbol{q}_i), \phi(\boldsymbol{k}_j)\rangle is the dot product defined through the kernel function. For thoughts in this area, one can refer to the paper "Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel". We will not extend further here.

Clever Use of Softmax

An earlier article, "Efficient Attention: Attention with Linear Complexities", provides a more interesting choice. It notes that in \boldsymbol{Q}\boldsymbol{K}^{\top}, where \boldsymbol{Q}, \boldsymbol{K} \in\mathbb{R}^{n\times d}, if "\boldsymbol{Q} is normalized along the d dimension and \boldsymbol{K} is normalized along the n dimension," then \boldsymbol{Q}\boldsymbol{K}^{\top} automatically satisfies normalization. Thus, the choice it provides is: Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax_2\left(\boldsymbol{Q}\right)softmax_1(\boldsymbol{K})^{\top}\boldsymbol{V} where softmax_1 and softmax_2 refer to Softmax operations performed along the first (n) and second (d) dimensions, respectively. In other words, we apply Softmax to \boldsymbol{Q} and \boldsymbol{K} individually, rather than applying it after calculating \boldsymbol{Q}\boldsymbol{K}^{\top}.

If we directly take \phi(\boldsymbol{q}_i)=softmax(\boldsymbol{q}_i) and \varphi(\boldsymbol{k}_j)=softmax(\boldsymbol{k}_j), then this form is clearly a special case of equation [eq:gen-att-2]. Additionally, this design has appeared more than once in CV, for example, A^2-Nets also includes the same approach.

My Own Conception

Here, I present a conception of my own. The starting point for this idea is no longer equation [eq:gen-att-2], but rather an approximation of the original definition [eq:std-att-2]. From the Taylor expansion, we have: e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j} \approx 1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j If \boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1, then the non-negativity of the right side can be guaranteed, allowing us to let \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)=1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j. At this point, readers might have already realized that to ensure \boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1, one only needs to perform l_2 normalization on \boldsymbol{q}_i and \boldsymbol{k}_j respectively. Therefore, the scheme I finally propose is: \text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = 1 + \left( \frac{\boldsymbol{q}_i}{\Vert \boldsymbol{q}_i\Vert}\right)^{\top}\left(\frac{\boldsymbol{k}_j}{\Vert \boldsymbol{k}_j\Vert}\right) This differs from the form in equation [eq:gen-att-2], but theoretically, it is closer to the original Scaled-Dot Attention.

Sparse Attention

We have previously introduced OpenAI’s Sparse Attention, which reduces the computational load of Attention by "retaining values only within small regions and forcing most attention to be zero." After special design, most elements of the Attention matrix are 0, so theoretically, it can save memory usage and computation. Subsequent similar works include "Explicit Sparse Transformer: Concentrated Attention Through Explicit Selection" and "Longformer: The Long-Document Transformer".

However, this approach clearly has two shortcomings:

1. How to choose the attention regions to retain is determined subjectively by humans, which carries a significant lack of intelligence;

2. It requires specific design and optimization in programming to achieve an efficient implementation, making it difficult to generalize.

Reformer

Reformer is also a representative improvement work, which reduces the complexity of Attention to \mathcal{O}(n\log n). In a sense, Reformer is also a type of Sparse Attention, but its sparsity pattern is not pre-specified. Instead, it uses LSH (Locality Sensitive Hashing) technology to (approximately) find the largest few Attention values quickly and then only calculates those values. Furthermore, Reformer reduces memory usage by constructing a reversible form of the FFN (Feedforward Network) to replace the original FFN and redesigning the backpropagation process.

Therefore, compared to the aforementioned Sparse Attention, Reformer solves the first shortcoming but still suffers from the second: high implementation complexity. Implementing LSH-based Attention is much more complex than standard Attention, and rewriting the backpropagation process for reversible networks is even further out of reach for ordinary readers.

Linformer

A work very similar to the Linear Attention introduced in this article is Facebook’s recently released Linformer. It still retains the original Scaled-Dot Attention form but, before performing Attention, uses two m\times n matrices \boldsymbol{E} and \boldsymbol{F} to project \boldsymbol{K} and \boldsymbol{V} respectively: Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax\left(\boldsymbol{Q}(\boldsymbol{E}\boldsymbol{K})^{\top}\right)\boldsymbol{F}\boldsymbol{V} In this way, \boldsymbol{Q}(\boldsymbol{E}\boldsymbol{K})^{\top} is only an n\times m matrix. The authors claim that even for very large sequence lengths n, m can be kept as a moderate constant, making this Attention linear as well. A similar idea to Linformer also appeared in an earlier CV paper "Asymmetric Non-local Neural Networks for Semantic Segmentation".

However, I believe the conclusion that "for ultra-long sequences, m can remain constant" is questionable. For long sequences, the original paper only performed MLM tasks, and it is clear that MLM does not require long-range dependencies as much, so this experiment is not very convincing. Therefore, whether Linformer is truly Linear remains to be discussed.

Autoregressive Generation

Another drawback of Linformer is that the operations \boldsymbol{E}\boldsymbol{K} and \boldsymbol{F}\boldsymbol{V} directly "mix" the information of the entire sequence, so it cannot easily mask future information (Causal Masking). Consequently, it cannot perform autoregressive generation tasks such as language modeling or Seq2Seq, which is also why the original authors only performed MLM tasks. In contrast, the several types of Linear Attention introduced in this article can achieve this. Taking equations [eq:gen-att] and [eq:gen-att-2] as examples, if we want to mask future information, we only need to change the summation \sum\limits_{j=1}^n to \sum\limits_{j=1}^i: Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i = \frac{\sum\limits_{j=1}^i \left(\phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)\right)\boldsymbol{v}_j}{\sum\limits_{j=1}^i \phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)}=\frac{ \phi(\boldsymbol{q}_i)^{\top} \sum\limits_{j=1}^i\varphi(\boldsymbol{k}_j)\boldsymbol{v}_j^{\top}}{ \phi(\boldsymbol{q}_i)^{\top} \sum\limits_{j=1}^i\varphi(\boldsymbol{k}_j)} There are two ways to implement the above equation: the first way is to let \boldsymbol{S}_i=\sum\limits_{j=1}^i\varphi(\boldsymbol{k}_j)\boldsymbol{v}_j^{\top} and \boldsymbol{z}_i=\sum\limits_{j=1}^i\varphi(\boldsymbol{k}_j), giving us: Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})_i =\frac{ \phi(\boldsymbol{q}_i)^{\top} \boldsymbol{S}_i}{ \phi(\boldsymbol{q}_i)^{\top} \boldsymbol{z}_i},\quad \begin{aligned}&\boldsymbol{S}_i=\boldsymbol{S}_{i-1}+\varphi(\boldsymbol{k}_i)\boldsymbol{v}_i^{\top}\\ &\boldsymbol{z}_i=\boldsymbol{z}_{i-1}+\varphi(\boldsymbol{k}_i) \end{aligned} This shows that this type of Attention can be implemented as an RNN model in a recursive manner. Its spatial complexity is the lowest, but it requires sequential calculation, making it suitable for use during inference decoding. The second way is to directly take the outer product of \varphi(\boldsymbol{K}) and \boldsymbol{V}\in\mathbb{R}^{n\times d} to get an n\times d\times d matrix, and then perform a cumsum operation along the n dimension. This yields \boldsymbol{S}_1, \boldsymbol{S}_2, \dots, \boldsymbol{S}_n all at once. This method is the fastest but has the largest memory footprint, making it suitable for training. However, since d^2 \gg n is often true, it is generally difficult to bear this spatial complexity during training, so the RNN form is mostly used.

Downsampling Techniques

From the results, Linformer’s \boldsymbol{E}\boldsymbol{K} and \boldsymbol{F}\boldsymbol{V} essentially shorten the sequence (downsampling). The most primitive method for shortening a sequence is Pooling, so I previously tried introducing Pooling technology into Transformers. Similar works have emerged recently, such as IBM’s "PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination" and Google’s "Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing". Besides Pooling, there are other downsampling techniques, such as implementation via 1D convolution with stride > 1. Based on this idea, perhaps we could replace the Position-Wise fully connected layers in the FFN with 1D convolutions with stride > 1? In short, many variations can be played with in this area, but like Linformer, it becomes difficult to perform autoregressive generation after such mixing.