Chasing Sparsity in Vision Transformers: An End-to-End Exploration

Abstract

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.

arXiv commentNeurIPS 2021

References in corpus (38)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Longformer: The Long-Document Transformer
• VisualBERT: A Simple and Performant Baseline for Vision and Language
• Transformer in Transformer
• Linformer: Self-Attention with Linear Complexity
• A Survey of Model Compression and Acceleration for Deep Neural Networks
• VL-BERT: Pre-training of Generic Visual-Linguistic Representations
• The State of Sparsity in Deep Neural Networks
• Axial Attention in Multidimensional Transformers
• DeepViT: Towards Deeper Vision Transformer
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up
• LXMERT: Learning Cross-Modality Encoder Representations from Transformers
• Long Range Arena: A Benchmark for Efficient Transformers
• Sparse Networks from Scratch: Faster Training without Losing Performance
• Picking Winning Tickets Before Training by Preserving Gradient Flow
• Stabilizing the Lottery Ticket Hypothesis
• Pre-Trained Image Processing Transformer
• Parameter Efficient Training of Deep Convolutional Neural Networks by Dynamic Sparse Reparameterization
• Rethinking Attention with Performers
• End-to-End Object Detection with Adaptive Clustering Transformer
• A topological insight into restricted Boltzmann machines
• Structured Pruning of a BERT-based Question Answering Model
• End-to-End Video Instance Segmentation with Transformers
• Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets
• CMT: Convolutional Neural Networks Meet Vision Transformers
• Compressive Transformers for Long-Range Sequence Modelling
• Reweighted Proximal Pruning for Large-Scale Language Representation
• The Lottery Tickets Hypothesis for Supervised and Self-supervised Pre-training in Computer Vision Models
• The Difficulty of Training Sparse Neural Networks
• A Unified Lottery Ticket Hypothesis for Graph Neural Networks
• Top-KAST: Top-K Always Sparse Training
• Do We Actually Need Dense Over-Parameterization? In-Time Over-Parameterization in Sparse Training
• EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets
• Selfish Sparse RNN Training
• Patch Slimming for Efficient Vision Transformers
• Sparse Weight Activation Training
• Efficient Lottery Ticket Finding: Less Data is More

Cited by in corpus (8)

• Recent Progress in Transformer-based Medical Image Analysis
• IA-RED${}^2$: Interpretability-Aware Redundancy Reduction for Vision Transformers
• Vision Transformers on the Edge: A Comprehensive Survey of Model Compression and Acceleration Strategies
• S${}^2$-MLPv2: Improved Spatial-Shift MLP Architecture for Vision
• Pruning Self-attentions into Convolutional Layers in Single Path
• OnDev-LCT: On-Device Lightweight Convolutional Transformers towards federated learning
• Deep Ensembling with No Overhead for either Training or Testing: The All-Round Blessings of Dynamic Sparsity
• Rethinking Token-Mixing MLP for MLP-based Vision Backbone