150億のパラメータを持つ、史上最大のビジュアルモデル「V-MoE」の全コードをGoogleがオープンソース化

昨年 6 月に Google Brain チームが発表した 43 ページの論文「Scaling Vision with Sparse Mixture of Experts」を覚えていますか?彼らは史上最大の視覚モデルである V-MoE を発表し、SOTA に近いトップ 1 の精度を達成しました。現在、Google Brain はモデルのトレーニングと微調整のためのすべてのコードをオープンソース化しています。

過去数十年にわたり、ディープラーニングの進歩は、少数のシンプルで柔軟なメカニズム、大規模なデータセット、より特殊なハードウェア構成など、いくつかの重要な要素によって推進されてきました。これらの進歩により、ニューラル ネットワークは、画像分類、機械翻訳、タンパク質予測などのタスクで素晴らしい結果を達成できるようになりました。

ただし、大規模なモデルやデータセットを使用すると、計算要件が増大します。最近の研究では、モデルの一般化と堅牢性を高めるには、大規模なモデルのサポートが必要であることが示されています。したがって、大規模なモデルをトレーニングする際には、トレーニング リソースの制限と調整することが非常に重要です。考えられるアプローチの 1 つは、条件付き計算を活用することです。これは、単一の入力に対してネットワーク全体をアクティブ化するのではなく、異なる入力に基づいてモデルのさまざまな部分をアクティブ化する方法です。このパラダイムは、Google の Pathways ビジョンや大規模言語モデルに関する最近の研究で強調されていますが、コンピューター ビジョンでは十分に調査されていません。

スパースゲートエキスパート混合 (MoE) ネットワークは、自然言語処理において優れたスケーラビリティを実証しています。しかし、コンピューター ビジョンでは、ほぼすべての高性能ネットワークは高密度であり、つまり、各入力が処理用のパラメーターに変換されます。

昨年 6 月、Google Brain の研究者らは、専門家の疎混合に基づく新しい視覚アーキテクチャである V-MoE (Vision MoE) を提案しました。 V-MoE を画像認識に適用すると、推論時に半分の計算量で最先端のネットワークのパフォーマンスを達成できます。さらに、この研究では、バッチ全体内の各入力のサブセットに優先順位を付けることができるルーティング アルゴリズムの拡張を提案し、それによって適応的な画像計算を実現します。これにより、V-MoE はテスト時にパフォーマンスとスムーズな計算をトレードオフできます。最後に、この研究では、V-MoE が視覚モデルをスケールアップする可能性を実証し、ImageNet で 90.35% を達成した 150 億のパラメータ モデルをトレーニングしました。

論文アドレス: https://arxiv.org/pdf/2106.05974.pdf

コードアドレス: https://github.com/google-research/vmoe

V-MoE

Google Brain は、ViT のさまざまなバリアント (ViT-S (mall)、ViT-B (ase)、ViT-L (arge)、ViTH (uge)) に基づいて、次のハイパーパラメータを使用して V-MoE を構築します。

ViT は転移学習設定で適切にスケーリングされ、事前トレーニングの計算が少なくても CNN よりも高い精度を達成できることが示されています。 ViT は、画像を一連のパッチとして処理します。入力画像は最初に同じサイズのパッチに分割され、Transformer の隠れ層に線形投影されます。位置埋め込みの後、パッチ埋め込み (トークン) は、主に交互に繰り返される自己注意層と MLP 層で構成される Transformer によって処理されます。 MLP には 2 つのレイヤーと GeLU 非線形性があります。 Vision MoE の場合、研究ではそのサブセットを MoE レイヤーに置き換えました。各エキスパートは MLP であり、次の図に示されています。

ビジョン モデルを大規模に拡張するために、ViT アーキテクチャ内の一部の密なフィードフォワード レイヤー (FFN) を、独立した FFN (エキスパートと呼ばれる) のスパースな混合に置き換えます。学習可能なルーティング レイヤーは、個々のトークンごとに対応するエキスパートを選択します。つまり、同じ画像からの異なるトークンが異なるエキスパートにルーティングされる可能性があります。各トークンは、合計 E 人のエキスパート (E は通常 32) のうち最大 K 人 (通常は 1 人または 2 人) のエキスパートにルーティングできます。これにより、トークンごとの計算を一定に保ちながら、モデルのサイズをスケーリングできます。次の図は、V-MoE エンコーダ ブロックの構造をより詳細に示しています。

V-MoE トランスフォーマー エンコーダー ブロック

実験結果

Google Brain はまず、大規模な画像データセット JFT-300M でモデルを事前トレーニングします。

下の図は、すべてのサイズ（小さい s/32 から巨大な H/14 まで）でのモデルの事前トレーニング結果を示しています。次に、新しいヘッド (モデルの最後のレイヤー) を使用して、モデルを新しい下流タスク (ImageNet など) に転送します。彼らは、新しいタスクの利用可能なすべての例でモデル全体を微調整するか、事前トレーニング済みのネットワークをフリーズし、少数の例を使用して新しいヘッドのみを微調整する（いわゆる少数ショット転送）という 2 つの転送設定を検討しました。

下の図 (右) は、モデルを ImageNet に移行した場合の効果をまとめたものです。ここでは、各画像カテゴリが 5 枚の画像のみでトレーニングされています (5 ショット転送と呼ばれます)。

左側は JFT-300M データセットの Precision@1 曲線、右側は ImageNet 5 ショット精度曲線です。

どちらの場合も、Google Brain は、与えられた量のトレーニング計算に対して、スパース モデルが密なモデルよりも大幅に優れているか、同様のパフォーマンスをより速く達成することを発見しました。視覚モデルの限界を探るために、研究者らは、JFT-300M 拡張データセット上で 24 個の MoE レイヤー (48 ブロック) を持つ 150 億個のパラメータ モデルをトレーニングしました。このこれまでで最大のビジョン モデルは、ImageNet で 90.35% のトップ 1 精度を達成しています。

優先ルーティング

実際には、動的にサイズが調整されるバッファを使用することはハードウェアの制限により非効率的であるため、モデルでは各エキスパートに対して事前定義されたバッファ容量が使用されることが多いです。エキスパートが「いっぱい」になると、この容量を超えて割り当てられたトークンは破棄され、処理されなくなります。したがって、容量が大きいほど精度は高くなりますが、計算コストも高くなります。

Google Brain はこの実装制約を利用して、推論時の V-MoE を高速化します。処理されるトークンの数よりも合計バッファ容量を減らすと、ネットワークはエキスパート レイヤーで一部のトークンの処理をスキップせざるを得なくなります。以前の研究のように、何らかの恣意的な方法でスキップするトークンを選択する代わりに、モデルは重要度スコアに従ってトークンをランク付けすることを学習します。これにより、多くの計算を節約しながら、高品質の予測を維持できます。この方法はバッチ優先ルーティング (BPR) と呼ばれ、動的な図は次のようになります。

ボリュームが大きい場合、バニラ ルーティングと優先ルーティングの両方ですべてのパッチが適切に処理されます。ただし、計算を節約するためにバッファ サイズを縮小すると、Vanilla ルーティングは任意のパッチの処理を選択するため、予測精度が低下することがよくあります。BPR は重要なパッチの処理をインテリジェントに優先するため、計算コストを抑えながら予測精度が向上します。

トークンを適切に削除することが、高品質でより効率的な推論予測を提供するために重要であることがわかりました。エキスパートの能力が低下すると、Vanilla ルーティング メカニズムのパフォーマンスが急速に低下します。対照的に、BPR は低ボリュームに対してより堅牢です。

全体的に、V-MoE は推論時に非常に柔軟であることがわかりました。たとえば、モデルの重みをさらにトレーニングすることなく、トークンごとに選択されるエキスパートの数を減らして、時間と計算を節約できます。

V-MoEを探索する

スパースネットワークの内部の仕組みについてはまだ解明すべきことがたくさんあるため、Google Brain は V-MoE のルーティング パターンも調査しました。 1 つの仮説は、ルーターが何らかの意味的コンテキスト (「車」の専門家、「動物」の専門家など) に基づいてトークンを区別して専門家に割り当てることを学習するというものです。

これをテストするために、彼らは 2 つの異なる MoE レイヤー (1 つは非常に初期のもので、もう 1 つは頭部に近いもの) の以下の画像を示しています。 x 軸は 32 人のエキスパートそれぞれに対応し、y 軸は画像クラスの ID (1 ～ 1000) を示します。図の各エントリは、特定の画像クラスに対応するトークンに対してエキスパートが選択された頻度を示しており、色が濃いほど頻度が高いことを示します。

結果は、初期の層では相関がほとんどない一方で、ネットワークの後半では各エキスパートが少数のカテゴリからのトークンのみを受け取って処理することを示しています。したがって、パッチの何らかの意味的クラスタリングはネットワークのより深い層に現れると結論付けることができます。

より高いルーティング決定は、画像カテゴリに関連付けられます。

Google Brain は、これがコンピューター ビジョンの大規模な条件付き計算の始まりに過ぎないと考えています。異種エキスパート アーキテクチャと条件付き可変長ルーティングも潜在的な研究方向です。スパース モデルは、大規模なビデオ モデリングなどのデータが豊富な領域で特に役立ちます。彼らは、オープンソースのコードとモデルによって、より多くの研究者がこの分野に注目するようになることを期待しています。