[論文レビュー] Adaptive Token Sampling For Efficient Vision Transformers

Mohsen Fayyaz, Soroush Abbasi Koohpayegani+6 more|arXiv (Cornell University)|Nov 30, 2021

Cell Image Analysis Techniques被引用数 4

ひとこと要約

本稿では、入力画像の内容に基づいて動的にトークン数を削減する微分可能でパラメータフリーのモジュール、Adaptive Token Sampling (ATS) を提案する。分類トークンからのアテンション重みを用いて、逆変換サンプリングにより顕著なパッチをスコア化・サンプリングすることで、ImageNet、Kinetics-400、Kinetics-600 において精度を落とさずにGFLOPsを最大2倍まで削減可能であり、事前学習済みモデルに即座に統合可能な効率性向上が実現される。

ABSTRACT

While state-of-the-art vision transformer models achieve promising results in image classification, they are computationally expensive and require many GFLOPs. Although the GFLOPs of a vision transformer can be decreased by reducing the number of tokens in the network, there is no setting that is optimal for all input images. In this work, we therefore introduce a differentiable parameter-free Adaptive Token Sampler (ATS) module, which can be plugged into any existing vision transformer architecture. ATS empowers vision transformers by scoring and adaptively sampling significant tokens. As a result, the number of tokens is not constant anymore and varies for each input image. By integrating ATS as an additional layer within the current transformer blocks, we can convert them into much more efficient vision transformers with an adaptive number of tokens. Since ATS is a parameter-free module, it can be added to the off-the-shelf pre-trained vision transformers as a plug and play module, thus reducing their GFLOPs without any additional training. Moreover, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate the efficiency of our module in both image and video classification tasks by adding it to multiple SOTA vision transformers. Our proposed module improves the SOTA by reducing their computational costs (GFLOPs) by 2X, while preserving their accuracy on the ImageNet, Kinetics-400, and Kinetics-600 datasets.

研究の動機と目的

• ビジョントランスフォーマーの計算コストがシーケンス長に比例して四乗的に増加するという問題に対処すること。
• 固定比率のプルーニングを避けるために、入力画像ごとに異なる内容に応じた動的でコンテンツに依存するトークン削減を実現すること。
• 微分可能でパラメータフリーのモジュールを設計し、再訓練なしに事前学習済みモデルに統合可能にすること。
• 画像および動画分類ベンチマークにおいて、精度を維持または向上させながらGFLOPsを削減すること。
• アーキテクチャの再トレーニングなしにエッジデバイスでの効率的推論を可能にするプラグアンドプレイソリューションを提供すること。

提案手法

• ATSは分類トークンからのアテンション重みを用いて、すべての入力トークンの重要度スコアを計算する。
• これらのスコアの累積分布関数に対して逆変換サンプリングを適用し、顕著なトークンのサブセットを選択する。
• 選択されたトークンは、アテンション重みと値の重み付き組み合わせを用いてソフトにダウンサンプリングされ、特徴表現が保持される。
• このモジュールは微分可能でパラメータフリーであるため、追加のトレーニングを必要とせず、任意のビジョントランスフォーマーブロックに統合可能である。
• 事前学習済みモデルに即座に統合可能なプラグアンドプレイレイヤーとして使用可能であり、バックボーンと共同でトレーニングすることも可能である。
• 複数段階にわたり段階的なトークンプルーニングが可能であり、画像の複雑さに応じてトークン数が変動する。

実験結果

リサーチクエスチョン

• RQ1 パラメータフリーで微分可能なモジュールが、入力コンテンツに応じてビジョントランスフォーマーのトークン数を動的に削減できるか？
• RQ2 固定比率プルーニングと比較して、適応的トークンサンプリングが計算コストを削減しながら精度を維持できるか？
• RQ3 ATSが微調整なしに事前学習済みビジョントランスフォーマーにプラグアンドプレイとして適用可能か？
• RQ4 複数段階にわたり段階的なトークンサンプリングを適用した場合、効率性と精度にどのような影響を与えるか？
• RQ5 学習可能なスコアリングネットワークと比較して、アテンションベースの重要度スコアリングが、効率性および一般化性能において優れているか？

主な発見

• ATSはDeiT-Sおよびその他の最先端ビジョントランスフォーマーにおいて、GFLOPsを最大2倍まで削減しながら、ImageNetにおけるトップ-1精度を維持または向上させた。
• DeiT-S+ATSモデルは、わずか2.9 GFLOPsで79.7%のトップ-1精度を達成し、同じFLOPレベルにおけるベースラインDeiT-S（78.9%）を上回った。
• 後段階（例：ステージ3）にATSを統合すると、前段階（73.1%）よりも高い精度（78.5%）が得られた。これは、より信頼性の高いアテンションスコアが得られるためである。
• ステージ3～11にATSを複数段階で統合した場合、2.6 GFLOPsで79.2%のトップ-1精度を達成し、段階的な効率性向上を示した。
• EViT-DeiT-SにATSを追加した場合、GFLOPsを3.0から2.5に削減しながら精度に変化がなく、静的プルーニングより優れた互換性と優位性を示した。
• 可視化結果から、複雑でごちゃついた画像ではより多くのトークンがサンプリングされ、単純な画像では少ないことが確認され、顕著な領域に集中した注視がなされている。