Efficient Transformers generally refer to all works aimed at improving the efficiency of the Transformer architecture. I have been following this field quite early; my earliest blog post can be traced back to 2019, "Born for Savings: From Standard Attention to Sparse Attention", at a time when there was very little work in this area. Later, such works gradually increased, and I followed some of them, such as Linear Attention, Performer, and Nyströmformer. I even conducted some explorations myself, such as the previous "Road to Transformer Upgrade" series. Eventually, as related works became more numerous and mostly uninteresting, I stopped paying much attention.
Feeling much like "rain after a long drought," an interesting efficient Transformer work has finally appeared—"Transformer Quality in Linear Time" from Google. After a careful reading, I believe the paper is truly "full of surprises."
What is the Surprise?
What kind of results deserve the description "surprising"? Is it an exaggeration? Let’s first look at what the paper achieves:
1. It proposes a new Transformer variant that still has quadratic complexity but is faster, has lower memory usage, and achieves better results compared to the standard Transformer.
2. It proposes a new linearization scheme for Transformers that not only improves the performance of original Linear Attention but also maintains the possibility of being used as a Decoder, while keeping efficient training parallelism during decoding.
To be honest, I think achieving either of the above points is very rare, and this paper manages both at once, which is why I describe it as "full of surprises." More importantly, the improvements in the paper are generally natural and elegant, unlike many similar works that feel forced. Furthermore, I have conducted simple replication experiments, and the results suggest that the paper’s reproducibility is quite good. It truly feels like "the standard Transformer is in danger."
Gated Attention
Without further ado, let’s get to the main topic. We know that the standard Transformer is constructed by alternating Attention layers and FFN (Feed-Forward Network) layers. The core of this paper is a new design called GAU (Gated Attention Unit) that merges the two. It is the key to making the new model faster, more efficient, and better. Additionally, it results in a model with only one type of layer, which is more elegant.
Initial Power
How do we merge Attention and FFN? First, a standard FFN is a two-layer MLP: \boldsymbol{O}=\phi(\boldsymbol{X}\boldsymbol{W}_u)\boldsymbol{W}_o where \boldsymbol{X}\in\mathbb{R}^{n\times d}, \boldsymbol{W}_u\in\mathbb{R}^{d\times e}, \boldsymbol{W}_o\in\mathbb{R}^{e\times d}, and \phi is an activation function. Later, "GLU Variants Improve Transformer" found that FFNs using GLU (Gated Linear Unit) perform better, which was adopted by mT5. Its form is: \boldsymbol{O}=(\boldsymbol{U}\odot\boldsymbol{V})\boldsymbol{W}_o,\quad \boldsymbol{U}=\phi_u(\boldsymbol{X}\boldsymbol{W}_u),\quad\boldsymbol{V}=\phi_v(\boldsymbol{X}\boldsymbol{W}_v) where \boldsymbol{W}_u, \boldsymbol{W}_v\in\mathbb{R}^{d\times e} and \odot is the element-wise product (Hadamard product). It is not surprising that GLU is more effective; it played a key role as early as Facebook’s 2017 paper "Convolutional Sequence to Sequence Learning". Additionally, my previous research on DGCNN also confirmed the effectiveness of GLU.
In general GLU, \boldsymbol{U} has no activation while \boldsymbol{V} uses Sigmoid. However, in this paper, both \boldsymbol{U} and \boldsymbol{V} use the Swish activation (also known as SiLU, Sigmoid Linear Unit). This can be found in the source code in the appendix; it differs slightly from mainstream GLU usage.
Strong Combination
Since GLU-style FFNs are more effective, we modify them. Note that FFNs cannot replace Attention because there is no interaction between tokens—each row of matrices \boldsymbol{U} and \boldsymbol{V} is calculated independently. To supplement this deficiency, a natural idea is to add token interactions to \boldsymbol{U} and \boldsymbol{V}. To reflect the combination with Attention, a natural design is: \boldsymbol{O}=(\boldsymbol{U}\odot\boldsymbol{A}\boldsymbol{V})\boldsymbol{W}_o \label{eq:mix} where \boldsymbol{A}\in\mathbb{R}^{n\times n} is the Attention matrix responsible for fusing information between tokens. The resulting \boldsymbol{O} contains token interactions and can, in principle, replace Attention. As for how to calculate \boldsymbol{A}, we will discuss that shortly.
In equation [eq:mix], if \boldsymbol{A} is the identity matrix \boldsymbol{I}, it is a GLU-style FFN. If \boldsymbol{U} is an all-ones matrix, it is standard Attention. Thus, [eq:mix] is a simple and natural fusion of Attention and FFN. We expect it to replace both while performing better.
Weak Attention
As mentioned, GLU itself is very strong (otherwise Facebook couldn’t have achieved SOTA in Seq2Seq using CNN+GLU). Since GLU is so strong, one might guess that it weakens the dependence on Attention. That is, while \boldsymbol{A} is indispensable in equation [eq:mix], we might be able to simplify its form. Indeed, the original paper uses the following simplified Attention matrix: \boldsymbol{A}=\frac{1}{n}\text{relu}^2\left(\frac{\mathcal{Q}(\boldsymbol{Z})\mathcal{K}(\boldsymbol{Z})^{\top}}{\sqrt{s}}\right)=\frac{1}{ns}\text{relu}^2\left(\mathcal{Q}(\boldsymbol{Z})\mathcal{K}(\boldsymbol{Z})^{\top}\right),\quad \boldsymbol{Z}=\phi_z(\boldsymbol{X}\boldsymbol{W}_z) \label{eq:relu-att} where \boldsymbol{W}_z\in\mathbb{R}^{d\times s}, s is the attention head size (the paper uses s=128), \mathcal{Q}, \mathcal{K} are simple affine transformations (like multiplying by \gamma and adding \beta in Layer Norm), and \text{relu}^2 is \text{relu} followed by squaring.
Similar to standard Scaled-Dot Self Attention, the attention matrix here is derived from the dot product of \boldsymbol{Q} and \boldsymbol{K} divided by the square root of the dimension. The complexity is still \mathcal{O}(n^2). The difference is the simplified transformation for \boldsymbol{Q}, \boldsymbol{K} and the use of \text{relu}^2. This activation function might be unfamiliar; it was found via NAS in the authors’ previous paper "Primer: Searching for Efficient Transformers for Language Modeling". The final 1/n is a simple normalization factor to eliminate the influence of length. This design’s success also indicates that softmax is not mandatory in attention mechanisms and can be replaced by a regular activation function plus simple normalization.
Note that according to the reference code in the paper’s appendix, the simplified scaling factor is actually 1/n^2 rather than 1/ns. I believe 1/ns is more reasonable; otherwise, when n is large enough, each attention term becomes too small. Moreover, compared to the softmax used in standard attention, the denominator is only of order \mathcal{O}(n), so n^2 feels unscientific. I have performed a simple comparison and found that at length 512, the 1/ns version is slightly better, so I will present it according to my intuition.
One for All
Now, please pay close attention—the real "heavyweight" is about to appear! Perhaps GLU is so strong that the dependence on Attention is so weak that the authors discovered: Only one head is enough!
We know that standard Transformers use multi-head attention, which generates matrices of size bhn^2 during computation, where b is batch size and h is the number of heads. Imagine when n=1000, 2000, or larger; n^2 is already bad enough, and multiplying it by h is "adding insult to injury" for both time and space complexity. Now, with GAU using only one head, it can achieve the same or even better results, increasing calculation speed and reducing memory usage—it’s almost a "free lunch."
When GAU has only one head, the parameter count for \boldsymbol{W}_z is very small. The main parameters are in \boldsymbol{W}_u, \boldsymbol{W}_v, \boldsymbol{W}_o, so the parameter count of GAU is approximately 3de. In a standard Transformer, Attention has 4d^2 parameters and FFN has 8d^2 (assuming e=4d in standard FFN), totaling 12d^2. Therefore, in terms of parameters, when e=2d, two GAU layers are roughly equivalent to the original Attention+FFN.
Thus, in GAU experiments, the authors fixed e=2d. A standard Transformer with "n layers of Attention + n layers of FFN" corresponds to a new model with "2n layers of GAU," which we call FLASH-Quad. "Quad" is short for "Quadratic," indicating the complexity is still quadratic. The meaning of FLASH will be discussed later.
Efficient Linear
FLASH-Quad is already an excellent replacement for the standard Transformer, but the authors were not satisfied with its quadratic complexity and proposed FLASH (Fast Linear Attention with a Single Head) with linear complexity. To this end, the authors proposed a "Mixed Chunk Attention" scheme, which can be used not only in GAU but also in standard Attention—a general linearization trick.
Existing Methods
Mainstream efficient Transformer works generally follow two paths: "Sparsification" and "Linearization."
The work mentioned at the beginning, "From Standard Attention to Sparse Attention", is one of the "Sparsification" works. Later works like Reformer also fall into this category, as do pooling-based works like Linformer. These works introduce inductive biases to force most attention to zero, theoretically reducing computation. However, they often require specialized programming optimization for speedup or are difficult to use as Decoders (pooling works). Furthermore, their performance depends heavily on the introduced inductive bias, which can feel unnatural.
As for "Linearization," I introduced it in "Exploration of Linear Attention". Many people are researching this, including Performer, Nyströmformer, and recent works like cosFormer and Flowformer. Simply put, these works change standard Attention \phi(\boldsymbol{Q}\boldsymbol{K}^{\top})\boldsymbol{V} to (\phi_q(\boldsymbol{Q})\phi_k(\boldsymbol{K})^{\top})\boldsymbol{V}=\phi_q(\boldsymbol{Q})(\phi_k(\boldsymbol{K})^{\top}\boldsymbol{V}), achieving linear complexity. These methods are easy to implement but have two main issues: first, low-rankness leads to significantly worse performance (see "Road to Transformer Upgrade: 3"); second, when used as a Decoder (Causal), they sacrifice training parallelism because they must be converted to RNN-style calculation, or they don’t sacrifice parallelism but require bhns^2 space complexity. Compared to bhn^2 in standard Attention, this only has an advantage when n \gg s^2. Even with s=64, n must be \gg 4096, which is unrealistic in most cases.
Mixed Chunk
FLASH adopts a "Local-Global" mixed chunking approach, combining the advantages of "Sparsification" and "Linearization." First, for an input sequence of length n, we divide it into n/c non-overlapping chunks of length c (assuming c divides n; the paper uses c=256). Let \boldsymbol{U}_g, \boldsymbol{V}_g \in \mathbb{R}^{c \times e}, \boldsymbol{Z}_g \in \mathbb{R}^{c \times s} be the g-th chunk. As in equation [eq:relu-att], we obtain \boldsymbol{Q}_g^{\text{quad}}, \boldsymbol{K}_g^{\text{quad}}, \boldsymbol{Q}_g^{\text{lin}}, \boldsymbol{K}_g^{\text{lin}} via four simple affine transformations of \boldsymbol{Z}_g.
We use \boldsymbol{Q}_g^{\text{quad}}, \boldsymbol{K}_g^{\text{quad}} to calculate intra-chunk self-attention: \hat{\boldsymbol{V}}_g^{\text{quad}}=\frac{1}{cs}\text{relu}^2\left(\boldsymbol{Q}_g^{\text{quad}}{\boldsymbol{K}_g^{\text{quad}}}^{\top}\right)\boldsymbol{V}_g This represents tokens within each chunk interacting with each other, which is a form of "Sparsification." Its complexity is roughly \mathcal{O}(n/c \times c^2) = \mathcal{O}(nc), which is proportional to n. In implementation, this is equivalent to multi-head attention with n/c heads and sequence length c, which can be fully parallelized. To use it as a Decoder, simply mask the upper triangle of the attention matrix.
The remaining \boldsymbol{Q}_g^{\text{lin}}, \boldsymbol{K}_g^{\text{lin}} are used for global attention via the linear attention method: \hat{\boldsymbol{V}}_g^{\text{lin}}=\frac{1}{n}\boldsymbol{Q}_g^{\text{lin}}\sum_{h=1}^{n/c} {\boldsymbol{K}_h^{\text{lin}}}^{\top}\boldsymbol{V}_h Note that this operation is equivalent to performing linear attention with the full matrices \boldsymbol{Q}^{\text{lin}}, \boldsymbol{K}^{\text{lin}} \in \mathbb{R}^{n \times s} and \boldsymbol{V}. Writing it this way just highlights the connection to chunking. For a Decoder, to prevent future information leakage, it is changed to a cumsum form: \hat{\boldsymbol{V}}_g^{\text{lin}}=\frac{1}{(g-1)n/c}\boldsymbol{Q}_g^{\text{lin}}\sum_{h=1}^{g-1} {\boldsymbol{K}_h^{\text{lin}}}^{\top}\boldsymbol{V}_h In this case, to maintain parallelism, we only need b(n/c)se space complexity. Without chunking, linear attention would require bns^2 (or bhns^2 with multi-head). Under current settings, e/c \ll s, so it saves more memory.
Finally, the two attention results are combined and integrated into the GAU to obtain the linear version of GAU: \boldsymbol{O}_g=\left[\boldsymbol{U}_g\odot\left(\hat{\boldsymbol{V}}_g^{\text{quad}} + \hat{\boldsymbol{V}}_g^{\text{lin}}\right)\right]\boldsymbol{W}_o The Transformer model built on the linear version of GAU is the FLASH model.
Discussion
I believe the reason for this "Local-Global" mixed attention, besides reducing computational cost, is that it yields an attention distribution more consistent with reality. Based on our empirical understanding of NLP, associations in natural language are mainly concentrated locally. While global, extremely long-range associations exist, they are not dominant. Thus, this mixed design helps the model highlight local associations without discarding long-range ones. The original paper also conducted ablation experiments showing that local attention is relatively more important than global attention, and the mixed version performs best.
Furthermore, some readers might worry if non-overlapping chunking hinders the prediction of boundary words. The paper mentions that introducing more complex overlapping local attention does improve performance but adds extra computational cost. Given the same computational budget, the gain from overlapping local attention is less than simply adding more layers of non-overlapping GAU. Thus, non-overlapping chunks balance speed and performance well.
Finally, this "Mixed Chunk" linearization scheme is essentially universal. It can be used not only in GAU but also in standard Transformers (retaining the Attention+FFN combination but linearizing Attention via Mixed Chunk). The paper calls this "MC-TFM" and compares it, showing that GAU is also more advantageous in linearization.
Experimental Analysis
Regarding the experimental results of GAU and FLASH, I believe two are most noteworthy.
The first is the comparison between the new GAU and standard Multi-Head Self-Attention (MHSA), which is essentially a comparison between FLASH-Quad and standard Transformer:
Note that the horizontal axis is speed and the vertical axis is performance. Points closer to the top-right are more ideal. The figure shows that for models of any size, GAU is more advantageous than the corresponding multi-head attention model.
The second is the experimental table for the FLASH model:
The table directly shows:
1. Although both FLASH-Quad and Transformer have quadratic complexity, FLASH-Quad performs better and is faster.
2. When the sequence is long enough, the linear-complexity FLASH is faster than FLASH-Quad with similar performance.
To be honest, even the speed increase of FLASH-Quad (which is still quadratic) is something many "linear complexity" works fail to achieve. The power of GAU is evident. Also, the paper specifically points out that the Rotary Position Embedding (RoPE) I proposed significantly improves the performance of both Transformer and FLASH. Thus, the Transformer+, Transformer++, FLASH-Quad, and FLASH in the experiments all use RoPE. I’ll take a little pride in that.
Additionally, the table above does not compare memory usage. In my tests, at the base scale with a sequence length of 1024, the maximum usable batch size for FLASH-Quad is nearly double that of the Transformer, meaning FLASH-Quad significantly reduces memory consumption. I also tried a small version of FLASH-Quad for Chinese pre-training and found it performed even better than RoFormer (RoPE+Transformer). The results reported in the paper are indeed credible. However, due to limited GPU resources recently, I haven’t been able to conduct deeper tests. I will share new results if I have them.
Extended Thoughts
The introduction to GAU and FLASH is basically complete. As of this blog post, the authors have not released the full source code on GitHub, but the appendix includes the key source code (TensorFlow version) which can be adapted easily. Implementation should not be difficult.
Now for some "nitpicking" regarding parts of the paper I find less than perfect.
First, I think FLASH-Quad and FLASH are not decoupled well enough. As stated at the beginning, both are "heavyweight" results. For me, FLASH-Quad is even more valuable because the quadratic complexity of self-attention provides enough degrees of freedom for tricks like UniLM. FLASH-Quad should be a very independent and noteworthy model, but in the paper, it feels like a transitional product for FLASH. Fortunately, the authors separated the GAU concept, which mitigates this.
Second, GAU can replace both Attention and FFN. By design, it aims to replace Self-Attention. The authors don’t seem concerned with its ability to replace Cross-Attention, and there are no experiments for it. Can GAU replace Cross-Attention? Theoretically yes, but it’s unclear if it can still use only one head, which is the highlight of GAU. Furthermore, the paper only conducts language modeling experiments (LM and MLM) and lacks "pre-training + fine-tuning" experiments, so GAU’s transfer performance is uncertain.
Finally, there is one thing I don’t quite understand: GAU/FLASH-Quad/FLASH use additive absolute, additive relative, and RoPE position encodings simultaneously. Theoretically, one should suffice. My GAU experiments used only RoPE and worked well. Is there a specific reason for using all three? Also, from the appendix code, the authors didn’t handle padding carefully, and the normalization factor for the Decoder recursion was not written perfectly (the sum of the first t items should be divided by t, not n). These are minor details for improvement. Of course, the authors’ original code might be correct, and the appendix might just be simplified for readability.
Summary
This article introduced a new efficient Transformer work from Google, which merges Attention and FFN into a new GAU layer to create the FLASH-Quad variant. The authors further proposed a "Mixed Chunk" linearization scheme to create FLASH with linear complexity. Experimental results show that both FLASH-Quad and FLASH are faster, more efficient, and better than the standard Transformer. Perhaps soon, "All You Need" will no longer be Attention, but GAU.
Reprinting: Please include the original link: https://kexue.fm/archives/8934