[論文レビュー] A Unified Pruning Framework for Vision Transformers

Hao Yu, Jianxin Wu|arXiv (Cornell University)|Nov 30, 2021

Advanced Neural Network Applications参考文献 43被引用数 5

ひとこと要約

本稿では、モデルアーキテクチャを変更せずに、自己注意機構、フィードフォワードネットワーク、正規化層、畳み込み層のすべてのコンponentに同時にチャネルをプルーニングする統合的構造的プルーニングフレームワークであるUP-ViTsを提案する。効率的な重要度スコアモジュールと段階的プルーニング戦略を用いることで、UP-ViTsは、圧縮モデルにおいて最先端の精度を達成し、例えばUP-DeiT-Tは、元のDeiT-Tと同等のFLOPsでImageNetで75.79%のトップ1精度を達成する。

ABSTRACT

Recently, vision transformer (ViT) and its variants have achieved promising performances in various computer vision tasks. Yet the high computational costs and training data requirements of ViTs limit their application in resource-constrained settings. Model compression is an effective method to speed up deep learning models, but the research of compressing ViTs has been less explored. Many previous works concentrate on reducing the number of tokens. However, this line of attack breaks down the spatial structure of ViTs and is hard to be generalized into downstream tasks. In this paper, we design a unified framework for structural pruning of both ViTs and its variants, namely UP-ViTs. Our method focuses on pruning all ViTs components while maintaining the consistency of the model structure. Abundant experimental results show that our method can achieve high accuracy on compressed ViTs and variants, e.g., UP-DeiT-T achieves 75.79% accuracy on ImageNet, which outperforms the vanilla DeiT-T by 3.59% with the same computational cost. UP-PVTv2-B0 improves the accuracy of PVTv2-B0 by 4.83% for ImageNet classification. Meanwhile, UP-ViTs maintains the consistency of the token representation and gains consistent improvements on object detection tasks.

研究の動機と目的

• 構造的一致性を保ちながら、効果的で一般化可能なビジョントランスフォーマー向けのプルーニング手法の不足に対処する。
• トークンレベルのプルーニングの限界を克服し、空間的構造を損なわず、下流タスクへの転送性を損なわないようにする。
• 従来のViTとその変種（畳み込み型や階層的設計を含む）の両方に適用可能な統合フレームワークを開発する。
• 高い精度と低計算量を実現する圧縮ViTを実現し、オブジェクト検出などの密予測タスクにうまく一般化できるようにする。
• モデルの整合性を保ちながら、すべてのコンponentを同時にプルーニングすることで、既存のプルーニング手法を上回る性能を達成する。

提案手法

• 勾配ベースのマグニチュード推定を用いて、事前に学習されたViTコンponent内の各フィルタの重要度スコアを計算する効率的な評価モジュールを設計する。
• MHSAs、FFNs、正規化層、畳み込みを含むすべてのコンponentにおいて、同時にかつ均一に、不要なチャネルを削除することで構造的プルーニングを適用する。
• 自己注意ヘッドの構造を保持し、表現の一貫性を維持する新しいチャネル削除戦略を提案する。
• 最も重要度の低いブロックを段階的に削除し、性能を維持するためにハイブリッドブロックに置き換える段階的ブロックプルーニング法を導入する。
• 微調整中に知識蒸留を用いて、プルーニング後の精度を保持する。教師モデルは元のアーキテクチャと一致させる。
• DeiT、PVTv2、Swin Transformerを含む複数のViT変種にわたる統一されたプルーニングパイプラインを適用し、圧縮後のアーキテクチャの一貫性を保証する。

実験結果

リサーチクエスチョン

• RQ1 モデルの構造的設計を変更せずに、ビジョントランスフォーマーのすべてのコンponentを同時に圧縮できる統合的プルーニングフレームワークを設計できるか？
• RQ2 FFN、MSA、正規化、畳み込みを含むすべてのレイヤーで構造的プルーニングを適用した場合、トークンレベルのプルーニングと比較して、精度と推論効率にどのような影響を与えるか？
• RQ3 プルーニング中に元のトークン表現を維持することで、オブジェクト検出などの下流タスクへの転送性が向上するか？
• RQ4 一度にプルーニングする方法と段階的プルーニングの方法とを比較した場合、ViTの幅と深さの圧縮において、どちらが相対的に効果的か？
• RQ5 さまざまなアーキテクチャとタスクにおいて、既存の最先端の圧縮ViTよりも優れた一般化性能を示せるか？

主な発見

• UP-DeiT-Tは、FLOPsを元のDeiT-Tと同等に保ったまま、ImageNetで75.79%のトップ1精度を達成し、元のDeiT-Tより3.59%高い。
• UP-PVTv2-B0は、計算コストが同一の条件下で、元のPVTv2-B0よりもImageNet分類タスクで4.83%高い精度を達成した。
• UP-DeiT-Tを10ブロックに圧縮する際、段階的ブロックプルーニングでは73.78%の精度を達成したのに対し、一度にブロックを削除する方法では73.20%にとどまった。
• コンponent 1（幅）を一度にプルーニングした場合、66.95%の精度を達成したが、これは段階的プルーニング（65.05%）を上回った。
• プルーニング後のUP-ViTsモデルは、保存されたトークン表現のおかげで、オブジェクト検出などの下流タスクにおいて一貫した性能向上を示した。
• フレームワークはDeiT、PVTv2、Swin Transformerを含むViTのさまざまな変種に効果的に一般化され、広範な適用可能性と頑健性を示した。