SPViT: Enabling Faster Vision Transformers via Soft Token Pruning

Zhenglun Kong; Peiyan Dong; Xiaolong Ma; Xin Meng; Mengshu Sun; Wei Niu; Xuan Shen; Geng Yuan; Bin Ren; Minghai Qin; Hao Tang; Yanzhi Wang

Source-linked AI summary

SPViT: Enabling Faster Vision Transformers via Soft Token Pruning

Zhenglun Kong, Peiyan Dong, Xiaolong Ma, Xin Meng, Mengshu Sun, Wei Niu, Xuan Shen, Geng Yuan, Bin Ren, Minghai Qin, Hao Tang, Yanzhi Wang

arXiv:2112.13890v2Dec 27, 2021cs.CVcs.AIcs.ARcs.LG

TL;DR

ViT pruning remains difficult to align with actual device acceleration, while existing dynamic token pruning can discard informative token content. SPViT combines multi-head token selection, soft token packaging, and latency-aware training to balance accuracy with hardware constraints, reducing ViT latency and enabling real-time mobile and FPGA deployment.

Problem
ViT pruning frameworks have a gap between actual device deployment and acceleration, while dynamic token pruning may lose informative content by discarding less informative tokens.
Method
SPViT uses an attention-based multi-head token selector, token packaging, and latency-aware training for hardware-constrained, per-image adaptive soft token pruning.
Results
SPViT achieves comparable accuracy with higher pruning rates, reduces DeiT-S and DeiT-T inference latency by 40%-60% within 0.5% accuracy loss, and enables real-time mobile and FPGA deployment.
Takeaways & Limitations
SPViT enables ViT models to meet real-time requirements on mobile devices and FPGA while balancing accuracy against hardware constraints.
Takeaways & Limitations
The authors identify large data movement from multiple blocks and intermediate results as a memory pressure that requires further hardware-deployment optimization.

Abstract

from arXiv · show

Recently, Vision Transformer (ViT) has continuously established new milestones in the computer vision field, while the high computation and memory cost makes its propagation in industrial production difficult. Pruning, a traditional model compression paradigm for hardware efficiency, has been widely applied in various DNN structures. Nevertheless, it stays ambiguous on how to perform exclusive pruning on the ViT structure. Considering three key points: the structural characteristics, the internal data pattern of ViTs, and the related edge device deployment, we leverage the input token sparsity and propose a computation-aware soft pruning framework, which can be set up on vanilla Transformers of both flatten and CNN-type structures, such as Pooling-based ViT (PiT). More concretely, we design a dynamic attention-based multi-head token selector, which is a lightweight module for adaptive instance-wise token selection. We further introduce a soft pruning technique, which integrates the less informative tokens generated by the selector module into a package token that will participate in subsequent calculations rather than being completely discarded. Our framework is bound to the trade-off between accuracy and computation constraints of specific edge devices through our proposed computation-aware training strategy. Experimental results show that our framework significantly reduces the computation cost of ViTs while maintaining comparable performance on image classification. Moreover, our framework can guarantee the identified model to meet resource specifications of mobile devices and FPGA, and even achieve the real-time execution of DeiT-T on mobile platforms. For example, our method reduces the latency of DeiT-T to 26 ms (26%$\sim $41% superior to existing works) on the mobile device with 0.25%$\sim $4% higher top-1 accuracy on ImageNet.

1 Introduction

SPViT addresses limitations in Vision Transformer pruning by combining adaptive token selection, soft token packaging, and latency-aware training for edge deployment. It reduces latency while preserving comparable accuracy and supports real-time mobile and FPGA execution.

Motivation: ViT pruning faces deployment gaps because head pruning alleviates only part of computation, while static pruning ignores image-specific information variation.Dynamic pruning can adapt per image, but careless selection may limit pruning or reduce accuracy.
Method: SPViT uses an attention-based multi-head token selector to score tokens from their characteristics across attention heads.The selector produces token importance scores for adaptive keep/prune decisions.
Method: Less informative tokens are compressed into a package token and concatenated with informative tokens for later transformer blocks.The package token retains residual information that may correct early-layer pruning errors and preserve useful background features.
Method: Latency-aware training links pruning rates to device specifications and progressively searches selector placement under a target latency budget.The strategy inserts selectors progressively, groups blocks with similar pruning rates, and removes redundant selectors while maintaining accuracy.
Results: 40%-60% inference-latency reductions for DeiT-S and DeiT-T occur within 0.5% accuracy loss, while PiTs and Swins incur negligible performance drops.The reported lightweight-model results establish SPViT as particularly effective for compressing smaller models.

2 Related Work

Related work adapts pruning and token reduction to Vision Transformers, but existing approaches differ in which tokens or dimensions they remove and how they preserve image information.

Vision Transformers: Vision Transformers enable global positional interaction, unlike CNNs’ restricted fixed-window spatial interactions.Variants such as DeiT, T2T-ViT, Mixer, PiT, LV-ViT, and PS-ViT address efficiency, training, or representation design.
Efficient ViT: Efficient ViT methods prune tokens, attention heads, or patches, estimate token importance, evolve tokens, or generate adaptive token vectors.Examples include S2ViTE, DynamicViT, Evo-ViT, TokenLearner, and PATCHMERGER.

3 Computational Complexity Analysis

The analysis distinguishes token, channel, and attention-head pruning by how dimensions propagate through Transformer operations. Because token count N is transitive, token pruning can reduce computation across operations, including quadratic attention terms, whereas channel and head pruning have narrower effects.

The corrected ViT complexity expression separates channel, attention, and feed-forward dimensions: 4NDchDattn + 2N^2Dattn + 8NDchDfc.
Pruning branches: Token-channel pruning reduces computation only in the current matrix multiplication because Dch is non-transmissible.A mask layer is needed before linear layers to propagate the reduction across layers.
In the Transformer block, N×Dch inputs produce Q, K, and V matrices of size N×Dattn through three Dch×Dattn transformations.
Pruning branches: Token pruning reduces all operations linearly or quadratically because N is transitive and appears as N^2 in attention operations.
Pruning branches: Attention-head pruning primarily reduces Dattn computation in the MSA module, which usually accounts for less than 40% of total ViT computation.

4 Latency-Aware Soft Pruning

SPViT combines adaptive multi-head token selection with soft token packaging and latency-aware training to reduce ViT computation under edge-device constraints. The framework uses progressive, hardware-aware pruning while retaining information from less informative tokens for later computation.

Framework Overview: SPViT inserts hierarchical token selectors between Transformer blocks, classifying tokens as informative or less informative before subsequent calculations.The selector updates pruning decisions across stages during hierarchical pruning.
Multi-Head Token Selector: Each attention head produces token scores from local and global features, while head-wise attention scores weight their contributions to the overall token score.The selector uses per-head feature processing and an attention-based branch to capture dependencies among heads.
Multi-Head Token Selector: Gumbel-Softmax converts token probability scores into differentiable binary keep/prune decisions during training.The resulting decisions are propagated to later layers and combined with subsequent decisions using a Hadamard product.
Token Packaging Technique: Less informative tokens are merged into a package token rather than discarded, allowing residual information to participate in subsequent blocks and help correct scoring mistakes.The package token is formed from the less informative tokens using their keep probabilities and is concatenated with informative tokens.
Hardware Efficiency: The selector adds less than 1% of total model GFLOPs, and its MLP, Softmax, Pooling, and Sigmoid operations are supported on edge platforms.The design avoids unsupported Top-k-style ordering and reuses GEMM hardware through fully connected layers.
Latency-Aware Training Strategy: Latency-aware training maps pruning rates to measured block latency and constrains total inference latency to the target device limit.The training objective combines classification, distillation, and latency-aware sparsity losses, while progressive layer-to-phase training determines selector placement and pruning rates.

5 Experiments

Experiments evaluate SPViT across flat and hierarchical ViTs, pruning configurations, accuracy–computation trade-offs, and deployment on mobile and FPGA hardware. The results show substantial computation and latency reductions with small accuracy losses, while progressive training identifies effective selector placement.

Main Results: SPViT reduces computation cost by 31%–43% across various backbones with only 0.1%–0.5% accuracy degradation.For DeiT-T, it reduces GFLOPs by 31% with top-1 accuracy of 72.10% versus 72.20%.
Results on Hierarchical Architectures: SPViT reduces computation cost by 23%–27% for Swin and 18%–24% for PiT, with accuracy degradation of 0.4%–0.5% and 0.2%–0.6%, respectively.These results extend soft token pruning to lightweight hierarchical architectures.
Deployment on Edge Devices: 26 ms per inference enables DeiT-T to meet the real-time requirement on Samsung Galaxy S20 mobile CPUs.For DeiT-S, latency decreases from 113 ms to 60 ms with only a 0.46% accuracy decrease.
Deployment on Edge Devices: DeiT-T and DeiT-S achieve 1.57× and 1.69× acceleration, respectively, on the Xilinx ZCU102 FPGA.The FPGA evaluation uses 16-bit fixed-point parameters and activation data.
Token Pruning Visualization: SPViT progressively removes less informative tokens while preserving representative object regions through adaptive per-image pruning.The visualized pruning phases move from background removal to object-focused and informative-feature localization.
Ablation Analysis: The 3-6-9 selector placement gives the highest accuracy and lowest computation cost, while each selector and package token adds 1.67 ms on Samsung Galaxy S20.Earlier blocks are difficult pruning locations because token representations are insufficiently encoded.

6 Conclusion

SPViT combines dynamic soft token pruning, token packaging, and latency-aware training to balance accuracy with hardware constraints. The framework is deployed on mobile and FPGA platforms while meeting real-time requirements.

Conclusion: SPViT combines an attention-based multi-head token selector, token packaging, and latency-aware training to balance accuracy with hardware constraints.The framework is deployed on mobile and FPGA platforms.

A Analysis and Discussion

SPViT provides token pruning as an alternative to changing ViT channel dimensions for model scaling. Comparisons with scaled DeiT models show better complexity–accuracy trade-offs under comparable computation.

Model Scaling Comparison: SPViT offers token pruning as an alternative to scaling ViTs by changing the number of channels.The comparison varies DeiT embedding dimensions from 192 to 384 and evaluates models under comparable GFLOPs.
Model Scaling Comparison: SPViT models achieve better complexity–accuracy trade-offs than the compared scaled DeiT models under comparable GFLOPs.Figure 8 illustrates the comparison between SPViT and model scaling.

A.2 Progressive Training Sparsity for Each Layer

Progressive training reveals layer groups with similar sparsity and accuracy behavior, supporting their consolidation into pruning phases. The reported layer structure places selectors at the fronts of these phases.

Progressive Training Sparsity for Each Layer: Layers 6–8 and 9–11 show similar accuracy and sparsity, supporting their combination into separate pruning phases.The selector is inserted before the layer indicated by the layer index.
Progressive Training Sparsity for Each Layer: Progressive training adds token selectors sequentially from the 11th layer toward the 1st layer to measure layer-wise sparsity and accuracy.The resulting measurements are reported in Table 7.

A.3 Model Latency on Hardware

SPViT reduces latency on mobile hardware while preserving competitive accuracy, including substantial gains for both lightweight and larger DeiT models. Its token packaging design and selector variants are further examined through ablations and model analyses.

Hardware latency: 38ms vs. 44ms: SPViT outperforms DeiT-T by up to 4.8% with smaller latency on Samsung Galaxy S20.The comparison uses Galaxy S20 measurements with a Snapdragon 865 processor.
Hardware latency: 60ms vs. 113ms: SPViT reduces DeiT-S latency by up to 47% with only 0.46% accuracy decrease.
Hardware latency: SPViT models outperform DynamicViT on both latency and accuracy for LV-ViT-S/M.
Token packaging: Token packaging passes package tokens between pruning modules either by addition, yielding 1 additional token, or concatenation, yielding 3.
Ablations: Token packaging improves DeiT performance by 0.12% with a class token, from 72.0% to 72.1%.
Ablations: SPViT accuracy shows a slight boost as batch size increases, while attention analysis indicates that pooling-layer redundancy can be recognized precisely.

B.1 Token Pruning Visualization

Visualizations show that token packaging changes pruning behavior across phases and helps focus retained information on objects rather than background. Head-specific heatmaps support scoring tokens individually because attention heads encode different image features.

Phase-wise pruning: Token packaging helps lock onto objects instead of background, especially in complex and multi-object images.
Phase-wise pruning: Figure 10 compares pruning visualizations with and without token packaging across rows and pruning phases.
Head-wise attention: Each DeiT-S attention head focuses on different image features and visual receptive fields, producing head-specific token importance.

C Main Results

SPViT is evaluated against representative vision transformer and pruning methods on ImageNet-1K using top-1 accuracy and GFLOPs under comparable computational settings. The comparison includes reproduced models matched to similar GFLOPs where indicated.

ImageNet-1K comparison: SPViT is compared with representative ViT pruning and architecture methods using top-1 accuracy and GFLOPs.
ImageNet-1K comparison: Table 13 reports comparisons under comparable GFLOPs and parameter counts, with reproduced results marked by an asterisk.Negative FLOPs-reduction values indicate FLOP increases; embedding dimensions and SPViT scaling variants are also identified.