PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Siyuan Wei, Tianzhu Ye, Shen Zhang, Yao Tang, Jiajun Liang

arXiv:2304.10716v1cs.CV

TL;DR

Vision transformers require substantial computation, and aggressive token pruning can lose information that affects prediction. TPS jointly prunes tokens and squeezes pruned-token information into reserved tokens, outperforming prior approaches across pruning intensities and improving accuracy by 1%-6% at 35% DeiT-small/tiny budgets. Experiments across transformer types also show flexibility and robustness to token-selection errors.

  • Problem

    Vision transformers have high computational costs, while token-pruning errors can discard information that affects performance during aggressive compression.

  • Method

    TPS divides tokens into reserved and pruned subsets, then transfers pruned-token information into reserved tokens through nearest-neighbor matching and similarity-based fusion.

  • Results

    TPS outperforms baselines by 1%-6% accuracy when DeiT-small and DeiT-tiny are compressed to 35% of their GFLOPs, with broader experiments showing flexibility and robustness.

  • Takeaways & Limitations

    Conserving information from pruned tokens enables more aggressive vision-transformer compression while maintaining stronger performance across architectures and pruning conditions.

  • Takeaways & Limitations

    Structured spatial operations constrain straightforward token-pruning integration in hybrid ViTs, and fine-tuning pretrained models increases total training time.

Abstract

from arXiv · show

Although vision transformers (ViTs) have shown promising results in various computer vision tasks recently, their high computational cost limits their practical applications. Previous approaches that prune redundant tokens have demonstrated a good trade-off between performance and computation costs. Nevertheless, errors caused by pruning strategies can lead to significant information loss. Our quantitative experiments reveal that the impact of pruned tokens on performance should be noticeable. To address this issue, we propose a novel joint Token Pruning & Squeezing module (TPS) for compressing vision transformers with higher efficiency. Firstly, TPS adopts pruning to get the reserved and pruned subsets. Secondly, TPS squeezes the information of pruned tokens into partial reserved tokens via the unidirectional nearest-neighbor matching and similarity-based fusing steps. Compared to state-of-the-art methods, our approach outperforms them under all token pruning intensities. Especially while shrinking DeiT-tiny&small computational budgets to 35%, it improves the accuracy by 1%-6% compared with baselines on ImageNet classification. The proposed method can accelerate the throughput of DeiT-small beyond DeiT-tiny, while its accuracy surpasses DeiT-tiny by 4.78%. Experiments on various transformers demonstrate the effectiveness of our method, while analysis experiments prove our higher robustness to the errors of the token pruning policy. Code is available at https://github.com/megvii-research/TPS-CVPR2023.

1. Introduction

Vision transformers are computationally expensive because token length drives quadratic attention costs, while aggressive token pruning can discard essential subject and context information. TPS addresses this by squeezing pruned-token information into reserved tokens, enabling more aggressive compression with stronger accuracy and robustness across settings.

  • Motivation: Token pruning reduces computation but can lose essential subject and background context, causing errors when the reserved-token count becomes very small.Examples include losing lawn context needed to identify a lawn mower and subject information needed to distinguish a baseball from a rugby ball.
  • Motivation: A reversed-pruning toy experiment shows that information in pruned tokens can provide complementary accuracy, increasingly so as pruning continues.The experiment uses DynamicViT [25] and defines the added performance as bonus accuracy.
  • Proposed method: TPS divides tokens into reserved and pruned subsets, then uses feature dispatch to squeeze essential pruned-token information into reserved tokens.This avoids simply discarding pruned tokens or collapsing them into a single token.
  • Experimental evidence: 1%-6% accuracy improvements over baselines are achieved when compressing DeiT-small and DeiT-tiny to 35% of their GFLOPs.The experiments compare TPS with prior token-pruning and token-reorganization approaches under aggressive compression.
  • Experimental evidence: Experiments across vanilla and hybrid ViTs show TPS flexibility, while random-token evaluations indicate greater robustness than token pruning and token reorganization.The evaluation covers ImageNet1K and iNaturalist 2019 [29].

2. Related Work

Prior ViT compression methods reduce computation by pruning or reorganizing tokens, but discarded information limits aggressive compression. TPS instead matches pruned tokens to reserved host tokens and fuses their features, with variants supporting different insertion points and scoring policies.

  • Token Pruning: Token pruning reduces computation by discarding redundant tokens, but pruned information loss causes accuracy drops that limit higher pruning ratios.Token reorganization preserves some background context by collapsing pruned tokens into one, yet noticeable performance loss remains at higher pruning ratios.
  • Motivation: A reversed-policy experiment shows that pruned tokens retain complementary predictive information, with bonus accuracy increasing as pruning becomes more aggressive.This observation motivates assimilating pruned-token information into reserved tokens rather than discarding it.
  • TPS Variants: The method includes dTPS and eTPS, which support inter-block and intra-block compression while adopting DynamicViT [25] and EViT [16] scoring methods for fair comparison.Both variants use Top-k selection with a fixed token reduction ratio and maintain constant tensor shape for inference optimization.
  • Token Pruning & Squeezing: TPS preserves discarded information by splitting tokens into reserved and pruned subsets, then matching each pruned token to a similar reserved host for feature fusion.Matching is unidirectional, so multiple pruned tokens may share one host; the resulting fusion processes only as many tokens as the reserved subset.
  • Token Squeezing: TPS uses similarity-based weighting rather than simple averaging, expanding closer tokens’ influence while limiting effects from imperfect token scoring.Cosine similarity between reserved and pruned features can outperform the attention map and introduces no extra matching parameters.

4. Experiment

Experiments show that TPS improves accuracy–computation trade-offs under aggressive pruning, generalizes across transformer backbones and datasets, and is robust to imperfect token-selection policies.

  • Comparison to baselines: At 35% of DeiT-small&tiny computational budgets, TPS improves accuracy by 1%-6% over DynamicViT and EViT, while achieving 1745 images/s on DeiT-small and exceeding DeiT-tiny accuracy by 4.78%.Under all tested pruning settings, TPS outperforms both baselines; the throughput comparison uses a single NVIDIA RTX 2080Ti with batch size 32.
  • Comparison to states of the art: TPS achieves a competitive accuracy–computation trade-off against numerous token-pruning, vanilla, and hybrid vision transformers on ImageNet1K after 30 fine-tuning epochs.The comparison integrates TPS into DeiT-small&tiny and LV-ViT-small&tiny.
  • Extension on more backbones: TPS generalizes across vanilla and hybrid ViTs, surpassing several token-pruning methods and improving PVT-tiny accuracy by 0.1% while reducing its GFLOPs by 13%.On iNaturalist 2019, dTPS improves DeiT-tiny accuracy by 0.3% and DeiT-small accuracy by 0.2% over DynamicViT after 30 epochs.
  • Epochs of training: Both TPS variants benefit from longer training and surpass pre-trained DeiT-small&tiny at 65% GFLOPs, with dTPS gaining more from extended training on DeiT-small.eTPS performs better after 30 epochs, whereas dTPS benefits more from longer training because its learnable scoring has a higher performance upper bound.
  • Feature Type: Using full token embeddings for matching is more favorable because they combine content and positional information, while current-feature cosine similarity outperforms reused attention with a minor computational increase.The feature-type study defines content features as x_i − p_i, but reports the entire embedding as preferable.
  • Robustness Experiments: Under random token-selection policies, dTPS and eTPS experience smaller accuracy drops than DynamicViT and EViT, indicating greater robustness to pruning-policy errors.The evaluation averages results over 100 random policies using DeiT-small fine-tuned for 30 epochs.

5. Conclusions and Limitations

The paper presents TPS as a more aggressive vision-transformer compression module that conserves information and avoids significant performance drops. It reports flexibility and robustness across backbones and quantitative analyses, while noting integration and training limitations.

  • Conclusions: TPS compresses vision transformers more aggressively while conserving information and avoiding a significant performance drop compared with token pruning and reorganization.The method is also reported to have better efficiency than prior token-pruning methods and vision-transformer state of the arts.
  • Conclusions: Extensive experiments across various backbones and quantitative analyses show TPS flexibility and robustness.
  • Limitations: Structured spatial operations in hybrid vision transformers restrict straightforward integration of token pruning.
  • Limitations: Fine-tuning pretrained models could be replaced by pruning-aware training-from-scratch schemes to shorten total training time.
  • Limitations: The authors plan to improve adaptation to hybrid vision transformers and apply TPS to denser prediction tasks.

1. Overview

The supplemental materials provide implementation variants, hybrid-ViT coverage, experimental settings, larger-model evaluations, keep-ratio studies, and additional TPS ablations.

  • Supplementary contents: The supplement details two TPS variants and their application to hybrid vision transformers.
  • Supplementary contents: It includes detailed experiment settings and evaluations on larger models and larger input sizes.
  • Supplementary contents: Additional analyses examine TPS under different token keep ratios and provide more ablations of the TPS design.

2. Visualizations

Visualization cases show TPS recovering information lost by imperfect pruning, correcting predictions that dynamicViT gets wrong.

  • Visualizations: TPS correctly predicts ImageNet1K cases that dynamicViT-DeiT predicts incorrectly by squeezing pruned-token information into similar reserved tokens.Imperfect pruning can remove background context and parts of the subject, producing close but incorrect predictions.
  • Visualizations: The visualized errors involve lost background context and incomplete subjects that leave the model with confusing visual evidence.

3. Details of Two Variants

TPS has dTPS and eTPS variants with different scoring, insertion, and parameterization choices, and it is adapted to preserve spatial structure in hybrid ViTs. Its matching and fusing designs, plus visual comparisons, support robust token squeezing under aggressive pruning.

  • Variants: TPS provides dTPS and eTPS variants that support fair comparisons with dynamicViT and EViT and can theoretically accompany any token-pruning method.
  • Variants: dTPS is inserted before transformer blocks and adds a small parameter overhead, while eTPS is inserted after attention and is parameter-free.The insertion difference follows from their distinct token-scoring mechanisms.
  • Variants: Under 30 training epochs eTPS outperforms dTPS, while under 100 epochs dTPS performs better, indicating extra dTPS parameters provide a higher upper limit.
  • Hybrid ViTs: For hybrid ViTs, TPS preserves complete spatial structure through masking during training and padding during inference, though this reduces computational shrinkage.PVT uses modified spatial-reduction attention, while CvT pads dropped tokens in convolutional projection.
  • Visualizations: In the visualized comparison, dynamicViT leaves pruned tokens blank, whereas TPS masks groups of matched tokens under the same pruning setup.
  • Matching and fusing: N:1 matching assigns each pruned token to its nearest reserved token, and similarity-weighted fusing has a lower accuracy drop than average fusing under random squeezing.

5. Detailed Experiment Settings

The experiments evaluate TPS across multiple vision-transformer architectures, pruning locations, token-keeping ratios, and datasets using architecture-specific fine-tuning settings.

  • Training protocol: The experiments initialize models from pretrained weights and fine-tune them with different pruning locations and token-keeping ratios using AdamW and cosine learning-rate scheduling.The general setup follows DeiT2 data augmentations and pretrained-model initialization.
  • Architecture-specific settings: TPS is evaluated on DeiT, CvT, LV-ViT, PS-ViT, and PVT with architecture-specific pruning and training configurations.The settings vary pruning locations, token-keeping ratios, learning rates, and whether dTPS or eTPS is used.
  • DeiT settings: DeiT experiments compare dTPS and eTPS across multi-layer pruning locations and token-keeping ratios of 0.5 or 0.7.The listed pruning configurations include three multi-layer location combinations, with the ratio held constant across pruning stages.
  • LV-ViT settings: LV-ViT experiments use separate pruning-location configurations for tiny and small models, each tested with token-keeping ratios of 0.5 and 0.7.LV-ViT-T and LV-ViT-S use different three-option pruning-location sets.
  • Keep-ratio evaluation: DeiT-S performance under different keep ratios is presented in Figure 11.The figure is identified as the evaluation of TPS under different keep ratios on DeiT-S.
  • Additional dataset: The iNaturalist 2019 experiment retrains DeiT for 300 epochs, then fine-tunes dynamicViT-DeiT and dTPS-DeiT for 30 epochs at pruning locations [7] and ratio 0.5.The fine-tuning models are initialized from the retrained DeiT weights.

6. TPS on Larger Models and with Larger Input Size

The paper evaluates TPS on DeiT-B and at 384 × 384 input resolution to test compatibility with larger models and larger inputs.

  • Larger models: TPS is evaluated on DeiT-B to demonstrate compatibility with larger models.The DeiT-B experiment is shown in Fig. 9a.
  • Larger input size: TPS is evaluated with 384 × 384 inputs to assess performance at a larger input size.The larger-input experiment is shown in Fig. 9b.
  • Evaluation scope: The larger-model and larger-input evaluations are both presented as compatibility tests for TPS.The passage links the DeiT-B and 384 × 384 experiments to broader applicability of the module.

7. TPS under Different Keep Ratios

The paper studies TPS across different keep ratios and ablates its matching, fusing, and token-scoring choices to assess design contributions.

  • Keep-ratio study: TPS performance under different keep ratios on DeiT-S is reported in Figure 11.The experiment directly varies the token keep ratio.
  • Matching and fusing ablation: TPS improvement benefits from compressing pruned-token information while leaving unmatched reserved tokens unchanged.This conclusion is reported from the matching and fusing ablations in Table 7a.
  • Token-scoring ablation: Token scoring is necessary for squeezing, because random token division produces a significant performance drop.The result is reported in Table 7b and connected to the paper’s robustness analysis.
Loading 2304.10716v1…