Rethinking Depth Pruning for Vision Transformers: A Heterogeneity-Aware Perspective
This work addresses the challenge of depth pruning in Vision Transformers for efficient deployment, offering a practical method that significantly improves speedup without accuracy loss.
HetDPT introduces a heterogeneity-aware depth pruning method for Vision Transformers that avoids dimension mismatch, achieving 1.58× speedup for DeiT-B with maintained accuracy and 1.39× for DeiT-S with negligible degradation. Combined with width pruning, it sets a new SOTA in extreme ViT pruning, improving acceleration from 4.24× to 5.19× while preserving near-lossless accuracy.
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 https://github.com/Efficient-AI-for-All/HetDPT.