Expediting Large-Scale Vision Transformer for Dense Prediction without Fine-tuning
Abstract
Vision transformers have recently achieved competitive results across various vision tasks but still suffer from heavy computation costs when processing a large number of tokens.
Many advanced approaches have been developed to reduce the total number of tokens in large-scale vision transformers, especially for image classification tasks.
Typically, they select a small group of essential tokens according to their relevance with the [ class ] token, then fine-tune the weights of the vision transformer.
Such fine-tuning is less practical for dense prediction due to the much heavier computation and GPU memory cost than image classification.
In this paper, we focus on a more challenging problem, i.e. , accelerating large-scale vision transformers for dense prediction without any additional re-training or finetuning.
In response to the fact that high-resolution representations are necessary for dense prediction, we present two non-parametric operators, a token clustering layer to decrease the number of tokens and a token reconstruction layer to increase the number of tokens.
The following steps are performed to achieve this: (i) we use the token clustering layer to cluster the neighboring tokens together, resulting in low-resolution representations that maintain the spatial structures; (ii) we apply the following transformer layers only to these low-resolution representations or clustered tokens; and (iii) we use the token reconstruction layer to re-create the high-resolution representations from the refined low-resolution representations.
The results obtained by our method are promising on five dense prediction tasks, including object detection, semantic segmentation, panoptic segmentation, instance segmentation, and depth estimation.
Accordingly, our method accelerates 40% ↑
FPS and saves 30% ↓ GFLOPs of “Segmenter+ViT-L/ 16 ” while maintaining 99 .
5% of the performance on ADE 20 K without fine-tuning the official weights.
Introduction
Transformer [67] has made significant progress across various challenging vision tasks since pioneering efforts such as DETR [4], Vision Transformer (ViT)
[18], and Swin Transformer [47].
By removing the local inductive bias [19] from convolutional neural networks [28, 64, 60], vision transformers armed with global self-attention show superiority in scalability for large-scale models and billion-scale dataset [18, 84, 61], self-supervised learning [27, 76, 1], connecting vision and language [53, 34], etc.
We can find from recent developments of SOTA approaches that vision transformers have dominated various leader-boards, including but not limited to image classification [74, 14, 17, 84], object detection [86, 46, 38], semantic segmentation [32, 11, 3], pose estimation [77], image generation [85], and depth estimation [42].
Although vision transformers have achieved more accurate predictions in many vision tasks, largescale vision transformers are still burdened with heavy computational overhead, particularly when processing high-resolution inputs [20, 46], thus limiting their broader application to more resourceconstrained applications and attracting efforts on re-designing light-weight vision transformer architectures [10, 51, 87].
In addition to this, several recent efforts have investigated how to decrease the model complexity and accelerate vision transformers, especially for image classification, and introduced various advanced approaches to accelerate vision transformers.
Dynamic ViT [55] and
EViT [44], for example, propose two different dynamic token sparsification frameworks to reduce the redundant tokens progressively and select the most informative tokens according to the scores predicted with an extra trained prediction module or their relevance with the [ class ] token.
TokenLearner [58] learns to spatially attend over a subset of tokens and generates a set of clustered tokens adaptive to the input for video understanding tasks.
Most of these token reduction approaches are carefully designed for image classification tasks and require fine-tuning or retraining.
These approaches might not be suitable to tackle more challenging dense prediction tasks that need to process high-resolution input images, e.g., 1024 × 1024 , thus, resulting in heavy computation and
GPU memory cost brought.
We also demonstrate in the supplemental material the superiority of our method over several representative methods on dense prediction tasks.
Rather than proposing a new lightweight architecture for dense prediction or token reduction scheme for only image classification, we focus on how to expedite well-trained large-scale vision transformers and use them for various dense prediction tasks without fine-tuning or re-training.
Motivated by these two key observations including (i) the intermediate token representations of a well-trained vision transformer carry a heavy amount of local spatial redundancy and (ii) dense prediction tasks require high-resolution representations, we propose a simple yet effective scheme to convert the “high-resolution” path of the vision transformer to a “high-to-low-to-high resolution” path via two non-parametric layers including a token clustering layer and a token reconstruction layer.
Our method can produce a wide range of more efficient models without requiring further fine-tuning or re-training.
We apply our approach to expedite two main-stream vision transformer architectures, e.g., ViTs and
Swin Transformers, for five challenging dense prediction tasks, including object detection, semantic segmentation, panoptic segmentation, instance segmentation, and depth estimation.
We have achieved encouraging results across several evaluated benchmarks and Figure 1 illustrates some representative results on both semantic segmentation and depth estimation tasks.
Related work
Pruning Convolutional Neural Networks.
Convolutional neural network pruning [2, 30, 71] is a task that involves removing the redundant parameters to reduce the model complexity without a significant performance drop.
Pruning methods typically entail three steps: (i) training a large, over-parameterized model to convergence, (ii) pruning the trained large model according to a certain criterion, and (iii) fine-tuning the pruned model to regain the lost performance [49].
The key idea is to design an importance score function that is capable of pruning the less informative parameters.
We follow [7] to categorize the existing methods into two main paths: (i) unstructured pruning (also named weight pruning) and (ii) structured pruning.
Unstructured pruning methods explore the absolute value of each weight or the product of each weight and its gradient to estimate the importance scores.
Structured pruning methods, such as layer-level pruning [72], filter-level pruning[48, 79], and image-level pruning [25, 65], removes the model sub-structures.
Recent studies [5, 6, 31] further extend these pruning methods to vision transformer.
Unlike the previous pruning methods, we explore how to expedite vision transformers for dense prediction tasks by carefully reducing & increasing the number of tokens without removing or modifying the parameters.
Efficient Vision Transformer.
The success of vision transformers has incentivised many recent efforts [58, 62, 50, 29, 73, 55, 37, 24, 69, 9, 35, 44, 57] to exploit the spatial redundancies of intermediate token representations.
For example, TokenLearner [58] learns to attend over a subset of tokens and generates a set of clustered tokens adaptive to the input.
They empirically show that very few clustered tokens are sufficient for video understanding tasks.
Token Pooling [50] exploits a nonuniform data-aware down-sampling operator based on K-Means or K-medoids to cluster similar [Figure] together to reduce the number of
ViT [55] observes that the accurate image recognition with vision transformers mainly depends on a subset of the most informative tokens, and hence it develops a dynamic token sparsification framework for pruning the redundant tokens dynamically based on the input.
EViT (expediting vision transformers) [44] proposes to calculate the attentiveness of the [ class ] token with respect to each token and identify the top- k attentive tokens according to the attentiveness score.
Patch
Merger [56] uses a learnable attention matrix to merge and combine together the redundant tokens, therefore creating a much more practical and cheaper model with only a slight performance drop.
Refer to [66] for more details on efficient transformer architecture designs, such as Performer [12] and Reformer [36].
In contrast to these methods that require either retraining or fine-tuning the modified transformer architectures from scratch or the pre-trained weights, our approach can reuse the once-trained weights for free and produce lightweight models with a modest performance drop.
Vision Transformer for Dense Prediction.
In the wake of success of the representative pyramid vision transformers [47, 70] for object detection and semantic segmentation, more and more efforts have explored different advanced vision transformer architecture designs [39, 8, 40, 41, 21, 81, 78, 88, 23, 43, 75, 80, 16, 83, 82, 26] suitable for various dense prediction tasks.
For example, MViT [40] focuses more on multi-scale representation learning, while HRFormer [81] examines the benefits of combining multi-scale representation learning and high-resolution representation learning.
Instead of designing a novel vision transformer architecture for dense prediction, we focus on how to accelerate a well-trained vision transformer while maintaining the prediction performance as much as possible.
Our approach.
The contribution of our work lies in two main aspects: (i) we are the first to study how to accelerate state-of-the-art large-scale vision transformers for dense prediction tasks without fine-tuning (e.g., "Mask2Former + Swin-L" and "SwinV2-L + HTC++").
Besides, our approach also achieves much better accuracy and speedup trade-off when compared to the very recent ACT [1] which is based on a clustering attention scheme; (ii) our token clustering and reconstruction layers are capable of maintaining the semantic information encoded in the original high-resolution representations.
This is the very most important factor to avoid fine-tuning.
We design an effective combination of a token clustering function and a token reconstruction function to maximize the cosine similarity between the reconstructed high-resolution feature maps and the original ones without fine-tuning.
The design of our token reconstruction layer is the key and not straightforward essentially.
We also show that our token reconstruction layer can be used to adapt the very recent EViT [44] and
DynamicViT [55] for dense prediction tasks in the supplementary.
[Figure]
Our Approach
Preliminary.
The conventional Vision Transformer [18] first reshapes the input image X ∈ R H × W ×into a sequence of flatten patches X p ∈ R N ×
( P 2 3 ) , where ( P , P ) represents the resolution of each patch, ( H , W ) represents the resolution of the input image, N =
( H × W ) /
P 2 represents the number of resulting patches or tokens, i.e. , the input sequence length.
The Vision Transformer consists of alternating layers of multi-head self-attention ( MHSA ) and feed-forward network ( FFN ) accompanied with layer norm ( LN ) and residual connections:
[Equation]
(1) where l ∈ { 1 , . . .
, L } represents the layer index, Z l ∈ R N × C , and Z 0 is based on X p .
The computation cost O ( LNC ( N + C )) mainly depends on the number of layers L , the number of tokens N , and the channel dimension C .
Despite the great success of transformer, its computation cost increases significantly when handling high-resolution representations, which are critical for dense prediction tasks.
This paper attempts to resolve this issue by reducing the computation complexity during the inference stage, and presents a very simple solution for generating a large number of efficient vision transformer models directly from a single trained vision transformer, requiring no further training or fine-tuning.
We demonstrate how our approach could be applied to the existing standard Vision Transformer in
Figure 2.
The original Vision Transformer is modified using two non-parametric operations, namely a token clustering layer and a token reconstruction layer.
The proposed token clustering layer is utilized to convert the high-resolution representations to low-resolution representations by clustering the locally semantically similar tokens.
Then, we apply the following transformer layers on the low-resolution representations, which greatly accelerates the inference speed and saves computation resources.
Last, a token reconstruction layer is proposed to reconstruct the feature representations back to high-resolution.
Token Clustering Layer.
We construct the token clustering layer following the improved SLIC scheme [33], which performs local k-means clustering as follows: - Initial superpixel center : We
apply adaptive average pooling ( AAP
) over the high-resolution representations from the α -th layer to compute the h × w
initial cluster center
representations
[Equation] where S α ∈ R hw × C , Z α ∈ R N × C , and hw ≪ N .
- Iterative local clustering : (i) Expectation step: compute the normalized similarity between each pixel p and the surrounding superpixel
i (we only consider the neighboring λ positions), (ii) Maximization step: compute the new superpixel centers:
[Equation]
where we iterate the above Expectation step and Maximization step for κ times, τ is a temperature hyper-parameter, and i ∈ { 1 , 2 , · · · , λ } .
We apply the following β transformer layers on S α instead of Z α , thus results in S α + β and decreases the computation cost significantly.
Token Reconstruction Layer.
We implement the token reconstruction layer by exploiting the relations between the high-resolution representations and the low-resolution clustered representations:
[Equation]
) where τ is the same temperature hyper-parameter as in Equation 3.
k-NN ( Z α,p ) represents a set of the k nearest, a.k.a, most similar, superpixel representations for Z α,i .
We empirically find that choosing the same neighboring positions as in Equation 3 achieves close performance as the k-NN scheme while being more easy to implementation.
In summary, we estimate their semantic relations based on the representations before refinement with the following β transformer layers and then reconstruct the high-resolution representations from the refined low-resolution clustered representations accordingly.
Finally, we apply the remained γ transformer layers to the reconstructed high-resolution features and the task-specific head on the refined high-resolution features to predict the target results such as semantic segmentation maps or monocular depth maps.
Extension to Swin Transformer.
We further introduce the window token clustering layer and window token reconstruction layer , which are suitable for Swin Transformer [46, 47].
Figure 3 illustrates an example usage of the proposed window token clustering layer and window token reconstruction layer.
We first cluster the K × K window tokens into k × k window tokens and then reconstruct K × K window tokens according to the refined k × k window tokens.
We apply the swin transformer layer equipped with smaller window size k × k on the clustered representations, where we need to bi-linear interpolate the pre-trained weights of relative position embedding table from (2 K − 1) ×
(2 K − 1) to (2 k − 1) × (2 k − 1) when processing the clustered representations.
In summary, we can improve the efficiency of Swin Transformer by injecting the window token clustering layer and the window token reconstruction layer into the backbones seamlessly without fine-tuning the model weights. Why our approach can avoid fine-tuning?
The reasons include the following two aspects: (i) our token clustering/reconstruction layers are non-parametric, thus avoiding retraining any additional parameters, (ii) the reconstructed high-resolution representations maintain high semantic similarity with the original high-resolution representations.
We take Segmenter+ViT-L/ 16 (on ADEK, α = 10 ) as an example and analyze the semantic similarity between the reconstructed high-resolution feature (with our approach) and the original high-resolution feature (with the original ViT-L/ 16 ) in Table 1.
Accordingly, we can see that the cosine similarities are consistently high across different transformer layers between the reconstructed high-resolution feature (with our approach) and the original highresolution feature.
In other words, our approach well maintains the semantic information carried in the original high-resolution feature maps and thus is capable of avoiding fine-tuning.
[Table]  Experiment
We verify the effectiveness of our method across five challenging dense prediction tasks, including object detection, semantic segmentation, instance segmentation, panoptic segmentation, and monocular depth estimation.
We carefully choose the advanced SOTA methods that build the framework based on either the plain ViTs [18] or the Swin Transformers [46, 47].
We can integrate the proposed token clustering layer and token reconstruction layer seamlessly with the provided official trained checkpoints with no further fine-tuning required.
More experimental details are illustrated as follows.
4.1 Datasets
COCO [45].
This dataset consists of 123 K images with 896 K annotated bounding boxes belonging to 80 thing classes and 53 stuff classes, where the train set contains 118 K images and the val set contains 5 K images.
We report the object detection performance of SwinV2 + HTC++ [46] and the instance/panoptic segmentation performance of Mask 2 Former [11] on the val set.
ADE 20 K [90].
This dataset contains challenging scenes with fine-grained labels and is one of the most challenging semantic segmentation datasets.
The train set contains 20 , 210 images withsemantic classes.
The val set contains 2 , 000 images.
We report the segmentation results with
Segmenter [63] on the val set.
PASCAL-Context [52].
This dataset consists of 59 semantic classes plus a background class, where the train set contains 4 , 996 images with and the val set contains 5 , 104 images.
We report the segmentation results with Segmenter [63] on the val set.
Cityscapes [13].
This dataset is an urban scene understanding dataset with 30 classes while onlyclasses are used for parsing evaluation.
The train set and val set contains 2 , 975 and 500 images respectively.
We report the segmentation results with Segmenter [63] on the val set.
KITTI [22].
This dataset provides stereo, optical flow, visual odometry (SLAM), and 3D object detection of outdoor scenes captured by equipment mounted on a moving vehicle.
We choose
DPT [54] as our baseline to conduct experiments on the monocular depth prediction tasks, which consists of around 26 K images for train set and 698 images for val set, where only 653 images have the ground-truth depth maps and the image resolution is of 1 , 241 × 376 .
NYUv 2 [59].
This dataset consists of 1 , 449 RGBD images with resolution of 640 × 480 , which captures 464 diverse indoor scenes and contains rich detailed dense annotations such as surface normals, segmentation maps, depth, 3D planes, and so on.
We report the depth prediction results of
DPT [54] evaluated on 655 val images. [Table] [Table] [Table]
[Table]
[Table] [Figure] [Figure] 4.2 Evaluation Metrics
We report the numbers of AP (average precision), mask AP (mask average precision), PQ (panoptic quality), mIoU (mean intersection-over-union), and RMSE (root mean squared error) across object detection, instance segmentation, panoptic segmentation, semantic segmentation, and depth estimation tasks respectively.
Since the dense prediction tasks care less about the throughput used in image classification tasks [44], we report FPS to measure the latency and the number of GFLOPs to measure the model complexity during evaluation.
FPS is tested on a single V100 GPU with Pytorch 1.10 and
CUDA 10.2 by default.
More details are provided in the supplementary material.
4.3 Ablation Study Experiments
We conduct the following ablation experiments on ADE 20 K semantic segmentation benchmark with the official checkpoints of Segmenter+ViT-L/ 16 [63] 3 by default if not specified.
Hyper-parameters of token clustering/reconstruction layer .
We first study the influence of the hyper-parameters associated with the token clustering layer, i.e., the number of neighboring pixels λ used in Equation 3, the number of EM iterations κ , and the choice of the temperature τ in Table 2.
According to the results, we can see that our method is relatively less sensitive to the choice of both λ and κ compared to τ .
In summary, we choose λ as 5 × 5 , κ as 5 , and τ as 50 considering both performance and efficiency.
Next, we also study the influence of the hyper-parameters within the token clustering layer, i.e., the number of nearest neighbors k within k-NN.
We do not observe obvious differences and thus set k as 20 .
More details are provided in the supplementary material.
Influence of cluster size choices .
We study the influence of different cluster size h × w choices based on input feature map of size HP × WP =
× 40 ( N = 1 , 600 ) 4 in Table 3.
According to the results, we can see that choosing too small cluster sizes significantly harms the dense prediction performance, and setting h × w as 28 × 28 achieves the better trade-off between performance drop and model complexity.
Therefore, we choose 28 × 28 on Segmenter+ViT-L/ 16 by default.
We also empirically find that selecting the cluster size h × w around N / 4 ∼ N / 2 performs better on most of the other experiments.
We conduct the following ablation experiments under two typical settings, including 20 × 20 ( ∼ N / 4 ) and 28 × 28 ( ∼ N / 2 ).
Comparison with adaptive average pooling and bi-linear upsample .
We report the comparison results between our proposed token clustering/reconstruction scheme and adaptive average pooling/bilinear upsample scheme in Table 4 and Table 5 under two cluster size settings respectively.
We choose to compare with adaptive average pooling and bi-linear upsample instead of strided convolution or deconvolution as the previous ones are non-parametric and the later ones require re-training or fine-tuning, which are not the focus of this work.
Specifically, we keep the inserted position choices the same and only replace the token cluster or token reconstruction layer with adaptive average pooling or bi-linear upsampling under the same cluster size choices.
According to the results, we can see that our proposed token clustering and token reconstruction consistently outperform adaptive average pooling and bi-linear upsampling under different cluster size choices.
Influence of inserted position of token clustering/reconstruction layer .
We investigate the influence of the inserted position of both token clustering layer and token reconstruction layer and summarize the detailed results in Figure 4 and Figure 5 under two different cluster size choices.
According to the results shown in Figure 4, our method achieves better performance when choosing α larger than 10 , therefore, we choose α = 10 as it achieves a better trade-off between model complexity and segmentation accuracy.
Then we study the influence of the inserted positions of the token reconstruction layer by fixing α = 10 .
According to Figure 5, we can see that our method achieves the best performance when setting α + β = 24 , in other words, we insert the token reconstruction layer after the last transformer layer of ViT-L/ 16 .
We choose α = 10 , α + β = 24 , and γ = 0 for all ablation experiments on ADE 20 K by default if not specified.
Combination with lighter vision transformer architecture .
We report the results of applying our method to lighter vision transformer backbones such as ViT-B/ 16 , in Table 6.
Our approach consistently improves the efficiency of Segmenter+ViT-B/ 16 at the cost of a slight performance drop without fine-tuning.
Specifically speaking, our approach saves more than 26% ↓ GFLOPs of a trained “Segmenter+ViT-B/ 16 ” with only a slight performance drop from 48 .
48% to 48 .
40% , which verifies our method also generalizes to lighter vision transformer architectures.
Comparison with uniform downsampling .
We compare our method with the simple uniform downsampling scheme, which directly downsamples the input image into a lower resolution.
Figuresummarizes the detailed comparison results.
For example, on ADE 20 K, we downsample the input resolution from 640 × 640 to smaller resolutions ( e.g. , 592 × 592 , 576 × 576 , 560 × 560 , and× 544 ) and report their performance and GFLOPs in Figure 6.
We also plot the results with our method and we can see that our method consistently outperforms uniform sampling on both ADE 20 K and PASCAL-Context under multiple different GFLOPs budgets.
4.4 Object Detection
We use the recent SOTA object detection framework SwinV 2 -L + HTC++ [46] as our baseline.
We summarize the results of combining our method with SwinV 2 -L + HTC++ on COCO object detection and instance segmentation tasks in Figure 7.
Implementation details.
The original SwinV 2 -L consists of { 2 , 2 , 18 , 2 } shifted window transformer blocks across the four stages.
We only apply our method to the 3 -rd stage with 18 blocks considering it dominates the computation overhead, which we also follow in the “Mask 2 Former + Swin-L” experiments.
We insert the window token clustering/reconstruction layer after the 8 -th/ 18 -th block within the 3 -rd stage, which are based on the official checkpoints 5 of SwinV 2 -L + HTC++.
In other words, we set α = 12 and α + β = 22 for SwinV 2 -L. The default window size is K × K = 32 × 32 and we set the clustered window size as k × k =
× 23 .
We choose the values of other hyperparameters following the ablation experiments.
According to the results summarized in Figure 7, compared to
SwinV 2 -L + HTC++, our method improves the FPS by 21% ↑ and saves the GFLOPs by nearly 20% ↓ while maintaining around 98% of object detection & instance segmentation performance. [Figure] [Figure] [Figure]
Q
[Figure] [Figure] m I o U ( % ) / m a s k
A P ( % ) ↑
[Figure] [Figure] 4.5 Semantic/Instance/Panoptic Segmentation We first apply our method to a plain ViT-based segmentation framework Segmenter [63] and illustrate the semantic segmentation results across three benchmarks including ADE 20 K, PASCAL-Context, and Cityscapes on the first row of Figure 1.
Then, we apply our method to a very recent framework
Mask 2 Former [11] that is based on Swin Transformer and summarize the semantic segmentation, instance segmentation [Figure] and panoptic segmentation results on COCO in Figure 8.
Implementation details.
The original ViT-L first splits an image into a sequence of image patches of size 16 × 16 and applies a patch embedding layer to increase the channel dimensions to 1024 , then applies 24 consecutive transformer encoder layers for representation learning.
To apply our method to the ViT-L backbone of Segmenter, we use the official checkpoints 6 of “Segmenter + ViT-L/ 16 ” and insert the token clustering layers and token reconstruction layer into the ViT-L/ 16 backbone without fine-tuning.
For the Mask 2 Former built on Swin-L with window size as 12 × 12 , we use the official checkpointsof “Mask 2 Former + Swin-L” and insert the window token clustering layer and the window token reconstruction layer into the empirically chosen positions, which first cluster 12 × 12 tokens into 8 × 8 tokens and then reconstruct 12 × 12 tokens within each window.
Figure 8 summarizes the detailed comparison results.
Accordingly, we can see that our method significantly improves the FPS by more than 35% ↑ with a slight performance drop on COCO panoptic segmentation task. [Table] [Table]
4.6 Monocular Depth Estimation
To verify the generalization of our approach, we apply our method to depth estimation tasks that measure the distance of each pixel relative to the camera.
We choose the DPT (Dense Prediction
Transformer)
[54] that builds on the hybrid vision transformer, i.e., R 50 +ViT-B/ 16 , following [18].
Implementation details.
The original R 50 +ViT-B/ 16 8 consists of a ResNet 50 followed by a ViTB/ 16 , where the ViT-B/ 16 consists of 12 transformer encoder layers that process 16 × downsampled representations.
We insert the token clustering layer & token reconstruction layer into ViT-B/ 16 and summarize the results on both KITTI and NYUv 2 on the second row of Figure 1.
We also report their detailed depth estimation results in Table 7, where we can see that our method accelerates DPT by nearly 30% / 37% ↑ on KITTI/NYUv 2 , respectively.
4.7 ImageNet- 1 K Classification
Finally, we apply our method to the ImageNet- 1 K classification task and compare our method with a very recent SOTA method EViT [44].
The key idea of EViT is to identify and only keep the top- k tokens according to their attention scores relative to the [ class ] token.
We empirically find that applying EViT for dense prediction tasks directly suffers from significant performance drops.
More details are illustrated in the supplementary material.
Implementation details.
We choose the recent SWAG [61] as our baseline, which exploits 3 .
6 billion weakly labeled images associated with around 27 K categories (or hashtags) to pre-train the large-scale vision transformer models, i.e., ViT-L/ 16 and ViT-H/ 14 .
According to their official implementations 9 ,
SWAG + ViT-H/ 14 and SWAG + ViT-L/ 16 achieve 88 .
55% and 88 .
07% top- 1 accuracy on ImaegNet1K respectively.
We apply our approach and EViT to both baselines and summarize the comparison results in Figure 9.
According to the results, our method achieves comparable results as EViT while being more efficient, which further verifies that our method also generalizes to the image classification tasks without fine-tuning.
Conclusion
In this paper, we present a simple and effective mechanism to improve the efficiency of large-scale vision transformer models for dense prediction tasks.
In light of the relatively high costs associated with re-training or fine-tuning large vision transformer models on various dense prediction tasks, our study provides a very lightweight method for expediting the inference process while requiring no additional fine-tuning.
We hope our work could inspire further research efforts into exploring how to accelerate large-scale vision transformers for dense prediction tasks without fine-tuning.
Acknowledgement
This work is partially supported by the National Nature Science Foundation of China under Grant 62071013 and 61671027 , and National Key R&D Program of China under Grant 2018 AAA 0100300 .
We first illustrate the overall details of our token clustering layer and token reconstruction layer in Figure 10.
We then present the example implementation of token clustering layer and token reconstruction layer based on PyTorch in Listing 1 and Listing 2, respectively.
B. More Hyper-parameter Details
We summarize the detailed hyper-parameter settings for the dense prediction methods based on plain
ViTs and Swin Transformers in Table 8 and Table 9, respectively.
Table 8 summarizes the hyper-parameters, including the inserted positions α & α + β of token clustering layer & token reconstruction layer, the number of remaining transformer layers after the token reconstruction layer γ , the total number of transformer layers L , the number of tokens before clustering HP × WP , the number of tokens after clustering h × w , the number of neighboring pixels λ , the number of EM iterations κ , the temperature value τ , and the number of nearest neighbors k , for
Segmenter, DPT, and SWAG.
Table 9 summarizes the hyper-parameters, including the inserted positions α & α + β of the window token clustering layer & window token reconstruction layer, the number of remaining transformer layers after the token reconstruction layer γ , the total number of transformer layers L , the number of window tokens before clustering K × K , the number of window tokens after clustering k × k , the number of neighboring pixels λ
We illustrate the evaluation details used for measuring the GFLOPs and FPS of different methods in
Table 10.
We choose the input resolutions for different methods with different backbones according to their official implementations.
To illustrate the effectiveness of our method more accurately, we do not include the complexity and latency brought by the especially heavy detection heads or segmentation heads within Mask 2 Former and SwinV 2 -L + HTC++.
For example, the GFLOPs of
SwinV 2 -L backbone accounts for only 56 .
7% of the whole model, therefore, we only report the
GFLOPs and FPS improvements of our method over the backbone.
D. Comparison with EViT [44] on Dense Prediction
To demonstrate the advantage of our approach over the representative method that is originally designed for the image classification tasks, i.e., EViT [44], we report the detailed comparison results in Figure 12. The original EViT propose to identify and only keep the top ρ % tokens according to their attention scores relative to the [ class ] token. Specifically, we follow the official implementations to insert the token identification module into the 8 -th, 14 -th, and 20 -th layer of ViT-L/ 16 (withlayers in total) to decrease the number of tokens by ( 1 - ρ % ), respectively.
We report the results of
EViT by choosing ρ % = 60% / 70% / 80% / 90% in Figure 12. Accordingly, we can see that our method significantly outperforms EViT across various GFLOPs & FPS settings when evaluating without either re-training or fine-tuning.
The EViT can not be used for dense prediction directly, as it only keeps around 21 . 6% ∼ 72 . 9% of the tokens at last. To reconstruct the missed token representations over the abandoned positions, we apply two different strategies, including (i) reusing the representations before the corresponding token identification module, and (ii) using our token reconstruction layer to reconstruct the missed token representations according to Figure 11a.
We empirically find the first strategy achieves much worse results, thus choosing the second strategy by default. E. Adapting DynamicViT [55] for Dense Prediction
To adapt DynamicViT for dense prediction tasks, we propose to add multiple token reconstruction layers to reconstruct high-resolution representations from the selected low-resolution representations iteratively.
Figure 11 (b) presents more details of the overall framework.
We also report the comparison results in Table 11.
F. Comparison with Clustered Attention [68], ACT [89], and SMRF [15]
We illustrate the key differences between our approach and the existing clustered attention approaches [68, 89, 15] the following two aspects: (i)
These clustering attention methods perform clustering within each multi-head self-attention layer (MHSA) independently while our approach only performs clustering once with the token clustering layer and refines the clustered representations with the following transformer layers. Therefore, our approach introduces a much smaller additional overhead caused by the clustering operation.
(ii)
These clustering attention methods only reduce the computation cost of each MHSA layer equipped with clustering attention as they maintain the high-resolution representations outside the MHSA layers while Our approach can reduce the computation cost of both MHSA layers and feed-forward network (FFN) layers after the token clustering layer. We further summarize their detailed differences and the experimental comparison resultswith ACT [89](without retraining) in Table 12 and Table 13, respectively. According to the results in Table 13, we can see that (i) ACT also achieves strong performance without retraining, (ii) our approach is a better choice considering the trade-off between performance and FPS & GFLOPs, e.g., our method achieves close performance as ACT ( 51 .
32 vs. 51 . 38 ) while running 70% faster ( 9 . 1 vs. 5 . 3 ) and saving more than 35% GFLOPs ( 388 . 2 vs. 614 . 7 ).
G. Visualization
We first present the visual comparison results of our approach in Figure 13a, which shows three different configurations over Segmenter+ViT-L/ 16 achieve 32 . 32% when setting the cluster size h × w as 8 × 8 / 16 × 16 / 24 × 24 , respectively. Then, we visualize both the original feature maps and the clustering feature maps in Figure 13b. Accordingly, we can see that the clustering feature maps, based on our token clustering layer, well maintain the overall structure information carried in the original high-resolution feature maps.
Last, to verify the redundancy in the tokens of vision transformer, we visualize the attention maps of neighboring tokens in Figure 14. [Figure] [Figure] [Figure]
be
[Figure]
[Figure] cluster_features example of [Figure] [Figure]
return features
[Table]
[Table] [Table] [Table] [Table]
[Table] [Figure]
[Figure]