[論文レビュー] DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification

Yongming Rao, Wenliang Zhao+4 more|arXiv (Cornell University)|Jun 3, 2021

Advanced Neural Network Applications参考文献 31被引用数 310

ひとこと要約

DynamicViTは、軽量な予測モジュールとアテンションマスキングを用いて、視覚トランスフォーマーにおける階層的で入力依存のトークンプルーニングを導入し、実質的なFLOPs削減を、精度の低下を最小限に抑えて達成する。

ABSTRACT

Attention is sparse in vision transformers. We observe the final prediction in vision transformers is only based on a subset of most informative tokens, which is sufficient for accurate image recognition. Based on this observation, we propose a dynamic token sparsification framework to prune redundant tokens progressively and dynamically based on the input. Specifically, we devise a lightweight prediction module to estimate the importance score of each token given the current features. The module is added to different layers to prune redundant tokens hierarchically. To optimize the prediction module in an end-to-end manner, we propose an attention masking strategy to differentiably prune a token by blocking its interactions with other tokens. Benefiting from the nature of self-attention, the unstructured sparse tokens are still hardware friendly, which makes our framework easy to achieve actual speed-up. By hierarchically pruning 66% of the input tokens, our method greatly reduces 31%~37% FLOPs and improves the throughput by over 40% while the drop of accuracy is within 0.5% for various vision transformers. Equipped with the dynamic token sparsification framework, DynamicViT models can achieve very competitive complexity/accuracy trade-offs compared to state-of-the-art CNNs and vision transformers on ImageNet. Code is available at https://github.com/raoyongming/DynamicViT

研究の動機と目的

• 有意義な画像パッチの疎性を利用して、ビジョントランスフォーマーの加速を動機づける。
• 入力に依存して段階的にトークンを剪定する動的トークンサ sparsificationフレームワークを提案する。
• エンドツーエンドで訓練可能な予測モジュールを、Gumbel-Softmaxとアテンションマスキングによる微分可能な剪定と共に開発する。
• ImageNet上で複数のバックボーントランスフォーマーにおいて、実質的なFLOPs削減とスループット向上を示す。

提案手法

• 複数のトランスフォーマーブロックに軽量な予測モジュールを挿入し、トークンごとの重要度を推定する。
• トークン特徴量からローカル-グローバル埋め込みを計算し、各トークンの削除/保持確率を予測する。
• 訓練時の微分可能性を保ちながら、Gumbel-Softmaxを用いてバイナリの保持/削除マスクをサンプルする。
• 自分注意機構において剪定されたトークンを含む相互作用を除去するアテンションマスキングを適用し、訓練時の計算を均一に保つ。
• クロスエントロピー損失、教師バックボーンに対する蒸留損失、KL発散、比率制約付き剪定損失の組み合わせで訓練する。
• 推論時には、学習されたスコアに従って各ステージで固定数のトークンを剪定し、目標保持比を達成する。

実験結果

リサーチクエスチョン

• RQ1 有用性の低いトークンを剪定しても、視覚トランスフォーマーは大幅な精度低下なしに加速できるだろうか。
• RQ2 トランスフォーマーフレームワーク内で、動的トークン剪定機構をエンドツーエンドで訓練するにはどうすればよいか。
• RQ3 階層的で入力依存的なトークン剪定が、異なるバックボーンにおけるモデルの効率と精度に与える影響は何か。

主な発見

• 階層的トークン疎化は入力トークンの最大66%を剪定でき、FLOPsを31%–37%削減し、バックボーン全体でスループットを40%以上向上させ、精度低下は約0.5%程度に留まる。
• DynamicViTは、ImageNetにおいて最先端のCNNsおよびビジョンTransformerと比較して、計算量と精度の競争力のあるトレードオフを達成する。
• 動的剪定は合理的に振る舞い、画像中心や物体の周辺のトークンを保持し、周辺領域を剪定する。剪定が進むにつれて解釈可能性を示す。
• このアプローチは、モデルスケーリングのための幅スケーリングの実用的な代替手段を提供し、動的トークン疎化によって同等またはそれより良い効率を達成する。
• 大型モデル（DeiT-B および 384x384 入力）は DynamicViT の恩恵を受け、FLOPsの大幅削減と控えめな精度低下を示す。
• アブレーションから、動的で学習された剪定が、静的またはランダムなトークン除去戦略より効果的であることが示される。