Exploring Linear Attention: Does Attention Really Need a Softmax?

线性Attention的探索：Attention必须有个Softmax吗？ 苏剑林 · 2020-07-04 · Machine Learning model text generation attention

As is well known, although Transformer-style models built on the attention mechanism have good parallelism properties, both their space and time complexity are $\mathcal{O}(n^2)$, where $n$ is the sequence length, so when $n$ is large, the computational cost of Transformer models becomes prohibitive. Recently, quite a lot of work has focused on reducing the computational cost of Transformer models — for instance, model pruning, quantization, distillation and other compression techniques, or alternatively modifying the attention structure itself so that its complexity can be brought down to $\mathcal{O}(n\log n)$ or even $\mathcal{O}(n)$.

A few days ago I came across the paper Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention, which introduced me to this line of exploration known as Linear Attention. I then went on to read several related papers and gained some useful insights, which I have gathered together in this post to summarize my current understanding of linear attention.

Attention

The most popular attention mechanism at present is undoubtedly Scaled-Dot Attention, which takes the form

\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax\left(\boldsymbol{Q}\boldsymbol{K}^{\top}\right)\boldsymbol{V}\label{eq:std-att}\end{equation}

Here, $\boldsymbol{Q}\in\mathbb{R}^{n\times d_k}, \boldsymbol{K}\in\mathbb{R}^{m\times d_k}, \boldsymbol{V}\in\mathbb{R}^{m\times d_v}$, and for simplicity we haven't explicitly written out the scaling factor of attention. In this post we're mainly concerned with the self-attention scenario, so for convenience of presentation we uniformly set $\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V}\in\mathbb{R}^{n\times d}$; in general we have $n > d$ and even $n\gg d$ (in BERT base, $d=64$). For more background, see my earlier post A Brief Reading of "Attention is All You Need" (Introduction + Code), and for some improvements built on top of it, see Breaking Through the Bottleneck: Building a Stronger Transformer and Google's New Work Synthesizer: We Still Don't Fully Understand Self-Attention. I won't go into further depth here.

Removing the Softmax

Readers may not expect this, but the key factor limiting the performance of attention is actually the Softmax in its definition! In fact, a simple derivation makes this clear. $\boldsymbol{Q}\boldsymbol{K}^{\top}$ At this step we get an $n\times n$ matrix, and it's precisely this step that determines that the complexity of attention is $\mathcal{O}(n^2)$. If there were no Softmax, we would just have three matrices multiplied together, $\boldsymbol{Q}\boldsymbol{K}^{\top}\boldsymbol{V}$, and since matrix multiplication satisfies associativity, we could first compute $\boldsymbol{K}^{\top}\boldsymbol{V}$, obtaining an $d\times d$ matrix, and then left-multiply it by $\boldsymbol{Q}$. Since $d \ll n$, this way the overall complexity is roughly only $\mathcal{O}(n)$ (dominated by the step of left-multiplying by $\boldsymbol{Q}$).

In other words, once we remove the Softmax, the complexity of attention can be reduced to the ideal linear level, $\mathcal{O}(n)$! This is clearly the ultimate goal we're after: Linear Attention, i.e. attention with linear complexity. So the theme of this post is to explore what linear attention looks like once we remove the Softmax.

A General Definition

The question is: can we still call it attention if we simply remove the Softmax? Can it still achieve effects comparable to standard attention? To answer this, let's first rewrite the definition of Scaled-Dot Attention, $\eqref{eq:std-att}$, equivalently as (all vectors in this post are column vectors)

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

So Scaled-Dot Attention is essentially a weighted average of $\boldsymbol{v}_j$ with $e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$ as weights. This lets us propose a generalized definition of attention:

\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)}\label{eq:gen-att}\end{equation}

That is, we replace $e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$ with a general function $\boldsymbol{q}_i, \boldsymbol{k}_j$ of $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)$. To preserve the distribution-like properties of attention, we require that $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$ always holds. In other words, if we want to define a new form of attention, we need to keep the form of equation $\eqref{eq:gen-att}$ while satisfying $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$.

This kind of generalized attention is also known in computer vision as a Non-Local network, from the paper Non-local Neural Networks.

A Few Examples

If we simply remove the Softmax, we get $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$, but the problem is that an inner product cannot guarantee non-negativity, so this isn't yet a valid choice. Below we introduce a few viable alternatives.

It's worth pointing out that of the linear attention variants introduced below, the first two come from the CV literature, while the third is something I came up with myself — so none of these have actually been tested experimentally on NLP tasks yet. NLP researchers working on model improvements now have a direction for experimentation (^_^). Incidentally, the CV community has quite a few pieces of work improving attention (besides the ones introduced below, there's also EMANet, etc.), and much of this content is worth reading for those of us working in NLP.

Kernel Function Form

A natural idea is: if every element of $\boldsymbol{q}_i,\boldsymbol{k}_j$ is non-negative, then the inner product will naturally be non-negative too. To achieve this, we can apply an activation function $\phi,\varphi$ to $\boldsymbol{q}_i,\boldsymbol{k}_j$ separately, i.e.

\begin{equation}\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \phi(\boldsymbol{q}_i)^{\top} \varphi(\boldsymbol{k}_j)\label{eq:gen-att-2}\end{equation}

where $\phi(\cdot),\varphi(\cdot)$ is an activation function with non-negative range. The paper mentioned at the start of this post, Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention, chose $\phi(x)=\varphi(x)=\text{elu}(x)+1$.

If we want to tell a story about it, equation $\eqref{eq:gen-att-2}$ can be connected to "kernel methods" — in particular, when $\phi=\varphi$, $\phi$ acts like a kernel function, and $\langle \phi(\boldsymbol{q}_i), \phi(\boldsymbol{k}_j)\rangle$ is exactly the inner product defined via that kernel function. For more along these lines, see the paper Transformer Dissection: An Unified Understanding for Transformer's Attention via the Lens of Kernel; I won't dwell on this further here.

A Clever Use of Softmax

Another, earlier paper, Efficient Attention: Attention with Linear Complexities, offers an even more interesting choice. It notices that in $\boldsymbol{Q}\boldsymbol{K}^{\top}$, $\boldsymbol{Q}, \boldsymbol{K}, \in\mathbb{R}^{n\times d}$, and that if "$\boldsymbol{Q}$ is normalized along the $d$ dimension, and $\boldsymbol{K}$ is normalized along the $n$ dimension," then $\boldsymbol{Q}\boldsymbol{K}^{\top}$ is automatically normalized as well. So the choice it makes is:

\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax_2\left(\boldsymbol{Q}\right)softmax_1(\boldsymbol{K})^{\top}\boldsymbol{V}\end{equation}

where $softmax_1$ and $softmax_2$ denote Softmax applied along the first dimension ($n$) and second dimension ($d$) respectively. In other words, here we apply Softmax separately to $\boldsymbol{Q},\boldsymbol{K}$, rather than applying Softmax after computing $\boldsymbol{Q}\boldsymbol{K}^{\top}$.

If we take $\phi(\boldsymbol{q}_i)=softmax(\boldsymbol{q}_i),\varphi(\boldsymbol{k}_j)=softmax(\boldsymbol{k}_j)$ directly, it's clear that this form is a special case of equation $\eqref{eq:gen-att-2}$. Incidentally, this design has appeared more than once in the CV literature — for example, A2-Nets also uses the same trick.

My Own Idea

Here, I'll present a construction of my own. The starting point of this construction is no longer equation $\eqref{eq:gen-att-2}$, but instead comes from an approximation of the original definition $\eqref{eq:std-att-2}$. By Taylor expansion we have

\begin{equation}e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j} \approx 1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j\end{equation}

If $\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$, then the non-negativity of the right-hand side is guaranteed, which lets us set $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)=1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$. At this point readers may already see that, to guarantee $\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$, we just need to apply $l_2$ normalization to $\boldsymbol{q}_i,\boldsymbol{k}_j$ separately. So the scheme I ultimately propose is:

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

This differs from form $\eqref{eq:gen-att-2}$, but theoretically it's a closer approximation to the original Scaled-Dot Attention.

Related Work

There has been a great deal of work aimed at reducing the computational complexity of attention by modifying its form; here I'll briefly list a few of them.

Sparse Attention

We previously covered OpenAI's Sparse Attention, which reduces the computational cost of attention by "only keeping values within a small local region, forcing most attention weights to be zero." After special design, most elements of the attention matrix become zero, so in theory it can also save on GPU memory and computation. Subsequent similar work includes Explicit Sparse Transformer: Concentrated Attention Through Explicit Selection and Longformer: The Long-Document Transformer.

But clearly this approach has two shortcomings:

1. How to choose which attention regions to keep is decided manually and subjectively, which is far from being an intelligent solution;
2. It requires specific engineering optimizations to achieve an efficient implementation, so it doesn't generalize easily.

Reformer

Reformer is another representative piece of improved work, bringing the complexity of attention down to $\mathcal{O}(n\log n)$. In a sense, Reformer is also a form of sparse attention, except that its sparsity pattern isn't specified in advance — instead, it uses LSH (Locality Sensitive Hashing) to (approximately) quickly find the largest attention values, and then only computes those. In addition, Reformer replaces the original FFN (Feedforward Network) with a reversible construction and redesigns the backpropagation process accordingly, thereby reducing GPU memory usage.

So, compared to the sparse attention discussed above, Reformer solves the first shortcoming, but it still has the second: high implementation complexity. Implementing LSH-based attention is far more complex than standard attention, and rewriting the backpropagation process for a reversible network is even more out of reach for the average reader.

Linformer

A piece of work quite similar to the linear attention discussed in this post is Facebook's recently released Linformer, which still keeps the original Scaled-Dot Attention form, but before performing attention, it uses two $m\times n$ matrices $\boldsymbol{E},\boldsymbol{F}$ to project $\boldsymbol{K},\boldsymbol{V}$ respectively, turning it into

\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax\left(\boldsymbol{Q}(\boldsymbol{E}\boldsymbol{K})^{\top}\right)\boldsymbol{F}\boldsymbol{V}\end{equation}

This way, $\boldsymbol{Q}(\boldsymbol{E}\boldsymbol{K})^{\top}$ becomes just an $n\times m$ matrix, and the authors claim that even for very long sequence lengths $n$, $m$ can be kept as a moderate constant, making this form of attention linear as well. A similar idea appears in an earlier CV paper, Asymmetric Non-local Neural Networks for Semantic Segmentation.

However, I think the claim that "$m$ can remain constant for extremely long sequences" is questionable. For long sequences, the original paper only experimented with MLM tasks, and it's clear that MLM doesn't really require long-range dependencies, so this experiment isn't very convincing. So whether Linformer is truly linear remains open to debate.

Autoregressive Generation

Another drawback of Linformer is that the operations $\boldsymbol{E}\boldsymbol{K},\boldsymbol{F}\boldsymbol{V}$ directly "blend" information from the entire sequence together, so it can't easily mask out future information (causal masking), and hence can't be used for language modeling, Seq2Seq, or other autoregressive generation tasks — this is precisely why the original authors only tested MLM tasks. By contrast, all the linear attention variants introduced in this post can handle this. Taking equations $\eqref{eq:gen-att}$ and $\eqref{eq:gen-att-2}$ as examples, if we want to mask out future information, we just need to change the summation $\sum\limits_{j=1}^n$ to $\sum\limits_{j=1}^i$:

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

There are two ways to implement the equation above. The first is to set $\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

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

This shows that this form of attention can be implemented recursively as an RNN-like model, which has the lowest space complexity but must be computed serially — suitable for use during decoding/inference. The second way is to directly take the outer product of $\varphi(\boldsymbol{K}),\boldsymbol{V}\in\mathbb{R}^{n\times d}$, obtaining an $n\times d\times d$ matrix, and then perform a $\text{cumsum}$ operation along the $n$ dimension, which gives us $\boldsymbol{S}_1,\boldsymbol{S}_2,\dots,\boldsymbol{S}_n$ all at once. This is the fastest approach but has the largest memory footprint, making it suitable for training. However, since we often have $d^2\gg n$, in most cases this space complexity is too costly to bear during training, so the RNN form is generally used instead.

Downsampling Techniques

Looking at the outcome, Linformer's $\boldsymbol{E}\boldsymbol{K}, \boldsymbol{F}\boldsymbol{V}$ essentially shortens the sequence (downsampling), and the simplest way to shorten a sequence is pooling — so I previously tried introducing pooling techniques into the Transformer myself. Recently, similar work has appeared as well, 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 — for instance, using a 1D convolution with stride > 1. Building on this idea, perhaps we could replace the position-wise fully connected layer in the FFN with a stride > 1 1D convolution? In any case, there should be plenty of room to experiment in this direction, though as with Linformer, once things are blended together this way, autoregressive generation becomes difficult.

Summary

This post has introduced some work that modifies the structure of attention in order to reduce its computational complexity, the main idea being that removing the Softmax from standard attention allows the complexity of attention to degrade to the ideal $\mathcal{O}(n)$ level (Linear Attention). Compared to other similar structural improvements, this modification manages to bring the complexity down to $\mathcal{O}(n)$ while still preserving all "token-to-token" attention, and it also preserves the possibility of being used for autoregressive generation.