Joint Token Pruning and Squeezing Towards More Aggressive Compression of Vision Transformers

To quantitatively verify the discarded information of pruned tokens, we conduct a toy experiment on DynamicViT [25] as Fig. 3 shows. It is easy to agree that the performance of pruned model declines as the pruning becomes more aggressive. Nevertheless, by exchanging reserved and pruned tokens (dubbed as the reversed policy in Fig. 3), we find that the pruned tokens can still handle partial cases correctly. Furthermore, the bonus accuracy, which is dedicated by the cases that only the reversed policy predicts rightly, rises along with the token pruning intensity. It implies that the exclusive information from pruned tokens matters more while the token pruning intensity grows.

These fun facts motivate us to assimilate the pruned tokens into the reserved tokens to prevent information loss, as shown in Fig. 1. As shown in Fig. 2 (c), TPS employs two steps to compress ViTs, including token pruning and squeezing.

In this section, we briefly review the basic procedure of token pruning. Note that our TPS is compatible with any token pruning techniques. Here, we introduce two variants of TPS: dTPS and eTPS, to cover both intra-block and inter-block token compression shown in Fig. 4. They follow the pruning parts of two baselines for a fair comparison with two typical baselines [25, 16].

As shown in Fig. 4, dTPS adopts the learnable token score prediction head from dynamicViT [25] and samples the binary decision mask by Straight-Through Gumbel-Softmax [12] for differentiability; eTPS utilizes the class token attention values to measure tokens’ importance as EViT [16]. In the inference stage of both variants, based on token scores, we devise the token selection policy using the Top-k operation with a fixed given token reduction ratio ${\rho}$. Both variants ensure the constant shape to benefit from the inference optimization on the computation graph. The tokens are separated into two subsets, $S^{r}$ and $S^{p}$, where the reserved tokens are placed in $S^{r}$ and the pruned ones are placed in $S^{p}$. More implementation details can be found in our codes.

After reserved & pruned tokens are split, we introduce our token squeezing part. Considering that the reserved ones contribute the majority of correct predictions, we aim to design a procedure that retains most of the attentive tokens while compressing information from rest, preserving the model’s overall performance. To avoid generating extra tokens as  [16, 14], we inject pruned tokens into similar reserved tokens. So, we apply a unidirectional nearest-neighbor matching algorithm from $S^{p}$ to $S^{r}$ in a many-to-one manner. After that, we employ a similarity-based fusing method to assimilate information from pruned tokens into partial reserved tokens. We summarize the above process as two steps: matching and fusing.

Matching. Given the two subsets $S^{r}$ and $S^{p}$, $I^{r}$ and $I^{p}$ are the corresponding token indices of $S^{r}$ and $S^{p}$. A similarity matrix $c_{i,j}$ for all $i\in I^{p}$ and $j\in I^{r}$ represents the interactions between the tokens for matching. For each pruned token $\boldsymbol{x}_{i}\in S^{p}$, we find its nearest token $\boldsymbol{x}_{*}^{host}\in S^{r}$ from the reserved token set $S^{r}$ as its host token:

$$\boldsymbol{x}_{*}^{host}=\mathop{argmax}\limits_{\boldsymbol{x}_{j}\in S^{r}}{c_{i,j}}.$$

(1)

Note that since the token matching step is unidirectional from $S^{p}$ to $S^{r}$, multiple pruned tokens can share the same host token and not each reserved token can serve as a host token. We then record the matching results in a mask matrix $\boldsymbol{M}\in\mathbb{R}^{N^{p}\times N^{r}}$ and its values are decided by:

$$m_{i,j}=\begin{cases}1,&\boldsymbol{x}_{j}\ is\ the\ host\ token\ of\ \boldsymbol{x}_{i},\\ 0,&otherwise,\end{cases}$$

(2)

where $N^{p}$ and $N^{r}$ denote the token number of two subsets. The mask helps that the following fusing step can be conducted with regular matrix operations on $S^{r}$ and $S^{p}$ while excluding the influence from non-matched pairs.

Although the attention map is a natural and free choice to measure interactions among tokens, we can acquire higher performances with the cosine similarity between $S^{r}$ and $S^{p}$ as the ablation experiment in Section. 4.2 discusses. Therefore in all of our experiments, the similarity matrix is defined as:

$$c_{i,j}=\frac{{\boldsymbol{x}_{i}^{T}}{\boldsymbol{x}_{j}}}{\|\boldsymbol{x}_{i}\|\|\boldsymbol{x}_{j}\|},for\ i\in I^{p},j\in I^{r}.$$

(3)

Since the similarity matrix $c_{i,j}$ is generated directly from input features, no extra parameters are introduced in the matching step.

Fusing. Simply averaging tokens can lead to feature dispersion because of discrepancies among the different tokens. EViT [16] utilizes the token importance scores to re-weight the aggregated tokens. Separately, we use a similarity-based weighting scheme. It expands the influence of closer tokens to the host tokens while also avoiding potential flaws from imperfect token scoring. As previously mentioned, the fusing step encompasses all tokens from two subsets and is controlled by the mask $\boldsymbol{M}$ to ensure that only host tokens and pruned tokens are mixed. This introduces a few redundant computations but increases practical training & inference throughput due to the efficiency of regular matrix operations.

Specifically, the reserved token $\boldsymbol{x}_{j}$ is updated by fusing the original feature and pruned tokens’ features as follows:

$$\boldsymbol{y}_{j}={w}_{j}\boldsymbol{x}_{j}+\sum_{\boldsymbol{x}_{i}\in S^{p}}{w_{i}\boldsymbol{x}_{i}},$$

(4)

where $w_{i}$ is the weight of each pruned token $\boldsymbol{x}_{i}\in S^{p}$, ${w}_{j}$ is the weight of the reserved token itself, and $\boldsymbol{y}_{j}$ is the updated one. The fusing weight $w_{i}$ depends on the mask value $m_{i,j}$ and similarity $c_{i,j}$:

$$w_{i}=\frac{\exp(c_{i,j})m_{i,j}}{\sum_{\boldsymbol{x}_{i}\in S^{p}}\exp(c_{i,j})m_{i,j}+\mathrm{e}}.$$

(5)

The reserved token always has the largest fusing weight $w_{j}$, as the similarity between ${x}_{j}$ and itself equals to $1$:

$$w_{j}=\frac{\mathrm{e}}{\sum_{\boldsymbol{x}_{i}\in S^{p}}\exp(c_{i,j})m_{i,j}+\mathrm{e}}.$$

(6)

According to the above equations, the reserved tokens that have not been chosen as host tokens remain unchanged, while the pruned tokens are squeezed into host tokens and replace the original ones.

As can be seen, our matching and fusing steps ensure that the number of processed tokens equals the number of reserved tokens, thereby maintaining a constant shape for efficient inference.

To prove our flexibility and generalization across different transformers, we also conduct experiments in hybrid ViTs [31, 33]. For plain transformer blocks, our TPS modules can be easily inserted to reduce the token number and achieve a significant speedup. If the layer contains operations that require a complete spatial structured input: e.g., convolution or pooling, the operation of our TPS will be adjusted slightly. For example, in PVT  [31] models, the TPS module is inserted before the first block of each stage with token pruning applied and generates the decision policy $\boldsymbol{D}$. For the attention layer, we decrease the token dimension size of the input and consequent query $\boldsymbol{Q}$. If the spatial-reduction layer is employed inside, the dropped token features are complemented with zeros to maintain the structured spatial input. More details can be found in supplementary materials.

Comparison to baselines. As  Fig. 5 shows, we compare our method with the token pruning baseline: DynamicViT, and the token reorganization baseline: EViT, by replacing their original pruning modules with our dTPS & eTPS modules. The contrast experiments involve DeiT-small&tiny. All the models in this part are fine-tuned 30 epochs under multiple pruning settings. Under all the settings, our method outperforms DynamicViT and EViT. Both dynamicViT and EViT encounter a larger accuracy drop along with the progressively aggressive pruning. While shrinking the computational budgets of DeiT to 35%, our method can avoid 1%-6% accuracy decline compared with baselines. Equipped with our TPS, we can accelerate the throughput of DeiT-small to 1745 images/s, which is beyond that of DeiT-tiny: 1686 images/s, and surpass the accuracy of DeiT-tiny by 4.78%.

Visual Comparisons. We demonstrate the cases from ImageNet1K [4], which DeiT predicts correctly at first but gives the wrong predictions after being applied with token pruning. As Fig. 7 shows, the imperfect pruning policy brings the context information loss, which leads to a close but incorrect prediction. However, our TPS remedies these cases by saving the pruned tokens’ information.

Comparison to states of the art. In Fig. 6, we demonstrate our TPS performances compared with other state-of-the-art transformers, including token pruning methods [20, 8, 25, 16, 35, 20], vanilla ViTs [28, 13, 39, 38, 2, 23] and hybrid ViTs [17, 31, 40, 34]. By integrating DeiT-small&tiny and LV-ViT-small&tiny with TPS and fine-tuning them only 30 epochs, we can achieve a quite competitive performance among numerous vision transformers from the perspective of accuracy-computation trade-off.

Extension on more backbones. As shown in Tab. 1 and Tab. 2, we incorporate TPS into different vanilla ViTs [28, 39, 38] and hybrid ViTs [31, 40] to prove the flexibility and generalization. For vanilla ViTs, our TPS outperforms EViT [16], Evo-ViT [35], A-ViT [36], IA-RED² and SPViT [14] with equal or slightly increasing computation while using DeiT [28], LV-ViT [13] as backbones. DeiT-small&tiny with TPS applied can surpass the pre-trained models by 0.3% and 0.7% in accuracy under 100 fine-tuning epochs. For hybrid ViTs, we can compress the GFLOPS of PVT-tiny by 13% and improve its accuracy by 0.1%.

Fine-Grained Visual categorization. We compare our dTPS with DynamicViT by fine-tuning DeiT on iNaturalist 2019 [29] as shown in Tab. 3. See the supplementary materials for the training details on iNaturalist 2019 [29]. Compared with dynamicViT, our dTPS obtains 0.3% accuracy improvement in DeiT-tiny and 0.2% accuracy improvement in DeiT-small when fine-tuning 30 epochs, respectively. We further fine-tune dTPS 100 epochs and observe a significant improvement in both backbones. Notably, dTPS fine-tuned with 100 epochs is only 0.1% lower than Deit-small while shrinking the computational budgets of Deit-small to 65%.

Epochs of training. Fig. 8 shows that both variants can benefit from longer training epochs and surpass the DeiT-small&tiny with only 65% GFLOPS. However, the benefit of epochs varies slightly in two variants. Because the class-token attention scoring requires no extra optimization target, eTPS performs better than dTPS under 30 epochs. On the other hand, dTPS can benefit more from longer training epochs in DeiT-small, for its learnable scoring brings higher performance upper bound.

Feature Type. We show the effects of feature type used to establish the matching relationships. Supposing $\boldsymbol{x}_{i}$ is the full embedding of the token and the position feature $\boldsymbol{p}_{i}$ is the corresponding positional embedding, we define the content feature as $\boldsymbol{x}_{i}-\boldsymbol{p}_{i}$. As Tab. 4 illustrates, the entire feature is more favorable for it contains both the content and position information.

Similarity Matrix. Considering that the dot-product attention of query and key measures tokens’ relationships naturally, we have tried reusing the previous attention to replace the computations of cosine similarities in the matching step. We believe the previous attention is outdated to measure current tokens’ relations and Tab. 5 shows that calculating the cosine similarity of current features outperforms reusing the attention with only a minor computational increase.

We generate random token selection policies to construct manufactured policy errors that simulate the cases brought by sub-optimal token pruning strategies. All models are based on DeiT-small and fine-tuned 30 epochs with identical pruning setups. We run the experiments under random policies 100 times and report the average results. By comparing the performances of our method with dynamicViT [25] and EViT [16], the accuracy drop from the original to random policies denotes the robustness under incorrect policies. As shown in Tab. 6, our inter-block version dTPS and intra-block version eTPS have fewer accuracy drops than dynamicViT [25] and EViT [16].

Generally, We insert dTPS modules between the patch embedding layer and the subsequent basic block of each pruned stage. Unlike TPS on vanilla ViTs, we reserve attentive tokens from the whole tokens set in each pruned stage and utilize the masking or padding operations to maintain the complete spatial structure during training and inference, respectively.

Training. The spatial reduction layer of the basic block in PVT requires input with a complete spatial structure. For the basic block with a spatial reduction block, given the policy $M$, we maintain the complete spatial structure and mask the dropped tokens in key $K$ and in value $V$ with zeros before the spatial reduction layer as follows:

$$S\!R\!A^{*}(Q,K,V)=M\!H\!A(Q,S\!R(K\odot M),S\!R(V\odot M))$$

(7)

Here, $S\!R\!A^{*}$ is the modified spatial reduction attention, $M\!H\!A$ is the multi-head attention operation, and $SR$ is the spatial reduction layer. Moreover, we perform the same masking operation on dropped tokens before the patch embedding layer of next stage.

For the basic block without a spatial reduction layer, no masking operation is needed ,and we conduct the attention masking strategy from dynamicViT [25] to erase the effects of the dropped tokens.

Inference. The inference procedure of dTPS is adjusted slightly to practically accelerate the spatial reduction layer. The input tokens are pruned by a top-k selection operation based on the scoring results. For the block with a spatial reduction layer, we pad the previously dropped tokens with zero in the $S\!R\!A^{*}$ to maintain the complete spatial structure.

$$K^{\prime}=Pad(K,M)$$

(8)

$$V^{\prime}=Pad(V,M)$$

(9)

$$S\!R\!A^{*}(Q,K,V)=M\!H\!A(Q,S\!R(M^{\prime}),S\!R(V^{\prime}))$$

(10)

Also, the same padding operation is utilized before the patch embedding layer of the next stage. For the block without a spatial reduction layer, no padding operation is needed either. The requirement of complete spatial structure leads to less shrinkage of computations.

The last stage of CvT [33] contains most of its blocks; therefore we only modify the last stage with our dTPS. The operations remain the same for other stages as the original CvT [33].

Training. The convolutional projection operation in CvT requires the input with a complete spatial structure. Given the policy $M$, we mask the dropped tokens with zeros before the convolution projection:

$$Q,K,V=ConvolutionalProjection(X\odot M)$$

(11)

Inference. The input tokens are pruned by a top-k selection operation based on the scoring results. To maintain the complete spatial structure, we pad the previously dropped tokens with zeros in the convolutional projection layer.

$$X^{\prime}=Pad(X,M)$$

(12)

$$Q,K,V=ConvolutionalProjection(X^{\prime})$$

(13)

ATS [8] conducts experiments on CvT [33] as well. It takes a variant of CvT [33] as the pre-trained model without convolutional projection in stage 3. It only performs token pruning in stage 3 to avoid the extra operation to maintain the structured spatial input. Compared to ATS [8], our method utilizes masking and padding during training and inference to keep the spatial structure.

All experiments follow the same data augmentations used in DeiT²²2 https://github.com/facebookresearch/DeiT [28]. All the model is initialized with pre-trained models’ weights and fine-tuned with different token pruning location and token keeping ratio. We adopt the AdamW [18] as the optimizer and a cosine learning rate scheduler.

TPS on DeiT [28]. The experiment settings of dTPS-DeiT follows dynamicViT³³3 https://github.com/raoyongming/DynamicViT except for basic learning rate is set to $\frac{batchsize}{1024}\times 2.5\times 10^{-4}$ and no stage of fixing backbone weights. The experiment settings of eTPS-DeiT follow EViT⁴⁴4https://github.com/youweiliang/evit. The pruning settings include combinations of three multi-layer pruning settings: $prune\_locs\in\{[4,7,10],[3,5,7,9],[4,6,8,10]\}$, and two token keeping ratios :$\rho\in\{0.5,0.7\}$. The token keeping ratio remains the same in all pruning stages.

TPS on LV-ViT [13]. The experiments of dTPS and eTPS on LV-ViT [13] follow the same training settings of dTPS and eTPS on DeiT, except for the basic learning rate is set to $\frac{batchsize}{1024}\times 1.0\times 10^{-4}$ for the stable convergence. For LV-ViT-T, the pruning settings include combinations of three multi-layer pruning settings: $prune\_locs\in\{[4,7,10],[3,5,7,9],[4,6,8,10]\}$, and two token keeping ratios :$\rho\in\{0.5,0.7\}$. For LV-ViT-S, the pruning settings include combinations of three multi-layer pruning settings: $prune\_locs\in\{[5,9,13],[3,6,9,12],[4,7,10,13]\}$, and two token keeping ratios :$\rho\in\{0.5,0.7\}$.

TPS on PS-ViT [39]. The experiments of dTPS on PS-ViT [39] follow the same training settings as ATS [8] on PS-ViT [39]. The basic learning rate is set to $\frac{batchsize}{768}\times 5.0\times 10^{-4}$, $prune\_locs$ is set to [3,6,9] and token keeping ratio $\rho$ is 0.5.

TPS on PVT [31]. The pruning stages include stage 2, stage 3, and stage 4. The token keeping ratio for all dTPS modules is set to 0.7. Basic learning rate is set to $\frac{batchsize}{1024}\times 2.5\times 10^{-4}$.

TPS on CvT [33]. The basic learning rate is set to $\frac{batchsize}{1024}\times 5.0\times 10^{-5}$. The dTPS modules are only inserted into stage 3, and the pruning locations include [3,6,9] for CvT-13,[5,9,13] for CvT-21. The token keeping ratio for all TPS modules is set to 0.5.

TPS on DeiT [28]. For the experiment on iNaturalist 2019 Classification [29], we re-train DeiT and fine-tune the model with dynamicViT or dTPS applied.

In the training step, we initialize DeiT [28] with weights of ImagetNet1K pre-trained model and re-train them for $300$ epochs. The basic learning rate is set to $\frac{batchsize}{1024}\times 10^{-3}$. The other settings follow DeiT [28].

In the fine-tuning step, we initialize dynamicViT-DeiT and dTPS-DeiT with weights from the last step and fine-tune them for 30 epochs with the same pruning setup. The token pruning location is set to [4,7,10], and the token keeping ratio is 0.5. We follow the same fine-tuning settings as the experiments on ImageNet1K, except for no distillation loss.