Castling-ViT: Compressing Self-Attention via Switching Towards Linear-Angular Attention at Vision Transformer Inference

Introduction

Vision Transformers (ViTs) have made significant progress in image classification, object detection, and many other applications. It is well recognized that the superior performance achieved by ViTs is largely attributed to their self-attention modules that can better capture global context . Nevertheless, ViTs’ powerful self-attention module comes at the cost of quadratic complexity with the number of input tokens, causing a major efficiency bottleneck to ViTs’ achievable runtime (i.e., inference latency) . To mitigate this issue, linear attention designs have been developed to alleviate the vanilla ViT attention’s quadratic complexity. In particular, existing efforts can be categorized into two clusters: (1) ViTs with local attention by restricting the attention window size , sharing the attention queries, or representing the attention queries/keys with low rank matrices; and (2) ViTs with kernel-based linear attention, which approximate the non-linearity softmax function by decomposing it into separate kernel embeddings. This enables a change in the matrix computation order for a reduced computational complexity .

Despite their promise in alleviating ViTs’ complexity and thus inference runtime, both the local and linear attention compromise ViTs’ performance due to the lack of capabilities to capture global or local context. To marry the best of both worlds, we advocate training ViTs with both (1) efficient but less powerful linear attention, i.e., without the high-order residuals in angular kernel expansion, and (2) powerful yet costly softmax-based masked attention. The latter helps approximate high-order residuals at the early training stage while being dropped during inference, based on an assumption that the remaining networks can gradually learn the high-order components at the later training stage . This concept resembles the “castling” move in chess when two pieces are moved at once. While it sounds promising, there are still two challenges to achieve this. First, existing linear attention modules still underperform their vanilla softmax-based counterparts. Therefore, a better linear attention is crucial for the final performance. We find that angular kernels perform equally as softmax-based attentions in terms of similarity measurements. While they still suffer from a quadratic complexity, they can be divided into linear terms and high-order residuals. The challenge is how to construct ViTs with only the linear terms. Second, doing so would require that the trained ViTs merely rely on the linear terms towards the end of training, which would call for an approximation of the above high-order residuals. The challenge is that how we can resort to costly but powerful modules to approximate high-order residuals during training but does not incur extra inference cost.

In this work, we develop techniques to tackle those challenges, and make the following contributions:

We propose a framework called Castling-ViT, which trains ViTs using both linear-angular attention and masked softmax-based quadratic attention, but then switches to having only linear-angular attentions during ViT inference to save computational costs.

We develop a new linear-angular attention leveraging angular kernels to close the accuracy gap between linear attention and softmax-based attention. It expands angular kernels where linear terms are kept while complex high-order residuals are approximated.

We use two parameterized modules to approximate the high-order residuals above: a depthwise convolution and an auxiliary masked softmax-based attention, where the latter’s attention masks are regularized to gradually become zeros to avoid inference overhead.

We conduct extensive experiments to validate the effectiveness of the proposed Castling-ViT. Results on classification, detection, and segmentation tasks consistently demonstrate its superior performance ($\uparrow$1.8% top-1 accuracy or $\uparrow$1.2 mAP) or efficiency (40% MACs savings) over state-of-the-art (SOTA) CNNs and ViTs.

Related Works

Vision Transformers (ViTs). ViT beats CNNs with a simple encoder-only transformer architecture taking the splitted non-overlapped image patches as sequential inputs, but relies on costly pretraining on a huge JFT-300M dataset . Later, DeiT and T2T-ViT leverage an improved ViT training recipe or enhanced tokenization mechanism to achieve a comparable accuracy without the necessity of costly pretraining. To further improve ViTs’ achievable accuracy-efficiency tradeoffs, CrossViT , PiT , PVT , MViT and Swin-Transformer propose a pyramid-like architecture, which is commonly used in CNNs ; DynamicViT , A-ViT , ToME , and MIA-Former propose to adaptively identify and remove unnecessary input tokens for saving computational costs. With the goal of deploying ViTs in resource-constrained devices, various efficient ViT architectures have been proposed . For example, LeViT , CvT, and MobileViT adopt more efficient self-attention implementation or incorporate convolutional feature extraction blocks into their early layers; EfficientFormer further enables pure ViTs to run as fast as MobileNets. In contrast, our Castling-ViT explores whether ViTs can learn both global and local features while still being efficient at runtime. Also, we target a generic linear-angular attention that can serve as a drop-in replacement for all kinds of ViT architectures and thus is orthogonal to new ViT architecture designs.

Efficient ViT Variants. As commonly recognized, ViTs rely heavily on their self-attention module which is however costly due to its quadratic computational complexity with the total number of input tokens . To make the self-attention module more efficient, a surge of linear attention works have been proposed and can be roughly categorized into two groups: local attention or kernel-based linear attention . For kernel-based linear attention, common designs approximate the softmax function or the full self-attention matrix with orthogonal features or kernel embeddings, then the computation order can be changed from $\mathbf{(QK)V}$ to $\mathbf{Q(KV)}$. For example, and decompose the exponential terms in softmax-based attention into kernel functions and exchange the computation order. Despite their decent performance, currently kernel-based linear attention in general underperform the softmax-based attention. Recent works also unify low rank approximated and sparse attention (can also be dropped at inference) to improve ViTs’ accuracy-efficiency tradeoffs. Different from the above works, we explore from a new perspective by taking spectral angles into consideration when measuring the similarities among tokens, resulting in linear-angular attention that can achieve comparable or even better performance than softmax-based attention. More efficient ViT variants are supplied to the Appendix.

ViTs have also been used as the backbones for downstream tasks, e.g., detection and segmentation . We supply more literature review to Appendix.

The Proposed Methods

Self-Attention. Self-attention module is a core component of the Transformer , and usually consists of multiple heads. Each head captures global-context information by measuring pairwise correlations among all $N$ tokens ($N$ denotes the total number of tokens) as defined below:

Kernel-based Linear Attention. The core idea of linear attention is to decompose the similarity measurement function into separate kernel embeddings, i.e., $\text{Sim}(\mathbf{Q},\mathbf{K})\approx\phi(\mathbf{Q})\phi(\mathbf{K})^{T}$, so that we can change the computation order to $\phi(\mathbf{Q})(\phi(\mathbf{K})^{T}\mathbf{V})$ based on the associative property of matrix multiplication. In this way, the attention complexity is quadratic to the feature dimension $d$ instead of the token length $N$. One straightforward implementation of linear attention is to use Gaussian RBF kernels to measure the similarity, which can serve as an unbiased estimation of $\text{exp}(\langle\cdot,\cdot\rangle)$ in Eq. (1):

where $\langle\cdot,\cdot\rangle$ denotes the inner product operator. According to , we can induce a function $\phi(\cdot)$ to approximate the Gaussian RBF kernel, mapping the input space to the feature space. Assuming that both $\mathbf{Q}$ and $\mathbf{K}$ are normalized as unit row vectors along the feature dimension, then the attention formula can be approximated by:

While doing so can alleviate the attention complexity to become linear w.r.t. the token length $N$, it often causes a nontrivial accuracy drop as compared to the corresponding ViTs with vanilla softmax-based attention .

Revisit Attention Designs. Recent ViTs have achieved a low complexity for classification tasks (e.g., LeViT ) but still be costly for downstream tasks due to high input resolutions, i.e., the number of tokens (e.g., MViTv2 ). Therefore, one dilemma is the tradeoff between the model efficiency and generalizability. As illustrated in Fig. 2, LeViT shrinks the feature dimension to be more efficient on low-resolution tasks but cannot be well generalized to downstream tasks, due to the resulting (1) feature bottlenecks, i.e., insufficient feature dimensions, and (2) fixed attention biases. On the other hand, MViTv2 performs token pooling for better fitting high-resolution tasks but is less effective for low-resolution due to the caused token bottlenecks, i.e., insufficient number of tokens. Our ablation studies in Table 1 show that (1) for ImageNet classification, both token and feature pooling can lead to accuracy drops in ViTs; and (2) for downsampling layers, post-$\mathbf{Q}$ pooling (i.e., pooling after linear projection) together with residual connections in MViTv2 performs better than pre-$\mathbf{Q}$ pooling (i.e., pooling during linear projection) in LeViT. Therefore, we develop Castling-LeViTs on top of attention with merely post-$\mathbf{Q}$ pooling and residual connections.

2 The Proposed Castling-ViT Framework

Castling-ViT Overview. Fig. 4 illustrates an overview of the proposed Castling-ViT, which makes linear attention more powerful than previous designs while still being efficient during inference. In particular, we propose (1) a novel kernel-based linear-angular attention from the spectral angle perspective to close the accuracy gap between linear attention and softmax-based attention; and (2) a training augmentation method that leverages softmax-based attention as an auxiliary branch to assist the linear-angular attention only during ViT training. Note that a mask is applied to the auxiliary branch to manifest linear-angular attention and drop the auxiliary branch.

Angular Kernel. In addition to the previously adopted polynomial, exponential, or RBF kernel that focuses on spatial similarity measurements , we propose to consider measuring spectral similarity via angular kernel as an alternative to existing softmax-based attentions, leading to similar or better performance since it additionally takes into consideration the nature of spectral characteristics, e.g., the spectral angle as a distance measurement function . Such a spectral angle between two vectors is defined as:

where $\|\cdot\|$ is the Euclidean distance and $\langle\cdot,\cdot\rangle$ is the inner product. The output range of $\theta$ is $[0,\pi]$. Such an angle can be used as a distance. In our design, we define the angular kernel as a similarity measurement function between the queries $\mathbf{Q}$ and keys $\mathbf{K}$ as:

and the output range is thus $$. With more aligned$\mathbf{Q}_{i}$and$\mathbf{K}_{j}$,$\theta$is closer to 0 and thus the similarity is closer to 1; In contrast, if$\mathbf{Q}_{i}$and$\mathbf{K}_{j}$have opposite features,$\theta$is closer to$\pi$ and thus the similarity is closer to 0.

Properties of Angular Kernel and Its Feature Space. One property of our angular kernel is that replacing the similarity measurement in self-attention with such a kernel provides an efficient way to implicitly map the input data to a high (even infinite after expansion) dimensional feature space , where the distances/angles are calculated based on a rich feature structure. Let $\phi(\cdot)$ denotes the implicit map induced by this kernel, the norm of mapped input data is:

which means that all data in the input space are mapped onto the sphere of radius 1 in the feature space. Also, the distance between two input features is given by:

That is, the square (Euclid) distance and the spectral angle is positively correlated and the distance range is $$. As shown in Fig. 3, the angles in the input space is correlated to the feature distance after applying our angular kernel.

Expansion of Angular Kernel. A natural following question is how to incorporate the above angular kernel for designing linear attention given its quadratic complexity w.r.t. the input token length. Recalling from trigonometric identities and the expansion of the arccos function into an infinite series, we reformulate the similarity function as:

where $\mathbf{Q}_{i}\cdot\mathbf{K}_{j}^{T}$ denotes the normalized linear kernel function $\langle\mathbf{Q}_{i},\mathbf{K}_{j}\rangle/\|\mathbf{Q}_{i}\|\cdot\|\mathbf{K}_{j}\|$, which is equivalent to the inner product if $\mathbf{Q}_{i}$ and $\mathbf{K}_{j}$ are unit vectors. We see that the first linear-angular terms can be directly used as the similarity measurement in linear attention, while the remaining higher-order terms of an infinite series introduce a much higher complexity. As such, we propose to adopt a relaxation to approximate it.

In this way, our attention module can be formulated as:

where $1/\pi\cdot Q\cdot(K^{T}\cdot V)$ is the linear term with $\mathcal{O}(N)$ complexity, $\mathbf{M}_{DW}\cdot V$ is the matrix form of DWConv, and $\mathbf{M}_{SparseAttn}$ is the normalized sparse softmax attention. The overall complexity to compute Eq. (9) is linear to the input token length, where the MACs of DWConv is also negligible (e.g., $<$ 1% of the total MACs). Also, normalizations ($\nicefrac{{\mathbf{x}}}{{\|\mathbf{x}\|_{2}}}$) are inserted to the $\mathbf{Q}/\mathbf{K}$ and sparse attention branches to help the similarity measurement following .

2.2 Switch Towards Linear-Angular Attention

Recall that we add a sparse softmax-based attention as an auxiliary branch in the linear-angular attention to help approximate the high-order residuals. Such a costly attention can be potentially dropped without hurting the inference accuracy, drawing inspiration from that the remaining network can gradually learn the high-order/frequency components at the later training stage. Next, we explain how it is constructed and dropped.

Sparse Training Augmentation. As illustrated in Fig. 4, we adopt a masked softmax-based attention as an auxiliary branch to augment ViT training. In particular, we first use a comparator with a predefined threshold $\epsilon$ to generate a binary mask, where attentions greater than $\epsilon$ are set to 1 and 0 otherwise. These masks are then applied to the attention maps to generate masked attention maps that will be summed up with our linear-angular attention together as the final attention. The resulting sparse attention is given by:

where $\tt{Mask}_{\epsilon}(x)=x\text{ if }x>\epsilon\text{ else }0$, acting as an element-wise threshold function. As such, the sparse attention captures the higher attention scores (i.e., strong local features), and can potentially complement our linear-angular attention by supplementing the missing higher-order terms. Such an assumption aligns well with recent findings that (1) low-rank and sparse approximations complement each other ; and (2) linear attention lacks local feature extraction capabilities over its softmax-based counterpart .

Castling During ViT Inference. As we are targeting efficient ViT inference, it is desired to reduce or completely remove the costly softmax-based attentions while only keeping our linear-angular attention at runtime, i.e., performing castling. In our experiments, we found that under the sparsity regularization above, the softmax attention naturally converges to all zeros as the training progresses. We consider both fixed and dynamic schedules for $\epsilon$ and find that our Castling-ViT is not sensitive to neither the threshold value nor the threshold schedule for a given task. Given a fixed mask threshold (e.g., $\epsilon=0.02$ in image classification experiments), the masks become all zeros at latter training stages and thus the auxiliary branch can be removed without hurting the model accuracy. We supply the visualization of the mask-evolving trajectory and our conjecture for understanding such a phenomenon in Sec. 4.5.

Experiments

Tasks, Datasets, and Models. Tasks and Datasets. We consider three benchmark datasets and three representative vision tasks to demonstrate the superiority of the proposed Castling-ViT, including image classification on ImageNet dataset with 1.2 million training and 50K validation images; Object detection on COCO dataset with 118K training and 5K validation images; Semantic segmentation on ADE20K dataset with 20K/2K/3K images for training, validation, and testing, respectively. Models. We apply our proposed Castling-ViT idea on top of various models. For the classification task, we consider LeViT , MViTv2 , and DeiT ; For the detection task, we consider models with efficient ViT backbones (e.g., PicoDet with modified ESNet and LCNet backbones with transformer blocks); For the segmentation task, we consider Mask2former with ViT-Base backbone.

Training Settings. For the classification task, we use a SGD optimizer with 0.9 momentum and $2\times 10^{-5}$ weight decay to train ViTs for 1000 epochs using 64 V100 GPUs, with each card having 64 (LeViT) or 32 (MViTv2/DeiT) batch sizes. The learning rate is 2.0 with first 11 epochs warm-up starting from 0.01 and decays by a factor of 0.9875 per epoch . Also, we use the distillation based on a teacher model with a 85.5% accuracy. For the detection task, we adopt SGD optimizer with momentum 0.9 and weight decay 4e-5 to train models on COCO. All models are trained on 8 V100 GPUs with each card having 80 batch sizes following PicoDet’s training recipe . Also, we follow LeViT’s training recipe to pretrain backbones on ImageNet; For the segmentation task, we follow Mask2former’s training recipe to train models on ADE20K, where ViT backbones are pretrained following MAE if specified.

Baselines and Evaluation Metrics. Baselines. For the classification task, we compare the proposed Castling-ViT with LeViT , MviTv2 , DeiT , Swin , CSWin , PVT , etc. For the detection task, we compare with FBNetV5 , YOLOX , YOLOv5, MobileDet , and EfficientDet . For the segmentation task, we compare with Mask2former with ViT backbones. Evaluation Metrics. We evaluate the Castling-ViT and all baselines in terms of accuracy-efficiency tradeoffs. Specifically, the accuracy metrics refer to top-1/5 accuracy for the classification task; AP, AP50, AP75 for the detection task (AP: average precision); mIoU, mAcc, and pAcc for the segmentation task (mIoU: mean intersection over union). For efficiency metrics, we compare the number of model parameters or inference FLOPs (or MACs).

2 Castling-ViT over SOTA Baslines

Image Classification. To evaluate the effectiveness of our proposed techniques on the image classification task, we apply the proposed Castling-ViT idea to three typical or SOTA ViT architectures: DeiT as typical ViTs, LeViT as representative efficient ViTs, and MViTv2 as representative hierarchical ViTs for downstream applications of high input resolutions, and compare their performance over baselines on ImageNet. As shown in Tab. 2, the comparison across a large MACs (or FLOPs) range starting from 0.4G to 17G. We categorize ViT models into four regimes: $<$1G, 1$\sim$3G, 3$\sim$10G, and $>$10G, and select baselines in each regime to benchmark separately for the clarity purpose. Castling-ViTs consistently perform better than all baselines across various MACs ranges in terms of the accuracy-efficiency tradeoff. For example, Castling-LeViT achieves 82.6% top-1 accuracy with only 1.40G MACs while LeViT requires 2.35G FLOPs instead, i.e., $\downarrow$40% MACs; On the other hand, under comparable MACs, Castling-MViTv2 achieves 84.1% accuracy vs. MViTv2 with 82.3% accuracy instead, i.e., $\uparrow$1.8% top-1 accuracy. Overall, Castling-ViT achieves an improved accuracy of $\uparrow$0.5% $\sim$ $\uparrow$6.6%, $\uparrow$1.0% $\sim$ $\uparrow$8.1%, $\uparrow$1.0% $\sim$ $\uparrow$4.1%, and $\uparrow$0.6% $\sim$ $\uparrow$2.6% over baselines under$<$1G MACs, 1$\sim$3G MACs, 3$\sim$10G MACs, and $>$10G MACs, respectively. Note that we calculate improvements under comparable MACs. In addition to the overall comparison, we also visualize the apple-to-apple benchmark, e.g., Castling-LeViT vs. LeViT, to validate the effectiveness of proposed techniques. As shown in Fig. 5, Castling-LeViT/MViTv2 achieves $\downarrow$25.7% $\sim$ $\downarrow$55.3% MACs reductions under comparable accuracies or offers a comparable or better accuracy ($\uparrow$0.6% $\sim$ $\uparrow$1.8%) under comparable MACs over corresponding LeViT/MViTv2 baselines.

Object Detection. We also extend the Castling-ViT to the downstream object detection task and compare it with previous efficient detectors on COCO dataset to evaluate its efficacy. Specifically, we construct the detector with modified ESNet or LCNet (replace the last one or two stages with transformer blocks) as backbones and follow PicoDet ’s detection head design as well as their training recipe. Fig. 1 and Tab. 3 show the overall comparison between the proposed Castling-ViT and other baselines. We can see that our Castling-ViT consistently achieves better accuracy-efficiency tradeoffs, leading to $\uparrow$6.0, $\uparrow$2.2 $\sim$ $\uparrow$2.3, $\uparrow$4.0 $\sim$ $\uparrow$5.9, $\uparrow$3.1 $\sim$ $\uparrow$4.0 mAP improvements as compared to YOLOv5, YOLOX , MobileDet , and FBNetv5 , respectively, under comparable or even less MACs. As for EfficientDet , our method achieves comparable accuracy-efficiency trade-offs. Apart from the overall comparison with baselines, we also provide the apple-to-apple comparison with detectors with softmax-based attention, we supply the comparison results and Castling-ViT’s breakdown analysis to Sec. 4.4. This set of experiments validate the effectiveness of the proposed Castling-ViT for servering as efficient detector backbones (e.g., $\leq$3G MACs) in the object detection task.

Semantic Segmentation. We further extend Castling-ViT to the semantic segmentation task to evaluate its effectiveness. Specifically, we use ViT-Base as the backbone in Mask2former framework to serve as our baseline and testbed. Then we build the Castling-ViT-Base as the backbone to benchmark on the ADE20K dataset. As shown in Tab. 4, Mask2former with Castling-ViT backbone achieves 15% total MACs reductions and 19% backbone MACs reduction under comparable or slightly better mIoU, i.e., $\uparrow$0.13% and $\uparrow$0.52% without or with MAE pretraining on ImageNet . This set of experiments validate that our proposed techniques could be well generalized to various downstream tasks that require large input resolutions.

3 Linear-Angular Attention over SOTA Baselines

We also conduct ablation studies among various kinds of kernels used in linear attention to evaluate the superiority of our proposed linear-angular kernel. Also shown in Tab. 5, we compare the angular kernel with five other commonly adopted kernels, results on three model and resolution settings consistently demonstrate that our proposed angular kernel helps achieve better mAP, e.g., $\uparrow$4.0% $\sim$ $\uparrow$4.2% over MC , $\uparrow$1.0% $\sim$ $\uparrow$1.2% over Softmax , $\uparrow$2.6% $\sim$ $\uparrow$3.2% over Cosine , $\uparrow$3.1% $\sim$ $\uparrow$4.6% and $\uparrow$0.2% $\sim$ $\uparrow$1.2% over ReLU-S and ReLU-E , respectively. This set of experiments validate the superiority of the proposed linear-angular kernels.

4 Ablation Studies of Castling-ViT

We conduct ablation studies on Castling-ViT’s linear-angular attention, added DWConv, as well as auxiliary masked softmax-based attention, as shown in Tab. 6 and Tab. 7, where the experiments are performed on Castling-ViT-S-320⋆/M-416⋆, respectively, without pretraining on ImageNet. Note that here ⋆ means that we use LCNet-ViT, i.e., replacing all layers in last two stages with transformer blocks, as backbones. Results in three tables show detailed performance breakdown and consistently demonstrate that all components in our proposed Castling-ViT contribute to the final performance. Specifically, linear-angular attention itself already achieves $\downarrow$0.1% $\sim$ $\uparrow$0.5% mAP improvements while leading to 10.1% $\sim$ 15.3% MACs reductions simultaneously. Adding DWConv in the middle of MLP layers further leads to $\downarrow$0.1% $\sim$ $\uparrow$0.6% mAP improvements while only incur negligible MACs overhead. Adding auxiliary masked softmax-based attention further increases the mAP by 0.2% without incurring any overhead since it only assists the training while being removed during inference. Note that although LCNet-ViT backbone leads to better mAP, it also introduces more MACs as compared to efficient convolutions. As such, LCNet-ViT still slightly underperforms LCNet backbones in terms of accuracy-efficiency trade-offs, since we directly replace the CNN backbones with transformer blocks, whose architecture could not be optimal for ViTs. It remains to an open problem to build ViT based models that achieves higher accuracy-efficiency trade-off than pure ConvNet models.

5 Discussion on the Auxiliary Branch

Trajectory of Nonzeros in Masks. To further understand the effect of the auxiliary branch, we visualize the trajectories of nonzeros in masks of auxiliary softmax-based attention. As shown in Fig. 6, we count the nonzeros in the first attention layer’s masks throughout the training of both Castling-ViT-S-416 and Castling-ViT-M-416 with the LCNet-ViT backbone on COCO. we observe that the introduced auxiliary attention will only assist the training in the early or middle training stages and will gradually vanish towards all zeros in the later training stage, which is well aligned with the assumption in our Castling-ViT. This set of experiments validate the idea of performing “castling” in ViTs, i.e., drop the auxiliary branch without sacrificing performance.

Conjecture of Castling. We conduct a synthetic experiment to give an analogy for explaining the castling phenomenon. In Fig. 7 (Left), we visualize the curve of (1) angular similarity function (denoted in Red), (2) linear-angular terms (denoted in Green), and (3) two-layer DNN’s approximation. We see that only keeping linear-angular terms leads to distortion, while DNN is capable of learning the missing high-frequency parts, whose loss trajectory is shown in Fig. 7 (Middle). Also, the learning trajectory of frequency components is visualized in Fig. 7 (Right), following F-Principle , where x-axis is training steps, y-axis categorizes both low and high frequency components, blue/red colors refer to large/small difference between the learned frequency components and target frequency components, respectively. We observe that DNNs fit target functions from low to high frequencies during training. It indicates that DWConv itself is not sufficient for approximating the high-order residual at early training stages. Costly attention that contains high-order residuals is then desired to help training while being dropped at inference since the remaining networks can gradually learn the high-frequency components at later training stages .

Conclusion

We present Castling-ViT that trains ViTs with both linear-angular and softmax-based quadratic attention but switches to only having the former during inference. Castling-ViT leverages angular kernels to measure the similarities between the queries and keys via spectral angles and highlights two enablers: (1) a new linear-angular attention mechanism: we decompose angular kernels to linear terms and high-order residuals, and keep only the former for inference; and (2) we approximate the high-order residuals using a depthwise convolution and an auxiliary masked softmax attention whose masks gradually become zeros during training without incurring inference overhead. Extensive experiments consistently validate Castling-ViT’s advantages.

Acknowledgment

The work is supported in part by the National Science Foundation (NSF) RTML program (Award number: 1937592) and the CoCoSys, one of the seven centers in JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA.

References

Appendix A More Literature Review

Efficient ViTs. As previously mentioned in Sec. 2, efficient attention can be roughly categorized into two groups: local attention or kernel-based linear attention . For local attention, Swin restricts the window size of self-attention, so that only neighboring tokens will perform similarity measurements each other instead of all tokens; MaxViT also adopts block attention within windows but additionally takes dilated global attention into account for learning both local and global information; QnA shares the attention queries among all tokens; Linformer approximates the queries and keys with low-rank factorization to reduce their vector length. For kernel-based linear attention, XCiT proposes a “transposed” version of self-attention that operates across feature channels rather than tokens, resulting in linear complexity with the number of tokens; Linformer explores a low-rank matrix to approximate the self-attention; Reformer replaces self-attention by one that uses locality-sensitive hashing, changing its complexity from $\mathcal{O}(n^{2})$ to $\mathcal{O}(n\log(n))$ where $n$ denotes the number of tokens; Longformer combines a windowed local-context self-attention and a task-motivated global attention that encodes inductive bias about that task; Nystromformer adapts the Nystrom method to approximate standard self-attention with $\mathcal{O}(n)$ complexity; Scatterbrain unifies both low-rank approximation and sparse attention to improve accuracy-efficiency tradeoffs. Different from all above works, we explore from a new perspective by taking spectral angles into consideration when measuring the similarity among tokens, resulting in linear-angular attention with sparse training techniques that can achieve comparable or even better performance than softmax-based attention.

ViTs for Downstream Tasks. Apart from image classification tasks, ViTs have also been leveraged to serve as backbones for downstream tasks, such as object detection and semantic segmentation . For example, DETR directly detects and predicts objects by combining a common CNN with a transformer architecture; Maskformer proposes to use a simple mask classification model to predict a set of binary masks, each associated with a single global class label prediction. One big difference is that ViTs can beat CNNs on classification tasks that have a lower image resolution while are still less efficient than lightweight CNNs on downstream tasks that heavily rely on multi-scale resolution features. Therefore, there have been various debates on designing powerful ViT backbones: (1) plain ViTs (e.g., ViTDet ) or hierarchical ViTs (e.g., MViTv2 , or Swin )? Plain ViTs win in terms of simplicity but could be hard to scale down to lower resolution and computation regimes; Hierarchical ViTs seamlessly match with feature pyramid networks (FPNs) for extracting multi-scale feature maps but have more design factors to be considered or searched over. (2) pure ViTs or hybrid CNN-ViTs? Pure ViTs are compatible with self-supervised masked autoencoder (MAE) pretraining ; Hybrid CNN-ViTs can suffer from the information leakage problem when adopting MAE pretraining, while being more efficient especially for feature extractions in early layers. Our proposed method does not fall into the aforementioned debates. Instead, it is compatible with all ViT variants relying on the softmax-based attention.

Appendix B More Results and Clarification

Improvement from Our Training Recipe. Recall that in Sec. 4, we conduct experiments on three classical computer vision tasks. For object detection and semantic segmentation, we follow the baseline’s training recipe for a fair and direct comparison. For the image classification, our training recipe has a minor difference due to the increased batch size and training epochs with more GPU nodes. As such, we further provide the detailed improvement breakdown here. Specifically, our adopted training recipe leads to $\uparrow$0.2% $\sim$ $\uparrow$1.6% top-1 accuracy improvements and our Caslting-ViT further reduces up to $\downarrow$40% MACs and increases $\uparrow$0.1% $\sim$ $\uparrow$1.2% top-1 accuracy simultaneously.

Ablation Studies on Image Classification. Our ablation studies are mostly done on the detection task as shown in Sec. 4.4 because of its less training time as compared to training ImageNet. Note that for these ablation studies, we do not adopt pretraining on ImageNet as specified in Sec. 4.4. After finishing the trial-and-error and when it comes to comparing with SOTA works, we then pretrain final models with the training recipe the same as LeViT , resulting in final results in Tab. 3. According to our experiments, training ImageNet takes nearly one week, while training COCO without pretraining on ImageNet takes only one day. In fact, ablation results on the classification task are consistent. To deliver more comprehensive ablation studies, we train Castling-LeViT-256 on ImageNet afterwards and find that: (1) + Lin.: 81.5%; (2) + Lin. & DWConv: 82.4%; (3) + Lin. & DWConv & SparseAttn: 82.6%, those results are consistent with our observation on detection experiments.

Conjecture of Why Linear-Angular Attention Sometimes Beats the Original Self-Attention. To better understand why the result of our Castling-ViT is even better than softmax-based ViTs. We summarize three differences between our method and previous linear attentions: (1) In addition to linear attention, we also take DWConv and sparse softmax-based attention into the training process; (2) We use a SGD optimizer instead of Adam, which is not common for training ViTs. Although Adam optimizer leads to faster convergence, we find that SGD optimizer helps to deliver better results if being trained sufficiently converged, e.g., we train 1000 epochs on ImageNet; (3) After revisiting the attention design, we remove token/feature pooling and adopt post-$\mathbf{Q}$ pooling and residual connections in our attention blocks. All above three differences contribute to the the final accuracy apart from the improvement of using linear-angular attention. We also show the breakdown analysis for each of these three points, see Sec. 4.4, Sec. B, and Sec. 3.1 for detailed analysis, respectively.

Actual Latency, Throughputs, and Memory Measurements. Our final models are dense and thus well compatible with GPUs. We measure and report the latency ($\downarrow$55%), throughputs ($\uparrow$16 $\sim$ $\uparrow$100%), and GPU memory ($\downarrow$16 $\sim$ $\downarrow$61%) for both classification and detection tasks, as shown in Tab. 8/10. For throughputs, we measure both our Castling-ViT and baselines under their maximum allowed batch sizes (bs), i.e., bs=16/2/1, for different input resolutions $r$=512/1024/1536 in a fair and consistent V100 environment. Note that when the input resolution $r$=224, our models cannot beat the baseline in terms of throughputs because of (1) the newly added DWConv; (2) the removal of token/feature pooling. However, in terms of accuracy-efficiency tradeoffs, our Castling-ViT consistently beats all baselines as shown in Sec. 4. For memory, we record the peak memory per image. For latency, we benchmark with SOTA CNN-based detectors. Our model achieves 37.3mAP at 3.9ms latency on a V100, while YOLOv5-S only achieves 36.7mAP at 8.7ms latency). Moreover, Castling-ViT wins more throughputs (up to $\uparrow$6.7$\times$) for smaller patch sizes and/or larger input resolutions, as shown in Tab. 9 and 8. Note that we record CUDA latency following the literature . All reported results are averaged among three runs.

Compare with ViT-based Baselines on Detection. We benchmark with SOTA CNN-based detectors under 6G MACs in Sec. 4 because that ViT-based detectors are too expensive. For example, our Castling-ViT achieves 37.3mAP at 5.3G MACs, while RetinaNet+PVT-Tiny only achieves 36.7mAP at even 221G MACs , as shown in 10.

Advantages of Angular Kernels? Angular kernels take into account extra spectral characteristics and enjoy good properties, e.g., positive semi-definite function $\rightarrow$ inner product in a high-dimensional and rich feature space, as analyzed in Sec. 3.2. It also achieves comparable accuracy with vanilla attention as validated by Sec. 4.

Large-Scale Ablation Studies on Attention Design. We use small ViTs for idea validation in Tab. 5 and the conclusion generalizes to larger ones. Here we add another ablation study on a larger model LeViT-384 as shown in Tab. 11, from which we see that the attention design insights consistently generalize from small models to larger models, further validating our design insights.

Why More Parameters Than Others in Low MACs? ViTs tend to have more parameters than CNNs under small MACs, e.g., LeViT and Efficient-ViT . For the LeViT, it features more layers with gradually downsampled input resolutions. For example, LeViT-256 requires 18.9M parameters at only 1.1G MACs, LeViT-384 requires 39.1M parameters at 2.4G MACs. Since we adopt LeViT-like structure to construct our Castling-ViT on image classification tasks, the parameter looks higher than other else baselines. Also, as indicated in Sec. 3.1, Castling-LeViT uses merely post-$\mathbf{Q}$ pooling, causing slightly higher hidden dimensions for $\mathbf{Q}$/$\mathbf{K}$ than LeViT. In this work, we focus more on the FLOPs/latency instead of parameters since storage is not a major concern in modern hardware .

Will Auxiliary Attention and DWConv Work for Existing Linear Attentions? Yes, we train a DeiT-Tiny (w/o distill.; Acc.: 72.2%) w/ linear attention for 300 epochs and observe that: (1) + Lin.: 68.3%; (2) + Lin. & DWConv: 71.7%; (3) + Lin. & SparseAttn: 70.2%; (4) + Lin. & DWConv & SparseAttn: 72.4%.

Clarify ReLU-S vs. ReLU-E in Tab. 5. During approximation, i.e., $\text{Sim}(\mathbf{Q},\mathbf{K})\approx\phi(\mathbf{Q})\phi(\mathbf{K})^{T}$, both of them use ReLU as $\phi(\cdot)$, but ReLU-S takes the whole Softmax as $\text{Sim}(\cdot)$, while ReLU-E takes the $\text{Exp}(\cdot)$ as $\text{Sim}(\cdot)$, e.g., Efficient-ViT, resulting in additional divisions.