Pruning Self-attentions into Convolutional Layers in Single Path

Abstract

Vision Transformers (ViTs) have achieved impressive performance over various computer vision tasks. However, modeling global correlations with multi-head self-attention (MSA) layers leads to two widely recognized issues: the massive computational resource consumption and the lack of intrinsic inductive bias for modeling local visual patterns. To solve both issues, we devise a simple yet effective method named Single-Path Vision Transformer pruning (SPViT), to efficiently and automatically compress the pre-trained ViTs into compact models with proper locality added. Specifically, we first propose a novel weight-sharing scheme between MSA and convolutional operations, delivering a single-path space to encode all candidate operations. In this way, we cast the operation search problem as finding which subset of parameters to use in each MSA layer, which significantly reduces the computational cost and optimization difficulty, and the convolution kernels can be well initialized using pre-trained MSA parameters. Relying on the single-path space, we introduce learnable binary gates to encode the operation choices in MSA layers. Similarly, we further employ learnable gates to encode the fine-grained MLP expansion ratios of FFN layers. In this way, our SPViT optimizes the learnable gates to automatically explore from a vast and unified search space and flexibly adjust the MSA-FFN pruning proportions for each individual dense model. We conduct extensive experiments on two representative ViTs showing that our SPViT achieves a new SOTA for pruning on ImageNet-1k. For example, our SPViT can trim 52.0% FLOPs for DeiT-B and get an impressive 0.6% top-1 accuracy gain simultaneously. The source code is available at https://github.com/ziplab/SPViT.

arXiv commentAccepted by TPAMI 2024

References in corpus (28)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
• Gaussian Error Linear Units (GELUs)
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• AMC: AutoML for Model Compression and Acceleration on Mobile Devices
• ConViT: Improving Vision Transformers with Soft Convolutional Inductive Biases
• ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness
• Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
• DeepViT: Towards Deeper Vision Transformer
• DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification
• Intriguing Properties of Vision Transformers
• LocalViT: Analyzing Locality in Vision Transformers
• ViTAE: Vision Transformer Advanced by Exploring Intrinsic Inductive Bias
• Approximating CNNs with Bag-of-local-Features models works surprisingly well on ImageNet
• AutoSlim: Towards One-Shot Architecture Search for Channel Numbers
• Do Vision Transformers See Like Convolutional Neural Networks?
• Chasing Sparsity in Vision Transformers: An End-to-End Exploration
• Large Memory Layers with Product Keys
• CMT: Convolutional Neural Networks Meet Vision Transformers
• Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
• NAS-BERT: Task-Agnostic and Adaptive-Size BERT Compression with Neural Architecture Search
• PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination
• The Generalization-Stability Tradeoff In Neural Network Pruning
• X-volution: On the unification of convolution and self-attention
• Searching for Efficient Multi-Stage Vision Transformers
• Scalable Vision Transformers with Hierarchical Pooling
• Transformed CNNs: recasting pre-trained convolutional layers with self-attention
• Rapid Elastic Architecture Search under Specialized Classes and Resource Constraints