[Paper Review] Expediting Large-Scale Vision Transformer for Dense Prediction without Fine-tuning

Weicong Liang, Yuhui Yuan+7 more|arXiv (Cornell University)|Oct 3, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes a non-parametric, fine-tuning-free method to accelerate large-scale Vision Transformers for dense prediction tasks by inserting token clustering and reconstruction layers into pre-trained models. The approach reduces FLOPs by 30% and boosts inference speed by 40% while preserving 99.5% of the original performance on ADE20K, demonstrating strong generalization across object detection, segmentation, and depth estimation without retraining.

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 fine-tuning. 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.

Motivation & Objective

• To address the high computational cost of large-scale Vision Transformers in dense prediction tasks like segmentation and detection.
• To eliminate the need for fine-tuning or retraining when accelerating pre-trained Vision Transformers.
• To maintain high performance while significantly reducing inference time and FLOPs on high-resolution inputs.
• To generalize across diverse dense prediction tasks including instance segmentation, depth estimation, and panoptic segmentation.
• To provide a lightweight, plug-and-play solution compatible with existing Vision Transformer architectures such as ViT and Swin Transformer.

Proposed method

• Introduces a token clustering layer that groups neighboring tokens into lower-resolution representations while preserving spatial structure.
• Applies subsequent transformer layers only on the clustered, low-resolution tokens to reduce computation.
• Inserts a token reconstruction layer that restores high-resolution representations from the refined low-resolution features.
• Uses non-parametric operators that do not require additional training or weight updates.
• Applies the method to ViT-L/16 and Swin-L by inserting clustering and reconstruction layers at empirically chosen stages.
• Maintains model performance by leveraging the rich intermediate representations of pre-trained models without fine-tuning.

(a) mIoU@Segmenter

Experimental results

Research questions

• RQ1 Can Vision Transformers be accelerated for dense prediction without fine-tuning or retraining?
• RQ2 How can high-resolution representations be preserved while reducing computational cost in large-scale Vision Transformers?
• RQ3 What non-parametric operations can effectively reduce the number of tokens without introducing learnable parameters?
• RQ4 How does the proposed method generalize across diverse dense prediction tasks such as object detection and depth estimation?
• RQ5 To what extent can the method reduce FLOPs and improve inference speed while maintaining high accuracy?

Key findings

• The method accelerates inference by 40% and reduces FLOPs by 30% on "Segmenter+ViT-L/16" while maintaining 99.5% of the original performance on ADE20K without fine-tuning.
• On COCO panoptic segmentation, the method improves FPS by over 35% with minimal performance drop.
• For depth estimation using DPT, the method achieves a 30–37% speedup on KITTI and NYUv2 benchmarks with only minor degradation in metrics.
• The approach generalizes to image classification, achieving comparable performance to EViT but with better efficiency on ImageNet-1K.
• The method maintains strong performance across five dense prediction tasks: object detection, semantic segmentation, panoptic segmentation, instance segmentation, and depth estimation.

(b) FPS@Segmenter