Linear Transformers Are Probably Not the Model You're Waiting For

线性Transformer应该不是你要等的那个模型 苏剑林 · 2021-08-09 · Machine Learning model matrix attention

In this blog, we've discussed linear attention several times already. The usual narrative for introducing linear attention goes something like this: standard attention has $\mathcal{O}(n^2)$ quadratic complexity, which is one of its main "pain points", so we $\mathcal{O}(n)$ an improved model with reduced complexity — namely, linear attention. After reading such introductions, some readers have been eagerly waiting for us to release a pretrained model based on linear attention, hoping it would relieve the suffering they've endured from BERT's computational cost.

However, what this post has to say is: readers with such hopes are likely to be disappointed. The gap between standard attention and linear attention falls far short of what you'd expect, and the reason BERT is so slow isn't actually the quadratic complexity of standard attention.

Rethinking BERT

Intuitively, shouldn't switching from quadratic complexity to linear complexity be a huge leap forward? Why would it instead "fall far short of expectations"? The reason this puzzle arises is that we've never carefully evaluated the overall computational cost of a conventional Transformer model like BERT.

As many readers already know, a Transformer's architecture is roughly an embedding layer plus several Transformer layers. The embedding layer's computational cost is negligible, so we mainly care about the Transformer layers. Ignoring lightweight components like residual connections and layer normalization, each Transformer layer mainly consists of two sublayers: Self-Attention (SA) and the FeedForward Network (FFN). Although the seminal Transformer paper claims that "Attention is all you need", plenty of subsequent work has demonstrated the necessity of modules like residual connections and FFNs — see for instance Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth.

Now let me ask you a question:

Do you think SA has a higher computational cost, or FFN?

Estimating the Computational Cost

There's no doubt that SA has complexity $\mathcal{O}(n^2)$ while FFN has complexity $\mathcal{O}(n)$. If you take this at face value and conclude that SA costs more than FFN, you'd be wrong!

We know addition is much cheaper than multiplication, so when estimating computational cost we mainly count the number of multiplications. In neural networks, the dominant operation is matrix multiplication. It's easy to see that multiplying a $a\times b$ matrix by a $b\times c$ matrix, done by definition, requires $abc$ multiplications, so $abc$ is the complexity of multiplying two matrices — this is the basis for our estimate of Transformer complexity.

Let $n$ be the sequence length, $d$ the head size (64 for the base version), and $h$ the number of heads (12 for the base version), so that $hd$ is what we usually call the "hidden size" (768 for the base version). For SA, we start with the $Q,K,V$ projection, i.e. a $n\times hd$ matrix multiplied by a $hd\times hd$ matrix, done three times, giving a cost of $3n(hd)^2$. Then there's the computation for $h$ attention heads: for each head, a $n\times d$ matrix of $Q$ is multiplied by a $d\times n$ matrix of $K^{\top}$ to get a $n\times n$ attention matrix (ignoring the cost of softmax and normalization for now), and then a $n\times n$ matrix is multiplied by a $n\times d$ matrix of $V$ to get a $n\times d$ matrix; both of these steps cost $n^2 d$, so the total cost is $h(n^2 d + n^2 d)$. Finally there's an output projection, which is again a $n\times hd$ matrix multiplied by a $hd\times hd$ matrix, costing $n(hd)^2$. So the total cost of SA is

\begin{equation}3n(hd)^2 + h(n^2 d + n^2 d) + n(hd)^2 = 4nh^2 d^2 + 2n^2 hd\end{equation}

As for FFN, it's simpler — just two fully connected layers, i.e. two matrix transformations (again ignoring the cost of the activation function). Typically, the first layer is a $n\times hd$ matrix multiplied by a $hd\times 4hd$ matrix, and the second layer is a $n\times 4hd$ matrix multiplied by a $4hd\times hd$ matrix. So the total cost is

\begin{equation}n\times hd\times 4hd + n\times 4hd\times hd = 8nh^2 d^2\end{equation}

Given this, for SA's cost to exceed FFN's, we'd need

\begin{equation}4nh^2 d^2 + 2n^2 hd > 8nh^2 d^2\quad\Leftrightarrow\quad n > 2hd\end{equation}

For the base version, this means $n > 1536$! In other words, only when the sequence length exceeds 1536 does SA's computational cost surpass FFN's — before that point, it's the linear-complexity FFN that dominates!

There's more. From the above results, we can derive the total computational cost of a Transformer layer:

\begin{equation}4nh^2 d^2 + 2n^2 hd + 8nh^2 d^2 = 12nh^2 d^2 + 2n^2 hd\end{equation}

This is a sum of a linear term and a quadratic term in $n$. When $n$ is large enough, the complexity is naturally $\mathcal{O}(n^2)$, but the condition for the quadratic term to dominate is

\begin{equation}2n^2 hd > 12nh^2 d^2\quad\Leftrightarrow\quad n > 6hd\end{equation}

For the base version, this means $n > 4608$! In other words, the Transformer's complexity only truly manifests as quadratic once the sequence length approaches 5000!

Putting It Together

Combining these results, we arrive at the following conclusion: for the base version, when the sequence length is no more than 1536, the Transformer's complexity is nearly linear; once the sequence length exceeds 1536, attention gradually becomes the dominant cost and the complexity slowly tends toward quadratic, but it's only once the length exceeds roughly 4608 that the quadratic term truly dominates. Of course this boundary is just an estimate and the actual situation may vary — treat it as a way to get a feel for the rough range and order of magnitude.

I've previously suggested to many readers that for "long text" tasks with a sequence length under 2000, they should just try an unrestricted-length model like NEZHA or RoFormer directly, without overthinking the tricks — and this is exactly why. No matter how clever your tricks are, at best they bring you down to linear complexity, and within this length range, the model is already nearly linear anyway, so the tricks don't save you much.

For readers who honestly just use BERT base, with maxlen typically no more than 512 — far below the thresholds above — please stop complaining about attention's quadratic complexity being a resource hog. The truth is:

BERT is slow mainly because it's genuinely big, not because of attention's quadratic complexity.

What "Linear" Really Means

The other reason people find linear attention "far short of expectations" is that they never actually analyzed the real computational cost of linear attention, leading to overly high expectations.

For an introduction to linear attention, see Exploring Linear Attention: Does Attention Really Need a Softmax?; I won't repeat it here. In short, linear attention computes attention in the order $Q(K^{\top} V)$. So, following the same estimation method as before, the per-head cost of linear attention is $2nd^2$, while standard attention is $2n^2 d$. Therefore, if $n > d$, linear attention is cheaper than standard attention. (Note: this isn't the only approach for achieving linear efficiency in attention, but the complexities involved are broadly similar, so the conclusion below is representative.)

For the base version, that condition is $n > 64$, which is quite easy to satisfy, so some readers might think "every bit of savings helps" or "why not use it if it's free." However, this conclusion assumes that standard attention and linear attention both use the same $d$. Readers who've carefully studied Performer: Linearizing Attention Complexity via Random Projections and The Path to the Upgraded Transformer: 3. From Performer to Linear Attention know that linear attention suffers from a much more severe "low-rank bottleneck" than standard attention. So if you switch to linear attention while keeping the same $d$, its performance will drop noticeably; to preserve roughly the same performance, linear attention needs a larger $d$ (typically around 4 times larger).

Given this, the actual computational cost of linear attention should be $2n(4d)^2$. For linear attention to be faster than standard attention, we'd need $n > 16d$, which for the base version comes out to $n > 1024$ — again beyond the range most readers would ever use. Moreover, even after switching to linear attention, the earlier conclusion about SA versus FFN cost still holds: across most sequence lengths, it's still FFN and other linear operations that dominate the computational cost, so switching to linear attention won't produce any noticeable speedup. So, in short:

Unless you're dealing with sequence lengths in the thousands, don't bother switching to linear attention.

Flipping Through the Papers Again

In fact, even without the analysis above, anyone who has carefully read work on improving attention efficiency can reach a similar conclusion just from certain figures in these papers: so-called "efficient" attention mechanisms generally only pay off at sequence lengths in the thousands — only in that regime do you see a clear performance improvement.

For example, in the earlier work Sparse Transformers, there's a figure showing that the sequence lengths handled are all 3000+:

The sequence lengths handled by Sparse Transformer are all 3000+

Or take the famous Reformer, whose performance demonstrations use sequence lengths measured in thousands (K):

The sequence lengths used to demonstrate Reformer's performance are all in units of K

The widely praised Longformer is no different:

The sequence lengths used to demonstrate Longformer's performance are in the thousands or even tens of thousands

And there's Google's classic work on linear attention, Performer, which shows that even at a sequence length of $2^{12}=4096$, the gap between Performer and the standard Transformer isn't particularly dramatic:

Performer's performance curve

Finally, the more recent work Luna provides a fairly comprehensive comparison table, which likewise supports our conclusion:

Performance comparison of various improved attention mechanisms in Luna

Looking across these existing works on efficient attention, we can conclude that the sequence lengths these improvements care about are mainly measured in thousands; sequence lengths with a clear efficiency gain are essentially always at least several thousand. Of course, our discussion above has mainly focused on time complexity — for space complexity, i.e. GPU memory usage, the reduction is generally larger than the improvement in time complexity, but overall, it's only at long sequence lengths that these methods pay off.

Time to Reset Your Expectations

So, if your sequence length is only a couple hundred tokens, don't expect improvements to attention itself to help at all — just honestly switch to a smaller model. You can reasonably hope that in the future, smaller models will be able to achieve the same level of performance. But don't expect that a model of the same size can be made more efficient just by modifying attention, because frankly, even if you removed attention entirely, you wouldn't gain much in performance.