[論文レビュー] Unified Visual Transformer Compression

Shixing Yu, Tianlong Chen+6 more|arXiv (Cornell University)|Mar 15, 2022

Advanced Neural Network Applications被引用数 28

ひとこと要約

UVCはViTにおける剪定、ブロックスキップ、知識蒸留を予算の下で共同最適化し、ImageNet上でDeiTとT2T-ViTのバックボーン全体で大幅なFLOPs削減を実現し、精度損失は最小限です。

ABSTRACT

Vision transformers (ViTs) have gained popularity recently. Even without customized image operators such as convolutions, ViTs can yield competitive performance when properly trained on massive data. However, the computational overhead of ViTs remains prohibitive, due to stacking multi-head self-attention modules and else. Compared to the vast literature and prevailing success in compressing convolutional neural networks, the study of Vision Transformer compression has also just emerged, and existing works focused on one or two aspects of compression. This paper proposes a unified ViT compression framework that seamlessly assembles three effective techniques: pruning, layer skipping, and knowledge distillation. We formulate a budget-constrained, end-to-end optimization framework, targeting jointly learning model weights, layer-wise pruning ratios/masks, and skip configurations, under a distillation loss. The optimization problem is then solved using the primal-dual algorithm. Experiments are conducted with several ViT variants, e.g. DeiT and T2T-ViT backbones on the ImageNet dataset, and our approach consistently outperforms recent competitors. For example, DeiT-Tiny can be trimmed down to 50\% of the original FLOPs almost without losing accuracy. Codes are available online:~\url{https://github.com/VITA-Group/UVC}.

研究の動機と目的

• Vision Transformersの計算コストを削減しつつ精度を損なわず、効率的な展開を促進する。
• 単一のエンドツーエンド最適化で複数の圧縮技術を組み合わせる統一フレームワークを提案する。
• グローバルなリソース予算の下で剪定、スキップ、蒸留の自動的なバランスを可能にする。
• ImageNet上でDeiTとT2T-ViTという人気のあるViT変種を評価し、最先端の剪定手法と比較する。

提案手法

• ViT圧縮を制約付きのエンドツーエンド最適化問題として定式化する。
• 蒸留損失を用いてネットワーク重み、層ごとの剪定マスク/比率、スキップ構成を共同学習する。
• 各トランスフォーマーブロック内のアテンションヘッドとMLPニューロンを二層レベルのグループスパース性で剪定する。
• ブロック削除を可能にするため、トランスフォーマーブロック全体のスキップ構成のためのビナリゲーティング変数を導入する。
• primal-dualアルゴリズムを用いてミニ最大化問題を解く。
• DeiTおよびT2T-ViTバックボーンでのリファレンス実装と実験結果を提供する。

実験結果

リサーチクエスチョン

• RQ1 リソース予算の下でViTブロック内およびブロック間のスキップ構成を共同最適化できるか。
• RQ2 剪定、スキップ、知識蒸留を統合することは、単一手法のViT圧縮よりも優れているか。
• RQ3 さまざまなFLOPs予算でDeiTとT2T-ViTが達成できる精度/遅延のトレードオフはどうなるか。
• RQ4 ImageNet上でUVCは既存のパッチ削減および剪定法とどのように比較されるか。

主な発見

モデル	方法	Top-1 精度（%）	FLOPs（G）	FLOPs残存率（%）
DeiT-Tiny	Baseline	72.2	1.3	100.0
DeiT-Tiny	SViTE	70.12	0.99	76.31
DeiT-Tiny	PatchSlimming	72.0	0.7	53.8
DeiT-Tiny	UVC	71.8	0.69	53.1
DeiT-Tiny	HVT	69.7	0.64	49.23
DeiT-Tiny	UVC (distilled)	71.3	0.64	49.23
DeiT-Tiny	UVC (other setting)	70.6	0.51	39.12
DeiT-Small	Baseline	79.8	4.6	100.0
DeiT-Small	SViTE	79.22	3.14	68.36
DeiT-Small	PatchSlimming	79.4	2.6	56.5
DeiT-Small	HVT	78.0	2.40	52.2
DeiT-Small	UVC	78.82	2.32	50.41
DeiT-Base	Baseline	81.8	17.6	100.0
DeiT-Base	SViTE	82.22	11.87	66.87
DeiT-Base	VTP	80.7	10.0	56.8
DeiT-Base	PatchSlimming	81.5	9.8	55.7
DeiT-Base	UVC	80.57	8.0	45.50
T2T-ViT-14	Baseline	81.5	4.8	100.0
T2T-ViT-14	PoWER	79.9	3.5	72.9
T2T-ViT-14	UVC	80.4	2.90	60.4
T2T-ViT-14	UVC (lower)	79.6	2.47	51.5
T2T-ViT-14	UVC (higher)	78.9	2.11	44.0

• UVCはDeiT-TinyでFLOPsを約50%削減でき、蒸留/非蒸留設定での精度低下は0.3%/0.9%と最小限である。
• UVCは類似またはより大きなFLOPs削減で、いくつかの現代的なViT圧縮法（例: SViTE、PatchSlimming、HVT、IA-RED 2）よりも精度の点で優れている。
• DeiT-Baseでは約45%のFLOPs削減を実現し、許容できる精度低下を伴い、安定した学習が難しいベースラインよりも競争力のある結果を示す。
• T2T-ViT-14に適用すると、44–60%のFLOPs削減など顕著な節約を達成し、いくつかのパッチベース削減を上回る。
• アブレーションにより、スキップ、剪定、蒸留を共同最適化する方が、いずれかの単一要素や逐次適用よりも良いトレードオフを生み出す。
• UVCのコードは公開されている。
• （本文の注記）UVCはImageNet上の複数のViT系で高い性能を示している。