[論文レビュー] Not All Patches are What You Need: Expediting Vision Transformers via Token Reorganizations

Youwei Liang, Chongjian Ge+4 more|arXiv (Cornell University)|Feb 16, 2022

Advanced Neural Network Applications被引用数 95

ひとこと要約

EViT は訓練中に attentive tokens を特定し inattentive tokens を統合することで Vision Transformer の推論を高速化し、パラメータを追加せず、同じコストで効率または精度を改善します。

ABSTRACT

Vision Transformers (ViTs) take all the image patches as tokens and construct multi-head self-attention (MHSA) among them. Complete leverage of these image tokens brings redundant computations since not all the tokens are attentive in MHSA. Examples include that tokens containing semantically meaningless or distractive image backgrounds do not positively contribute to the ViT predictions. In this work, we propose to reorganize image tokens during the feed-forward process of ViT models, which is integrated into ViT during training. For each forward inference, we identify the attentive image tokens between MHSA and FFN (i.e., feed-forward network) modules, which is guided by the corresponding class token attention. Then, we reorganize image tokens by preserving attentive image tokens and fusing inattentive ones to expedite subsequent MHSA and FFN computations. To this end, our method EViT improves ViTs from two perspectives. First, under the same amount of input image tokens, our method reduces MHSA and FFN computation for efficient inference. For instance, the inference speed of DeiT-S is increased by 50% while its recognition accuracy is decreased by only 0.3% for ImageNet classification. Second, by maintaining the same computational cost, our method empowers ViTs to take more image tokens as input for recognition accuracy improvement, where the image tokens are from higher resolution images. An example is that we improve the recognition accuracy of DeiT-S by 1% for ImageNet classification at the same computational cost of a vanilla DeiT-S. Meanwhile, our method does not introduce more parameters to ViTs. Experiments on the standard benchmarks show the effectiveness of our method. The code is available at https://github.com/youweiliang/evit

研究の動機と目的

• MHSA におけるトークンレベルの冗長性を特定して Vision Transformers (ViTs) の加速を動機づける。
• attentive tokens を保存し inattentive tokens を融合する訓練時のトークン再編成を提案する。
• 推論時に MHSA と FFN の計算を削減し、追加パラメータなしで EViT を示す。
• 同じ計算予算でより多くのトークン（高解像度）入力を許可することで精度を向上させる。
• oracle を用いてトークンの関連性を案内させる効果を検討し、既存の加速手法と比較する。

提案手法

• MHSA ヘッドごとに各画像トークンへのクラス・トークンの平均的な注意度を計算する。
• 上位の k 個の attentive トークンを保持し、attention の低いトークンを 1 つの融合トークンに統合する。
• 低い注意度のトークンを、それらの attentive を重みとして用いた重み付き平均で融合する（x_fused = sum_{i in N} a_i x_i）。
• 選択した層で ViT の訓練にトークン再編成を組み込み、keep rate に対してコサインスケジュールを適用する。
• 重要なトークンを識別する oracle ViT を用いた訓練を任意で行い、oracle 重みで EViT を初期化する。
• 同じ計算コストでより多くのトークンを入力として高解像度訓練を実証し、ImageNet 実験で検証する。

実験結果

リサーチクエスチョン

• RQ1 ViT の訓練中にトークン再編成を行うと推論コストを削減しつつ精度を維持できるか？
• RQ2 低注意度トークンの融合は単純なトークン削除と比べて情報をより多く保持し、訓練を安定させるか？
• RQ3 固定計算量の下での EViT の性能はどうなるか、入力解像度を上げた場合は？
• RQ4 oracle ViT を用いてトークン選択を導くことが精度と効率に及ぼす影響は？

主な発見

• EViT は ImageNet で約 0.3% の精度低下のみで DeiT-S 推論を約 50% 速くできる。
• EViT は同じ MACs でより高いスループットを達成し、より高解像度の入力を使用する場合に精度を維持または向上できる（例: DeiT-S は同じ計算量で top-1 が 1% 向上）。
• inattentive トークンの融合は情報を保持し、トークン剪定のみより訓練の安定性と精度を改善する。
• oracle を用いた訓練はさらなる精度向上につながる（例えば DeiT-S が oracle 設定で 79.8% → 80.7% に向上）一方で計算量を維持または削減できる。
• DynamicViT と比較して、同じ計算量でより少ないパラメータでより良い精度を提供し、長期的な訓練でも追加の利益を示す。
• EViT は DeiT および LV-ViT など異なる ViT バリアントに適用でき、設定間での精度-スループットのトレードオフが良好になる。