Rethinking Depth Pruning for Vision Transformers: A Heterogeneity-Aware Perspective

While prior studies have successfully compressed vision Transformers (ViTs) through various pruning techniques, most have concentrated on width pruning to achieve significant reductions in model size. Depth pruning, which removes entire layers from a ViT, is notoriously difficult for accuracy recovery despite its potential to deliver higher speedups, limiting the acceleration achieved by existing joint width-and-depth pruning methods. In this work, we reveal that the failure of existing depth pruning methods lies in their neglect of heterogeneity between different layers, and we introduce HetDPT, a heterogeneity-aware depth pruning method that avoids dimension mismatch. Comprehensive experiments on ImageNet-1K, CIFAR-100, COCO, and ADE20K validate our method: HetDPT achieves a 1.58$\times$ speedup for DeiT-B while maintaining accuracy and a 1.39$\times$ speedup for DeiT-S with nearly no accuracy degradation. Furthermore, when combined with width pruning, HetDPT+ sets a new state-of-the-art record in extreme ViT pruning, enhancing the acceleration ratio from 4.24$\times$ to 5.19$\times$ for the Isomorphic-Pruning-2.6G configuration while maintaining near-lossless accuracy; our code is available at \url{https://github.com/Efficient-AI-for-All/HetDPT}.

Modern vision AI models can recognize images very accurately, but they are often large and slow to run. This makes them expensive to use and difficult to deploy on devices with limited computing power. A common way to make such models faster is to remove unnecessary parts, but removing whole layers from vision Transformers often causes a large drop in accuracy. In this paper, we study why this happens and find that different types of layers behave very differently when removed. Treating all layers in the same way can therefore lead to poor pruning decisions. We propose HetDPT, a method that carefully accounts for these differences when deciding which layers to remove. Our method also avoids structural issues that would otherwise make the pruned model invalid. Across several image recognition, object detection, and segmentation benchmarks, HetDPT makes vision Transformer models faster while keeping their accuracy nearly unchanged. When combined with existing compression methods, it achieves even stronger speedups. This work can help make powerful vision AI models cheaper, faster, and easier to deploy in real-world applications.