[論文レビュー] Chasing Sparsity in Vision Transformers: An End-to-End Exploration

Tianlong Chen, Yu Cheng+4 more|arXiv (Cornell University)|Jun 8, 2021

Advanced Neural Network Applications参考文献 95被引用数 85

ひとこと要約

この論文は、視覚トランスフォーマー（ViT）のエンドツーエンドのスパース学習を導入し、動的スパースサブネットワーク、構造化スパース性、データとアーキテクチャの共同スパース性によって、精度を維持しつつ学習メモリと推論コストを削減します。

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.

研究の動機と目的

• ViTにおけるエンドツーエンドのスパース性を動機づけ、学習メモリと推論コストを削減する。
• 固定されたパラメータ予算の下で精度を維持または向上させるスパースViT学習法を開発する。
• ハードウェア効率のために、スパース性を非構造的な形態から構造化された形態へ拡張する。
• データのスパース性とアーキテクチャのスパース性を共同探索して追加の効率を得る。
• DeiTバックボボーンを横断してImageNetで有意なFLOPsとレイテンシ削減を示す。

提案手法

• 固定パラメータ予算でスパースViTを訓練する Sparse Vision Transformer Exploration (SViTE) を導入する。
• ハードウェアに優しいスパース性のために自己注意ヘッドの剪定と成長を導くことで、SViTEを Structured Sparse ViT Exploration (S2ViTE) に拡張する。
• 有益な入力トークンを共同選択しモデルをスパース化することにより Sparse Vision Transformer Co-Exploration (SViTE+) を追加する。
• ヘッドの重要度を評価するためにテイラー展開に基づくプロキシと、剪定のためにMLPニューロンのL1ノルムを用いる。
• Gumbel-Softmaxとストレートスルーの工夫を用いた学習可能なトークン選択子を組み込み、上位k個の有益なパッチを選択してデータスパースを有効にする。

実験結果

リサーチクエスチョン

• RQ1 エンドツーエンドのスパース学習は、精度を犠牲にせずにViTのFLOPsとパラメータ削減を大幅に実現できるのか？
• RQ2 構造化スパース性（例: 注意ヘッドの剪定）は、非構造化スパース性と比較してハードウェアにやさしい利得をもたらすのか？
• RQ3 データスパース性（トークン選択）とアーキテクチャスパース性の共同探索は、性能を損なうことなく追加の効率向上をもたらすのか？
• RQ4 暗黙の正則化によってスパースViTモデルは一般化性能が高く、特定のスパース性の下で精度が向上する可能性があるのか？
• RQ5 ImageNetでDeiT-Tiny/Small/Baseのバックボーンに対して、スパース戦略はどのように性能を発揮するのか？

主な発見

• SViTEは、バックボーンとスパース性に応じて約25.56%～57.50%のFLOPs削減を伴うスパースな DeiT をもたらし、精度低下はほとんどなく（通常は0.5%以内）。
• SViTE+はトークン選択を用いて、DeiT-Smallでデータスパース性50%、モデルスパース性5%で最大4.40%の精度向上を達成し、FLOPsを49.32%削減、実行時間も4.40%短縮。
• 構造化スパース性を用いるS2ViTEは非構造化のバリアントと同等以上で、密なベースラインを上回りつつ実行時間削減を大幅に実現する（最大で24.70%の削減が報告されている）。
• データスパースは正則化として機能する。SViTE+-Smallは最大10%のトークン削減とそれに対応する実行時間・FLOPs削減を実現し、時には精度向上も見られる。
• 構造化スパース性（S2ViTE）はいくつかの設定でSSPを上回ることがあり、40%の構造化スパース性を持つS2ViTE-BaseはFLOPsを約34%削減しつつ、Dense DeiT-Baseを最大で1.24%の精度向上で上回る。
• ImageNet-1KのDeiT-Tiny/Small/Base全体で、提案手法は一貫した効率向上を提供し、精度は競争力があるか改善され、より小さなスパースネットワークでもより大きな密なネットワークを上回ることがある。