SOFT: Softmax-free Transformer with Linear Complexity

Introduction

Recently the step change brought by Transformers in natural language processing (NLP) seems to have arrived in vision . Indeed, with less inductive bias in its architecture design than Convolution neural networks (CNNs), pure Vision Transformer (ViT) and its variants have shown to be able to outperform CNNs on various vision tasks . However, there is a bottleneck in any Transformer based model, namely its quadratic complexity in both computation and memory usage. This is intrinsic to the self-attention mechanism: given a sequence of tokens (e.g., words or image patches) as input, the self-attention module iteratively learns the feature representations by relating one token to all other tokens. This results in a quadratic complexity $O(n^{2})$ with the token sequence length $n$ in both computation (time) and memory (space) since an $n\times n$ sized attention matrix needs to be computed and saved during inference. This problem is particularly acute in vision: a 2D image after tokenization will produce a far longer sequence than those in NLP even with a moderate spatial resolution. This quadratic complexity thus prevents a ViT model from modeling images at high spatial resolutions, which are often crucial for visual recognition tasks.

A natural solution is to reduce the complexity of self-attention computation via approximation. Indeed, there have been a number of attempts in NLP . For example, takes a naive approach by shortening the length of Key and Value via learnable projections. Such a coarse approximation would inevitably cause performance degradation. In contrast, both leverage the kernel mechanism to approximate softmax normalization to linearize the computation in self-attention. instead adopts a hashing strategy to selectively compute the most similar pairs. Recently, uses Nyström matrix decomposition to reconstruct the full attention matrix with polynomial iteration for approximating the pseudo-inverse of the landmark matrix. Nonetheless, softmax normalization is simply duplicated across the matrix decomposition process, which is theoretically unsound. We empirically found that none of these methods are effective when applied to vision (see Sec. 4.2).

In this work, we identify that the limitations of existing efficient Transformers are caused by the use of softmax self-attention, and for the first time propose a softmax-free Transformer. More specifically, in all existing Transformers (with or without linearization), a softmax normalization is needed on top of scaled dot-product between token feature vectors . Keeping this softmax operation challenges any subsequent linearization efforts. To overcome this obstacle, we introduce a novel softmax-free self-attention mechanism, named as SOFT, with linear complexity $O(n)$ in both space and time. Specifically, SOFT uses Gaussian kernel to define the similarity (self-attention) function without the need for subsequent softmax normalization. With this softmax-free attention matrix, we further introduce a novel low-rank matrix decomposition algorithm for approximation. The robustness of the approximation is theoretically guaranteed by employing a Newton-Raphson method for reliably computing the Moore-Penrose inverse of the matrix.

We make the following contributions. (I) We introduce a novel softmax-free Transformer with linear space and time complexity. (II) Our attention matrix approximation is achieved through a novel matrix decomposition algorithm with theoretical guarantee. (III) To evaluate our method for visual recognition tasks, we design a family of generic backbone architectures with varying capacities using SOFT as the core self-attention component. Extensive experiments show that with a linear complexity (Figure 1(b)), our SOFT models can take in as input much longer image token sequences. As a result, with the same model size, our SOFT outperforms the state-of-the-art CNNs and ViT variants on ImageNet classification in the accuracy/complexity trade-off (Figure 1(a)).

Related work

There is a surge of research interests recently in exploiting Transformers for visual recognition tasks , inspired by their remarkable success in NLP . Core to these NLP and vision transformers is the same self-attention mechanism that computes a self-attention matrix by exhaustively comparing token pairs. This means a quadratic complexity with the sequence length in both space and time, which thus limits the scalability of Transformers in dealing with long sequences. This limitation is more serious in vision than NLP: To process an image with at least thousands of pixels, patch-wise tokenization is a must for Transformers to control the computational cost. Given higher resolution images, the patch size also needs to be enlarged proportionally sacrificing the spatial resolution. This limits the capability of Transformers, e.g., learning fine-grained feature representation as required in many visual recognition tasks.

Recently, there have been a number of linear/efficient variants of Transformers in NLP. For example, learns to shrink the length of Key and Value based on a low-rank assumption. adopts a hashing strategy to selective the most similar pairs and only compute attention among them. utilize different kernel functions for approximating softmax-based self-attention matrix. applies random feature mapping on the sequences to approach the original softmax function. decreases the time and memory consumption of the attention matrix by replacing the softmax function with its linear-complexity recurrent alternative. When applied to visual recognition tasks, however, we show that these models have considerable performance degradation compared to the standard Transformers (see Sec. 4.2).

The most related work to SOFT is which uses the Nyström matrix decomposition to avoid computing the full attention matrix. However, this method suffers from several theoretical defects: (1) As the standard self-attention needs to apply row-wise softmax normalization on the full attention matrix, a direct application of matrix decomposition is infeasible. As a workaround, softmax is simply applied to all the ingredient matrices in . Such an approximation is not guaranteed theoretically. (2) With a polynomial iteration method, it is not guaranteed that the generalized attention matrix inverse can be computed when the matrix is a nearly singular one in practice. In contrast to all the above methods, in this paper we propose a softmax-free self-attention mechanism that facilitates matrix decomposition for complexity minimization with theoretical guarantees.

Method

Formally, $X$ is first linearly projected into three $d_{e}$-dimensional spaces (query, key, and values) as:

Whilst this softmax self-attention has been the de facto choice and seldomly questioned, as discussed earlier it is not necessarily suited for linearization. To facilitate the design of linear self-attention, we introduce a softmax-free self-attention function with the dot-product replaced by a Gaussian kernel as:

To preserve the symmetric property of attention matrix as in Eq (3), we set the project matrices $W_{q}$ and $W_{k}$ in Eq (1) identical (i.e., $Q=K$). Our self-attention matrix is then written as:

For notation simplicity, we define the matrix formulation as: $S=\text{exp}\left(Q\ominus K\right)$.

Our self-attention matrix $S$ has three important properties: (1) It is symmetric; (2) All the elements lie in a unit range of $$; (3) All diagonal elements hold the largest value$1$(self-reinforced), with the bottom ones (corresponding to most dissimilar token pairs) being close to . As Gaussian kernel is a positive definite kernel ,$S$is deemed a Gram matrix. However, we find that when using our kernel-based self-attention matrix$S$ without linearization, the training of a transformer fails to converge. This might explain why softmax dot-product based self-attention is so popular in vanilla transformers.

2 Low-rank regularization via matrix decomposition with linear complexity

To solve the convergence and quadratic complexity problems, we leverage matrix decomposition as a unified solution with low-rank regularization. In particular, we consider Nyström , which is originally a low-rank matrix approximation algorithm. This enables our model’s complexity to be reduced significantly without computing the full self-attention matrix $S$.

We make this choice because our $S$ is positive semi-definite (i.e., a Gram matrix) without follow-up normalization which are all necessary conditions for Nyström. In contrast, totally ignores these requirements, leading to theoretical flaw in its approximation.

To define the Nyström method formally, let us express $S=\text{exp}\left(Q\ominus K\right)$ as a block matrix:

and $A^{\dagger}$ is the Moore-Penrose (a generalized) inverse of $A$.

As $K$ is identical to $Q$, we have $\widetilde{K}=\widetilde{Q}$. Given these $m$ tokens, we then compute $A$ and $P$ as:

We finally obtain the regularized self-attention matrix $\hat{S}$ of SOFT as:

leading to Algorithm 1. The low-rank regularization is conducted as follows. For computing the attention score between any two tokens, we first correlate each of them with sampled tokens using our self-attention function (Eq (5)); With this correlation representation we then compute their similarity under the modulation of the generalized inverse of $\widetilde{Q}$’s correlation matrix. Similar as standard Nyström, our design associates the input tokens w.r.t. a small space spanned by sampled tokens, giving a proper estimation of the original attention relationships subject to a low-rank constraint. The correctness of this method is proved in Appendix A.1.

We now prove that $A_{k}$ finally converges to Moore-Penrose inverse of $A_{m\times m}$, if $\alpha$ is sufficiently small .

When $\alpha$ is sufficiently small, $A_{k+1}=2A_{k}-A_{k}AA_{k}$, $A_{k}$ converges to $A^{\dagger}$.

Though $\alpha=2/\|A\|_{1}^{2}$ which ensures good convergence behavior in Algorithm 2 (see more details in Appendix A.2.1), in practice, we find that using an alternative form gives more stable training and faster convergence. Specifically, in $\|I-A\frac{2\beta^{n}}{\|A\|_{1}^{2}}\|_{1}\leq 1$ where $\beta$ equals to $0.5$, we find the smallest $n_{i}$ that holds this inequality. Then, we initialize $\alpha$ as $\alpha=\frac{2\beta^{n_{i}}}{\|A\|_{1}^{2}}$.

The following proposition comes with the proof of Theorem 1:

$\|AA_{k}A-A\|$ and $\|A_{k}-A^{\dagger}\|$ decreases to monotonously, if $\alpha$ is sufficiently small.

The detail of proposition 1 is shown in Appendix A.2.2. This ensures that our estimated inverse is sufficiently accurate for matrix decomposition, subject to that our SOFT attention is regularized.

We summarize the complexity of SOFT in space and time. For time complexity, it involves: (1) Sampling: $\mathcal{O}(nd_{e})$. (2) Calculating three decomposed matrices: $\mathcal{O}(nmd_{e}+mnd_{e}+m^{2}d_{e})=\mathcal{O}(2mnd_{e}+m^{2}d_{e})$; (3) Moore-Penrose inverse: $\mathcal{O}(\mathcal{T}\times m^{3})=\mathcal{O}(\mathcal{T}m^{3})$, where $\mathcal{T}$ is the iteration steps. (4) All matrix multiplication: $\mathcal{O}(nm^{2}+mnd_{e}+mnd_{e})=\mathcal{O}(nm^{2}+2mnd_{e})$. The total time complexity is $\mathcal{O}((d_{e}+4md_{e}+m^{2})n+\mathcal{T}m^{3}+d_{e}m^{2})$. The space complexity is decided by four decomposed matrices with $\mathcal{O}(n\times m)+\mathcal{O}(m\times m)+\mathcal{O}(m\times n)+\mathcal{O}(n\times d_{e})=\mathcal{O}((2m+d_{e})n+m^{2})$. As we keep $m$ ($m\ll n$) a fixed constant in our model, both time and space complexity are $\mathcal{O}(n)$, making SOFT a linear self-attention.

3 Instantiations

Figure 2 shows how our proposed softmax-free self-attention block (SOFT block) can be implemented in a neural network. We replace the self-attention block with our SOFT block in the traditional Transformer, that is, we stack a SOFT block with a feed forward residual block to form a softmax-free Transformer layer (SOFT layer).

Focusing on the general image recognition tasks, we integrate our SOFT layer into the recent pyramidal Transformer architecture to form our final model SOFT. Further, several improvements are introduced in patch embedding (i.e., tokenization). Specifically, unlike that uses a combination of non-overlapping convolution and layer normalization , we adopt a stack of overlapping convolutions, batch normalization and ReLU non-linearity. Concretely, the $\tt STEM$ is implemented by 3 units of $\tt 3x3\;Conv{\to}BN{\to}ReLU$, with the stride of 2, 1, 2 respectively. Then, one such unit is applied to each of three following down-sampling operations with stride of 2 in the multi-stage architecture.

The architecture hyper-parameters of SOFT are: $d$: the input channel dimension of SOFT layer. $d_{e}$: the embedding dimension of tokens in SOFT block. In practice, we set $d_{e}=d$. $h$: the head number of SOFT block. $d_{h}$: the channel dimension of each head and $d_{h}=d_{e}/h$. $n$: the input token sequence length of a SOFT block. $m$: the bottleneck token sequence length of SOFT block. $sp$: the sampling ratio of token sequence length sampling, which is the ratio between input token sequence length and the bottleneck token sequence length. $e$: the expansion ratio of the 2-layer feed forward block. In SOFT, for all the stages we set $d_{h}=32$, $e=4$ and $m=49$, $sp$ varies in each stage according to the input token sequence length. Table 2 details the family of our SOFT configurations with varying capacities (depth and width).

Experiments

Dataset: We evaluate the proposed SOFT on the ILSVRC-2012 ImageNet-1K dataset with 1.28M training images and 50K validation images from 1,000 classes. Following the common practice, we train a model on the training set and evaluate on the validation set. Metrics: For model performance, the top-1 accuracy on a single crop is reported. To assess the cost-effectiveness, we also report the model size and floating point operations (i.e., FLOPs). Implementation details: We use the code base with the default setting to train and test all the models. Specifically, we use weight decay of 0.05 and 10 epochs of linear warm-up. We conduct 300 epochs training with an $\tt AdamW$ optimizer and decreasing learning rate with the cosine annealing schedule. During training, random flipping, mixup and cutmix are adopted for data augmentation. Label smoothing is used for loss calculation. All our variants are trained with a batch size of 1024 on 32G NVIDIA V100 GPUs. We also implement our method using the Mindspore .

2 Comparison with existing linear Transformers

We compare our method with three existing linear Transformer models: Linformer , Performer , Nyströmformer in terms of model complexity and accuracy.

Two experimental settings are adopted. Under the first setting, for all methods we use the same $\tt Tiny$ (Table 2) architecture for a fair comparison. That is, we replace the core self-attention block in SOFT with each baseline’s own attention block with the rest of the architecture unchanged. Note that the spatial reduction module of is a special case of Linformer . We set the reduction ratio to be identical to ours. With the same uniform sampling idea, we replace the 1D window averaging of Nyströmformer (for NLP tasks) with 2D average pooling (for images). The downsampling ratio remains identical to ours. It is also worth mentioning that there is no official code released for Reformer and the local Sensitive Hash (LSH) module has strict requirements on the length of input tokens. We thus do not include this method in our comparison.

From Table 1 we can make the following observations: (i) Linear Transformer methods substantially reduce the memory and FLOPs while maintain similar parameter size comparing to the Transformer on the $\tt Tiny$ architecture; (ii) Our approach SOFT achieves the best classification accuracy among all the linearization methods. (iii) Our inference speed is on-par with other compared linear Transformers and our training speed is slightly slower than Nystromformer and both are slower than Performer and Linformer. Note that the slow training speed of our model is mostly due to the Newton-Raphson iteration which can only be applied sequentially for ensuring the accuracy of Moore-Penrose inverse. In summary, due to the on-par inference speed we consider the training cost increase is a price worth paying for our superior accuracy.

Under the second setting, we focus on the memory efficiency of SOFT against the baselines. Here we follow the ViT network structure, stacking 12 attention layers with hidden dimension $d=384$, heads $h=12$, bottleneck token sequence length $m=49$. Different attention blocks from the three linearized Transformer variants, Linformer , Performer , and Nyströmformer are studied. For each Transformer variant, we adjust its token sequence length $n$ in a linear increment. Specifically, we use a token sequence length of $784\times p$ where $p=1,2,3,4,5,6,7,8$ and set batch size 1 to verify whether the memory consumption increases “quadratically” or “linearly”. Figure 1(b) shows all compared transformer variants including our SOFT indeed have a linear memory usage complexity. This is in contrast with the standard Transformer which cannot cope with long token sequences with a quadratic complexity.

3 Comparison with state-of-the-art CNNs and ViTs

We compare with state-of-the-art alternatives and report the top-1 accuracy on the ImageNet-1K validation set. FLOPs are calculated at batch size 1. From Figure 1(a) and Table 3, the following observations are made: (i) Overall, ViT and its variants yield better classification accuracy over CNNs. (ii) We achieve the best performance among the recent pure vision Transformer based methods including ViT and DeiT , as well as the state-of-the-art CNN RegNet . (iii) Our SOFT outperforms the most similar (in architecture configuration) Transformer counterparts PVT at all variants. Since the attention module is the main difference, this validates directly the effectiveness of our model. (iv) We can also beat the latest ViT variants Twins which is designed to address the efficiency limitation of ViT. We have done so with less parameters and fewer float point computation.

To gain some insights into how attention is learned using our SOFT and the alternatives, Figure 3 shows the attention masks of various compared models. For each model, we show the output from the first two attention heads. It is evident that SOFT exhibits robustness and versatility in capturing local and long distance relations among pixels. It is interesting to note that, although SOFT is trained on an object categorization dataset in ImageNet , it seems to be able to learn both semantic concepts shared across instances in the same category and instance specific features. For instance, in the bottom-right example of a bird class, one attention head focuses on the black bird only, while the other attend to both birds in the image. More examples are shown in Appendix A.4.

4 Ablation studies

Unlike the earlier non-pyramidal vision Transformers (e.g., ViT ), most recent pyramidal (multi-scale) Transformers (e.g., PVT ) use convolution layers to reduce the spatial resolution (i.e., token sequence length) between stages. In this study, we ablate SOFT with a pyramidal architecture (our default SOFT-$\tt Small$), SOFT w/o a pyramidal architecture and DeiT-S (no pyramidal architecture either). We replace the Transformer layer with a SOFT layer to get SOFT w/o a pyramidal architecture. Note all three variants have similar parameters and FLOPs. Table 5a shows that the conv-based pyramidal architecture is clearly superior to a non-pyramidal design, and our non-pyramidal counterpart is even slightly better than DeiT-S whilst enjoying linear complexity.

In this study, we examine how the bottleneck token sequence length $m$, sampled from $n$ tokens, influences the model’s performance. We change the bottleneck token sequence length in all stages to $36,49,64,81$. Table 4a shows that longer bottleneck token would increase the memory cost and the computational overhead. $m=49$ seems to give the best trade-off between the performance and computational overhead. The memory usage is measured with the batch size of 128.

We ablate SOFT with overlapped convolution (our default choice, same as many recent works) and SOFT with non-overlapped convolution in our $\tt Tiny$ configuration. Table 5b shows that SOFT with overlapped convolution performs better than SOFT without overlapped convolution. Our non-overlapped convolution variant still outperforms the PVT which also has the same non-overlapped convolution by a clear margin.

Newton-Raphson’s convergence: We study how many iterations the Newton-Raphson method needs to converge when computing the Moore-Penrose inverse $A^{\dagger}$. We use $\|AA_{k}A-A\|_{p}/\|A\|_{p}$ with $p=2$ (see Proposition 1) as the convergence metric to quantify the difference between $A_{k}$ and $A^{\dagger}$. Figure 4 shows that our approximation converges within 20 iterations across all stages.

5 Additional experiments on NLP tasks

In this section, we evaluate our method against other linear counterparts on four tasks of the Long Range Arena (LRA) benchmark e.g., Listops , byte-level IMDb reviews text classification , byte-level document retrieval , and image classification on sequences of pixels .

Implementations. We use the Pytorch version of LRA benchmark, provided by . For the evaluation protocol, we strictly follow . We omit the Pathfinder(1K) task as we cannot replicate the result of Nyströmformer . For our SOFT, we simply use the average pooling with window size 4, stride 4 to sample the bottlenecks. We follow the configurations of , with 2 layers, 64 and 128 hidden dimension respectively, and 2 attention heads. The results in Table 6 shows that our SOFT outperforms both the standard and alternative efficient Transformers on three out of four tasks, as well as the average result.

Conclusions

We have introduced a novel softmax-free self-attention (SOFT) mechanism for linearizing Transformer’s complexity in space and time. Unlike existing linear Transformers that aim to approximate the conventional softmax based self-attention, SOFT employs a Gaussian kernel based attention which eliminates the need for softmax normalization. This design enables a full self-attention matrix to be approximated via a low-rank matrix decomposition. The robustness of the approximation is achieved by calculating its Moore-Penrose inverse using a Newton-Raphson method. Extensive experiments show that SOFT yields superior trade-off in accuracy and complexity.

This work was funded in part by Shanghai Municipal Science and Technology Major Projects (No.2018SHZDZX01 and No.2021SHZDZX0103), Mindspore, National Science Foundation of China under Grant No.11690013, 71991471 and the scientific-technological innovation plan program of Universities guided by the Ministry of Education, China.

References

Appendix A Appendix

Nyström method aims to calculate a low-rank approximation for a Gram matrix. For Transformers, the self-attention matrix can be viewed as a Gram matrix $S$ with a Gaussian kernel $k$ applied to the query $Q$, with each element $S_{ij}$ expressed as:

$k(x,y)$ means operating Gaussian kernel $k$ to $(x,y)$, which can be written in the feature space as:

$n$ is the dimension of a feature space, $\lambda_{i}$ denotes the eigenvalue and $\phi_{i}$ denotes the eigenfunction of kernel $k$. According to the eigenfunction’s definition, we can get:

where $p(x)$ is the probability distribution of $x$. And {$\phi_{i}(x)$} are $p$-orthogonal:

$\delta_{ij}$ is when $i\neq j$, $1$ when $i=j$. To get an approximation of the eigenfunctions, we sample $\{x_{1},x_{2},\cdots,x_{q}\}$ from $p(x)$, then:

This inspires us to approximate the Gram matrix $S$. Let $S^{(m)}$ be a submatrix of $S$, consisting of $m\times m$ elements from $S$. Gram matrix is a symmetric positive semi-definite matrix, so it has a spectral decomposition:

where $U^{(m)}$ is column orthogonal and $\Lambda^{(m)}$ is a diagonal matrix with the diagonal elements as the eigenvalues of $S^{(m)}$. Substituting the $y$ to $x_{j}$ and applying the approximation above to $S$, we can get:

A.2 Newton method

$A$ is a symmetric positive semi-definite matrix and $A_{ij}\leq 1$, $\forall 1\leq i,j\leq n$, $A_{ii}=1,\quad 1\leq i\leq n$ in our case. $A_{0}$ is chosen to be $\alpha A$, so the $A_{k}$ can be written as $A_{k}=C_{k}A=AD_{k}$ for some matrix $C_{k},D_{k}$, leading to the fact that

This is because $A_{k+1}=A_{k}(2I_{n}-AA_{k})=(2I_{n}-A_{k}A)A_{k}$ and $A_{0}=\alpha A$. We make a difference between $A^{\dagger}$ and $A_{k+1}$:

We norm both sides of the equation above:

And we left multiply $A$ on the both sides of (26), then norm the equation:

We choose $\alpha$ sufficiently small so that the initial value satisfy $\|AA^{\dagger}-AA_{0}\|<1$. We set $\alpha=\frac{2}{\|A\|_{1}^{2}}$ to ensure it is small enough . Then the $\|AA^{\dagger}-AA_{k}\|\rightarrow 0$, when $k\rightarrow\infty$. The inequality (27) implies that $A_{k}\rightarrow A^{\dagger},\ k\rightarrow\infty$.

A.2.2 Proof of proposition 1

Note that when we multiply $A$ on both sides of (26), the equation turns to be:

Similarly norm both sides of (29), considering that $\|AA^{\dagger}-AA_{k}\|\rightarrow 0$ and $\|AA^{\dagger}-AA_{k}\|<1$ always holds, $\|A-AA_{k}A\|$ monotonically decreases to . The inequality (27) implies that $\|A_{k}-A^{\dagger}\|$ decreases to monotonously .

Note that although $\|A-AA_{k}A\|$ monotonically decreases to , $\|A_{k}AA_{k}-A_{k}\|$ cannot be proved to monotonically decrease to 0.

A.3 Non-linearized gaussian kernel attention

In our formulation, instead of directly calculating the Gaussian kernel weights, they are approximated. More specifically, the relation between any two tokens is reconstructed via sampled bottleneck tokens. As the number $m$ (e.g., 49), of bottleneck tokens is much smaller than the token sequence length, our attention matrix is of low-rank. This has two favorable consequences: (I) The model now focuses the attentive learning on latent salient information captured by the $m$ bottleneck tokens. (II) The model becomes more robust against the underlying token noise due to the auto-encoder style reconstruction .

This explains why the model with an approximated gram matrix performs better than the one with a directly estimated matrix. Further, exact Gaussian kernel attention computation leads to training difficulties. We first hypothesized that this might be due to lacking normalization (as normalization often helps with training stability and convergence), and tested a variant with softmax on top of an exact Gaussian kernel attention matrix. However, it turns out to suffer from a similar failure. We cannot find a solid hypothesis so far and will keep investigate this problem.

A.4 Attention visualization

Figure 5 shows more visualization of the attention masks by various Transformers and our SOFT. For each model, we show the output from the first two attention heads (up and down row). It is noteworthy that SOFT exhibits better semantic diversity of the multi-head mechanism than other methods. Moreover, when we sample the patch at the boundary of multiple objects, SOFT is able to more precisely capture all these objects.