AiluRus: A Scalable ViT Framework for Dense Prediction

Vision transformers (ViTs) have emerged as a prevalent architecture for vision tasks owing to their impressive performance. However, when it comes to handling long token sequences, especially in dense prediction tasks that require high-resolution input, the complexity of ViTs increases significantly. Notably, dense prediction tasks, such as semantic segmentation or object detection, emphasize more on the contours or shapes of objects, while the texture inside objects is less informative. Motivated by this observation, we propose to apply adaptive resolution for different regions in the image according to their importance. Specifically, at the intermediate layer of the ViT, we utilize a spatial-aware density-based clustering algorithm to select representative tokens from the token sequence. Once the representative tokens are determined, we proceed to merge other tokens into their closest representative token. Consequently, semantic similar tokens are merged together to form low-resolution regions, while semantic irrelevant tokens are preserved independently as high-resolution regions. This strategy effectively reduces the number of tokens, allowing subsequent layers to handle a reduced token sequence and achieve acceleration. We evaluate our proposed method on three different datasets and observe promising performance. For example, the "Segmenter ViT-L" model can be accelerated by 48% FPS without fine-tuning, while maintaining the performance. Additionally, our method can be applied to accelerate fine-tuning as well. Experimental results demonstrate that we can save 52% training time while accelerating 2.46 times FPS with only a 0.09% performance drop. The code is available at https://github.com/caddyless/ailurus/tree/main.

Paper

References (38)

Scroll for more · 26 remaining

Similar papers

Peer review

Reviewer viNG4/10 · confidence 4/52023-06-30

Summary

ViTs process input images as a sequence of tokens, and as the token size increases, the computational complexity escalates. The paper aims to accelerate the training and inference of ViTs by reducing the number of tokens. To achieve this, the authors introduce an adaptive resolution strategy, leveraging the observation that dense prediction tasks primarily emphasize object shape and contour. They present a density-based clustering algorithm, spatial-aware DPC, which effectively prunes tokens. The proposed method is extensively evaluated on three diverse datasets, demonstrating its successful acceleration of ViT-based models while preserving performance.

Strengths

- This paper presents an efficient ViT that reduces the complexity of the self-attention operation with a novel dense-based clustering algorithm. - The paper is well written, and the motivation and technical details are clearly presented. - The experiments show good qualitative and quantitative results compared to baselines.

Weaknesses

- The idea of merging tokens to speed up ViT is not new (STViT, PaCa-ViT, etc.), and the proposed algorithm relies on DPC, which also reduces the novelty. - Comparisons are only conducted with Expedite, while there are many other efficient ViTs as introduced in the related works. Though some of them requires fine-tuning, but as least can be compared with AiluRus in table 4.

Questions

- Can the proposed algorithm be applied to classification tasks? - In Table. 2, the mIoU of AiluRus with 900 clusters outperforms the baseline, I’m interested in the explanation for this.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes.

Reviewer 2MFv6/10 · confidence 3/52023-07-05

Summary

This paper first points out that when handling long token sequences, ViTs particularly suffer in terms of computation. Motivated by an observation that dense prediction tasks focus more on the contours or shapes of object, the authors propose to leverage adaptive resolution strategy for different regions depending on each regions' importance. Specifically, regions where repetitive patterns are present, e.g., walls, they are merged to predefined anchors to form low resolution regions, while other "important" regions are preserved for further processing. Thanks to this design, the number of tokens have been greatly reduced, and their shapes are recovered via unfolding operation at the output end. This method has shown to its effectiveness using three different datasets. Especially, for the variant that uses Segmenter ViT-L, almost 48% FPS acceleration is accomplished, which can also benefit to reduce the training time required to learn the models.

Strengths

1. I really like the motivation and idea of this paper. It is a well known issue that the number of tokens fed to ViTs will inevitably require expensive computations and memory footprints, and to me, this paper nicely addresses this issue with a novel approach. 2. This paper is very easy to read and understand. 3. Significant improvements are made when this method is applied.

Weaknesses

1. Figure 2 visualization looks good. I also found the authors saying in L177 that some regions looking like they are assigned to a same group (same color) is not actually in the same group. I believe this visualization in Fig 2. is one of the most important one in this paper, which I would recommend looking for different visualization to make this figure look more clear and easy to catch the idea without reading the texts. Maybe scale the values then visualize? or change the colors? 2. Adding ViT-B results in Table 1 would make this table more informative. 3. What could be the reason for performance improvements in Table 2 last row (AiluRus clusters 900) ? Minor comments : typo in L200.

Questions

See above weaknesses. I do not have any major issue.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

"We attribute this to the complex designs used in object detection, such as the FPN and the complicated decoder. AiluRus can be improved by mitigating the performance drop caused by these complex design" If this was also addressed in this paper, it would have made this paper much stronger. However, this paper already shows good results and analysis in main task (segmentation).

Reviewer PvMg4/10 · confidence 4/52023-07-06

Summary

This paper proposed a spatial-aware density-based clustering algorithm for ViT-based dense prediction tasks. The proposed method can merge the original tokens into fewer representative tokens to accelerate the transformer inference without finetuning. The experiments on different datasets show that the proposed method can accelerate the ViT inference significantly with a slight performance drop.

Strengths

1. This paper focuses on an important topic of ViT inference acceleration, which may facilitate further applications. 2. The experiments on different datasets prove that the proposed method can speed up both training and inference significantly. The experiments also show that the proposed method can outperform the Expedite method, especially at shallow layers. 3. The visualization results are also reasonable and meaningful.

Weaknesses

1. The technical contribution is incremental compared with the previous method. The DPC-based token clustering module has been proposed in the previous method[29] for transformer-based dense prediction tasks. The only difference between the proposed method in this paper and the previous method is introducing spatial information into the DPC module. 2. The experiments are not sufficient to justify the algorithm design. a. The experiment results in Table 3 show that the different settings of α only result in a negligible accuracy variance. The results suggest that the proposed spatial information is incremental to the original DPC. b. The DPC module may also introduce computation overhead for each transformer block. It is mandatory to present the break-down time consumption for each block, e.g. the DPC module, the attention module, the feature map restore module, etc. 3. The DPC module will generate representative tokens adaptively for each image. It is not clear whether it will affect the batch inference performance. It is better to show a comparison in the batch inference mode. 4. Some typos and format errors a. In L.204, Fig.5 and Fig.4 should be Fig.3 and Fig.4. b. In L.217, Fig. 7a and Fig. 7 should be Fig. 7a and Fig. 7b. c. The line space between L.154 and L.155 is not sufficient.

Questions

1. Why does the proposed method perform better than existing methods when applied to shallower layers? 2. How does the proposed method recover the output of each transformer block?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have adequately addressed the limitations.

Reviewer ZFoB5/10 · confidence 3/52023-07-09

Summary

This paper proposes a new approach to improving the efficiency of Vision Transformers (ViTs) for dense prediction tasks such as semantic segmentation. Recognizing that these tasks emphasize the contours or shapes of objects more than the textures within, the authors introduce a method to apply adaptive resolution to tokens in ViTs through density-based clustering. At the intermediate layer of the ViT, a spatial-aware density-based clustering algorithm is used to select anchors from the token sequence. Tokens near these anchors are merged to form low-resolution regions, while others are kept as high-resolution, thereby significantly reducing the number of tokens processed in subsequent layers. This strategy aims to accelerate processing without loss of critical information. The proposed method showcased promising results on different semantic segmentation datasets compared to a couple of previous works, including a baseline method [13] which also conducts token clustering. The proposed method provides higher speed up with less performance drop on both well-trained and fine-tuning scenarios.

Strengths

- The paper presents promising empirical results compared to a closely related baseline [13], especially in reducing fine-tuning training time. - The paper is well-intuitive and easy to follow. - The paper provides extensive analysis and ablation on parameters, cluster numbers, and clustering layers. - The paper provides visual analysis to demonstrate that the clustering of tokens towards the segmentation task is reasonable.

Weaknesses

- The manuscript suffers from a substantial number of typographical errors. Some are actually critical. E.g., the abusing usage of delta and sigma (Ln138-142). - The paper can be improved by better structure and writing. E.g., The figures and Tables are all over the place instead of being roughly around the page when they are mentioned. - While I recognize the promising performance of the paper, I am concerned that the paper might not contain enough novelty and theory contribution for publication in the underlying conference. The idea to decrease the number of tokens has been proposed in related works, as mentioned in the paper itself. The clustering idea is also proposed by [13]. The critical contribution of the paper is adapting a different existing density-based clustering method to the semantic segmentation task. In addition, I don't think the paper provides enough concrete analysis on why the proposed clustering method is better than the one in [13], despite empirical results. - The paper refers to the text's segmentation and object detection, while only semantic segmentation results are presented in the main experiment section. For the sake of fairness, I would recommend removing the misleading reference to the object detection task or using it for future work. - The visualization example doesn't contribute much to the paper as the idea of decreasing the token for tasks that do not need local fine-gran information has been proposed and demonstrated in related works. Plus, the figures are not quite self-evident. For E.g. the clustering validity can depend on ontology.

Questions

- Does clustering location refers to layer number? If so, please declare. - Any idea why the new clustering method provides more robust clustering than [13]? - The ablation study seems to indicate that alpha makes little difference, any thoughts on that?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

I don't think there are any outstanding concerns about any negative societal impacts in this work.

Reviewer gX5G6/10 · confidence 4/52023-07-10

Summary

This paper studies the problem of accelerating vision transformers for dense prediction without fine-tuning and proposes a spatial-aware density-based clustering algorithm to select anchors from the token sequence. The proposed method is evaluated across three different datasets and demonstrates promising performance.

Strengths

+ The idea is reasonable and interesting, and the performance is good. + The abundant ablation studies are provided to verify the effectiveness of the proposed method.

Weaknesses

- As said in the title, the proposed method is for dense prediction, of course, segmentation is typical, But it is better to include another dense prediction task for further verifying the effectiveness of the proposed method, for example, instance segmentation, depth estimation. - It is better to calculate the reconstruction error between the reconstructed feature map and the original feature map. Meanwhile, make a comparison to the other SOTA method, e.g. Expedite.

Questions

Please refer to Weaknesses

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

None

Area Chair XNw72023-08-19

Has the rebuttal addressed your concerns?

Dear Reviewer gX5G, Could you please read the author rebuttal and acknowledge if your concerns have been addressed? The discussion period will end very soon on Monday, August 21. Thank you for your time in reviewing this submission! Best, AC

Reviewer ZFoB2023-08-21

I appreciate the rebuttal from the author. I read it th The answers addressed some of my concerns. I would recommend adding more discussion on `why` the proposed method is better than [3], in additional to experimental results. I increased my rating accordingly.

Authorsrebuttal2023-08-21

We sincerely appreciate your response. Our general response analysis provides a detailed analysis of the advantage of Ailurus, and the results indicate that Ailurus only aggregates tokens with high intra-similarity and thus could preserve their relations going through transformer blocks. At the same time, Expedite [13] suffers from severe relation distortion during forward. As a result, Ailurus could recover the feature map at the output layer with little distortion and performs much better than Expedite [13]. We also recognize that this is a critical contribution of Ailurus and will add this discussion to our later revisions. Sincerely appreciate for your suggestion!

Area Chair XNw72023-08-19

Has the rebuttal addressed your concerns?

Dear Reviewer ZFoB, Could you please read the author rebuttal and acknowledge if your concerns have been addressed? The discussion period will end very soon on Monday, August 21. Thank you for your time in reviewing this submission! Best, AC

Area Chair XNw72023-08-19

Has the rebuttal addressed your concerns?

Dear Reviewer PvMg, Could you please read the author rebuttal and acknowledge if your concerns have been addressed? The discussion period will end very soon on Monday, August 21. Thank you for your time in reviewing this submission! Best, AC

Area Chair XNw72023-08-19

Has the rebuttal addressed your concerns?

Dear Reviewer 2MFv, Could you please read the author rebuttal and acknowledge if your concerns have been addressed? The discussion period will end very soon on Monday, August 21. Thank you for your time in reviewing this submission! Best, AC

Area Chair XNw72023-08-19

Has the rebuttal addressed your concerns?

Dear Reviewer viNG, Could you please read the author rebuttal and acknowledge if your concerns have been addressed? The discussion period will end very soon on Monday, August 21. Thank you for your time in reviewing this submission! Best, AC

Reviewer 2MFv2023-08-20

Thanks for the thorough responses. I've read the rebuttal, and I am still leaning toward acceptance. Therefore, I am keeping my initial rating.

Reviewer PvMg2023-08-21

Feedback after the rebuttal

Thanks for the authors' careful feedback. I have read the rebuttal and some of the concerns are addressed, e.g., time consumption, output recovery, etc. However, some of my questions still exist, i.e. the novelty, the impact of $\alpha$, and the batch inference. 1. The difference between the TCFormer and the proposed method is not significant. It is better to compare the TCFormer and the Ailurus in different dense prediction tasks. 2. It is better to present the inference speed of Ailurus with batch 16 for a fair comparison. Above all, I will keep the original rating.

Authorsrebuttal2023-08-21

We appreciate your thoughtful feedback and would like to address your remaining concerns as follows: **Comparison with TCFormer** 1. It's important to highlight that TCFormer relies on fine-tuning, making it inaccessible for well-trained models. On the other hand, our proposed model, Ailurus, can be integrated into well-trained models to instantly accelerate inference. For instance, we utilized the recent text-to-video generation model, ControlVideo\[1\], as a case study to exemplify Ailurus' versatility. We incorporated Ailurus into ControlVideo, which resulted in halving the number of tokens during inference. The accordingly generated results are provided in the supplementary PDF file. Specifically, Ailurus maintains a comparable quality to the original model while halving the time required. Its flexibility and parameter-free characteristics allow Ailurus to be applied in a broader range of scenarios where TCFormer cannot be utilized. 2. As mentioned in our general response, replacing the DPC algorithm with a naive strategy in the Ailurus framework still yields satisfactory performance. This demonstrates that the DPC algorithm isn't a prerequisite for Ailurus to function effectively. The key to Ailurus' robustness and superiority lies in its framework, which manages to preserve the relationships between the remaining and merged tokens with little distortion. 3. As TCFormer is designed for human-centric tasks, its application to dense prediction tasks presents numerous challenges. Specifically, TCFormer relies on specific architectures and training, limiting its usability for direct comparison with Ailurus in well-trained models. Additionally, the unique architecture of TCFormer makes it difficult to leverage the benefits of extensive pre-trained models within the ViT-family. Consequently, comparing these two methods becomes challenging and less meaningful due to their distinct application scenarios. Nevertheless, we are diligently working to adapt TCFormer for compatibility with dense prediction tasks and conduct a comparative analysis against Ailurus. As the discussion period is drawing to a close, this task is challenging, but we will invest our best efforts to address it. If time permits, we will promptly report our results. **Performance Comparison with Batch Size 16** Below we present the performance comparison using a batch size of 16. As you will see, Ailurus still demonstrates outstanding performance. | Method | Num clusters | Batch size | FPS | Throughput | mIoU | | :------- | :----------: | :--------: | :--------------------: | :---------------------: | :---: | | Baseline | 1024 | 16 | 32.2 | 35.6 | 48.48 | | Ailurus | 324 | 16 | 61.2 ($\uparrow$ 90%) | 76.1 ($\uparrow$ 114%) | 46.45 | | | 400 | 16 | 53.7 ($\uparrow$ 67%) | 67.4 ($\uparrow$ 89%) | 47.12 | We are open to further questions and sincerely look forward to your response. \[1\]. Zhang, Yabo, et al. "ControlVideo: Training-free Controllable Text-to-Video Generation." arXiv preprint arXiv:2305.13077 (2023).

Authorsrebuttal2023-08-21

**TCFormer for dense prediction tasks** We follow your suggestion to apply TCFormer to dense prediction tasks. Initially, we intend to integrate TCFormer into the mmseg framework for semantic segmentation. We carefully implement the TCFormer under the mmseg framework according to the officially released code. However, we find that TCFormer is an extremely computationally intensive backbone. Even with its lightest configuration, the FPS for inferring ADE20K images at a resolution of 512x512 under the Segmenter framework is only 2.03. In comparison, a typical ViT-Base model achieves an FPS of 32.2, while Ailurus can surpass 50 without sacrificing performance. Upon analyzing the TCFormer code, we discover that it heavily relies on sparse matrix multiplication. Although these operations may not have high FLOPs, they introduce significant latency. Therefore, it is nearly impossible and impractical to apply TCFormer to dense prediction tasks given its high computational complexity. Furthermore, we attempt to integrate the core TCFormer design, the CTM module, into ViT for acceleration. This would allow us to compare the performance of TCFormer with Ailurus using the same backbone and initialization. However, we encountered the same issue of excessive complexity for the CTM module. In a similar setup, the time required for the CTM module to execute is even higher than the inference time of the model itself. The details are presented below: | Method | Total (ms) | blocks (ms) | extra (ms) | mIoU | |:------| :---------:|:----------:|:--------:|:-----:| |TCFormer | 33.8 | 8.6 | 25.2 | 0.14 | |Percents | - | 25.0% | 75.0% | - | |Ailurus | 10.0 | 8.7 | 1.3 | 47.35 | |Percents | - | 87.0% | 13.0% | - | Despite these issues, we proceed to train the ViT model with the integrated CTM module in the Segmenter framework for comparison with Ailurus. We followed the instructions provided in the TCFormer paper to set up the CTM module and trained ViT-CTM with the same optimization parameters and initialization as Ailurus. However, we find that even with only the integration of CTM into ViT, it raises a significant negative impact on training and causes very poor performance. Thus, it is untrivial to integrate the CTM module to ViTs for acceleration. |Method | Num clusters | FPS | mIoU | |:------| :---------:|:----------:|:--------:| |Baseline | 1024 | 32.2 | 49.60 | |ViT-CTM | 410 (1024 * 0.4) | 29.1 ($\downarrow$ 10%) | 33.64 ($\downarrow$ 15.96) | |Ailurus | 400 | 53.7 ($\uparrow$ 67%) | 49.04 ($\downarrow$ 0.56) | In conclusion, TCFormer is specifically designed for human-centric tasks, and adapting it to dense prediction tasks poses significant challenges. Additionally, TCFormer exhibits high computational complexity, rendering it unsuitable for accelerating ViTs. In contrast, Ailurus can be seamlessly integrated into well-trained models, providing immediate acceleration or expediting training without the need for additional hyper-parameter adjustments. These two methods differ significantly in terms of motivation, implementation, and application. Ailurus stands out for its flexibility and lightweight nature, enabling its deployment in various tasks that TCFormer can't accomplish. We sincerely appreciate your diligent review of this paper, and we would be grateful to receive any response or feedback from you.

Authorsrebuttal2023-08-22

Dear reviewer PvMg: In the past day, we made our best efforts to conduct a comparison between TCFormer and Ailurus for dense prediction tasks. The results indicate that TCFormer is not suitable for dense prediction tasks and is unable to surpass Ailurus. For detailed information, please refer to our response. We would greatly appreciate it if our response could effectively address any of your concerns, and we would be grateful for any feedback or response from you.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC