[論文レビュー] Vision Transformer Architecture Search

Xiu Su, Shan You+7 more|arXiv (Cornell University)|Jun 25, 2021

Advanced Neural Network Applications被引用数 10

ひとこと要約

この論文では、共有重みとプライベートクラストークンを用いたスーパitransformerを採用することで、多様なアーキテクチャを効率的に探索できる、ビジョントランスフォーマー向けのニューラルアーキテクチャサーチ手法ViTASを提案する。ハードウェア制約の下で深さ、幅、シーケンス長、アテンションヘッド数を同時に探索することで、1.3G FLOPsの制約下でImageNetで74.7%のトップ1精度を達成し、ベースラインのViTよりも2.5%高い性能を発揮した。

ABSTRACT

Recently, transformers have shown great superiority in solving computer vision tasks by modeling images as a sequence of manually-split patches with self-attention mechanism. However, current architectures of vision transformers (ViTs) are simply inherited from natural language processing (NLP) tasks and have not been sufficiently investigated and optimized. In this paper, we make a further step by examining the intrinsic structure of transformers for vision tasks and propose an architecture search method, dubbed ViTAS, to search for the optimal architecture with similar hardware budgets. Concretely, we design a new effective yet efficient weight sharing paradigm for ViTs, such that architectures with different token embedding, sequence size, number of heads, width, and depth can be derived from a single super-transformer. Moreover, to cater for the variance of distinct architectures, we introduce extit{private} class token and self-attention maps in the super-transformer. In addition, to adapt the searching for different budgets, we propose to search the sampling probability of identity operation. Experimental results show that our ViTAS attains excellent results compared to existing pure transformer architectures. For example, with $1.3$G FLOPs budget, our searched architecture achieves $74.7\%$ top-$1$ accuracy on ImageNet and is $2.5\%$ superior than the current baseline ViT architecture. Code is available at \url{this https URL}.

研究の動機と目的

• NLPモデルに由来するビジョントランスフォーマー（ViTs）の最適でない設計を改善すること。
• 固定されたハードウェア予算の下で、深さ、幅、シーケンス長、アテンションヘッド数といった複数のアーキテクチャ次元を効率的に探索できること。
• 1つのスーパitransformerが多様なViTアーキテクチャを表現できる重み共有メカニズムを設計すること。
• スーパーネット内のアーキテクチャのばらつきに対応するため、プライベートクラストークンとアテンションマップを導入すること。
• アイデンティティ演算のサンプリング確率を学習することで、異なるFLOP予算に適応した探索を可能にすること。

提案手法

• すべてのアーキテクチャバリアント間でアテンションおよびフィードフォワード重みを共有するスーパitransformerを設計する。
• 各アーキテクチャバリアント用にプライベートクラストークンと自己アテンションマップを導入し、アーキテクチャの多様性に対処する。
• アイデンティティ演算のための学習可能なサンプリング確率を用いた微分可能アーキテクチャサーチ戦略を実装する。
• トークン埋め込みサイズ、シーケンス長、アテンションヘッド数、幅、深さの変化を含む探索空間を設定する。
• 精度とFLOP予算のバランスを取る微分可能な探索目的関数を用いてスーパーネットを最適化する。
• 性能を維持したまま探索コストを削減するため、早期停止とプルーニングを適用する。

実験結果

リサーチクエスチョン

• RQ1 深さ、幅、シーケンス長が異なる多様なビジョントランスフォーマーのアーキテクチャを統一されたスーパーネットが効果的に表現できるか？
• RQ2 プライベートクラストークンとアテンションマップの導入が、共有アーキテクチャサーチスペース内での性能向上にどのように寄与するか？
• RQ3 学習可能なサンプリング確率を用いた微分可能アーキテクチャサーチは、固定されたFLOP予算下でViTの精度をどの程度向上できるか？
• RQ4 ViTASは、ImageNetにおける既存のViTアーキテクチャと比較して、精度と効率の両面で優れているか？
• RQ5 提案された重み共有方式は、性能を損なわせることなく、効率的かつ効果的なアーキテクチャサーチを可能にするか？

主な発見

• ViTASは1.3G FLOPsの予算下でImageNetで74.7%のトップ1精度を達成し、ベースラインのViTよりも2.5ポイント高い性能を発揮した。
• 同じハードウェア制約下で、既存の純粋なトランスフォーマー・モデルと比較して、探索されたアーキテクチャは優れた精度を示した。
• プライベートクラストークンとアテンションマップの使用により、スーパーネット内での多様なアーキテクチャ間の干渉が顕著に低減され、性能が向上した。
• 適応的サンプリング確率を用いた微分可能探索により、最小限の計算オーバーヘッドでアーキテクチャ空間の効率的探索が可能になった。
• 重み共有メカニズムにより、1つのスーパーネット内で深さ、幅、シーケンス長、アテンションヘッド数の全次元にわたる効果的なアーキテクチャサーチが実現された。
• 指定されたFLOP予算下で、純粋なトランスフォーマー・モデルの中で、ViTASは最先端の性能を達成した。