AdaTape 適応計算とダイナミックな読み書きを持つ基礎モデル

AdaTape 基礎モデルの適応計算とダイナミックな読み書き

Googleの研究インターンであるFuzhao Xueと研究科学者であるMostafa Dehghaniによって投稿されました。

適応的計算とは、機械学習システムが環境の変化に応じてその動作を調整する能力を指します。従来のニューラルネットワークは固定の機能と計算能力を持っており、すべての入力に対して同じ数のFLOPを使用して処理します。一方、適応的かつ動的な計算を持つモデルは、入力の複雑さに応じて各入力の処理に割り当てる計算予算を変調します。

ニューラルネットワークにおける適応的計算は、2つの主要な理由で魅力的です。まず、適応性を導入するメカニズムは、いくつかの難解なタスクを解決する上で重要な役割を果たす帰納的バイアスを提供します。たとえば、異なる入力に対して異なる計算ステップ数を許可することは、異なる深さの階層をモデリングする必要がある算術問題の解決に重要です。次に、動的な計算によって提供される柔軟性によって、推論のコストを調整する能力が実践者に与えられます。これらのモデルは新しい入力の処理により多くのFLOPを使用するように調整できるためです。

ニューラルネットワークは、さまざまな入力に対して異なる関数または計算予算を使用することで適応的になります。ディープニューラルネットワークは、入力とパラメータに基づいて結果を出力する関数として考えることができます。適応的関数タイプを実装するために、一部のパラメータは入力に基づいて選択的に活性化されます。このプロセスは条件付き計算と呼ばれます。関数タイプに基づく適応性は、ミクスチャーオブエキスパートの研究で探求され、各入力サンプルの希薄に活性化されたパラメータはルーティングによって決定されます。

• 「ダウンフォール」の欠陥が世代を超えたインテルチップで貴重なデータを公開します
• パンダのコピー・オン・ライトモードの詳細な調査：パートI
• PageRankによる大規模グラフの分析

適応的計算の研究のもう一つの領域は、動的な計算予算です。T5、GPT-3、PaLM、ViTなどの標準的なニューラルネットワークとは異なり、最近の研究では、変動的な計算予算がトランスフォーマーの課題でのパフォーマンス向上につながることが示されています。これらの研究の多くは、トランスフォーマーレイヤーの数に基づいて計算予算を割り当てることによって適応性を実現しています。たとえば、Adaptive Computation Time（ACT）アルゴリズムは、リカレントニューラルネットワークに対して適応的な計算予算を提供するために提案されました。ユニバーサルトランスフォーマーは、ACTアルゴリズムをトランスフォーマーに拡張し、各入力例またはトークンに使用されるトランスフォーマーレイヤーの数に応じて計算予算を割り当てます。PonderNetなどの最近の研究は、同様のアプローチを採用しながら、動的な停止メカニズムを改善しています。

「Elastic Input Sequenceを使用した適応的計算」という論文では、適応的計算を利用する新しいモデル、AdaTapeを紹介しています。このモデルは、適応的計算の前のアプローチと比較して、独自の視点で柔軟な入力シーケンスを作成するためのトランスフォーマーベースのアーキテクチャです。AdaTapeは、入力の複雑さに基づいて追加される可変サイズのテープトークンのシーケンスを動的に選択するために、適応的なテープ読み取りメカニズムを使用します。AdaTapeは実装が非常に簡単であり、必要に応じて精度を向上させるための効果的なツマミを提供しますが、モデルの深さではなく入力シーケンスに適応性を直接注入するため、他の適応基準と比較して効率的です。最後に、AdaTapeはイメージ分類などの標準的なタスクだけでなく、アルゴリズムのタスクでも優れたパフォーマンスを提供し、品質とコストのトレードオフを維持します。

適応的計算トランスフォーマーと弾性入力シーケンス

AdaTapeは、適応的な関数タイプと動的な計算予算の両方を使用します。具体的には、トークン化後の入力シーケンスのバッチ（たとえば、ビジョントランスフォーマーのイメージからの非重複パッチの線形射影）に対して、AdaTapeは各入力を表すベクトルを使用して可変サイズのテープトークンのシーケンスを動的に選択します。

AdaTapeは、「テープバンク」と呼ばれるトークンのバンクを使用して、モデルと適応的なテープ読み取りメカニズムを介して相互作用するすべての候補テープトークンを保存します。テープバンクを作成するための2つの異なる方法を探求しています：入力駆動型バンクと学習可能なバンク。

入力駆動型バンクの一般的なアイデアは、元のモデルのトークナイザーとは異なるアプローチを使用して、入力からトークンのバンクを抽出することです。これにより、異なる解像度の画像や異なる抽象度の情報など、異なる視点から入力からの情報への動的なオンデマンドアクセスが可能になります。

一部の場合、異なる抽象度のトークン化は不可能であり、入力駆動型テープバンクは実現不可能です。これに対処するために、AdaTapeは学習可能なベクトルセットを使用してテープバンクを生成するより一般的なアプローチを提供します。このアプローチは学習可能なバンクと呼ばれ、モデルは入力例の複雑さに基づいてトークンを動的に取得する埋め込み層と見なすことができます。学習可能なバンクにより、AdaTapeはより柔軟なテープバンクを生成し、各入力例の複雑さに基づいて計算予算を動的に調整する能力を提供します。つまり、より複雑な例はバンクからより多くのトークンを取得することができるため、モデルはバンクに格納された知識を利用するだけでなく、それを処理するためにより多くのFLOPを使用することができます。

最後に、選択されたテープトークンが元の入力に追加され、次のトランスフォーマーレイヤーに供給されます。各トランスフォーマーレイヤーでは、すべての入力トークンとテープトークンに対して同じマルチヘッドアテンションが使用されます。ただし、2つの異なるフィードフォワードネットワーク（FFN）が使用されます。1つは元の入力のすべてのトークン用であり、もう1つはすべてのテープトークン用です。入力トークンとテープトークン用に別々のフィードフォワードネットワークを使用することで、わずかに良い品質が得られることが観察されました。

AdaTapeの概要。異なるサンプルに対して、テープバンクから異なる数のトークンを選択します。テープバンクは、入力から追加の詳細情報を抽出することなどで駆動することができます。またはトレーニング可能なベクトルのセットであることもあります。適応的なテープ読み取りは、異なる入力に対して、可変長の異なるテープトークンのシーケンスを再帰的に選択するために使用されます。これらのトークンは、単純に入力に追加され、トランスフォーマーエンコーダに供給されます。

AdaTapeは有用な帰納バイアスを提供します

AdaTapeを標準のトランスフォーマーと比較するために、標準のトランスフォーマーには解けない非カウンターフリーまたは周期的な正規言語であるパリティタスクを用いて、AdaTapeを評価します。パリティタスクでは、1と0、-1のシーケンスが与えられた場合、モデルはシーケンス内の1の数の奇数または偶数を予測する必要があります。パリティは最も単純な非カウンターフリーまたは周期的な正規言語ですが、標準のトランスフォーマーではこのタスクは解けません。

パリティタスクの評価。標準のトランスフォーマーとユニバーサルトランスフォーマーはこのタスクを実行できず、ランダムな推測ベースラインと同等のパフォーマンスを示しました。

短くてシンプルなシーケンスで評価されたにもかかわらず、標準のトランスフォーマーやユニバーサルトランスフォーマーはパリティタスクを実行できず、モデル内でカウンタを維持することができません。しかし、AdaTapeは、入力選択メカニズム内に軽量の再帰を組み込んでいるため、カウンタの暗黙的な維持を可能にする帰納バイアスを提供し、すべてのベースラインを上回ります。これは、標準のトランスフォーマーでは不可能です。

画像分類の評価

また、画像分類タスクでAdaTapeも評価します。そのために、ImageNet-1KでAdaTapeをスクラッチからトレーニングしました。以下の図は、AdaTapeとA-ViT、Universal Transformer ViT（UViTとU2T）を含むベースライン手法の正確性と速度（各コードが処理する画像の数で測定）を示しています。品質とコストのトレードオフに関しては、AdaTapeは他の適応型トランスフォーマーベースラインよりも優れたパフォーマンスを発揮します。効率性に関しては、パラメータ数で見た場合、より大きなAdaTapeモデルの方が小さなベースラインよりも速いです。これらの結果は、TPUなどの多くのアクセラレータには適していないことを示す以前の研究の結果と一致しています。

私たちはImageNetをスクラッチからトレーニングすることでAdaTapeを評価しました。A-ViTについては、論文からの結果だけでなく、スクラッチからトレーニングしてA-ViT（Ours）を再実装しました。

AdaTapeの動作の研究

パリティタスクとImageNet-1Kのパフォーマンスに加えて、JFT-300Mの検証セットで入力駆動型バンクを使用してAdaTapeのトークン選択の振る舞いも評価しました。モデルの振る舞いをよりよく理解するために、入力駆動型バンクでのトークン選択結果をヒートマップとして可視化しました。明るい色ほどその位置がより頻繁に選択されていることを示します。ヒートマップは、AdaTapeがより頻繁に中央のパッチを選択することを明らかにしています。これは、中央のパッチが通常よりも情報量が多いことに合致しています。特に、主要なオブジェクトが画像の中央にある自然画像のデータセットの文脈では、この結果はAdaTapeの知能を示しており、より情報量の多いパッチを効果的に特定して優先することでパフォーマンスを向上させることができることを強調しています。

AdaTape-B/32（左）とAdaTape-B/16（右）のテープトークン選択ヒートマップを可視化しました。より熱い/明るい色は、この位置のパッチがより頻繁に選択されていることを意味します。

結論

AdaTapeは、適応的なテープ読み取りメカニズムによって生成される弾性的なシーケンス長が特徴です。これにより、AdaTapeは標準のトランスフォーマーと既存の適応型トランスフォーマーの両方にとって困難なタスクを解決する可能性を持つ新しい帰納的バイアスを導入します。画像認識のベンチマークで包括的な実験を行うことにより、計算が一定の場合においてAdaTapeが標準のトランスフォーマーや適応的アーキテクチャのトランスフォーマーよりも優れたパフォーマンスを発揮することを示しました。

謝辞

この投稿の著者の一人、Mostafa Dehghaniは現在Google DeepMindに在籍しています。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!