トレーニングは不要、自動的にスケーラブルなビジュアルトランスフォーマーが登場

現在、ビジョントランスフォーマー (ViT) の分野には 2 つの大きな問題点があります。1. ViT を設計および拡張するための効果的な方法が不足している。2. ViT のトレーニングの計算コストが畳み込みネットワークよりもはるかに高い。

これら 2 つの問題に対処するために、テキサス大学オースティン校、シドニー工科大学、Google の研究者らは、効率的かつ原則的な方法で ViT を自動的に設計および拡張する、トレーニング不要の ViT 自動拡張フレームワークである As-ViT (Auto-scaling Vision Transformers) を提案しました。

論文リンク: https://arxiv.org/abs/2202.11921

具体的には、研究者らはまず、トレーニング不要の検索プロセスを使用して ViT トポロジの「シード」を設計しました。この極めて高速な検索は、ViT ネットワークの複雑さを包括的に研究することで実現され、真の精度との強いケンドール・タウ相関がもたらされました。次に、「シード」トポロジから始めて、さまざまな ViT レイヤーの幅/深さを増やすことで ViT の拡張ルールを自動化し、1 回の実行でさまざまな数のパラメーターを持つさまざまなアーキテクチャを実現します。最後に、ViT はトレーニングの初期段階では粗粒度のトークン化を許容できるという経験に基づいて、本研究では、ViT をより迅速かつ経済的にトレーニングするための漸進的なトークン化戦略を提案します。

統合フレームワークである As-ViT は、ViT アーキテクチャを手動で調整したり拡張したりすることなく、分類 (ImageNet-1k で 83.5% トップ 1) と検出 (COCO で 52.7% mAP) タスクで優れたパフォーマンスを実現します。エンドツーエンドのモデル設計と拡張プロセスは、V100 GPU でわずか 12 時間しかかかりません。

ネットワークの複雑さを考慮したViTの自動設計と拡張

ViT 設計をスピードアップし、面倒な手作業を回避するために、この研究では、効率的で自動化された原則的な ViT 検索と拡張を目指しています。具体的には、解決する必要がある問題が 2 つあります。1) トレーニング コストを最小限に抑えて、またはゼロにして、最適な ViT アーキテクチャ トポロジを効率的に見つけるにはどうすればよいでしょうか。 2) モデル サイズのさまざまな要件を満たすために、ViT トポロジの深さと幅をどのように拡張しますか?

ViTの位相空間の拡張

設計と拡張の前に、まず As-ViT のトポロジカル検索空間を拡張します。入力画像は最初に 1/4 スケールの解像度のブロックに埋め込まれ、段階的な空間縮小とチャネル倍増の戦略が採用されます。これは、マルチスケール機能を必要とする検出などの高密度予測タスクを容易にするためです。

多様体伝播による初期化時の ViT 複雑性の評価

ViT のトレーニングは遅いため、トレーニングされたモデルの精度を評価してアーキテクチャ検索を実行すると、非常にコストがかかります。最近では、ローカル線形グラフ（Mellor et al.、2020）、勾配感度（Abdelfattah et al.、2021）、線形領域の数（Chen et al.、2021e;f）、ネットワークトポロジー（Bhardwaj et al.、2021）などの方法を使用して、ReLUベースのCNNを使用したトレーニングフリーのニューラルアーキテクチャ検索方法が数多くあります。

ただし、ViT には、自己注意、ソフトマックス、GeLU などのより複雑な非線形関数が備わっています。したがって、より一般的な方法で学習能力を測定する必要があります。新しい研究では、研究者らは、ViT を介した多様体の伝播の複雑さを測定し、ViT によって複雑な関数をどの程度近似できるかを推定することを検討しました。直感的に言えば、複雑なネットワークは、単純な入力を出力層で複雑な多様体に伝播できるため、強力な学習能力を持つ可能性があります。 UT オースティンの研究では、ViT を介して単純な円形入力の複数の複雑性をマッピングしました: h(θ) = √ N [u^0 cos(θ) + u^1 sin(θ)]。ここで、N は ViT 入力の次元 (たとえば、ImageNet 画像の場合、N = 3 × 224 × 224) であり、u^0 と u^1 は円が存在する R^N の 2D サブスペースの正規直交基底を形成します。

ViTトポロジー報酬を検索

研究者らは、L^E (アルゴリズム 1) に基づくトレーニング不要の検索を提案しました。ほとんどの NAS (ニューラル アーキテクチャ検索) 手法では、代理推論として単一パスまたはスーパー ネットワークの精度または損失値を評価します。このトレーニングベースの検索を ViT に適用すると、より多くの計算コストが必要になります。サンプリングされた各アーキテクチャについて、ViT をトレーニングする代わりに、L^E が計算され、検索プロセスをガイドするための報酬として考慮されます。

L^Eに加えて、ViTの訓練可能性を示すためにNTK条件数κΘ=λ_max/λ_minも含まれています（Chen et al.、2021e; Xiao et al.、2019; Yang、2020; Hron et al.、2020）。 λ_maxとλ_minはNTK行列Θの最大固有値と最小固有値です。

この検索で​​は強化学習法が使用され、戦略は結合分類分布として定義され、ポリシー勾配によって更新されます。この研究では、戦略を 500 ステップ更新し、戦略を収束させるのに十分であることを確認しました (エントロピーは 15.3 から 5.7 に低下しました)。検索プロセスは非常に高速です。ViT トレーニングをバイパスする L^E の単純な計算により、ImageNet-1k データセットではわずか 7 GPU 時間 (V100) しかかかりません。 L^EとκΘのサイズの違いを考慮するために、この研究ではそれらを相対的な値の範囲で正規化します（アルゴリズム1の5行目）。

表 3 は、新しい検索方法の ViT トポロジ統計をまとめたものです。 L^EとκΘは、(1)オーバーラップのあるトークン（K_1〜K_4はどちらもストライドより大きい）、および(2)より深いレイヤーでのFFN拡張率が大きい（E_1 < E_2 < E_3 < E_4）ことを非常に好むことがわかります。注意の分割と前頭葉の数において、L^E と κΘ に有意な好みは見られませんでした。

ViT独自の原則的な拡大

最適なトポロジを取得した後、次に解決すべき問題は、ネットワークの深さと幅のバランスをどのように取るかということです。

現在、ViT 拡張機能についてはそのような経験則はありません。最近の研究では、さまざまなリソース制約を満たすために、さまざまなサイズの畳み込みネットワークをスケールアップまたは拡大する試みがなされてきました (Liu et al.、2019a; Tan & Le、2019)。しかし、原理的な展開ルールを自動的に見つけるように ViT をトレーニングすると、膨大な計算コストがかかります。異なる ViT バリアントを検索することも可能です (セクション 3.3 で説明) が、これには複数回の実行が必要です。対照的に、「スケールアップ」は、単一の実験で複数のモデルバリアントを生成するより自然な方法です。したがって、この研究では、検索された基本的な「シード」ViTを、トレーニング不要で原理的に効率的な方法で、より大きなモデルに拡張することを試みます。この自動拡張方法はアルゴリズム 2 で説明されています。

Inception アーキテクチャでは、各ステージに 1 つのアテンション ブロックがあり、初期の隠し次元は C = 32 です。各反復では、上方へのさらなる拡張に最適な深さと幅が見つかります。深さについては、どの段階を深くするか（つまり、どの段階で注意ブロックを追加するか）を研究し、幅については、最適な拡張率（つまり、チャネルの数をどの程度拡張するか）を研究します。

拡張された軌道を下の図 3 に示します。研究者らは、自律的な拡張とランダムな拡張を比較し、拡張原理では、より浅いがより広いネットワークを使用して、より広い幅と引き換えに深さを放棄することを好むことを発見した。この拡張は、Zhai et al. (2021) によって開発されたルールに似ています。対照的に、ResNet と Swin Transformer (Liu et al.、2021) は、より狭く、より深くすることを選択します。

段階的に柔軟な再トークン化による効率的なViTトレーニング

この研究では、段階的に柔軟になる再トークン化トレーニング戦略を提案することで肯定的な答えを提供します。線形投影の重みの形状に影響を与えずにトレーニング中にトークンの数を更新するために、この研究では最初の線形投影層で異なるサンプリング粒度を採用しています。最初の投影カーネル K_1 = 4 とストライド = 4 を例に挙げます。トレーニング中、研究者は重みの形状と構造を変更せずに、最初の投影カーネルの (ストライド、拡張) ペアを (16, 5)、(8, 2)、(4, 1) に徐々に変更しました。

この再トークン化戦略は、ViT のカリキュラム学習を促進します。トレーニングの開始時に粗いサンプリングが導入され、トークンの数が大幅に削減されます。言い換えれば、As-ViT は、トレーニングの初期段階で、非常に低い計算コスト (フル解像度トレーニングのわずか 13.2% FLOP) で、画像から大まかな情報を素早く学習します。トレーニングの後の段階では、徐々にきめ細かいサンプリングに切り替え、競争力のある精度を維持しながら完全なトークン解像度を回復します。図 4 に示すように、ViT をトレーニングの初期段階で粗いサンプリングでトレーニングすると、計算コストが非常に低く抑えられながらも高い精度を達成できます。異なるサンプリング粒度間の遷移によりパフォーマンスが急激に上昇しますが、ネットワークは最終的に競争力のある最終パフォーマンスを回復します。

図 4 に示すように、ViT をトレーニングの初期段階で粗いサンプリングを使用してトレーニングすると、計算コストが非常に低く抑えられながらも高い精度を達成できます。異なるサンプリング粒度間の遷移によりパフォーマンスが急激に上昇しますが、ネットワークは最終的に競争力のある最終パフォーマンスを回復します。

実験

AS-VIT: 自動延長VIT

この研究では、検索された As-ViT トポロジーを表 4 に示します。このアーキテクチャは、最初の投影（トークン化）ステップと 3 つの再埋め込みステップでトークン間の強力なオーバーラップを促進します。 FFN の拡張率は、最初は狭くなり、その後、深い層に行くほど広くなります。少数の注意分割を利用して、グローバル情報をより適切に集約します。

画像分類

以下の表 5 は As-ViT と他のモデルの比較を示しています。従来のトランスフォーマーベースおよび CNN ベースのアーキテクチャと比較すると、As-ViT は同等の数のパラメーターと FLOP で SOTA パフォーマンスを実現します。

効率的なトレーニング

研究者らは、表 6 の各トークンの削減フェーズの期間を調整し、その結果を表 6 に示しました。標準的なトレーニングには 42.8 TPU 日が必要ですが、効率的なトレーニングでは、高い精度を維持しながら、トレーニング FLOP を最大 56.2%、トレーニング TPU 日を最大 41.1% 節約できます。

トポロジーと拡張への貢献

検索型トポロジーと拡張ルールの貢献をよりよく検証するために、本研究ではさらに多くのアブレーション研究を実施しました（表7）。まず、検索されたトポロジーは拡張前に直接トレーニングされます。この研究で検索されたシード トポロジは、図 2 の 87 個のランダム トポロジの中で最良のトポロジよりも優れています。

第二に、この研究では複雑性に基づくルールと「ランダム拡張 + As-ViT トポロジ」を比較しました。この研究における自動拡張は、さまざまな拡張の下でのランダム拡張よりも優れています。

COCOデータセット上のオブジェクト検出

この研究では、As-ViT を標準的な CNN および以前の Transformer ネットワークと比較しました。比較は、他の設定は変更せずにトランクのみを変更して行いました。以下の表 8 の結果から、As-ViT は ImageNet に基づいて設計されており、その複雑さは分類のために測定されているにもかかわらず、マルチスケールの特徴をキャプチャし、最先端の検出パフォーマンスを実現できることがわかります。

​​