Expediting Large-Scale Vision Transformer for Dense Prediction without Fine-tuning

Introduction

Transformer has made significant progress across various challenging vision tasks since pioneering efforts such as DETR , Vision Transformer (ViT) , and Swin Transformer . By removing the local inductive bias from convolutional neural networks , vision transformers armed with global self-attention show superiority in scalability for large-scale models and billion-scale dataset , self-supervised learning , connecting vision and language , 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 , object detection , semantic segmentation , pose estimation , image generation , and depth estimation .

Although vision transformers have achieved more accurate predictions in many vision tasks, large-scale vision transformers are still burdened with heavy computational overhead, particularly when processing high-resolution inputs , thus limiting their broader application to more resource-constrained applications and attracting efforts on re-designing light-weight vision transformer architectures . 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 and EViT , 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 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\times 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

Convolutional neural network pruning 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 . The key idea is to design an importance score function that is capable of pruning the less informative parameters. We follow 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 , filter-level pruning, and image-level pruning , removes the model sub-structures. Recent studies 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 to exploit the spatial redundancies of intermediate token representations. For example, TokenLearner 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 exploits a nonuniform data-aware down-sampling operator based on K-Means or K-medoids to cluster similar tokens together to reduce the number of tokens while minimizing the reconstruction error. Dynamic ViT 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) 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 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 for more details on efficient transformer architecture designs, such as Performer and Reformer . 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 for object detection and semantic segmentation, more and more efforts have explored different advanced vision transformer architecture designs suitable for various dense prediction tasks. For example, MViT focuses more on multi-scale representation learning, while HRFormer 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 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 and DynamicViT for dense prediction tasks in the supplementary.

Our Approach

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 , which performs local k-means clustering as follows:

-Initial superpixel center: We apply adaptive average pooling ($\operatorname{AAP}$) over the high-resolution representations from the $\alpha$-th layer to compute the $\mathsf{h}\times\mathsf{w}$ initial cluster center representations:

-Iterative local clustering: (i) Expectation step: compute the normalized similarity between each pixel $p$ and the surrounding superpixel $i$ (we only consider the neighboring $\lambda$ positions), (ii) Maximization step: compute the new superpixel centers:

where we iterate the above Expectation step and Maximization step for $\kappa$ times, $\tau$ is a temperature hyper-parameter, and $i\in\{1,2,\cdots,\lambda\}$. We apply the following $\beta$ transformer layers on $\mathbf{S}_{\alpha}$ instead of $\mathbf{Z}_{\alpha}$, thus results in $\mathbf{S}_{\alpha+\beta}$ 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:

where $\tau$ is the same temperature hyper-parameter as in Equation 3. $\textrm{k-NN}(\mathbf{Z}_{\alpha,p})$ represents a set of the k nearest, a.k.a, most similar, superpixel representations for $\mathbf{Z}_{\alpha,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 $\beta$ transformer layers and then reconstruct the high-resolution representations from the refined low-resolution clustered representations accordingly.

Finally, we apply the remained $\gamma$ 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 . Figure 3 illustrates an example usage of the proposed window token clustering layer and window token reconstruction layer. We first cluster the $\mathsf{K}\times\mathsf{K}$ window tokens into $\mathsf{k}\times\mathsf{k}$ window tokens and then reconstruct $\mathsf{K}\times\mathsf{K}$ window tokens according to the refined $\mathsf{k}\times\mathsf{k}$ window tokens. We apply the swin transformer layer equipped with smaller window size $\mathsf{k}\times\mathsf{k}$ on the clustered representations, where we need to bi-linear interpolate the pre-trained weights of relative position embedding table from $(2\mathsf{K}-1)<span class="katex-display"><span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.6667em;vertical-align:-0.0833em;"></span><span class="mord">×</span></span></span></span></span>(2\mathsf{K}-1)$ to $(2\mathsf{k}-1)<span class="katex-display"><span class="katex"><span class="katex-mathml"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.6667em;vertical-align:-0.0833em;"></span><span class="mord">×</span></span></span></span></span>(2\mathsf{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, $\alpha$=$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 high-resolution 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.

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 or the Swin Transformers . 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.

COCO . 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++ and the instance/panoptic segmentation performance of Mask$2$Former on the val set.

ADE$20$K . 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 with $150$ semantic classes. The val set contains $2,000$ images. We report the segmentation results with Segmenter on the val set.

PASCAL-Context . 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 on the val set.

Cityscapes . This dataset is an urban scene understanding dataset with $30$ classes while only $19$ classes 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 on the val set.

KITTI . 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 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\times 376$.

NYUv$2$ . This dataset consists of $1,449$ RGBD images with resolution of $640\times 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 evaluated on $655$ val images.

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 , 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.

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$ {https://github.com/rstrudel/segmenter##ade20k}, MIT License 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 $\lambda$ used in Equation 3, the number of EM iterations $\kappa$, and the choice of the temperature $\tau$ in Table 6. According to the results, we can see that our method is relatively less sensitive to the choice of both $\lambda$ and $\kappa$ compared to $\tau$. In summary, we choose $\lambda$ as $5\times 5$, $\kappa$ as $5$, and $\tau$ 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 $\mathsf{h}\times\mathsf{w}$ choices based on input feature map of size $\frac{\mathsf{H}}{\mathsf{P}}\times\frac{\mathsf{W}}{\mathsf{P}}=40\times 40$ ($\mathsf{N}=1,600$)We choose $\mathsf{H}\times\mathsf{W}=640\times 640$ and $\mathsf{P}=16$, thus, $\frac{\mathsf{H}}{\mathsf{P}}\times\frac{\mathsf{W}}{\mathsf{P}}=40\times 40$ or $\mathsf{N}=1,600$, on ADE$20$K. in Table 6. According to the results, we can see that choosing too small cluster sizes significantly harms the dense prediction performance, and setting $\mathsf{h}\times\mathsf{w}$ as $28\times 28$ achieves the better trade-off between performance drop and model complexity. Therefore, we choose $28\times 28$ on Segmenter+ViT-L/$16$ by default. We also empirically find that selecting the cluster size $\mathsf{h}\times\mathsf{w}$ around $\mathsf{N}/4\sim\mathsf{N}/2$ performs better on most of the other experiments. We conduct the following ablation experiments under two typical settings, including $20\times 20$ ($\sim\mathsf{N}/4$) and $28\times 28$ ($\sim\mathsf{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/bi-linear upsample scheme in Table 6 and Table 6 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 5 and Figure 5 under two different cluster size choices. According to the results shown in Figure 5, our method achieves better performance when choosing $\alpha$ larger than $10$, therefore, we choose $\alpha=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 $\alpha=10$. According to Figure 5, we can see that our method achieves the best performance when setting $\alpha+\beta=24$, in other words, we insert the token reconstruction layer after the last transformer layer of ViT-L/$16$. We choose $\alpha=10$, $\alpha+\beta=24$, and $\gamma=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\%\downarrow$ 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. Figure 9 summarizes the detailed comparison results. For example, on ADE$20$K, we downsample the input resolution from $640\times 640$ to smaller resolutions (e.g., $592\times 592$, $576\times 576$, $560\times 560$, and $544\times 544$) and report their performance and GFLOPs in Figure 9. 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 Object Detection

We use the recent SOTA object detection framework SwinV$2$-L + HTC++ 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 9.

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 https://github.com/microsoft/Swin-Transformer, MIT License of SwinV$2$-L + HTC++. In other words, we set $\alpha=12$ and $\alpha+\beta=22$ for SwinV$2$-L. The default window size is $\mathsf{K}\times\mathsf{K}$=$32\times 32$ and we set the clustered window size as $\mathsf{k}\times\mathsf{k}$=$23\times 23$. We choose the values of other hyperparameters following the ablation experiments. According to the results summarized in Figure 9, compared to SwinV$2$-L + HTC++, our method improves the FPS by $21\%\uparrow$ and saves the GFLOPs by nearly $20\%\downarrow$ while maintaining around $98\%$ of object detection & instance segmentation performance.

5 Semantic/Instance/Panoptic Segmentation

We first apply our method to a plain ViT-based segmentation framework Segmenter 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 that is based on Swin Transformer and summarize the semantic segmentation, instance segmentation, and panoptic segmentation results on COCO in Figure 9.

Implementation details. The original ViT-L first splits an image into a sequence of image patches of size $16\times 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 https://github.com/rstrudel/segmenter##model-zoo, MIT License 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\times 12$, we use the official checkpoints https://github.com/facebookresearch/Mask2Former/blob/main/MODEL_ZOO.md, CC-BY-NC $4.0$ of “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\times 12$ tokens into $8\times 8$ tokens and then reconstruct $12\times 12$ tokens within each window. Figure 9 summarizes the detailed comparison results. Accordingly, we can see that our method significantly improves the FPS by more than $35\%\uparrow$ with a slight performance drop on COCO panoptic segmentation task.

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) that builds on the hybrid vision transformer, i.e., R$50$+ViT-B/$16$, following .

Implementation details. The original R$50$+ViT-B/$16$ https://github.com/isl-org/DPT, MIT License consists of a ResNet$50$ followed by a ViT-B/$16$, where the ViT-B/$16$ consists of $12$ transformer encoder layers that process $16\times$ 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\%\uparrow$ on KITTI/NYUv$2$, respectively.

7 ImageNet-111K Classification

Finally, we apply our method to the ImageNet-$1$K classification task and compare our method with a very recent SOTA method EViT . 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 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 https://github.com/facebookresearch/SWAG, CC-BY-NC $4.0$, SWAG + ViT-H/$14$ and SWAG + ViT-L/$16$ achieve $88.55\%$ and $88.07\%$ top-$1$ accuracy on ImaegNet-1K 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$.

References

Checklist

Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]

Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]

Did you discuss any potential negative societal impacts of your work? [N/A]

Did you describe the limitations of your work? [N/A]

If you are including theoretical results…

Did you state the full set of assumptions of all theoretical results? [N/A]

Did you include complete proofs of all theoretical results? [N/A]

Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [TODO]We will release the code soon.

Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]

Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [TODO]

Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]

If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…

If your work uses existing assets, did you cite the creators? [Yes] We use the official checkpoints provided by:

{https://github.com/rstrudel/segmenter#model-zoo}, MIT License

https://github.com/facebookresearch/Mask2Former/blob/main/MODEL_ZOO.md, CC-BY-NC $4.0$

https://github.com/microsoft/Swin-Transformer, MIT License

https://github.com/facebookresearch/SWAG, CC-BY-NC $4.0$

https://github.com/isl-org/DPT, MIT License.

Did you mention the license of the assets? [Yes] We mark the license of these assets as above.

Did you include any new assets either in the supplemental material or as a URL? [Yes] We add the URL of these assets in the footnote.

Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]

Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]

If you used crowdsourcing or conducted research with human subjects…

Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]

Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]

Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]

Appendix

A. Illustrating More Details of Our Approach

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 LABEL:lst:token_clustering and Listing LABEL:lst:token_reconstruct, 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 13 and Table 13, respectively.

Table 13 summarizes the hyper-parameters, including the inserted positions $\alpha$ & $\alpha$+$\beta$ of token clustering layer & token reconstruction layer, the number of remaining transformer layers after the token reconstruction layer $\gamma$, the total number of transformer layers $\mathsf{L}$, the number of tokens before clustering $\frac{\mathsf{H}}{\mathsf{P}}\times\frac{\mathsf{W}}{\mathsf{P}}$, the number of tokens after clustering $\mathsf{h}\times\mathsf{w}$, the number of neighboring pixels $\lambda$, the number of EM iterations $\kappa$, the temperature value $\tau$, and the number of nearest neighbors $\mathsf{k}$, for Segmenter, DPT, and SWAG.

Table 13 summarizes the hyper-parameters, including the inserted positions $\alpha$ & $\alpha$+$\beta$ of the window token clustering layer & window token reconstruction layer, the number of remaining transformer layers after the token reconstruction layer $\gamma$, the total number of transformer layers $\mathsf{L}$, the number of window tokens before clustering $\mathsf{K}\times\mathsf{K}$, the number of window tokens after clustering $\mathsf{k}\times\mathsf{k}$, the number of neighboring pixels $\lambda$, the number of EM iterations $\kappa$, the temperature value $\tau$, and the number of nearest neighbors $\mathsf{k}$, for Mask$2$Former and SwinV$2$-L + HTC++.

C. More Evaluation Details

We illustrate the evaluation details used for measuring the GFLOPs and FPS of different methods in Table 13. 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 , we report the detailed comparison results in Figure 12. The original EViT propose to identify and only keep the top $\rho\%$ 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$ (with $24$ layers in total) to decrease the number of tokens by ($1$-$\rho\%$), respectively. We report the results of EViT by choosing $\rho\%$=$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\%\sim 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 11(a). 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 12 (b) presents more details of the overall framework. We also report the comparison results in Table 13.

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 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 (without retraining) in Table 13 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 13(a), which shows three different configurations over Segmenter+ViT-L/$16$ achieve $32.13\%$/$48.21\%$/$51.32\%$ when setting the cluster size $\rm{h}\times\rm{w}$ as $8\times 8$/$16\times 16$/$24\times 24$, respectively.

Then, we visualize both the original feature maps and the clustering feature maps in Figure 13(b). 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.