QT-ViT: Improving Linear Attention in ViT with Quadratic Taylor Expansion

Authors: yixing  xu
chao  li
dong  li
xiao  sheng
fan  jiang
lu  tian
emad  barsoum
Venue: NeurIPS_cc_2024_Conference
Published: 2024-04-23 10:00:46

Abstract: Vision transformer model (ViT) is widely used and performs well in vision tasks due to its ability to capture long-range dependencies. However, the time complexity and memory consumption increase quadratically with the number of input patches which limits the usage of ViT in real-world applications. Previous methods have employed linear attention to mitigate the complexity of the original self-attention mechanism at the expense of effectiveness. In this paper, we propose QT-ViT models that improve the previous linear self-attention using quadratic Taylor expansion. Specifically, we substitute the softmax-based attention with second-order Taylor expansion, and then accelerate the quadratic expansion by reducing the time complexity with a fast approximation algorithm. The proposed method capitalizes on the property of quadratic expansion to achieve superior performance while employing linear approximation for fast inference. Compared to previous studies of linear attention, our approach does not necessitate knowledge distillation or high-order attention residuals to facilitate the training process. Extensive experiments demonstrate the efficiency and effectiveness of the proposed QT-ViTs, showcasing the state-of-the-art results. Particularly, the proposed QT-ViTs consistently surpass the previous SOTA EfficientViTs under different model sizes, and achieve a new Pareto-front in terms of accuracy and speed.

Main paper: 
 introduction: 
 Compared to convolutional neural networks (CNNs), vision transformers (ViTs) are getting more and more attention due to their strong performance across various computer vision tasks, such as image classification [33, 16, 8, 37, 38, 21] , object detection [4, 13] , semantic segmentation [36, 15] and low-level vision [20, 32, 31, 30] . The effectiveness of ViT comes from the multi-head selfattention (MHSA) operation that allows the model to capture long-range information by calculating the attention score between each pair of patches. However, this mechanism necessitates quadratic time and storage complexity O(n 2 ) related to the number of input patches n, and the original ViTs require significant computational and storage resources when applied to real-world applications.To overcome the aforementioned problem, previous researches focus on improving the original selfattention mechanism by using local attention such as window attention [23] , dilated attention [10] and random attention [27] . Another family of methods is to utilize linear attention [11, 2, 5, 3] that decomposes the original softmax function into two non-linear kernels so that the order of matrix multiplications in attention score calculation is changed to reduce the quadratic computational complexity into a linear one. Many papers focus on designing non-linear kernels and novel linearattention architectures for better approximation, e.g., Hydra-attention [2] uses the hydra trick to their multi-head attention by setting as many heads as features. Performer [5] uses fast attention via a positive orthogonal random features approach to approximate the softmax attention. EfficientViT [3] replaces softmax with ReLU non-linear activation and applies depthwise and group convolution to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). improve its performance. Flatten transformer [14] utilizes focused attention based on ReLU to force the attention operation to focus on more informative regions.Previous linear attention methods reduce the complexity of the attention mechanism from O(n 2 d) to O(nd 2 ) at the expense of the performance on visual tasks, where d is the patch dimensionality. Some of them necessitate the knowledge distillation method [7] or high-order attention residuals [39] to make up for the performance gap. However, the GPU memory consumption will severely increase which makes these methods unsuitable for training large transformer models.In this paper, we explore the utilization of second-order (quadratic) Taylor expansion to approximate the original softmax attention. We theoretically show that this approximation can be decomposed into two non-linear kernels through the utilization of the Kronecker product [34] . By employing this approach, the computational complexity can be changed from O(n 2 d) to O(nd 3 ). We then propose a fast approximation algorithm to accelerate the computation of the Kronecker product, thereby reducing the complexity to O(nd 2 ). In contrast to the first-order (linear) Taylor expansion [7] and other linear attention methods, we can utilize the high-order information within the softmax function to achieve superior performance while at the same time preserving the efficiency of linear attention. Experimental results on the ImageNet dataset show that the proposed QT-ViTs can achieve a superior accuracy-speed trade-off when compared to other state-of-the-art methods, as shown in Fig. 1 . Additionally, we conduct experiments on object detection and semantic segmentation tasks to further validate the effectiveness of our approach. 
preliminaries: 
 In this section, we first introduce the preliminaries of softmax attention and linear attention. Then, we provide an overview of various instantiations of the original linear attention method used in ViT and analyze their advantages and disadvantages.Softmax self-attention operation is the key component in the transformer model. Given an input matrix X ∈ R N ×d where N is the number of patches and d is the dimension of each patch, we first map the input matrix to the query, key and value embeddings by using the matrix multiplications:EQUATIONwhere W Q , W K and W V ∈ R d×d are learnable matrices. Then, the attention score is computed on each pair of patches to capture the global information as shown below:EQUATIONwhereSim(Q k , K i ) = exp(Q k K ⊤ i / √ d) is the similarity measurement function in the softmax attention, Q k , K i (K j ), V i , O kare the corresponding k-th, i-th (j-th), i-th, k-th row vectors of the query, key, value and output matrices, respectively. The inner product of the query-key pair is first computed to calculate the similarity between the pair, then a scale is applied for stability and a softmax function is used to transfer the similarity into probability. This probability is applied to the value matrix to get the final attention score output. The softmax attention computes the inner products of all the query-key pairs and results in a O(N 2 d) time complexity.The overhead of the computation of Eq. 2 mainly comes from the matrix multiplication. By decomposing the similarity function into two separate kernel embeddings, i.e.,Sim(Q k , K i ) = ϕ(Q k )ϕ(K i ) ⊤ ,and the original softmax attention function can be changed into linear attention by exchanging the order of matrix multiplication:EQUATIONwhere the complexity is changed from O(N 2 d) to O(N d 2 ). Since the patch dimension d is always smaller than the number of patches N in the popular ViT architectures, the computation overhead can thus be reduced.However, in order to losslessly decompose the similarity function in the softmax attention Sim(Q k , K i ) into the product of two kernel embeddings ϕ(Q k ) and ϕ(K i ), the dimensionality of the kernel function needs to be infinite which is unable to apply to real-world applications. Thus, a series of instantiations are proposed trying to compute ϕ(•) efficiently while preserving as much information of the original similarity function as possible.In the following, we use q ≜ Q k and k ≜ K i to represent row vectors in query matrix Q and key matrix K that do not belong to any specific row.Linear transformer was first proposed in [19] and ϕ(x) = elu(x) + 1 was used as the kernel function. EfficientViT [3] used ϕ(x) = ReLU(x) as the instantiation. Both methods ensure that ϕ(q)ϕ(k) ⊤ ≥ 0 which is consistent with the property of the similarity measurement function Sim(•). Flatten Transformer [14] argued that previous approximations smooth the distribution of linear attention which failed to focus on more informative regions, and proposed a focused function ϕ p (x) = ||ReLU(x)|| ||ReLU(x) * * p || ReLU(x) * * p where || • || represents the Euclidean norm and (•) * * p is elementwise power p of the input. Hydra attention [2] used cosine similarity as the kernel ϕ(x) = x/||x|| 2 , PolyNL [1] used mean kernel ϕ(x) = x/ √ N , and AFT-Simple [41] proposed different functions for query ϕ(q) = sigmoid(q) and key ϕ(k) = softmax(k), respectively. These methods suffered from the performance drop since they lacked sufficient expression ability to replicate the original softmax attention mechanism.Besides the aforementioned methods, some studies approximated the similarity function with kernel expansions such as angular kernel expansion [39] with Sim(q, k) = 1/2 + 1/π • (qk ⊤ ) + H r or first order Taylor expansion [7] with Sim(q, k) = 1 + qk ⊤ / √ d + H r where H r represents the high-order residuals. These methods necessitated the masked output of original softmax attention as H r and applied the knowledge distillation (KD) method to further enhance the performance which severely increased the GPU memory consumption and were unsuitable for training large transformer models.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: It takes lots of computational resources for training the models from scratch. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methods: 
 In this section, we propose to use second-order (quadratic) Taylor expansion to approximate the similarity measurement function Sim(•, •)in Eq. 2. Compared to the first-order (linear) Taylor expansion [7] , quadratic approximation contains less information in the high-order residuals. Therefore, we can directly ignore them and derive a good performance without utilizing masked softmax attention output or the KD method.However, it is non-trivial to decompose the quadratic Taylor expansion into separate kernel embeddings with linear time complexity. Thus, in the following we first give a theoretical derivation by using the Kronecker product to decompose the quadratic expansion. Then, a fast approximation algorithm is applied to accelerate the computation of the Kronecker product.The quadratic Taylor expansion of the similarity measurement function is expressed as:EQUATIONwhere < •, • > is the dot product and ϕ(x) = x 4 √ d , 1 is used for vectors q and k. However, since the quadratic term exists in Eq. 4, it is challenging to decompose the equation into two separate kernel embeddings. In the following, we show that this problem can be solved by using the Kronecker product.Given two vectors a = {a i } d i=1 and b = {b i } d i=1 , we can easily derive:EQUATIONThis is equal to first computing the Kronecker product of each vector and then applying dot product, i.e., given K r (x) = vec(x⊗x) where ⊗ represents the Kronecker product and vec(•) is the vectorized output, we have:EQUATIONThen, we can apply Eq. 6 to Eq. 4 and decompose the similarity function into two separate kernel embeddings:EQUATIONwhereEQUATIONis the kernel function applied to the query and key vectors. Note that given a vector x ∈ R d , Kronecker product gives an output vector with quadratic length K r (x) ∈ R d 2 . Thus, the time complexity of linear attention using the decomposed quadratic Taylor expansion is O(N d 3 ) . Compared to the original softmax attention with O(N 2 d) time complexity, the proposed method does not yield an advantage.Recall that the computational burden primarily arises from the Kronecker product that quadratically expands the input dimension. Thus, there are several simple ways to reduce the dimension. For example, a pooling function can be applied on the input vector y = pool(x) ∈ R d/p where p is the dimensionality reduction factor. The output dimension of the Kronecker product can be reduced toK r (y) ∈ R d 2 /p 2, and the corresponding time complexity of the linear attention is O(N d 3 /p 2 ). Another way is to divide the input vector into c chunksx = [x 1 , • • •, x c ] and compute the Kronecker product within each chunk K r (x i ) ∈ R d 2 /c 2, and finally concatenate them together to derive theoutput o = concat K r (x 1 ), • • •, K r (x c ) ∈ R d 2 /c. The time complexity of the linear attention using this method is O(N d 3 /c).Although methods mentioned above can decrease the computational load, they do not actually reduce the time complexity. In the following, we propose a fast approximation algorithm to accelerate the computation of the Kronecker product, and reduce the computational complexity from O(Nd 3 ) to O(N d 2 ).By rewriting the definition of K r (ϕ(x)) in Eq. 8 in its element-wise form, we can get:EQUATIONNote that the order of the elements in the above equation does not influence the result of the inner product < K r (ϕ(q), K r (ϕ(k)) > in Eq. 7 as long as K r (ϕ(q)) and K r (ϕ(k)) change the order of their elements in the same manner. Thus, Eq. 9 can be written as:EQUATIONwhich is divided into four terms. The first is the quadratic term that contains d 2 components representing the multiplication of each pair of elements in x (including self-multiplication), the second and third terms are the linear term with length d each, and the fourth term is the constant term. Since the computational load of the inner product in Eq. 10 mainly comes from the quadratic term, it is important to reduce the number of elements in this term. Randomly preserving d items from d 2 elements is an efficient approach but leads to poor results. Employing grouping techniques help selecting the most representative items at the cost of increasing the computational complexity compared to random selection. We empirically find that using the self-multiplication terms {x 2 i } d i=1 can effectively represent all quadratic terms, while at the same time maintaining high efficiency.Therefore, the Kronecker product in Eq. 10 can be replaced with a compact version:EQUATIONin which we merge items of the same kind and multiply them by the square root of the number of the same items so as not to affect the inner-product result in Eq. 7. Learnable scalar parameters α, β and γ are used as the trade-off parameters. Note that given a vector x ∈ R d , this compact version of the Kronecker product gives an output of length 2d + 1. Therefore, the time complexity of linear attention using the decomposed quadratic Taylor expansion is reduced fromO(N d 3 ) to O(N d 2 ).We further found that the linear term can be discarded without hurting the classification performance, thus we set β = 0 in the following experiments.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA]Answer: [No] Justification: We do not include code. Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). • The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 In this section, we apply our linear attention with quadratic Taylor expansion to vision transformers and propose a series of QT-ViT models. We empirically investigate the effectiveness and efficiency of the proposed models on the ImageNet-1k classification dataset. Additional results regarding the performance on object detection and semantic segmentation tasks are provided in the appendices. Compared methods and evaluation metrics. To verify the effectiveness of the proposed QT-ViTs, we compare our method with a series of competitors including (1) Vision transformers with linear attention such as ViTALiTy [7] , Castling-ViT [39] , EfficientViT [3] , FLatten Transformer [14] and Hydra Attention ViT [2] ; (2) Other vision transformers with sparse attention or hierarchical architectures such as Swin [23] , SwinV2 [22] , FasterViT [17] , PoolFormer [40] , MobileViT [25] , MobileViTV2 [26] and CSwin [9] ; (3) State-of-the-art CNN models and CNN-Transformer combined model architectures such as CoAtNet [6] , CMT [12] , ConvNeXt [24] , EfficientNet [28] and EfficientNetV2 [29] .The proposed QT-ViTs and other baseline models are evaluated based on the accuracy-speed tradeoffs as shown in Fig. 1 . Furthermore, we measure the classification performance with top-1/top-5 accuracy. The efficiency of the model is represented by the FLOPs and parameters. Finally, we evaluate the inference speed of the models on the AMD Instinct MI250 GPU in Fig. 1 .Experimental results. The effectiveness and efficiency of the proposed QT-ViTs are evaluated on the ImageNet-1k dataset by comparing them to other state-of-the-art baseline methods mentioned above. The results are shown in Tab. 1 and all methods are gathered by their FLOPs into five groups including: <1G, 1∼3G, 3∼5G, 5∼10G and >10G.As shown in the 22.0 8.40 83.9 -SwinV2-T [22] 28.0 6.60 82.8 -EfficientViT-L1 [3] 53.0 5.30 84.5 96.9 EfficientViT-L2 [3] 64.0 6.96 85.1 97.0 Castling-MViTv2-B [39] 51.9 9.82 85.0 97.2 FasterViT-2 [17] 75.9 [41] sigmoid & softmax σ(x) softmax(x) 78.9 Castling-ViT [39] angular kernel Sim(q, k)= 1 2 + 1 π • (qk ⊤ )79.1 ViTALiTy [7] 1st order Taylor expansion[x/ 4 √ d, 1] 78.5QT-ViT (ours) 2nd order Taylor expansion1 √ 2 Kr (ϕ(x)), 1 √ 279.6In this section, we conduct several ablation studies to further verify the superiority of our proposed quadratic Taylor expansion kernel.Results of using different kernels. We compare the results of using quadratic Taylor expansion kernel with other kernels used in various linear attention vision transformers including EfficientViT [3] which uses ReLU kernel, Hydra attention [2] that utilizes cosine kernel, PolyNL [1] with the mean kernel, AFT-Simple [41] that proposes different kernels for query and key, angular kernel expansion [39] and first order Taylor expansion [7] . All methods use the same training settings and network architecture except the kernel used for computing the linear self-attention. The baseline method uses the original self-attention operation with O(N 2 d) computational complexity and is used as a strong baseline for comparison.As the results shown in Tab. 2, the proposed quadratic (2nd order) Taylor expansion outperforms all other competitors which demonstrates the effectiveness of the proposed method. For example, we achieve a 1.1% better top-1 accuracy compared to the linear (1st order) Taylor expansion kernel, 0.2% better than ReLU, 0.5% better than cosine, 0.8% better than mean, 0.7% better than sigmoid & softmax and 0.5% better than the angular kernel.Ablation on reducing the time complexity of the Kronecker product. In Sec. 3.2, we mentioned several ways of reducing the computational burden of the Kronecker product including:Method 1: applying a pooling function on the input vector y= pool(x) ∈ R d/p ; The classification results of using different methods mentioned above are shown in Tab. 3, in which the baseline method computes the original Kronecker product and is used to compare with other efficient methods. We can see that reducing input dimension with the pooling function (method 1) or dividing it into chunks (method 2) are not efficient enough since p and c are small compared to the dimension d of the vector and has a higher time complexity. Besides, the classification performances are not satisfying because too much information is lost. Randomly selecting quadratic items (method 3) is computationally friendly but has sub-optimal performance. The proposed compact version of the original Kronecker product (method 4) performs best among all the methods which indicates that using the self-multiplication terms to represent quadratic terms is enough to preserve the information in the output of the Kronecker product. We plot the results of the selfattention maps from the last block given a specific query (column 1, marked as red on the original images) using different attention methods including first-order Taylor expansion [7] (column 2), ReLU non-linearity function [3] (column 3) and the quadratic Taylor expansion used in the proposed QT-ViT (column 4). We can see that the proposed QT-ViT can exhibit a more focused and sharper response on attention feature maps. Furthermore, given a query vector, QT-ViT captures reasonable features on the feature map more accurately compared to the competitors. For example, QT-ViT concentrates on both ears of the dog given the query on the left ear of the dog. We can intuitively observe the advantages of QT-ViT from Fig. 2 . 
conclusion: 
 In this paper, we propose a new linear attention method to approximate the usage of softmax selfattention in the original vision transformer models. By conducting quadratic Taylor expansion of the similarity measurement function with the help of the Kronecker product, we can successfully decompose the similarity function into the product of two kernel embeddings while reserving highorder information and maintaining the effectiveness of the original self-attention. Furthermore, we propose a fast approximation algorithm to accelerate the computation of the Kronecker product and reduce the time complexity from O(N d 3 ) to O(N d 2 ) without much loss of information. We conduct experiments on the proposed QT-ViT models using the benchmark dataset ImageNet-1k, and the results show that we can achieve a better accuracy-efficiency trade-off compared to other state-of-the-art transformers and CNNs. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We conduct experiments on the COCO 2017 dataset to further validate the effectiveness of the proposed QT-ViT models. The COCO 2017 dataset has 118K training images, 5K validation images and 20K test-dev images. We use the mask R-CNN [18] as the object detection framework, EfficientViTs [3] and the proposed QT-ViTs as the backbone models. The experimental results are shown in the following table. We can see that the proposed QT-ViT has better mAP than the state-of-the-art model EfficientViT under the same training setting. In the following table, we show the results of using absolute positional embedding (APE) on object detection. Note that APE has little impact on the latency, FLOPs, and top-1 accuracy for image classification tasks thus we do not show the corresponding results in the main section. We further verify the effectiveness of the proposed QT-ViT on the semantic segmentation task using the ADE20K dataset, which contains 20K training images from 150 semantic categories, 2K validation images and 3K test-dev images. UperNet [35] is used as the framework for the experiments.As shown in the table below, using QT-ViTs as the backbone models achieve better mIoU on the ADE20K dataset than using the EfficientViTs as backbones. The results show that the proposed QT-ViT performs well on the semantic segmentation task.The masked output of the original softmax attention has been shown to be useful in previous studies such as Vitality [7] and Castling-ViT [39] . It is also shown to be useful in our method and the experiments using QTViT-1 on the ImageNet dataset are shown below. However, it requires more GPU memory during training, which is not suitable for training large models. Thus, we do not use this strategy in our method. For the integrity of the paper, we still list the results of using original softmax in the following table, the experiments are conducted with QT-ViT-1 on ImageNet dataset. This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: The abstract and introduction accurately reflect the paper's contributions and scope.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper presents a method for improving the prior art of linear self-attention using quadratic Taylor expansion with a fast approximation algorithm. It has been reviewed by several knowledgable reviewers. The majority consensus from the reviewers (dQG7, LZeB and QZCU) was that the proposed fast approximation algorithm is novel and effective. The clarifications and new results presented during the reviewer-author discussion addressed their concerns raised in the initial comments. This meta-reviewer concurs, and recommends **acceptance**. 

The reviewer 5zgq had concerns mainly regarding the difference between the proposed method and some related work pointed out during the reviewer-author discussion. The authors and the reviewer 5zgq had very informative discussions. Although the authors did not finally convince the reviewer 5zgq, s/he did mentioned that the proposed concatenation of self-multiplication terms in the fast approximation algorithm is somewhat new. So, this meta-review leans towards looking at the positive aspects of the proposed method. 

The authors are encouraged to carefully revise their paper by,
- incorporating the clarifications, discussions, new results and new references; 
- providing comparisons on memory footprints that are missing in the submission and the discussions;
- providing more detailed analyses on effects between w/ APE and w/o APE in both image classification tasks and downstream tasks.

Reviews: review: paper_summary: paper_summary: This paper proposed a new linear complexity sequence modeling strategy for image modeling. To achieve this, the authors first replace the softmax attention with second-order Taylor expansion, then accelerate its computation with a fast approximation algorithm. The effectiveness of the proposed method is validated on the image classification task.
strength_weakness: strength_weakness: 1. It is unclear why directly using the first-order Taylor expansion is worse than using the second-order Taylor expansion with linear approximation. The author should provide a section to discuss this.

2. Experiment details are missing. For example, different image resolutions can largely impact image classification performance. The author fails to provide the image resolution in the experiment section. 

3. Some mathematical derivation can be moved to the supplemental section. For example, there are some redundant equations in Eq. 4, 6, 7, 9. 

4. There is no actual training/inference speed comparison. Counting flops can be misleading in some cases. I would like to know the actual speed instead of flops.

5. Missing linear attention vision backbone. For example, VVT (TPAMI 23).

6. For image modeling, object detection, and semantic segmentation are also important. I do not know why the author decided to move these two tasks into the supplemental section. However, the semantic segmentation results are extremely low when compared with other linear image backbones. The author do not provide any comments on this.
questions: questions: As above.
limitations: limitations: N/A
paper_summary: paper_summary: This paper introduces QT-ViT models, which enhance linear self-attention using quadratic Taylor expansion. The similarity function is decomposed into the product of two kernel embeddings via the Kronecker product. By employing a fast approximation algorithm, the computational cost is reduced while maintaining overall performance. Experiments on ImageNet Classification with various model sizes show consistent improvement over the baseline EfficientViT method. Additionally, visualizations demonstrate that the QT-ViT model has a more focused attention feature map.
strength_weakness: strength_weakness: - W1. The overall improvement is incremental. Compared to the EfficientViT baseline, the performance is almost the same under the same computation budget across different model sizes and tasks. For ImageNet classification, it would be helpful to include results on ImageNet-v2 and ImageNet-real to demonstrate the robustness and consistency of the improvement.

- W2. The image resolution is not listed in the experiment section. Performance on different image resolutions should also be reported and compared with other methods. When using a much larger image resolution (and more patch tokens), how are the computation cost and latency affected?

- W3. It would be better to include the detection and segmentation results in the main text to support the effectiveness of the proposed QT-ViT across different tasks.

- W4. The network architecture should be detailed, and additional components like absolute positional embedding should be ablated.

- W5. The mIoU performance of EfficientViT-B3 is reported as 49.0 on the ADE20k dataset in the original paper, which is inconsistent with the 38.0 reported in this paper.
questions: questions: - The author should provide more implementation details about the network, experimental setting, and comparison with other methods.
limitations: limitations: - The latency is tested on an AMD GPU, and I'm not sure if there is any difference on the implementation side. Will the performance gain be consistent on other type of devices?
paper_summary: paper_summary: This paper proposes a novel method to compute the kernel function in linear attention. They decompose the softmax attention with Tayler expansion and utilize the first two items to approximate the exponential function. The Kronecker product is used to decompose the quadratic Taylor expansion into two kernel functions, and the self-multiplication terms in the output of the kernel are used to replace the quadratic terms for fast inference. The experimental results on multiple CV tasks show the effectiveness and efficiency of the proposed method.
strength_weakness: strength_weakness: - What is the dimension of $\alpha$, $\beta$ and $\gamma$ used in Eq.11? Are there any ablation studies using different parameters?

- I notice that you replace $K_r(\phi(x))$ in Eq.8 with Eq.11, and you use $\gamma$ to represent the constant term. Then, could the constant term $1/\sqrt{2}$ in Eq.11 be merged with $\gamma$? What about the scaling factor $1/\sqrt{2}$ in Eq.8?

- There are 6 models with different model sizes in Tab.1. Why do you plot only 3 models in Fig.1?
questions: questions: - You have mentioned that you do not necessitate the masked output of the original softmax attention during training as the previous methods do. I wonder if this can be added back to get further performance gain on QTViT?

- You mentioned that QTViT can exhibit a more focused and sharper response around Fig.2. Is there any insight about this phenomenon?

- In Tab.1, the performance gains of QT-ViT 4$\sim$6 to the SOTA are more marginal compared to QT-ViT 1$\sim$3 to the SOTA. Are there any explanations for this phenomenon?

Overall, I would like to see the accuracy and speed trade of the other 3 models to make sure that the proposed method can surpass EfficientViT on all different model sizes.
limitations: limitations: N/A
paper_summary: paper_summary: In this paper, the authors propose QT-ViT models, which improve the traditional linear self-attention methods by using a second-order (quadratic) Taylor expansion to approximate the original softmax attention and then accelerate this process using a fast approximation algorithm reducing computational complexity from $O(n^2d)$ to $O(nd^3)$ and further to $O(nd^2)$. This method leverages the properties of quadratic expansion for better performance while maintaining the speed of linear approximation. Extensive experiments on image classification, as well as object detection and semantic segmentation tasks, demonstrate that QT-ViT models achieve state-of-the-art accuracy-speed trade-offs, surpassing previous methods across various model sizes.
strength_weakness: strength_weakness: + In Table 1, the top 1 acc of the proposed QT-ViT models over EfficientViT appears to diminish. This trend raises concerns about the scalability and robustness of QT-ViT models as they are scaled up. The paper should include a detailed analysis of this trend.

+ While Figure 1 effectively visualizes the latency of the models, the paper does not include this crucial metric in Table 1, The paper should include latency metrics in Table 1 to provide a comprehensive comparison of the models.

+ Table 2 presents the results of using different kernels, but it is unclear whether all these kernels have the same time complexity. The paper should clearly state whether the time complexities of each kernel in Table 2 are the same. If not, the time complexity of each method should be added in Table 2.

+ In Eq.11. the authors state that using the self-multiplication terms can effectively represent all quadratic terms. Could the authors provide more details for this finding?

+ It seems that the improvement in segmentation is larger than in image classification and object detection. It is better to provide some analysis.
questions: questions: Please see the weakness. The instability of accuracy is my major concern of this method.
limitations: limitations: N/A

Comments: title: Regret to hear that
Rebuttal to reviewer 5zgq
Are there any further questions?
Rebuttal to reviewer 5zgq
Rebuttal to reviewer 5zgq
Thank you
Let's engage in the reviewer-author disscussion
Thanks for the detailed feedback.
Are there any further questions?
Are there any further questions?
Respond to Reviewer dQG7
Thanks
Respond to Reviewer LZeB
comment: We I regret to hear that you wouldn’t change your score. In fact, simply concat(q^2,q) yield bad result, since the coefficients before each term and the constant term are important to derive our sota results.

We believe that the simplicity of our implementation is an advantage rather than drawback. The analysis and derivation to the final equation form is the most important part of our method, which is not discussed by other papers before.

After all, we appreciate your time for reviewing this paper.
Well, all linear methods share the same complexity of O(nd^2). 

The paper begins with a second-order Taylor expansion to reduce the complexity to O(nd^3), and then utilizes a feature map of concat(q^2, q) to further reduce the complexity to O(nd^2). 

However, the transformation between Taylor expansion to concat(q^2, q) is not smooth.

Using concat(q^2, q) is somewhat new. However, I would like to see stronger evidence for its efficacy. The results of other vision tasks are not compelling.

Thanks for the quick response and for discussing the related work. Considering the others' ratings, I would like to keep my initial rating as the avg 5.5 reflects the quality of this work.
The paper [a2] you mentioned does not solve any of the problems in [a1] that I mentioned above. In fact, they do not even cite the paper [a1]. This paper simplifies the second-order Taylor expansion to $\phi(x)=x^2$ as shown on page 5, and uses a linear transformation before applying the kernel function. Thus, their implementation is much simpler than the second-order Taylor expansion since they discard all other terms except for the quadratic term. There is no theoretical analysis to show the relationship between this method and the second-order Taylor expansion.

Thus, compared to this paper [a2], our contributions are concluded in the following:

1. We improve the linear self-attention using the **correct** quadratic Taylor expansion through rigorous theoretical analysis and derivations, while there is no such analysis in [1, a1, a2] and the conclusions they derive are either incorrect [1, a1] (see the dimension $d’$ I mentioned above) or has little connection to second-order Taylor expansion [a2].

2. We propose a fast approximation algorithm to reduce the time complexity of our method from $O(Nd^3)$ to $O(Nd^2)$, while they use methods such as reducing the dimension of input through the linear layer or using hardware-efficient algorithm based on CUDA, which have their own drawbacks as we analyzed previously.

3. These three papers [a1, a2, 1] are all published within 2-3 months before the submission deadline of NeurIPS 2024. Thus it is unfair to judge our novelty based on them according to the policy of the NeurIPS conference.

Besides, the contribution of the hardware implementation of different platforms is about the paper [1] you mentioned above rather than ours.

Hope the above analysis can address your concern.

[a2] Linear Transformers with Learnable Kernel Functions are Better In-Context Models. ArXiv, 2024.
The issue in [a1] has been discussed in the paper "Linear Transformers with Learnable Kernel Functions are Better In-Context Models". However, the distinction between this paper and your approach is that one uses "concat" while the other uses "add".

The contribution of the hardware implementation of different platforms is not enough for a NeurIPS paper.

Can you rephrase the main insight or contribution of your paper based on this information?
Dear reviewer 5zgq,

We appreciate you taking the time to review this paper. Since it is less than 10 hours before the deadline for the discussion phase, we hope that we have addressed all your concerns.

If you still have any questions, we are happy to discuss them with you. Thank you very much.

Best,

Authors of paper 275
First of all, this paper directly used the conclusion of 2nd Taylor expansion from paper [a1] at the start of page 7. However, the conclusion in paper [a1] has some problems. For example, in page 6 in paper [a1], they have $d'=1+d+d^2$. However, the dimension of $\varphi(x)$ should be $(d+1)^2+1=2+2d+d^2$ which can be rigorously derived from Eq.8 in our paper. Besides, paper [a1] does not have any theoretical analysis and derivations to explain the dimension $d'$. Thus, we believe that our conclusion is correct.

Secondly, the acceleration of this paper is to utilize the hardware-efficient algorithm based on CUDA, which limits the utilization of its method and can only be accelerated on NVIDIA GPUs and is not able to be accelerated with other devices. The time complexity of their method is still $O(Nd^3)$, and mapping the input to lower dimensions as they used will yield bad results (see method 1 and 2 in Table 3 in our paper). Our method provides a fast approximation algorithm to reduce the time complexity to $O(Nd^2)$ and can be used without limitation.

Thirdly, this paper was published on Arxiv only two months before the deadline of NeurIPS 2024, thus we do not cite this paper or discuss it. We will add these discussions in our final version.

Thank you very much.

[a1] The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry. ArXiv, 2024.
Thanks for your prompt reply.

I am digging into the approximation and found a paper [1] that is very similar to the proposed method, i.e., doing a second-order Taylor expansion of exp(x) and then writing it as inner products. I noticed that [1] was first published on arxiv on 28 February 2024 and accepted in ICLR 24. Here is an implementation as well: https://github.com/lucidrains/taylor-series-linear-attention.

Would you please provide a detailed discussion of the differences between these two methods? Also, the paper is missing a citation of [1].

[1] Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D., ... & Ré, C. (2024). Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668.
Thanks for your comments.

1. If I understand it correctly, you are asking why using our fast approximation algorithm can well represent the second-order Taylor expansion. In fact, as Table 3 shows, we have tried several different methods to reduce the number of quadratic terms, and using the self-multiplication terms yields the best result among all of the methods. Generally speaking, the quadratic terms can be viewed as a square matrix ${\bf M}$ in which each element ${\bf M}_{ij}$ represents the multiplication result of ${\bf x}_i$ and ${\bf x}_j$. Thus, the self-multiplication terms can be viewed as the elements on the diagonal of the matrix which can represent the whole matrix to some extent. Besides, we expand the self-multiplication results with respect to the number of quadratic terms and use a learnable parameter 
 to further adjust the ratio. Thus, the overall method can represent all quadratic terms.

2. The training speeds of EfficientViT and QTViT are shown below, in which we use a batch size of 1024. The inference speeds of different models on NVIDIA V100 GPU are shown in **Figure 1 in the rebuttal pdf** which is also used to answer Q7 from reviewer QZCU.

|model|training speed (ms/batch)|
|-|-|
|EfficientViT-B1|534.5|
|QTViT-1|538.0|
|EfficientViT-B2|586.6|
|QTViT-2|591.5|
|EfficientViT-B3|756.8|
|QTViT-3|763.4|

Hope these answers can address your concerns.
1. The author addressed why using second-order Taylor expansion is better than the first-order Taylor expansion, but did not directly explain why using the first-order Taylor expansion is worse than using the second-order Taylor expansion with linear approximation. This does not resolve my concern in this case.

2. The actual training speed is still missing. Also, the inference speed is measured on an AMD Instinct MI250 GPU. I am not sure if all other methods have been fully optimized on AMD GPUs since other methods were originally implemented for NV GPUs.
Thank you very much for your positive decision.

Yes, the resolution tested in A2 is 288 rather than 228. Since the largest resolution used in EfficientViT-B1 is 288, we use the same setting as EfficientViT-B1 in A2. We further extend the resolution to 320 and 384. The results are shown below, but sorry for the limited discussion time left, we can only provide the FLOPs and latencies results.
|resolution|model|FLOPs (G)|Latency (ms)|
|-|-|-|-|
|320|QTViT-1|1.06|2.44|
|320|EfficientViT-B1|1.06| 2.42|
|384|QTViT-1|1.53|2.63|
|384|EfficientViT-B1|1.53| 2.59|

Best,

Authors of paper 275
I am satisfied with the authors' rebuttal. Since they have addressed all of my concerns, I decide to keep my score positive.
Thanks for the author's reply. I am glad to see that the author provides the accuracy-speed trade-off on all 6 different models, and the results show a consistent improvement. Other concerns are also addressed. Thus, I would like to increase my score from 6 to 7.
Dear Rev. 5zgq,

We look forward to seeing your comments on the authors' rebuttal, as well as any further clarifications as you may need. 

Thanks
Thank you for the detailed rebuttal and the additional results provided. I appreciate the effort in addressing the issues raised. The consistent improvement over the ImageNet dataset and its variants, along with the consistent latency improvements across different device types, demonstrates the robustness of your approach. Additionally, the extra experiments and ablation studies on detection and segmentation effectively highlight the method's potential.

However, I do have one minor concern. If I understand correctly, the resolution tested in A2 is 288 rather than 228. While the results are promising, 288 is still not considered a very large resolution where differences in latency might be more pronounced. Furthermore, the performance gain, though present, remains marginal compared to EfficientViT across various model sizes and computation costs.

The overall results suggest that the paper meets the acceptance threshold. Based on the improvements and additional evidence provided, I am willing to increase my score to 5.
Dear reviewer QZCU,

We appreciate your comments and suggestions in the review. All of your concerns in your original review have been addressed, including the overall improvement and results on ImageNetv2 and real, experiments on different image resolutions, ablations on positional embedding, experiments on semantic segmentation, and inference speed comparison on different devices.

We hope you can read the rebuttal and let us know if you have further questions. Thanks for your time.

Best, 

Authors of paper 275
Dear reviewer 5zgq,

We appreciate your comments and suggestions in the review. All of your concerns in your original review have been addressed, including the clarity of second-order Taylor expansion, experiments on different image resolutions, mathematical derivation, inference speed, comparison with VVT and experiments on semantic segmentation.

We hope you can read the rebuttal and let us know if you have further questions. Thanks for your time.

Best,

Authors of paper 275
Thanks for your question. 

We do not use KD method in the table in A4. In fact, adding KD method will largely increase the GPU memory (at least double the GPU memory usage). We use 8 GPUs during training. Each AMD MI250 GPU has a maximum GPU memory of 64GB. 

It is still enough for training QTViT-1 with the help of original softmax attention and KD method, but it will cause OOM issue when training larger models such as QTViT-3. Thus, it shows the advantage of our method that can derive a good performance without utilizing masked softmax attention output or the KD method.
Thanks for the author’s reply. I have some additional questions. Do you use KD method in the table in A4? How many GPUs do you use during training?
Thanks for your question. In fact, we have tested the latencies of Poolformer series and the results are shown below.

|model|latency (ms)|Top-1 Acc (%)|
|-|-|-|
|s12|3.78|77.2|
|s24|7.12|80.3|
|s36|9.91|81.4|
|m36|12.67|82.1|
|m48|16.56|82.5|

The results of Poolformer series are not as strong as other series and including them in Figure 1 will cause other models to be crowded together, making it difficult for readers to distinguish the strengths and weaknesses. Thus, we decide not to include Poolformer series in Figure 1. Other model series listed in Table 1 but not shown in Figure 1 are due to the similar reason.

We will add a 'Latency' column in Table 1 in the final version of our paper.
I appreciate your response. I noticed that some model series, such as Poolformer, are listed in Table 1 but are not shown in Figure 1. Could you please explain this?