[論文レビュー] PSViT: Better Vision Transformer via Token Pooling and Attention Sharing

Boyu Chen, Peixia Li+7 more|arXiv (Cornell University)|Aug 7, 2021

CCD and CMOS Imaging Sensors参考文献 44被引用数 20

ひとこと要約

PSViTは、空間的および層間の冗長性を低減するため、学習可能なトークンプールとアテンション共有を導入することで、ビジョントランスフォーマーを改善し、より優れた特徴表現と優れた速度-精度トレードオフを実現する。コンactな探索空間とAutoML（SPOS）を用いることで、ImageNet上でDeiTに比べて最高6.6%のトップ-1精度向上を達成し、下流の検出およびセグメンテーションタスクでもResNetsを上回る性能を発揮する。

ABSTRACT

In this paper, we observe two levels of redundancies when applying vision transformers (ViT) for image recognition. First, fixing the number of tokens through the whole network produces redundant features at the spatial level. Second, the attention maps among different transformer layers are redundant. Based on the observations above, we propose a PSViT: a ViT with token Pooling and attention Sharing to reduce the redundancy, effectively enhancing the feature representation ability, and achieving a better speed-accuracy trade-off. Specifically, in our PSViT, token pooling can be defined as the operation that decreases the number of tokens at the spatial level. Besides, attention sharing will be built between the neighboring transformer layers for reusing the attention maps having a strong correlation among adjacent layers. Then, a compact set of the possible combinations for different token pooling and attention sharing mechanisms are constructed. Based on the proposed compact set, the number of tokens in each layer and the choices of layers sharing attention can be treated as hyper-parameters that are learned from data automatically. Experimental results show that the proposed scheme can achieve up to 6.6% accuracy improvement in ImageNet classification compared with the DeiT.

研究の動機と目的

• 固定されたトークン数が各層にわたって与えられるために生じるビジョントランスフォーマーにおける空間的冗長性を解消すること。
• 隣接する層における類似したアテンションパターンが原因で生じる層間アテンションマップの冗長性を低減すること。
• 最適なトークン数とアテンション共有構成を学習することで、計算割り当ての柔軟性を高めること。
• 特徴表現を向上させ、ビジョントランスフォーマーのアーキテクチャにおいてより優れた速度-精度トレードオフを達成すること。
• AutoMLを用いた原理的でデータ駆動型のViTアーキテクチャ設計手法の開発。

提案手法

• 特徴の複雑さに応じて、ネットワークの異なる段階でトークン数を動的に削減するトークンプールを導入し、適応的に行う。
• 隣接するトランスフォーマーレイヤー間でアテンションを共有することで、類似したアテンションマップを再利用し、計算量を削減する。
• トークンプールの位置とアテンション共有の選択をハイパーパrameterとして組み合わせた、コンパクトな探索空間を構築する。
• SPOS（Single Path One-Shot）を用いてニューラルアーキテクチャサーチを実施し、同時にトークン数とアテンション共有パターンを最適化する。
• トークンプールとアテンション共有を学習可能なハイパーパrameterとして扱い、微分可能アーキテクチャサーチによりエンドツーエンドで学習する。
• 1次元および2次元のトークンプールのバリエーションを採用し、2次元プールが画像タスクにおいて空間構造をよりよく保持する。

実験結果

リサーチクエスチョン

• RQ1 層に跨る動的トークンプールは、ビジョントランスフォーマーにおける特徴表現を向上させ得るか？
• RQ2 隣接するViT層におけるアテンションマップはどの程度冗長性を示し、それを活用可能か？
• RQ3 トークンプールとアテンション共有のためのコンパクトな探索空間は、効果的なAutoMLベースのアーキテクチャサーチを可能にするか？
• RQ4 トークン数とアテンション共有の共同最適化は、より優れた精度と効率のトレードオフをもたらすか？
• RQ5 PSViTはオブジェクト検出やインスタンスセグメンテーションなどの下流ビジョンタスクにどの程度転送可能か？

主な発見

• PSViTは、DeiTに比べてImageNetで最高6.6%高いトップ-1精度を達成し、顕著な性能向上を示した。
• PSViT-2D-Tinyは、COCOオブジェクト検出で40.8%のmAP、インスタンスセグメンテーションで37.7%のmAPを達成し、ResNet18でさえも上回った。
• 2次元トークンプールのバリエーションが1次元バージョンを上回ったことから、空間構造を保持することで特徴品質が向上することが示された。
• 最後のトランスフォーマーレイヤーではアテンション共有を避けるべきである。なぜなら、最終分類に最も利益をもたらす独立した自己アテンションが有効だからである。
• ネットワークの最初の段階が、高い計算コストと特徴の滑らかさのため、アテンション共有から最も恩恵を受ける。
• 探索空間の設計により、効果的なAutoML最適化が可能となり、PSViT-2D-Tinyは最小限の人的介入で優れた性能を達成した。