Source-linked AI summary
Chasing Sparsity in Vision Transformers: An End-to-End Exploration
Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, Zhangyang Wang
TL;DR
ViTs are costly to train and deploy, while conventional pruning can add substantial retraining expense. The paper introduces end-to-end sparse training that dynamically explores subnetworks, structured attention heads, and selected input tokens. On ImageNet, the approach substantially reduces computation with nearly preserved accuracy and can sometimes improve it.
Problem
ViTs have large models and high training costs, while conventional post-training pruning can require additional retraining budgets.
Method
The paper dynamically trains sparse ViT subnetworks under a fixed parameter budget, extends sparsity to attention heads, and learns to select important input tokens.
Results
ImageNet experiments across DeiT backbones show substantial computation savings with nearly undamaged accuracy, including a 0.28% accuracy improvement with 49.32% FLOPs savings on DeiT-Small.
Takeaways & Limitations
End-to-end sparse (co-)training can reduce ViT training-memory and inference costs without necessarily sacrificing accuracy.
Abstract
from arXiv · showhide
Vision transformers (ViTs) have recently received explosive popularity, but their enormous model sizes and training costs remain daunting. Conventional post-training pruning often incurs higher training budgets. In contrast, this paper aims to trim down both the training memory overhead and the inference complexity, without sacrificing the achievable accuracy. We carry out the first-of-its-kind comprehensive exploration, on taking a unified approach of integrating sparsity in ViTs "from end to end". Specifically, instead of training full ViTs, we dynamically extract and train sparse subnetworks, while sticking to a fixed small parameter budget. Our approach jointly optimizes model parameters and explores connectivity throughout training, ending up with one sparse network as the final output. The approach is seamlessly extended from unstructured to structured sparsity, the latter by considering to guide the prune-and-grow of self-attention heads inside ViTs. We further co-explore data and architecture sparsity for additional efficiency gains by plugging in a novel learnable token selector to adaptively determine the currently most vital patches. Extensive results on ImageNet with diverse ViT backbones validate the effectiveness of our proposals which obtain significantly reduced computational cost and almost unimpaired generalization. Perhaps most surprisingly, we find that the proposed sparse (co-)training can sometimes improve the ViT accuracy rather than compromising it, making sparsity a tantalizing "free lunch". For example, our sparsified DeiT-Small at (5%, 50%) sparsity for (data, architecture), improves 0.28% top-1 accuracy, and meanwhile enjoys 49.32% FLOPs and 4.40% running time savings. Our codes are available at https://github.com/VITA-Group/SViTE.
1 Introduction
ViTs deliver strong vision performance but impose substantial training and inference costs. This paper explores end-to-end dynamic sparsity to reduce those costs while preserving accuracy across model, structure, and input data.
- ViTs achieve strong results across multiple vision tasks but have enormous parameter counts and rising training costs.
- Post-training pruning reduces inference complexity but often requires costly retraining, making it especially burdensome for ViTs.
- Sparse training instead identifies and trains subnetworks from scratch, motivated by evidence that sparse winning tickets can reach dense-network accuracy.
- The paper unifies end-to-end dynamic sparsity to reduce both inference costs and training memory overhead in ViTs.
- SViTE jointly optimizes parameters and connectivity under a fixed small parameter budget, while S2ViTE guides prune-and-grow over self-attention heads.
- SViTE+ adds a learnable token selector that identifies vital patch embeddings for joint data and architecture sparsity.
- 49.32% FLOPs and 4.40% running time savings accompany a 0.28% accuracy improvement for SViTE+ on DeiT-Small at 50% model and 5% data sparsity.
2 Related Work
Related work covers ViT architectures, compression, efficient attention, pruning, and sparse training. The paper positions its approach as distinct from post-training compression and attention approximations by jointly exploring sparse weights, structures, and tokens during training.
- ViT encodes images as patch-token sequences processed by transformer encoders, with variants improving training-data efficiency.
- Existing ViTs remain heavy to train and massive after training, motivating compression methods that reduce both training and inference costs.
- Initial ViT compression studies prune intermediate features or tokens after training but do not jointly address weight pruning and efficient training.
- Efficient attention methods reduce self-attention complexity through approximations such as low-rank decomposition, representing an orthogonal direction to this paper’s techniques.
- Pruning includes unstructured weight-element removal and structured removal of channels or attention heads, with structured sparsity generally better aligned with hardware efficiency.
- Sparse training adaptively identifies and optimizes sparse subnetworks and has included evolutionary connectivity updates and weight reallocation.
3 Methodology
The methodology dynamically explores sparse ViT subnetworks during training, supports unstructured and structured sparsity, and jointly selects informative input tokens for additional efficiency. It uses scheduled prune-and-grow updates, importance criteria for heads and neurons, and differentiable top-k token selection.
- Sparse ViT Exploration: SViTE trains a sparse ViT while jointly optimizing parameters and exploring connectivity under a fixed sparsity budget.Sparse topology updates repeatedly prune existing connections and grow new ones during training.
- Sparse ViT Exploration: Erdős-Rényi sparsity distributes remaining weights across layers, while cosine-decayed prune-and-grow updates continue until 80% of training.Weight magnitude determines pruning, and gradient magnitude determines newly grown connections.
- Structured Sparsity: Structured S2ViTE extends sparse exploration to hardware-friendly self-attention head pruning using first-order importance estimates.The method inherits the unstructured variant’s sparsity distribution and update schedule.
- Structured Sparsity: MLP neurons are pruned using the ℓ1-norms of their associated weight vectors and regrown using high-magnitude gradients.Newly added units are initialized to zero.
- Data and Architecture Sparsity: SViTE+ scores input token embeddings and differentiably selects the top-k informative tokens using Gumbel-Softmax and straight-through optimization.Reducing input tokens compresses intermediate features and yields additional efficiency gains.
4 Experiments
Experiments on ImageNet with DeiT backbones show that end-to-end sparse training reduces computation while preserving or improving accuracy across unstructured, structured, and data-plus-architecture sparsity settings.
- SViTE with Unstructured Sparsity: SViTE reduces DeiT FLOPs by 25.56%–57.50% across 30%–60% sparsity, with accuracy drops within 0.5%.The reductions span Tiny, Small, and Base backbones, while SViTE consistently surpasses OMP, GMP, TP, and Small-Dense baselines.
- S2ViTE with Structured Sparsity: S2ViTE achieves 23.79%–33.63% FLOPs savings and additional 10.57%–24.70% running-time reductions at 30%–40% structured sparsity.It maintains competitive top-1 accuracy and consistently outperforms the structured pruning baseline SSP.
- S2ViTE with Structured Sparsity: At 40% structured sparsity, S2ViTE-Base surpasses dense DeiT-Base by 0.42%–1.24% accuracy while reducing parameters by 34.41%, FLOPs by 33.13%, and running time by 24.70%.The authors associate this gain with implicit regularization and greater structural redundancy in larger ViTs.
- SViTE+ with Data and Architecture Sparsity Co-Exploration: SViTE+-Small removes 5%–10% of tokens while achieving 4.40%–7.63% running-time and 49.32%–52.38% FLOPs savings with improved top-1 accuracy.The benefit is attributed to reduced input and intermediate feature dimensions; adding token slimming to S2ViTE is less effective.
- Ablation and Generalization Study: SViTE and S2ViTE generalize beyond DeiT: on TNT-S, SViTE gains 0.13 accuracy points with 37.54% FLOPs savings, while S2ViTE cuts FLOPs by 32.96% and running time by 23.71%.S2ViTE-TNT-S retains almost unimpaired accuracy at 40% structured sparsity.
- Visualization: Visualization reveals structural redundancy in attention heads and emergent vertical patterns in unstructured masks, indicating potential hardware speedups.S2ViTE removes unnecessary heads, while SViTE-Base activates fewer heads than dense DeiT-Base.
- Visualization: The token selector preferentially removes patches around or in the background while preserving patches within objects of interest.These learned selection patterns support the effectiveness of SViTE+ token selection.
5 Conclusion and Discussion of Broader Impact
The paper introduces end-to-end sparse ViT exploration methods and reports ImageNet validation with reduced memory and computational costs. It also anticipates substantially lower energy costs and more economical deployment, such as on smartphones.
- SViTE, S2ViTE, and SViTE+ explore sparse patterns in ViT architecture and input token embeddings.The methods target training memory bottlenecks and inference efficiency, including running time and FLOPs.
- Comprehensive ImageNet experiments validate the effectiveness of the proposed sparse ViT methods.
- Sparse vision transformers are highly likely to reduce memory and energy costs substantially, supporting economic deployment in applications such as smartphones.
Checklist
The checklist reports reproducibility materials, training details, asset citations, and ethics disclosures. It also records that error bars were not reported because the ImageNet experiments were extremely resource-consuming.
- The checklist says the work describes its limitations and discusses potential negative societal impacts and ethics-review compliance.
- The paper provides publicly available data, detailed implementations or citations, training details, and code at the SViTE repository.
- The authors did not report error bars because each reported number required around 960 V100 GPU hours.They state that confidence intervals will be reported in future versions.
- The paper cites the creators of the publicly available ImageNet data but does not separately state its asset license.
- The authors state that they did not collect or curate new data, and mark participant-related disclosures as not applicable.
A1 More Implementation Details
The implementation details describe the computing resources and training-budget rationale, then compare SViTE with longer-trained DeiT baselines. The comparison reports better accuracy with fewer parameters and computations at 40% structured sparsity.
- Experiments use Tesla V100-SXM2-32GB GPUs, with each experiment running on 8 V100s for 4 ∼5 days.
- 600 SViTE training epochs approximately match the FLOPs of 300 dense DeiT epochs at 40% sparsity.The stated comparison is ∼0.95x FLOPs for SViTE versus 1x FLOPs for dense DeiT.
- SViTE remains able to achieve better accuracy with fewer parameters and fewer training and inference computations than longer-trained DeiT baselines.The authors report this specifically at 40% structured sparsity.
- The longer-epoch comparison uses DeiT-Small and DeiT-Base on ImageNet-1K.
A2 More Experimental Results
Additional experiments visualize explored sparse topologies, compare token and architecture sparsity, and examine layerwise attention-map sparsity. The results show emergent structural patterns and improved combined-sparsity performance in the reported ablation.
- Unstructured SViTE exploration produces emergent structural patterns, including completely pruned MLP neurons, from an initial random mask.The authors identify potential for real-world hardware acceleration.
- S2ViTE explores connectivity patterns distinct from the initial topology while pruning multi-attention heads and MLPs at 40% structural sparsity.
- 79.91 test accuracy results from combining 10% data sparsity, a token selector, and 50% unstructured SViTE sparsity.The token-selector-only setup reached 78.67, versus 79.90 for DeiT-Small.
- SViTE+ attention maps are denser than dense ViT in bottom layers but sparser in top layers.The analysis uses 10% data sparsity and 50% model sparsity, counting elements below 10^-4.