[論文レビュー] Scaling Vision Transformers to 22 Billion Parameters

Mostafa Dehghani, Josip Djolonga+40 more|arXiv (Cornell University)|Feb 10, 2023

Multimodal Machine Learning Applications被引用数 118

ひとこと要約

本論文は ViT-22B を提示する。22B パラメータのビジョントランスフォーマーで、安定性と効率性のためのアーキテクチャ変更を行い、分類、ゼロショット、密集予測、ビデオ、フェアネス/頑健性のベンチマークで最先端または競合的な結果を達成する。

ABSTRACT

The scaling of Transformers has driven breakthrough capabilities for language models. At present, the largest large language models (LLMs) contain upwards of 100B parameters. Vision Transformers (ViT) have introduced the same architecture to image and video modelling, but these have not yet been successfully scaled to nearly the same degree; the largest dense ViT contains 4B parameters (Chen et al., 2022). We present a recipe for highly efficient and stable training of a 22B-parameter ViT (ViT-22B) and perform a wide variety of experiments on the resulting model. When evaluated on downstream tasks (often with a lightweight linear model on frozen features), ViT-22B demonstrates increasing performance with scale. We further observe other interesting benefits of scale, including an improved tradeoff between fairness and performance, state-of-the-art alignment to human visual perception in terms of shape/texture bias, and improved robustness. ViT-22B demonstrates the potential for "LLM-like" scaling in vision, and provides key steps towards getting there.

研究の動機と目的

• 22B-parameter Vision Transformer (ViT-22B) を実現するスケーラブルな学習技術を示す。
• ViT-22B を多様なタスクで評価する: 画像分類、ゼロショット転移、密集予測、ビデオ。
• モデルのスケールが大きくなるにつれて、公平性、頑健性、キャリブレーション、そして人間適合性への影響を分析する。
• 大規模なViTsが小型バックボーンへの蒸留を通じて有効な教師になり得ることを示す。

提案手法

• パラレルレイヤー、QK正規化、バイアスの省略を導入し、スケール時の訓練を安定化・高速化。
• TPUv4 上での2Dメッシュシャーディングを用いた非同期モデル並列アプローチでスループットを最大化する（1.15k トークン/秒/コア）。
• 大規模モデルとバッチサイズに適合させるためにモデルパラメータと活性をシャードし、計算と通信を重ね合わせる。
• 4B 画像のJFT由来データセットで事前学習。画像あたり256トークン、177kステップのスケジュール。
• 線形プロービング、ロックドイメージチューニング、エンドツーエンドのファインチューニングを複数の下流タスクで評価する。

実験結果

リサーチクエスチョン

• RQ1 アーキテクチャの変更を伴うViT規模のトレーニングは、ViT-22Bに対して安定かつ効率的な訓練を生み出せるか？
• RQ2 従来のViTおよびLiTベースの手法と比較し、標準およびOODの画像分類タスクでViT-22Bはどの程度性能を発揮するか？
• RQ3 ViTのスケールアップはゼロショット、転移、ドメイン横断の頑健性、公平性、人間適合性の指標を改善するか？
• RQ4 固定バックボーンとして使用した場合、ViT-22Bは密集予測とビデオタスクに対して強力な特徴表現を提供できるか？

主な発見

• ViT-22B は凍結特徴抽出器としての強力な ImageNet パフォーマンスを達成（89.5% の精度）し、対応テキストタワーを用いたゼロショット ImageNet で 85.9%。
• ViT-22B を ViT-B/16 および ViT-L/16 に蒸留すると、これら小型モデルの ImageNet 精度が最先端となる（それぞれ 88.6%、89.6%）。
• ObjectNet におけるゼロショット結果はモデルサイズと共に改善され、この難解なデータセットで ViT-22B の新しいSOTAを樹立。
• ViT-22B は人間への形状バイアス適合を改善（87% の形状バイアス）し、サブグループとキャリブレーション指標全体で公正性/頑健性のトレードオフが向上。
• 密集予測転移（ADE20k 少数ショット）と単眼深度推定は ViT-22B の特徴から恩恵を受け、ViT-L および ViT-G のベースラインを上回る。
• 固定 ViT-22B バックボーンを用いた動画評価は、従来の4Bパラメータモデルに対して競争力のある結果を示し、フルファインチューニングでの向上余地がある。