Google の Transformer が NeurIPS 2021 で 4 つの SOTA 賞を受賞

[[440180]]

Google の Vision Transformer に関する新しい取り組みが NeurIPS 2021 に含まれました。この記事では、Google が TokenLearner メソッドを提案しました。このメソッドは、Vision Transformer で使用すると計算量を最大 8 倍削減し、さらに優れた分類パフォーマンスを実現できます。

現在、Transformer モデルは、物体検出やビデオ分類などのコンピューター ビジョン タスクにおいて最先端の結果を達成しています。

画像をピクセルごとに処理する標準的な畳み込み方法とは異なり、Vision Transformer (ViT) は、画像を一連のパッチ トークン (つまり、複数のピクセルで構成される画像の小さな部分) として扱います。

これは、ニューラル ネットワークの各層で、ViT モデルがマルチヘッド セルフ アテンションを使用して、各トークンのペア間の関係に基づいてパッチ トークンを処理することを意味します。

このようにして、ViT モデルは画像全体のグローバル表現を構築できます。

入力側では、画像は複数の部分に均等に分割され、トークンが形成されます。たとえば、512×512 ピクセルの画像は、16×16 ピクセルのパッチ トークンに分割されます。中間層では、前の層の出力が次の層のトークンになります。

ここに一文追加させてください。ビデオを処理する場合、16x16x2 ビデオ クリップ (16x16 画像の 2 フレーム) などのビデオ「パイプライン」がトークンになります。ビジュアル トークンの品質と量によって、Vision Transformer の全体的なパフォーマンスが決まります。

多くの Vision Transformer 構造が直面する主な課題は、妥当な結果を得るために必要なトークンが多すぎることが多いことです。

たとえば、16x16 パッチのトークン化の場合でも、1 つの 512x512 画像は 1024 個のトークンに相当します。フレーム数が多いビデオの場合、各レイヤーで数万個のトークンを処理する必要がある場合があります。

Transformer の計算量はトークンの数に応じて 2 乗的に増加するため、通常、Transformer が大きな画像や長いビデオを処理するのは困難になります。

ここで疑問が生じます。各レイヤーでこれほど多くのトークンを処理することが本当に必要なのでしょうか?

Google は、「TokenLearner: 学習した 8 つのトークンは画像や動画に何をもたらすのか?」で「適応」の概念について言及しました。この論文はNeurIPS 2021で発表される予定です。

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

プロジェクトアドレス:

https://github.com/google-research/scenic/tree/main/scenic/projects/token_learner

実験により、TokenLearner は、常に均一な画像分布によって形成されたトークンに依存するのではなく、より少数のトークンを適応的に生成できることが示され、これにより Vision Transformer の実行速度が向上し、パフォーマンスが向上します。

TokenLearner は、画像テンソル (つまり入力) を受け取り、小さなトークンのセットを生成する学習可能なモジュールです。このモジュールは Vision Transformer モデル内のさまざまな位置に配置できるため、後続のすべてのレイヤーで処理されるトークンの数が大幅に削減されます。

実験では、TokenLearner を使用すると分類パフォーマンスを低下させることなくメモリと計算を半分以上節約でき、入力に適応する能力により精度を向上させることもできることが示されています。

TokenLearnerとは何ですか?

TokenLearner は、実際には単純な空間注意法です。

各 TokenLearner が有用な情報を学習するには、まず重要な領域を強調表示する空間注意マップを計算する必要があります (畳み込み層または MLP を使用)。

次に、このような空間注意マップを使用して入力の各領域に重み付けし（不要な領域を破棄するため）、その結果を空間的にプールして最終的な学習トークンを生成します。

単一の画像に適用された TokenLearner モジュールの視覚化。

TokenLearner は、テンソルのピクセルのサブセットに対して空間処理を実行することを学習し、入力に適応したトークン ベクトルのセットを生成します。

この操作は並列に複数回繰り返され、元の入力から n 個 (約 10 個) のトークンが生成されます。

つまり、TokenLearner は、重み値に基づいてピクセル選択を実行し、その後グローバル平均化を実行していると見ることもできます。

注意マップを計算する関数は、さまざまな学習可能なパラメータによって制御され、エンドツーエンド方式でトレーニングされることは注目に値します。これにより、さまざまな入力における空間情報をキャプチャする際に注意機能を最適化できます。

実際には、モデルは複数の空間注意関数を学習し、それを入力に適用して、異なるトークン ベクトルを並行して生成します。

TokenLearner モジュールは、各出力トークンの空間注意マップを生成することを学習し、それを使用して入力トークンを抽象化します。

したがって、TokenLearner を使用すると、モデルは固定された均一なトークン化された入力を処理するのではなく、特定の認識タスクに関連する少数のトークンを処理できるようになります。

つまり、TokenLearner は適応型トークンを有効にして、入力に基づいてトークンを動的に選択できるようにすることで、トークンの総数を効果的に削減し、Transformer ネットワークの計算を大幅に削減します。

これらの動的かつ適応的に生成されたトークンは、画像分野の ViT やビデオ分野の ViViT (Video Vision Transformer) などの標準の Transformer アーキテクチャでも使用できます。

TokenLearnerはどこにありますか?

TokenLearner モジュールを構築した後、次のステップはそれをどこに配置するかを決定することです。

まず研究者たちは、224x224 の入力画像サイズを使用して、標準 ViT アーキテクチャ内のさまざまな位置に配置してみました。

TokenLearner によって生成されるトークンの数は 8 個と 16 個で、標準の ViT で使用される 196 個または 576 個のトークンよりもはるかに少なくなります。

下の図は、ViT B/16 内のさまざまな相対位置に TokenLearner を挿入したモデルの ImageNet 5 ショット分類精度と FLOP を示しています。ViT B/16 は、12 個のアテンション レイヤーを持つ基本モデルです。実行時には 16x16 パッチ トークンを使用します。

JFT 300M事前トレーニングによるImageNet 5ショット精度は、ViT B/16におけるTokenLearnerの相対的な位置に関係している。

位置 0 は、TokenLearner が Transformer レイヤーの前に配置されることを意味します。ベースラインは、標準の ViT B/16 ImageNet 5 ショット分類精度と FLOP です。

計算能力は数十億の浮動小数点演算（GFLOPS）で測定されます

ネットワークの最初の 4 分の 1 (1/4) に TokenLearner を挿入すると、計算量をベースラインの 3 分の 1 未満に削減しながら、ベースラインとほぼ同じ精度が達成されることがわかりました。

さらに、TokenLearner を後の層 (ネットワークの 3/4 以降) に配置すると、TokenLearner を使用しない場合よりもパフォーマンスが向上し、適応性により高速化も実現します。

TokenLearner の前後のトークンの数は非常に異なるため (たとえば、前は 196 個、後では 8 個)、TokenLearner モジュール後の相対的な計算量はほとんど無視できます。

TokenLearner VS ViT

TokenLearner を使用した ViT モデルを、ImageNet の few-shot で同じ設定を使用して通常の ViT モデルと比較します。

TokenLearner は、ネットワークの 1/2 や 3/4 など、各 ViT モデルの中央の異なる位置に配置されます。このうち、モデルはJFT 300Mによって事前トレーニングされています。

グラフから、TokenLearner モデルは精度と計算の複雑さの点で ViT よりも優れていることがわかります。

ImageNet 分類における ViT モデルのさまざまなバージョンのパフォーマンス

TokenLearner を、24 個のアテンション レイヤーと、初期トークンの L/10 および L/8 として 10x10 (または 8x8) パッチを持つような、より大きな ViT モデルに挿入します。

その後、これら 2 つのモデルを 48 層の ViT G/14 モデルと比較します。

ご覧のとおり、TokenLearner は G/14 モデルに匹敵するパフォーマンスを発揮しながら、非常に少ないパラメーターと計算しか必要としません。

左: ImageNetデータセットにおける大規模TokenLearnerモデルとViT G/14の分類精度の比較。右: パラメータ数とFLOPSの比較

高性能ビデオモデル

ビデオ理解はコンピューター ビジョンにおける重要な課題の 1 つであり、TokenLearner は複数のビデオ分類データセット ベンチマークで SOTA パフォーマンスを達成しました。

その中で、Kinetics-400 と Kinetics-600 でのパフォーマンスは従来の Transformer モデルを上回り、Charades と AViD では従来の CNN モデルも上回ります。

Video Vision Transformer (ViViT) と組み合わせることで、TokenLearner は各期間に 8 個 (または 16 個) のトークンを学習します。

左: ビデオ分類タスク、右: さまざまなモデルの比較

時間が経つにつれて、キャラクターがシーン内を移動すると、TokenLearner はさまざまな空間位置の変化に気づき、それをトークン化します。

TokenLearnerの空間注意マップの視覚化

結論は

Vision Transformer はコンピューター ビジョンの分野では強力なモデルですが、大量のトークンと膨大な計算量が、ViT を大きな画像や長いビデオに適用する際のボトルネックとなっていました。

この論文では、著者らは、これほど大量のトークンを保持し、レイヤーセット全体にわたって完全に処理する必要はないことを示しています。

さらに、著者らは、入力画像に基づいてトークンを適応的に抽出するモジュールを学習することで、計算を節約しながらより良いパフォーマンスを実現できることを実証しています。

最後に、複数の公開データセットでの検証により、TokenLearner がビデオ表現学習タスクで非常に優れたパフォーマンスを発揮することも示されました。