Linear Transformers are Probably Not the Models You Are Waiting For

Reflections on BERT

According to intuitive understanding, shouldn’t switching from quadratic complexity to linear complexity result in a "quantum leap" in performance? Why would it "fall far short of expectations"? The main reason for this confusion is that we have never carefully evaluated the overall computational workload of a conventional Transformer model (such as BERT).

Many readers already know that the structure of a Transformer is roughly an Embedding layer plus several Transformer layers. The computational load of the Embedding layer is minimal; we mainly care about the Transformer layers. Ignoring layers with relatively small computational loads like Residual connections and Layer Normalization, each Transformer layer consists of two main sub-layers: Self-Attention (SA) and the FeedForward Network (FFN). Although the seminal work on Transformers claimed that "Attention is all you need," many subsequent works have demonstrated the necessity of modules like residuals and FFN, such as "Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth."

Now, let me ask you a question:

Which do you think has a larger computational load: SA or FFN?

Evaluating Computational Workload

Undoubtedly, the complexity of SA is \mathcal{O}(n^2), while the complexity of FFN is \mathcal{O}(n). If you simply take this for granted and say that the computational load of SA is larger than that of FFN, you would be wrong!

We know that addition is much faster than multiplication, so when estimating computational load, we mainly calculate how many multiplications are required. In neural networks, the primary operation is matrix multiplication. It is easy to estimate that by definition, multiplying an a \times b matrix by a b \times c matrix requires abc multiplications. Therefore, abc is the complexity of multiplying two matrices. This is the basis for our estimation of Transformer complexity.

Let n be the sequence length, d be the head_size (64 in the base version), and h be the number of heads (12 in the base version). Then hd is what we usually call the hidden_size (768 in the base version). For SA, it starts with the projection transformations of Q, K, V, i.e., an n \times hd matrix multiplied by an hd \times hd matrix performed 3 times. Thus, the computational load is 3n(hd)^2. Next is the operation of the h attention heads. For each head, first, the n \times d matrix Q is multiplied by the d \times n matrix K^{\top} to get an n \times n Attention matrix (ignoring the computational load of softmax and normalization for now). Then, the n \times n matrix is multiplied by the n \times d matrix V to get an n \times d matrix. The computational load for both steps is n^2 d. So the total computational load for the heads is h(n^2 d + n^2 d). Finally, there is an output projection transformation, which is an n \times hd matrix multiplied by an hd \times hd matrix, with a computational load of n(hd)^2. Therefore, the total computational load of SA is: 3n(hd)^2 + h(n^2 d + n^2 d) + n(hd)^2 = 4nh^2 d^2 + 2n^2 hd As for FFN, it is simpler. It consists of two fully connected layers, which are two matrix transformations (ignoring the computational load of activation functions). The general parameter setting is: the first layer is an n \times hd matrix multiplied by an hd \times 4hd matrix, and the second layer is an n \times 4hd matrix multiplied by a 4hd \times hd matrix. So the total computational load is: n \times hd \times 4hd + n \times 4hd \times hd = 8nh^2 d^2 In this way, if the computational load of SA is larger than that of FFN, it means: 4nh^2 d^2 + 2n^2 hd > 8nh^2 d^2 \quad \Leftrightarrow \quad n > 2hd For the base version, this means n > 1536! In other words, only when the sequence length exceeds 1536 does the computational load of SA exceed that of FFN. Before that, the linear-complexity FFN dominates!

Furthermore, from the results above, we can obtain the total computational load of a Transformer layer as: 4nh^2 d^2 + 2n^2 hd + 8nh^2 d^2 = 12nh^2 d^2 + 2n^2 hd This is a sum of linear and quadratic terms with respect to n. When n is large enough, the complexity is naturally \mathcal{O}(n^2). However, the condition for the quadratic term to dominate is: 2n^2 hd > 12nh^2 d^2 \quad \Leftrightarrow \quad n > 6hd For the base version, this means n > 4608! That is to say, only when the sequence length approaches 5000 does the complexity of the Transformer truly manifest its quadratic nature!

Comprehensive Conclusion

Synthesizing the above results, we can conclude: For the base version, when the sequence length does not exceed 1536, the complexity of the Transformer is nearly linear; when the sequence length exceeds 1536, the computational load of the Transformer gradually becomes dominated by Attention, and the complexity slowly tends toward quadratic, until the length exceeds 4608, where the quadratic term truly dominates. Of course, these boundaries are estimates; actual conditions may vary slightly. Use these to perceive the range and order of magnitude.

I have previously suggested to many readers that for "long text" tasks with lengths not exceeding 2000, they should directly try models with no length limits like NEZHA or RoFormer without overthinking tricks. The reason is the same. No matter how many tricks you use, you can at most reduce it to linear complexity, and within this length range, the model itself is already nearly linear. Various tricks won’t save much.

For readers who honestly use BERT-base, the maxlen generally does not exceed 512, which is far below the aforementioned limits. Therefore, stop complaining about how the quadratic complexity of Attention wastes hardware, because the fact is:

BERT is slow mainly because it is truly large, not because of the quadratic complexity of Attention.

The Meaning of "Linear"

Another reason for the confusion regarding Linear Attention "falling far short of expectations" is the failure to analyze the computational load of Linear Attention from a practical perspective, leading to excessively high expectations.

For an introduction to Linear Attention, you can refer to "Exploration of Linear Attention: Does Attention Necessarily Need Softmax?". Briefly, Linear Attention calculates attention in the order of Q(K^{\top} V). Following the previous estimation method, the computational load for each head in Linear Attention is 2nd^2, while standard Attention is 2n^2 d. Thus, if n > d, Linear Attention saves more computational load than standard Attention. (Note: there is more than one way to achieve linear efficiency in Attention, but generally, the complexity is similar, so the following conclusions are representative.)

For the base version, that would be n > 64. This limit is easily reached, so some readers might think "saving a little is better than nothing" or "might as well use it." However, this assumes that both standard Attention and Linear Attention use the same d. Readers who have carefully studied "Performer: Linearizing Attention Complexity with Random Projections" and "Transformer Upgrade: 3. From Performer to Linear Attention" know that Linear Attention suffers from a more severe "low-rank bottleneck" than standard Attention. Therefore, if you switch to Linear Attention while keeping the same d, the performance will drop significantly. To maintain roughly the same effect, Linear Attention needs a larger d (usually about 4 times the original).

In this case, the computational load of Linear Attention would be 2n(4d)^2. For Linear Attention to be faster than standard Attention, we need n > 16d. For the base version, this means n > 1024, which also exceeds the range used by most readers. Moreover, after switching to Linear Attention, the previous conclusions regarding the computational load of SA and FFN still hold—that is, for most sequence lengths, the dominant computational load is still linear operations like FFN. After switching to Linear Attention, one cannot feel a significant speed increase. So, overall:

Unless your sequence length is in the thousands or tens of thousands, don’t think about switching to Linear Attention.

Reviewing the Papers

In fact, even without the above analysis, any reader who has carefully read works related to Attention efficiency improvements can draw similar conclusions from certain figures in the papers: so-called "efficient" Attention is generally only applicable to sequence lengths in the thousands or tens of thousands; only in these scenarios is there a significant performance improvement.

For example, the earlier work Sparse Transformers includes a figure showing that the sequence lengths processed are all 3000+:

For instance, the famous Reformer demonstrates performance with sequence lengths in units of K (thousands):

The widely acclaimed Longformer is the same:

There is also Google’s classic work on Linear Attention, Performer, which shows that even if the sequence length is 2^{12}=4096, the gap between Performer and Transformer is not particularly significant:

Finally, the relatively new work Luna provides a comprehensive comparison table, which also supports our conclusion:

From the existing works on efficient Attention, we can conclude: The sequence lengths these improvement works care about are mainly in the thousands. Sequence lengths with significant computational efficiency improvements are basically several thousands. Of course, our previous discussion mainly focused on time complexity; for space complexity (VRAM usage), the reduction is generally larger than the improvement in time complexity, but overall, it is only valuable for long sequences.

Change Your Expectations

Therefore, if your sequence length is only one or two hundred, do not expect any improvement from Attention itself; just honestly switch to a smaller model. You can hope that in the future there will be smaller models that achieve the same good results, but do not expect models of the same size to improve efficiency by modifying Attention. To put it bluntly, even if you completely remove Attention, performance won’t improve much.

Original address: https://kexue.fm/archives/8610