Intra-head pruning for vision transformers via inter-layer dimension relationship modeling.

Abstract

Transformer models have demonstrated good performance across a range of natural language processing and computer vision tasks. However, the huge computational cost imposed by transformer models poses a significant obstacle to their practical implementation on platforms with limited hardware. To address this challenge, recent academic studies have been focused on head pruning, a strategy that effectively eliminates unimportant components in transformer models. Although these pruning methods have shown significant improvements, they suffer from severe accuracy loss due to coarse pruning granularity and fail to consider the interdependence between layers when discarding zero-valued components. This is crucial for achieving a network architecture with efficient compression. Therefore, we propose a novel intra-head pruning (IHP) technique to sparsely train pruned vision transformers. Specifically, our method utilizes a trainable row parameter delicately designed to participate in the sparse training of the model. Furthermore, we introduce a relationship matrix which serves as the key to the grouping pruning process. The grouping policies ensures consistent and coherent elimination of redundant components, thereby maintaining the structural integrity and functional consistency of the pruned network. Experimental results on benchmark datasets (CIFAR-10/100, ImageNet-1K) show that this method can significantly reduce the computational cost of the mainstream vision transformers such as DeiT, Swin Transformer, and CCT, with a small decrease in accuracy. Especially on ILSVRC-12, under the same FLOPs reduction ratio of 46.20%, the Top-1 accuracy improves by 0.47% compared to advanced methods for DeiT-tiny.