Pruning Self-attentions into Convolutional Layers in Single Path

TL;DR

SPViT addresses the dual drawbacks of Vision Transformers: high computational cost from global self-attention and lack of local inductive bias. It introduces a weight-sharing scheme that expresses convolutional operations as subsets of MSA parameters, enabling a single-path NAS space to prune both MSA and FFN blocks efficiently. The method selects between MSA, bottleneck convolutions, and skip connections, and jointly prunes FFN expansion ratios under a target FLOPs, followed by fine-tuning with optional knowledge distillation. Empirically, SPViT achieves state-of-the-art pruning performance on ImageNet-1k for DeiT and Swin models, often with significant FLOPs reductions and small or positive accuracy gains, while revealing architectural tendencies such as favoring locality in shallow blocks and preserving global modeling in deep blocks.

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.

Figures (5)

• Figure 1: Multi-path vs. single-path search space. (a) Multi-path search space. MSA and different types of convolutional operations are added as separate trainable paths, leading to an expensive search cost. $\mathcal{P}$ and $g_p$ denote the ordered operation set and the learnable gate for the $p$-th operation, respectively. (b) Single-path search space. We propose to formulate searching as a subset selection problem where the outputs of the convolutional operations are directly indexed from the MSA intermediate results, thereby easing the optimization difficulty and reducing the search cost together with the number of trainable parameters.
• Figure 2: Overview of our weight-sharing scheme, which expresses the bottleneck convolution with MSA parameters. First, we restrict the value and output projections in the MSA layer to only process the features within local regions. Then, we learn to ensemble the value projections with a parameterized softmax $\sigma(\boldsymbol{z})$ to derive the convolutional kernel. After convolution, we project the convolutional output to the original dimensionality by learning to ensemble the output projections and deriving the bottleneck convolution output.
• Figure 3: The overview of our SPViT during the search. SPViT block (top left): We replace the original MSA and FFN layers with our UMSA and UFFN layers to search for a compact architecture, while keeping the other components intact. UFFN layer (top right): We search for important FFN hidden dimensions. Given the input features $\boldsymbol{X}$, we employ learnable binary gates for each of the $\alpha c_{in}$ transformed features $O_t(\boldsymbol{X})$, where $\alpha$ and $c_{in}$ represent the MLP expansion ratio and input dimensions. We optimize the learnable binary gates to search pruned FFN structures. UMSA layer (bottom): We search for MSA or convolutional layers in single path. In addition to the standard MSA operation (lower branch as described by Eq. (\ref{['eq:msa']})), we get the output of bottleneck convolutional layers (upper branch). The additional steps are in red lines where we index and scale the values $\boldsymbol{V}$ with the scaling factors $\sigma(\mathbf{z})$ to get the convolution output ${\rm Conv}{(\boldsymbol{X})}$, and then derive the bottleneck convolution output ${\rm BConv}(\boldsymbol{X})$ as described in Eq. (\ref{['eq:msa2conv7']}). For simplicity, we only show deriving the ${\rm Conv}{(\boldsymbol{X})}$ of one kernel size. In this way, we get both the MSA output and the bottleneck convolution outputs simultaneously and finally optimize the learnable gates to select operations.
• Figure 4: SPViT searched architectures under different target FLOPs on ImageNet-1k. (a)-(c), (d)-(f), and (g)-(i) are architectures searched by SPViT-DeiT-S, SPViT-DeiT-B, and SPViT-Swin-Ti, respectively. Dashed lines represent the operations and the number of hidden dimensions for FFN layers before pruning. Solid lines represent the two types of architecture configurations searched by SPViT.
• Figure 5: Attention probabilities of the pre-trained DeiT-B model averaged over 100 images. We visualize the attention maps for all twelve heads in the shallow blocks (Block 1 and 2) and deep blocks (Block 11 and 12). Each map shows the attention probabilities between a query feature embedding (green square) and the other feature embeddings. Darker color indicates higher attention probability and vice versa. Best viewed in color.