It is well known that although Transformer models based on attention mechanisms exhibit good parallel performance, their space and time complexity are both $\mathcal{O}(n^2)$, where $n$ is the sequence length. Consequently, when $n$ is large, the computational requirements of Transformer models become prohibitive. Recently, numerous efforts have been devoted to reducing Transformer computational costs, including model pruning, quantization, distillation, and other compression techniques, as well as modifying the attention structure to reduce complexity to $\mathcal{O}(n\log n)$ or even $\mathcal{O}(n)$.
Recently, I read the paper "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention", which introduced me to the exploration of linear attention. After reviewing several related papers, I gained valuable insights and have compiled my understanding of linear attention in this article.
Attention Mechanisms#
Currently, the most popular attention mechanism is Scaled-Dot Attention, formulated as:
Here, $\boldsymbol{Q}\in\mathbb{R}^{n\times d_k}$, $\boldsymbol{K}\in\mathbb{R}^{m\times d_k}$, and $\boldsymbol{V}\in\mathbb{R}^{m\times d_v}$. For simplicity, we omit the explicit scaling factor of attention. Since this article focuses on self-attention scenarios, we uniformly set $\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V}\in\mathbb{R}^{n\times d}$ for convenience in presentation. In general scenarios, we have $n > d$ or even $n\gg d$ (e.g., in BERT base, $d=64$). For related interpretations, readers can refer to my articles "A Brief Reading of 'Attention is All You Need' (Introduction + Code)", as well as some of its improvements in "Breaking Through the Bottleneck: Building a More Powerful Transformer" and "Google's New Work Synthesizer: We Still Don't Fully Understand Self-Attention", which will not be discussed in depth here.
Removing Softmax#
In fact, a simple derivation reveals this conclusion. The step $\boldsymbol{Q}\boldsymbol{K}^{\top}$ yields an $n\times n$ matrix, which determines that the complexity of attention is $\mathcal{O}(n^2)$. Without Softmax, we would have three matrix multiplications: $\boldsymbol{Q}\boldsymbol{K}^{\top}\boldsymbol{V}$. Since matrix multiplication satisfies associativity, we can first compute $\boldsymbol{K}^{\top}\boldsymbol{V}$, obtaining a $d\times d$ matrix, and then left-multiply it by $\boldsymbol{Q}$. Since $d \ll n$, the approximate complexity of this approach is only $\mathcal{O}(n)$ (with the left multiplication by $\boldsymbol{Q}$ dominating the computation).
In other words, removing Softmax from attention can reduce the complexity to the ideal linear level $\mathcal{O}(n)$! This is exactly our ultimate goal: Linear Attention, attention with linear complexity. Therefore, the theme of this article is to explore linear attention after removing Softmax.
Generalized Definition#
The question is: can direct removal of Softmax still be considered attention? Can it achieve the effects of standard attention? To answer this, we first equivalently rewrite the definition of Scaled-Dot Attention from equation (1) as (vectors in this article are column vectors):
Thus, Scaled-Dot Attention is essentially a weighted average of $\boldsymbol{v}_j$ with weights $e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$. Therefore, we can propose a generalized definition of attention:
That is, we replace $e^{\boldsymbol{q}_i^{\top}\boldsymbol{k}_j}$ with a general function $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)$ of $\boldsymbol{q}_i, \boldsymbol{k}_j$. To preserve the similar distribution characteristics of attention, we require $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$ to hold consistently. In other words, if we want to define a new form of attention, we should retain the form of equation (3) and satisfy $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)\geq 0$.
This general form of attention is also known as Non-Local Networks in computer vision, originating from the paper "Non-local Neural Networks".
Examples#
If we directly remove Softmax, then $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j) = \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$. However, the problem is that inner products cannot guarantee non-negativity, so this is not a reasonable choice. Below, we briefly introduce several feasible schemes.
Kernel Function Form#
A natural idea is: if each element of $\boldsymbol{q}_i,\boldsymbol{k}_j$ is non-negative, then the inner product will naturally be non-negative. To achieve this, we can apply activation functions $\phi,\varphi$ to $\boldsymbol{q}_i,\boldsymbol{k}_j$ respectively:
where $\phi(\cdot),\varphi(\cdot)$ are activation functions with non-negative ranges. The paper "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention" mentioned at the beginning of this article chooses $\phi(x)=\varphi(x)=\text{elu}(x)+1$.
If we must tell a story, equation (4) can be associated with "kernel methods", especially when $\phi=\varphi$, where $\phi$ acts as a kernel function, and $\langle \phi(\boldsymbol{q}_i), \phi(\boldsymbol{k}_j)\rangle$ is the inner product defined by the kernel function. Relevant considerations can be found in the paper "Transformer dissection: An unified understanding for transformer's attention via the lens of kernel", which will not be elaborated further here.
Ingenious Use of Softmax#
An earlier paper, "Efficient Attention: Attention with Linear Complexities", presents a more interesting alternative. 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, it proposes:
where $\text{softmax}_1$ and $\text{softmax}_2$ denote Softmax operations along the first ($n$) and second ($d$) dimensions, respectively. In other words, we apply Softmax to $\boldsymbol{Q},\boldsymbol{K}$ individually, rather than after computing $\boldsymbol{Q}\boldsymbol{K}^{\top}$.
If we directly take $\phi(\boldsymbol{q}_i)=\text{softmax}(\boldsymbol{q}_i),\varphi(\boldsymbol{k}_j)=\text{softmax}(\boldsymbol{k}_j)$, then clearly this form is also a special case of equation (4). Moreover, this design has appeared multiple times in computer vision, such as in A2-Nets, which also includes similar approaches.
Personal Conception#
Here, I present my own conception. The starting point of this idea is not equation (4), but rather an approximation of our original definition from equation (2). From Taylor expansion, we have:
If $\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$, then the right-hand side is guaranteed to be non-negative, allowing us to set $\text{sim}(\boldsymbol{q}_i, \boldsymbol{k}_j)=1 + \boldsymbol{q}_i^{\top}\boldsymbol{k}_j$. By now, readers might have realized that to ensure $\boldsymbol{q}_i^{\top}\boldsymbol{k}_j\geq -1$, we only need to normalize $\boldsymbol{q}_i,\boldsymbol{k}_j$ separately using $l_2$ normalization. Therefore, the final scheme I propose is:
Although this differs from the form in equation (4), it is theoretically closer to the original Scaled-Dot Attention.
Related Work#
There are many works focused on reducing the computational complexity of attention by modifying its form. Below, we briefly enumerate some of them.
Sparse Attention#
We previously introduced OpenAI's Sparse Attention, which reduces attention computation by "retaining only values in small regions and forcing most attention values to zero." After special design, most elements of the attention matrix become zero, thus theoretically saving 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 has two obvious shortcomings:
1. How to select the attention regions to retain is determined subjectively by humans, lacking intelligence.
2. It requires specific design optimizations in programming to achieve efficient implementation, making it difficult to generalize.
Reformer#
Reformer is another representative improvement work that reduces attention complexity to $\mathcal{O}(n\log n)$. In a sense, Reformer is also a form of sparse attention, but its sparse pattern is not predetermined. Instead, it uses LSH (Locality Sensitive Hashing) technology to (approximately) quickly find the largest several attention values and compute only those values. Additionally, Reformer replaces the original FFN (Feedforward Network) with a reversible form and redesigns the backpropagation process to reduce memory usage.
Thus, compared to the aforementioned sparse attention, Reformer addresses its first shortcoming but still suffers from the second: implementation complexity. Implementing LSH-based attention is much more complex than standard attention, and rewriting the backpropagation process for reversible networks is beyond the reach of most readers.
Linformer#
A work very similar to the linear attention discussed in this article is Facebook's recently released Linformer, which retains the original Scaled-Dot Attention form but projects $\boldsymbol{K},\boldsymbol{V}$ using two $m\times n$ matrices $\boldsymbol{E},\boldsymbol{F}$ before applying attention:
Thus, $\boldsymbol{Q}(\boldsymbol{E}\boldsymbol{K})^{\top}$ becomes only an $n\times m$ matrix. The authors claim that even for large sequence lengths $n$, $m$ can remain a moderate constant, making this attention linear as well. Similar ideas to Linformer appear in earlier CV papers such as "Asymmetric Non-local Neural Networks for Semantic Segmentation".
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, making it difficult to simply mask future information (Causal Masking). Consequently, it cannot perform language modeling, Seq2Seq, and other autoregressive generation tasks, which explains why the original authors only tested MLM tasks. In contrast, the linear attention methods introduced in this article can all achieve this. Taking equations (3) and (4) 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$:
There are two ways to implement this equation. The first method involves setting $\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)$:
This indicates that such attention can be implemented recursively as an RNN model, with minimal space complexity but requiring serial computation, making it suitable for predictive decoding. The second method involves directly computing the outer product of $\varphi(\boldsymbol{K}),\boldsymbol{V}\in\mathbb{R}^{n\times d}$, yielding an $n\times d\times d$ matrix, and then applying $\text{cumsum}$ along the $n$ dimension, thus obtaining $\boldsymbol{S}_1,\boldsymbol{S}_2,\dots,\boldsymbol{S}_n$ in one go. This approach is fastest but consumes the most memory, making it suitable for training. However, often $d^2\gg n$, and in many cases, the memory requirements during training are prohibitive. Therefore, the RNN form is often preferred.
Downsampling Techniques#
From the results, Linformer's $\boldsymbol{E}\boldsymbol{K}, \boldsymbol{F}\boldsymbol{V}$ essentially shorten (downsample) the sequence. The most straightforward method to shorten a sequence is Pooling, so I previously attempted to incorporate Pooling techniques into Transformers. Recent similar works include 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 using 1D convolutions with stride > 1. Based on this idea, perhaps we could replace the Position-Wise fully connected layers in FFN with 1D convolutions with stride > 1? In any case, many variations can be explored in this direction, but like Linformer, such blending makes autoregressive generation challenging.
Conclusion#
This article introduces several works that modify the structure of attention to reduce its computational complexity. The main idea is to remove the Softmax from standard attention, allowing attention complexity to degrade to the ideal $\mathcal{O}(n)$ level (Linear Attention). Compared to other similar improvement structures, this modification not only reduces complexity to $\mathcal{O}(n)$ but also retains all "token-token" attention relationships and preserves the possibility of autoregressive generation.
The key contributions of this exploration are:
- Complexity Reduction: Achieving $\mathcal{O}(n)$ complexity compared to the original $\mathcal{O}(n^2)$ of standard attention.
- Generalized Formulation: Proposing a generalized attention definition that allows for various similarity functions beyond exponential dot products.
- Practical Implementations: Presenting multiple practical approaches including kernel methods, dimension-wise normalization, and normalized inner products.
- Comparative Analysis: Contrasting linear attention with other efficiency-focused approaches like sparse attention, Reformer, and Linformer.
- Autoregressive Compatibility: Maintaining the ability to perform causal masking for autoregressive generation tasks.
These findings suggest that the Softmax operation, while integral to standard attention formulations, may not be strictly necessary for achieving effective attention mechanisms. The exploration of linear attention opens new avenues for efficient Transformer architectures capable of handling longer sequences while preserving modeling capabilities.
Future research directions include empirical validation of these linear attention variants on diverse NLP tasks, investigation of hybrid approaches combining different efficiency techniques, and exploration of theoretical guarantees for approximation quality compared to standard attention.
Original Article: Su Jianlin. Exploring Linear Attention: Must Attention Have a Softmax? Scientific Spaces.
How to cite this translation:
BibTeX: