[論文レビュー] Skip-Attention: Improving Vision Transformers by Paying Less Attention

Shashanka Venkataramanan, Amir Hossein Ghodrati+3 more|arXiv (Cornell University)|Jan 5, 2023

CCD and CMOS Imaging Sensors被引用数 6

ひとこと要約

Skip-Attention (SkipAt) は、より深い層での高コストな自己注意計算を、以前の層からの自己注意計算を再利用することで、ビジョントランスフォーマーの性能を向上させます。深さ方向の畳み込みにインspiredされた軽量なパラメトリック関数を導入することで、画像分類、セグメンテーション、ノイズ除去、動画処理のあらゆるタスクにおいて、ベースラインのViTよりも高いスループットと精度を達成し、ADE20Kでは最大40%の高速化、自己教師あり事前学習では26%の高速化を実現しました。

ABSTRACT

This work aims to improve the efficiency of vision transformers (ViT). While ViTs use computationally expensive self-attention operations in every layer, we identify that these operations are highly correlated across layers -- a key redundancy that causes unnecessary computations. Based on this observation, we propose SkipAt, a method to reuse self-attention computation from preceding layers to approximate attention at one or more subsequent layers. To ensure that reusing self-attention blocks across layers does not degrade the performance, we introduce a simple parametric function, which outperforms the baseline transformer's performance while running computationally faster. We show the effectiveness of our method in image classification and self-supervised learning on ImageNet-1K, semantic segmentation on ADE20K, image denoising on SIDD, and video denoising on DAVIS. We achieve improved throughput at the same-or-higher accuracy levels in all these tasks.

研究の動機と目的

• 自己注意の計算コストがシーケンス長に比例して四乗的に増加するビジョントランスフォーマー（ViTs）における高コスト問題に対処すること。
• 多頭注意（MSA）の出力が層をまたいで強く相関しており、特にCLSトークンにおいて冗長性が顕著であることを特定すること。
• より深い層における高コストなMSA計算を、以前の層からの表現を再利用することで置き換えるプラグインモジュールを開発すること。
• パフォーマンスを維持しながら効率性を向上させる軽量なパラメトリック関数を設計し、より高速な推論と学習を可能にすること。
• 画像分類、セマンティックセグメンテーション、画像および動画のノイズ除去、自己教師あり学習を含む多様なビジョンタスクへの汎用性を実証すること。

提案手法

• 深層のMSAブロックを、以前の層（例：層3〜8）からの注意特徴を再利用することで置き換えるプラグインモジュールであるSkipAtを提案する。
• 再利用された特徴をターゲット層の表現空間に適応させるために、学習可能なパラメトリック関数Φを導入する。
• 表現力と効率性のバランスを考慮し、デフォルトとして5×5の深さ方向畳み込み（DwC）とチャネル拡張比2をΦ関数として採用する。
• 分類タスクおよび密度予測タスクの両方（画像レベルおよびピクセルレベル出力）にスキップ機構を適用する。
• ベースラインViTと同一のデータおよび学習プロトコルを用いて、エンドツーエンドでモデルを訓練する。
• 任意のViTアーキテクチャ（小、ベース、大）に適用可能であり、多様なデータセットおよびタスクに適用可能である。

Figure 2 : Attention correlation . Mean of the attention heads from the CLS token of a pretrained ViT-T/16 at different layers from the validation set of ImageNet-1K. Numbers below each attention map indicates the cosine similarity of $A_{l}^{\texttt{[CLS]}}$ with $A_{l-1}^{\texttt{[CLS]}}$ .

実験結果

リサーチクエスチョン

• RQ1 パフォーマンスの低下を伴わずに、より深いViT層における自己注意計算を、以前の層からの特徴を再利用することでどの程度近似できるか？
• RQ2 単純なパラメトリック関数が、直接再利用やベースラインViTと比較して、再利用された注意特徴のパフォーマンスを効果的に向上させられるか？
• RQ3 SkipAtは、密度予測および動画理解を含む多様なビジョンタスクにおいて、推論速度とFLOPsにどのような影響を与えるか？
• RQ4 SkipAtは自己教師あり事前学習においてパフォーマンスを維持または向上させられるか？また、精度の損失なしに事前学習時間を短縮できるか？
• RQ5 精度とスループットのバランスを最適化するためのパラメトリック関数の最適な設定（カーネルサイズ、チャネル拡張比）は何か？

主な発見

• ADE20Kセマンティックセグメンテーションベンチマークにおいて、SkipAtはベースラインのViT-T/16と比較して、推論速度を40%高速化するとともに、トップ1精度を0.6%向上させた。
• ImageNet-1Kでは、SkipAtは67.7%のトップ1精度を達成し、1秒あたり6.9千枚（6.9×10³ img/sec）のスループットを実現。ベースラインViTと比較して、精度は1.9%向上、スループットは19%向上した。
• DAVISにおける動画ノイズ除去タスクでは、SOTA手法と同等のPSNR（35.16）を達成したが、FLOPsは17%少ない（77.1 GFLOPs vs. 93.2 for UniFormer）。
• ImageNet-1Kにおける自己教師あり事前学習では、SkipAtが26%の高速化を実現したが、下流タスクの精度に損失はなく、表現学習における効率性向上を示した。
• 最適なSkipAt設定は、5×5の深さ方向畳み込みとチャネル拡張比2を組み合わせたもので、ImageNet-1Kでは67.7%のトップ1精度と6.9K img/secのスループットを達成した。
• アブレーション実験では、特徴再利用に恒等写像を使用した場合、精度が4.7%低下したが、SkipAtのパラメトリック関数は、精度と効率の両面ですべてのベースラインを上回った。

Figure 3 : CKA analysis of $A^{[\texttt{CLS}]}$ and $Z^{\text{MSA}}$ across different layers of pretrained ViT-T/16 on the validation set of Imagenet-1K. Vanilla ViT-T/16 has high correlation across both attention maps (layer 3 to 10) and $Z^{\text{MSA}}$ (layer 2 to 8)