[論文レビュー] ClusTR: Exploring Efficient Self-attention via Clustering for Vision Transformers

Yutong Xie, Jianpeng Zhang+3 more|arXiv (Cornell University)|Aug 28, 2022

Multimodal Machine Learning Applications被引用数 6

ひとこと要約

ClusTRは、視覚変換器におけるコンテンツベースのスパース自己注意機構を提案し、意味的類似性に基づいてキーと値のトークンをクラスタリングおよび集約することで、計算複雑性を低減し、長距離モデリングを効率化する。この手法は、22.7MパラメータでImageNetで83.2%のTop-1精度を達成し、密度型およびグリッドベースのスパース自己注意法を上回る最先端の性能を発揮する。

ABSTRACT

Although Transformers have successfully transitioned from their language modelling origins to image-based applications, their quadratic computational complexity remains a challenge, particularly for dense prediction. In this paper we propose a content-based sparse attention method, as an alternative to dense self-attention, aiming to reduce the computation complexity while retaining the ability to model long-range dependencies. Specifically, we cluster and then aggregate key and value tokens, as a content-based method of reducing the total token count. The resulting clustered-token sequence retains the semantic diversity of the original signal, but can be processed at a lower computational cost. Besides, we further extend the clustering-guided attention from single-scale to multi-scale, which is conducive to dense prediction tasks. We label the proposed Transformer architecture ClusTR, and demonstrate that it achieves state-of-the-art performance on various vision tasks but at lower computational cost and with fewer parameters. For instance, our ClusTR small model with 22.7M parameters achieves 83.2\% Top-1 accuracy on ImageNet. Source code and ImageNet models will be made publicly available.

研究の動機と目的

• 密度型自己注意の2次関数的計算複雑性を、特に密度予測タスクにおいて解決すること。
• グリッドベースのダウンサンプリングの限界、例えば細粒度の詳細の損失や余分な背景トークンの含むことの是正。
• 意味的豊かさを保持しながらシーケンス長を短縮するコンテンツに依存したトークン集約法の開発。
• クラスタリングガイドド自己注意を多スケールモデリングに拡張し、密度予測タスクにおける性能向上を図ること。
• モデルサイズと計算コストを低減しつつ、最先端の性能を達成すること。

提案手法

• キーと値のトークンに対して、そのコンテンツ類似性に基づいてクラスタリングを適用し、集約によってトークン数を削減する。
• クエリトークンはクラスタリングされたキーと値トークンに注目するため、長距離依存性モデリングを維持するクラスタリングガイドド自己注意機構を形成する。
• クラスタ数は柔軟に調整可能であり、固定グリッドベースのダウンサンプリングとは異なり、シーケンス長と計算コストを動的に制御可能である。
• 異なる段階で異なるクラスタ数を適用することで、多スケールクラスタリングを導入し、階層的特徴学習を可能にする。
• この手法は、ClusTRと名付けられた新しい変換器アーキテクチャに統合され、スモール、ベーシック、タイニーの3つのバージョンで提供される。
• モデルは画像分類、セマンティックセグメンテーション、オブジェクト検出、ポーズ推定のベンチマークで訓練および評価される。

実験結果

リサーチクエスチョン

• RQ1 キーと値のコンテンツベースのクラスタリングは、性能を損なわずにVision Transformerにおける計算複雑性を低減できるか？
• RQ2 グリッドベースや畳み込みベースのダウンサンプリングと比較して、クラスタリングガイドド自己注意は意味情報を保持し、精度を向上させるか？
• RQ3 多スケールクラスタリングガイドド自己注意は、セマンティックセグメンテーションなどの密度予測タスクにおける性能を向上させるか？
• RQ4 クラスタリングに基づくスパース自己注意機構は、少ないパラメータとFLOPsで最先端の結果を達成できるか？
• RQ5 クラスタ数の柔軟性が、さまざまなビジョンタスクにおけるモデルの効率性と精度に与える影響は何か？

主な発見

• 22.7MパラメータのClusTR-Sは、ImageNetで83.2%のTop-1精度を達成し、FLOPsとパラメータが少ない競合モデルを上回る。
• クラスタリングガイドド法は、グリッドベースのトークン集約（77.2%対76.7%）と比較して、パラメータを18%削減し、Top-1精度を0.5ポイント向上させる。
• k-NNおよび空間的還元自己注意と比較して、ClusTRは0.4%高いTop-1精度を達成しながら、パラメータを18%削減し、FLOPsを5%削減する。
• 多スケールクラスタリングは、FLOPsとパラメータのわずかな増加にもかかわらず、ADE-20KセマンティックセグメンテーションのmIOUを1.4ポイント向上（42.6対41.2）させる。
• 40.3MパラメータでImageNetで84.1%のTop-1精度を達成し、強力なスケーラビリティを示す。
• アブレーションスタディにより、グリッドベースやk-NNスパース自己注意よりも、クラスタリングガイドド自己注意が精度と効率の両面で優れていることが確認される。