[論文レビュー] Auto-scaling Vision Transformers without Training

Wuyang Chen, Wei Huang+4 more|arXiv (Cornell University)|Feb 24, 2022

Advanced Neural Network Applications被引用数 7

ひとこと要約

As-ViTは、複雑さに基づく探索により最初に最適な「シード」アーキテクチャを発見し、その後体系的に深さと幅をバランスさせながら拡大することで、トレーニング不要で自動化されたビジョントランスフォーマー（ViTs）の設計およびスケーリングフレームワークを提案する。本手法は12時間のV100 GPU時間で運用され、ImageNet-1kで83.5％のトップ1精度、COCO検出で52.7％のmAPという最先端の性能を達成しており、訓練中のFLOPsを最大56.2％まで削減する画期的な段階的リトークナイゼーション戦略を採用している。

ABSTRACT

This work targets automated designing and scaling of Vision Transformers (ViTs). The motivation comes from two pain spots: 1) the lack of efficient and principled methods for designing and scaling ViTs; 2) the tremendous computational cost of training ViT that is much heavier than its convolution counterpart. To tackle these issues, we propose As-ViT, an auto-scaling framework for ViTs without training, which automatically discovers and scales up ViTs in an efficient and principled manner. Specifically, we first design a "seed" ViT topology by leveraging a training-free search process. This extremely fast search is fulfilled by a comprehensive study of ViT's network complexity, yielding a strong Kendall-tau correlation with ground-truth accuracies. Second, starting from the "seed" topology, we automate the scaling rule for ViTs by growing widths/depths to different ViT layers. This results in a series of architectures with different numbers of parameters in a single run. Finally, based on the observation that ViTs can tolerate coarse tokenization in early training stages, we propose a progressive tokenization strategy to train ViTs faster and cheaper. As a unified framework, As-ViT achieves strong performance on classification (83.5% top1 on ImageNet-1k) and detection (52.7% mAP on COCO) without any manual crafting nor scaling of ViT architectures: the end-to-end model design and scaling process cost only 12 hours on one V100 GPU. Our code is available at https://github.com/VITA-Group/AsViT.

研究の動機と目的

• トレーニングに高コストがかかる中で、原理的かつ効率的なビジョントランスフォーマー（ViTs）の設計およびスケーリング手法の欠如に対処すること。
• CNNと比較してはるかに重いViTのトレーニングにかかる計算コストを低減すること。
• 一元的でエンドツーエンドのフレームワーク内で、アーキテクチャ発見とスケーリングを両方自動化すること。
• 手動でのハイパーパrameterチューニングや広範なハイパーパrameterサーチを一切行わず、高性能なViTを実現すること。
• 実際のモデル精度と強く相関する、スケーラブルでトレーニング不要の探索手法を開発すること。

提案手法

• ネットワークの複雑さ指標、特に期待される長さ歪みを精度の代理指標として用いるトレーニング不要なアーキテクチャ探索により、ViTトポロジーの高速評価が可能になる。
• 過去の手動によるViT設計に基づく緩められた探索空間から、複雑さに基づくランク付けにより高価なトレーニングを回避して「シード」ViTトポロジーを発見する。
• 段階的に層ごとに深さと幅を増加させながら、各段階でネットワークの複雑さをバランスさせる原理的なスケーリング則を採用する。
• 段階的リトークナイゼーションは、トレーニング中にダイレータやストライドを動的に調整することで、トークン数を削減し、トレーニングを加速する。
• 単一のシードアーキテクチャから出発し、一度の実行でパラメータ数が異なる複数のViTバリアントを生成するフレームワークである。
• ViTは訓練の初期段階で粗いトークナイゼーションに耐えられることに着目し、効率的な段階的精錬を可能にする。

実験結果

リサーチクエスチョン

• RQ1 ネットワークの複雑さ指標に基づくトレーニング不要な探索は、実際のトレーニングなしにViT性能を信頼性高く予測できるか？
• RQ2 深さと幅をバランスさせつつ性能を維持するための最適なViTスケーリング戦略は何か？
• RQ3 特に大規模モデルにおいて、精度を損なわずViTのトレーニングをどのように高速化できるか？
• RQ4 単一のシードアーキテクチャを最小限の追加コストで複数のバリエーションに体系的にスケーリングできるか？
• RQ5 段階的リトークナイゼーションは、トレーニング効率を向上させつつモデル精度を保持できるか？

主な発見

• 期待される長さ歪み指標は、計算コストと実際のViT精度とのケンダール-τ相関の両立において最良のトレードオフを達成し、トレーニング不要な探索に最適である。
• As-ViTフレームワークは、7時間のV100 GPU時間で「シード」ViTトポロジーを発見し、さらに5時間でスケーリングを実行し、合計12時間で設計とスケーリングを完了する。
• As-ViTはImageNet-1kで83.5％のトップ1精度、COCOオブジェクト検出で52.7％のmAPを達成し、既存の手作業によるViTを上回る性能を示した。
• 段階的リトークナイゼーション戦略により、固定された粗いトークナイゼーションと比較して、トレーニングFLOPsを最大56.2％まで削減した。
• フレームワークは、単一のシードから、異なるパラメータ数を持つViTバリアントのフルファミリーを生成でき、ハードウェア制約に応じた効率的な展開が可能である。
• 本手法は、探索と設計の段階で高価なトレーニングループを回避できるように、複雑さベースの指標を用いたViTの効果的なスケーリングが可能であることを示した。