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

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

Advanced Neural Network Applications参考文献 32被引用数 8

ひとこと要約

本稿では、入力画像の複雑さに応じてビジョントランスフォーマーのトークン数を動的に削減する、パラメータフリーで微分可能なモジュールであるアダプティブトークンサンプリング（ATS）を提案する。各画像ごとに重要度をスコア化し、顕著なトークンをサンプリングすることで、ImageNet上でGFLOPsを37％削減しながら精度を維持する。既存モデルに即座に統合可能な効率性向上ソリューションを実現する。

ABSTRACT

While state-of-the-art vision transformer models achieve promising results for image classification, they are computationally very 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 Sampling (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 anymore static but it varies for each input image. By integrating ATS as an additional layer within 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 off-the-shelf pretrained vision transformers as a plug-and-play module, thus reducing their GFLOPs without any additional training. However, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate our module on the ImageNet dataset by adding it to multiple state-of-the-art vision transformers. Our evaluations show that the proposed module improves the state-of-the-art by reducing the computational cost (GFLOPs) by 37% while preserving the accuracy.

研究の動機と目的

• ビジョントランスフォーマーの高い計算コスト（推論に多数のGFLOPsを要する）に対処すること。
• 多様な入力画像に対して最適でない固定トークン数の制限を克服すること。
• 入力画像ごとに適応的にトークン削減を可能にする、パラメータフリーで微分可能なモジュールの開発。
• 再訓練を必要とせず既存のビジョントランスフォーマーに統合可能であり、同時にエンド・ツー・エンドの微調整をサポートすること。
• 分類精度を損なわずに顕著な計算効率の向上を達成すること。

提案手法

• ATSは、クラストークンのアテンションマップに基づき、トークンの重要度を順位付けする微分可能なスコアリング機構を導入する。
• スコアに基づいて、最も顕著なトークンのみを適応的にサンプリングすることで、1画像あたりのトークン総数を削減する。
• ATSは標準的なトランスフォーマーブロック内に追加レイヤーとして挿入され、固定トークン処理パイプラインに置き換える。
• パラメータフリーであるため、微調整なしに事前学習済みビジョントランスフォーマーに直接適用可能である。
• 微分可能であるため、必要に応じてエンド・ツー・エンドの学習が可能で、トークンサンプリングの最適化が可能になる。
• 本手法はあらゆるビジョントランスフォーマー・アーキテクチャと互換性があり、プラグアンドプレイモジュールとして適用可能である。

実験結果

リサーチクエスチョン

• RQ1 パラメータフリーで微分可能なモジュールが、入力画像の複雑さに応じてビジョントランスフォーマーのトークン数を動的に削減できるか？
• RQ2 適応的トークンサンプリングは、精度の低下を伴わずに顕著な計算コストの削減を実現できるか？
• RQ3 提案手法のモジュールは、再訓練を伴わず既存の事前学習済みビジョントランスフォーマーに適用可能か？
• RQ4 ATSを搭載したビジョントランスフォーマーの性能は、元のバージョンと比較してGFLOPsと精度の面でどのように異なるか？
• RQ5 ATSを含むフルモデルのエンド・ツー・エンド学習により、さらに効率性と精度が向上するか？

主な発見

• ATSはImageNetデータセットにおいて、ビジョントランスフォーマーの計算コストをGFLOPsの観点から37％削減した。
• ATSを搭載したビジョントランスフォーマーの精度は、元のバージョンと比べて変化がなかった。
• 本モジュールは、追加のトレーニングを要せず、既存の事前学習済みビジョントランスフォーマーに即座にプラグアンドプレイで適用可能である。
• ATSの微分可能設計により、エンド・ツー・エンドの微調整が可能となり、サンプリングプロセスのさらなる最適化が可能になった。
• 適応的トークンサンプリング機構は、各画像に対して最も情報量の多いトークンのみを的確に特定・保持し、効率性を向上させた。
• 本手法は、顕著に推論複雑性を低減しながらも、最先端の性能を維持した。