[論文レビュー] Pruning Self-attentions into Convolutional Layers in Single Path

Haoyu He, Jianfei Cai+5 more|arXiv (Cornell University)|Nov 23, 2021

Advanced Neural Network Applications参考文献 64被引用数 16

ひとこと要約

SPViTは、マルチヘッド自己注意（MSA）とフィードフォワードネットワーク（FFN）の両方の剪定を一括で最適化する単一パスの剪定フレームワークを提案する。この手法では、MSAと畳み込み演算の間で重みを共有することで、計算コストが低く、自動的なアーキテクチャ探索が可能になる。SOTAの剪定性能を達成し、DeiT-Bでは52.0%のFLOPs削減を実現しながら、ImageNet-1kでのトップ-1精度を0.6%向上させた。

ABSTRACT

Vision Transformers (ViTs) have achieved impressive performance over various computer vision tasks. However, modeling global correlations with multi-head self-attention (MSA) layers leads to two widely recognized issues: the massive computational resource consumption and the lack of intrinsic inductive bias for modeling local visual patterns. To solve both issues, we devise a simple yet effective method named Single-Path Vision Transformer pruning (SPViT), to efficiently and automatically compress the pre-trained ViTs into compact models with proper locality added. Specifically, we first propose a novel weight-sharing scheme between MSA and convolutional operations, delivering a single-path space to encode all candidate operations. In this way, we cast the operation search problem as finding which subset of parameters to use in each MSA layer, which significantly reduces the computational cost and optimization difficulty, and the convolution kernels can be well initialized using pre-trained MSA parameters. Relying on the single-path space, we introduce learnable binary gates to encode the operation choices in MSA layers. Similarly, we further employ learnable gates to encode the fine-grained MLP expansion ratios of FFN layers. In this way, our SPViT optimizes the learnable gates to automatically explore from a vast and unified search space and flexibly adjust the MSA-FFN pruning proportions for each individual dense model. We conduct extensive experiments on two representative ViTs showing that our SPViT achieves a new SOTA for pruning on ImageNet-1k. For example, our SPViT can trim 52.0% FLOPs for DeiT-B and get an impressive 0.6% top-1 accuracy gain simultaneously. The source code is available at https://github.com/ziplab/SPViT.

研究の動機と目的

• マルチヘッド自己注意層に起因する高い計算コストと、インダクティブバイアスの欠如というVision Transformerの課題に取り組むこと。
• 導入された局所的なインダクティブバイアスを備えた、効率的で自動的な事前学習済みViTの剪定を可能にすること。
• NASベースの剪定における探索コストと最適化の難易度を軽減するため、単一パスの探索空間でMSAと畳み込み演算を統合すること。
• 目標FLOPs制約下で、MSA層とFFN拡張比の両方を同時に最適化すること。
• 探索されたアーキテクチャの実証的分析を通じて、ViTの構造的重要な要素を解明すること。

提案手法

• 畳み込み演算をMSAパラメータの部分集合として表現する重み共有スキームを提案し、注目と畳み込みの間でパラメータを共有可能にする。
• すべての候補演算（MSAと畳み込み）を1層あたりの1つのMSAに符号化する単一パスの探索空間を構築し、探索の複雑さを低減する。
• MSA層に学習可能なバイナリゲートを導入し、どのMSAパラメータのサブセットを使用するかを選択することで、畳み込み演算への有効な剪定を実現する。
• 学習可能なゲートをFFN層に拡張し、隠れ次元の剪定とMLP拡張比の最適化を同時に実行する。
• 学習可能なゲートをエンドツーエンドで最適化することで、目標FLOPs制約下で自動的に効率的なアーキテクチャを発見する。
• 事前学習済みのMSAパラメータを用いて畳み込みカーネルを初期化し、効果的で安定した学習を可能にする。

Figure 1: Multi-path vs. single-path search space. (a) Multi-path search space. MSA and different types of convolutional operations are added as separate trainable paths, leading to an expensive search cost. $\mathcal{P}$ and $g_{p}$ denote the ordered operation set and the learnable gate for the $p

実験結果

リサーチクエスチョン

• RQ1 統合された単一パス探索空間は、NASベースのViT剪定における計算コストと最適化の難易度を低減できるか？
• RQ2 MSAと畳み込み演算の間で重みを共有することで、畳み込みカーネルの効果的で安定した初期化が可能になるか？
• RQ3 統合された探索空間内でMSAとFFN層を同時に剪定することで、個別に剪定する場合よりも優れた効率-精度トレードオフが達成できるか？
• RQ4 探索されたアーキテクチャにおける構造的選択（例：浅いブロックでMSAを畳み込みに剪定するか）は、モデルの機能的挙動とどのように関係するか？
• RQ5 FFNの容量は、さまざまなViTブロックで果たす役割は何か？また、標準的および階層的ViTアーキテクチャにおいて、その役割はどのように変化するか？

主な発見

• SPViTは、DeiT-Bで52.0%のFLOPs削減を実現しながら、ImageNet-1kでのトップ-1精度を0.6%向上させ、トレーニング後の剪定において新たなSOTAを樹立した。
• この手法は、浅いMSA層を畳み込み層やスキップ接続に剪定することに成功し、浅いブロックで局所的注目パターンを示す傾向と整合的である。
• すべてのモデルで深いMSA層が保持されていることから、グローバル表現学習におけるその重要性が裏付けられた。
• 標準的なViTの深いブロックにおけるFFN層は、より高い冗長性を示し、優先的に剪定される傾向にあり、後段の段階で低い容量で十分であることが示唆された。
• スウィン-Tiのような階層的ViTでは、各ステージの浅い層でより多くの隠れ次元が保持されていることから、トークンマージ後に高い容量が必要であることが示された。
• 単一パスの探索空間により、マルチパスNASと比較して探索コストと最適化の難易度が低減され、MSAとFFNの剪定を効率的に統合的に探索可能になった。

Figure 2: Overview of our weight-sharing scheme, which expresses the bottleneck convolution with MSA parameters. First, we restrict the value and output projections in the MSA layer to only process the features within local regions. Then, we learn to ensemble the value projections with a parameteriz