Source-linked AI summary
Vision Transformer Pruning
Mingjian Zhu, Yehui Tang, Kai Han
TL;DR
Vision transformers’ storage, memory, and computational demands hinder deployment on mobile devices. The paper prunes transformer dimensions using learned importance scores and sparsity regularization. On ImageNet-1K, accuracy decreases by 1.1% when 40% of dimensions are pruned, while the authors report reduced computation and parameters with high accuracy maintained.
Problem
Vision transformers’ storage, run-time memory, and computational demands hinder deployment on mobile devices.
Source: §1 · p. 1Method
VTP learns dimension-importance scores, applies L1 regularization to encourage sparsity, prunes corresponding linear projections across MHSA and MLP blocks, and fine-tunes the compact model.
Source: §1 · p. 1, §2.2 · pp. 1–2Results
On ImageNet-1K, accuracy decreases by 1.1% when 40% of dimensions are pruned and by 0.5% when 20% are pruned, relative to the base model DeiT-B.
Source: §3.3 · p. 3Takeaways & Limitations
Experiments on Imagenet-100 and ImageNet-1K show that VTP reduces computation costs and model parameters while maintaining high accuracy of the original vision transformers.
Source: §4 · p. 3
Abstract
from arXiv · showhide
Vision transformer has achieved competitive performance on a variety of computer vision applications. However, their storage, run-time memory, and computational demands are hindering the deployment to mobile devices. Here we present a vision transformer pruning approach, which identifies the impacts of dimensions in each layer of transformer and then executes pruning accordingly. By encouraging dimension-wise sparsity in the transformer, important dimensions automatically emerge. A great number of dimensions with small importance scores can be discarded to achieve a high pruning ratio without significantly compromising accuracy. The pipeline for vision transformer pruning is as follows: 1) training with sparsity regularization; 2) pruning dimensions of linear projections; 3) fine-tuning. The reduced parameters and FLOPs ratios of the proposed algorithm are well evaluated and analyzed on ImageNet dataset to demonstrate the effectiveness of our proposed method.
1 INTRODUCTION
Vision transformers have broad computer-vision applications, but their storage, memory, and computation demands impede deployment on edge devices. The paper proposes pruning vision-transformer dimensions using learned importance scores.
1 INTRODUCTION
- Vision transformers’ storage, run-time memory, and computational demands impede deployment on edge devices such as robotics and mobile phones.The introduction notes that CNN compression techniques exist but efficient vision transformers remain an urgent need.Source: §1 · p. 1
- Earlier transformer compression and pruning methods described here focus on natural-language-processing tasks, while vision transformers create a need for efficient computer-vision models.The passage mentions compression through parameter decomposition, parameter sharing, sparsification, distillation, and pruning.Source: §1 · p. 1
- VTP learns dimension-importance scores, applies L1 regularization to sparsify them, and prunes dimensions with smaller scores.The authors report that the resulting compact network can compress and accelerate original ViT (DeiT) models.Source: §1 · p. 1
2 APPROACH
The approach targets the fully connected projections in transformer attention and MLP blocks, which contribute to computation and storage burden. It learns dimension-importance scores, sparsifies them, prunes corresponding projections across the blocks, and fine-tunes the compact model.
2.1 Transformer
- Fully connected layers in the transformer’s attention and MLP architecture contribute to computation and storage burden.The described architecture includes MHSA, MLP, layer normalization, activation, and shortcut connections.Source: §2.1 · p. 1
2.2 Vision Transformer Pruning
- VTP learns importance scores for linear-projection dimensions to retain important features and remove less useful ones.Discrete binary scores are relaxed to real values for optimization; L1 regularization encourages sparsity, after which a threshold sets scores to zero or one according to the pruning rate.Source: §2.2 · pp. 1–2
- Pruning is applied across all MHSA and MLP blocks, replacing their linear projections with projections corresponding to the pruned features.The compact transformer is fine-tuned to diminish the accuracy drop.Source: §2.2 · p. 2
3 EXPERIMENTS
The experiments evaluate VTP on ImageNet-100 and ImageNet-1K, using sparsity regularization, dimension pruning, and fine-tuning. Results report parameter and FLOPs savings alongside accuracy changes.
3.1 Datasets
- ImageNet-1K is the ILSVRC2012 classification dataset with 1.2 million training images and 50,000 validation images across 1,000 classes; ImageNet-100 is a randomly sampled 100-class subset.The same data augmentation strategy is used for both datasets.Source: §3.1 · p. 2
3.3 Results and Analysis
- On ImageNet-100, pruning 40% of dimensions after training with a 0.0001 sparse rate saves 45.3% of parameters and 43.0% of FLOPs while accuracy is maintained.The reported savings track the pruning rate, and sparse rate does not highly influence pruning effectiveness.Source: Table 1 · p. 3, §3.3 · p. 3
- On ImageNet-100, pruning 40% of dimensions saves 45.3% of FLOPs and drops accuracy by 1.92%; pruning 20% saves 22.0% of FLOPs and drops accuracy by 0.96%.The comparison is between the baseline and VTP models.Source: Table 2 · p. 3, §3.3 · p. 3
- On ImageNet-1K, accuracy decreases by 1.1% with 40% of dimensions pruned and by 0.5% with 20% pruned, compared with DeiT-B.The authors state that these results demonstrate generalization to large-scale datasets.Source: Table 3 · p. 3, §3.3 · p. 3
3.2 Implementation Details
- The implementation trains with L1 regularization, selects a 0.0001 sparse rate on ImageNet-100, prunes using thresholds set by pruning rate, then fine-tunes without L1 regularization.Sparsity training uses 100 epochs; fine-tuning otherwise follows the same optimization setting.Source: §3.2 · p. 3
4 CONCLUSION
The conclusion presents VTP as a method for pruning vision-transformer dimensions through L1 regularization. Experiments on ImageNet-100 and ImageNet-1K report lower computation costs and parameter counts while maintaining high accuracy.
4 CONCLUSION
- Experiments on ImageNet-100 and ImageNet-1K show that VTP reduces computation costs and model parameters while maintaining high accuracy.The conclusion describes the method as simple and efficient.Source: §4 · p. 3
- L1 regularization encourages dimension-wise sparsity, allowing important transformer dimensions to emerge automatically.The method prunes dimensions with small importance scores.Source: §4 · p. 3
- The authors identify reducing the number of heads and layers as future directions for further vision-transformer compression.Source: §4 · p. 3
Related papers
This paper cites
- Attention Is All You Need2017
- An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale2020
- Swin Transformer: Hierarchical Vision Transformer using Shifted Windows2021
- FitNets: Hints for Thin Deep Nets2014
- Big Bird: Transformers for Longer Sequences2020
- TinyBERT: Distilling BERT for Natural Language Understanding2019