Google Brain Quoc、操作プリミティブから効率的なTransformerバリアントを検索するためのPrimerをリリース

[[426884]]

モデルのパフォーマンスを向上させるには、パラメータを調整し、活性化関数を変更するのが一般的です。最近、Google Brain の責任者である Quoc 氏は、効率的な Transformer バリアントを自動的に検索し、いくつかの効果的なモデル プライマーを見つけることができる検索フレームワークをリリースしました。その中で、ReLU と square を組み合わせると、パフォーマンスが最も向上します。

現在、自然言語処理の開発による利益は、大規模な Transformer ベースの言語モデルから得られていますが、これらの言語モデルのトレーニングと推論のコストは非常に高いため、民間の錬金術師の意欲を削いでいます。

モデルパラメータの数が一定数に達すると、研究者はモデルを縮小しながらパフォーマンスを維持する方法についても検討しています。

Google Brain チームは最近、より効率的な Transformer バリアントを見つけることでトレーニングと推論のコストを削減することを目指した論文を arxiv にアップロードしました。

提案された方法は、従来の方法と比較して、Tensorflow プログラムのプリミティブ上でトランスフォーマーを定義および検索し、より低いレベルで検索を実行します。そして、自己回帰言語モデリング用のオリジナルのTransformerやその他のモデルバリアントよりもトレーニングコストが低い、Primerと呼ばれるモデルアーキテクチャを提案しました。

https://arxiv.org/abs/2109.08668

この論文の著者は偉大な Quoc V. Le 氏で、スタンフォード大学での博士課程の指導教官は Andrew Ng 教授でした。彼は現在 Google の研究科学者であり、Google Brain の創設メンバーの 1 人です。また、seq2seq の著者の 1 人であり、ニューラル アーキテクチャなどの手法を提案した Google AutoML の創設者であり、EfficientNet などの著者でもあります。

研究者らは、TensorFlow (TF) の操作を使用して、Transformer バリアントの検索空間を構築しました。この検索空間では、各プログラムは自己回帰言語モデルのスタック可能なデコーダー ブロックを定義します。長さ n の入力テンソルと長さ d の埋め込みシーケンスが与えられると、プログラムは同じ形状のテンソルを返すことができます。

積み重ねられると、その出力は各シーケンス位置での次のトークンの予測された埋め込みを表し、プログラムはモデル アーキテクチャのみを指定し、それ以外は何も指定しません。つまり、入力および出力の埋め込み行列自体、および入力の前処理と重みの最適化は、このプログラムの範囲外です。

デコーダー モデル プログラム (DNA) は、自己回帰言語モデルを定義します。各 DNA にはサブルーチンのセットがあり、サブルーチン 0 は MAIN 関数のエントリ ポイントです。各サブルーチンは命令で構成され、TensorFlow コードの行に変換されます。

命令操作は、基本的な TensorFlow ライブラリ関数またはプリミティブ語彙内の親 DNA サブルーチンの 1 つにマッピングされます。プリミティブ語彙は、ADD、LOG、MATMUL などの単純なプリミティブ TF 関数で構成されていますが、自己注意などの高レベルのビルディング ブロックは検索空間内の操作ではありません。自己注意は、低レベルの操作から構築できます。

DNA のサブルーチン ライブラリは、命令を通じて関数として実行できる追加プログラムで構成されています。各サブルーチンは、サブルーチン ライブラリ内のより高いインデックスを持つサブルーチンのみを呼び出すことができるため、ループが発生する可能性が排除されます。

操作のパラメータを親命令のパラメータ セットで入力します。このパラメータ セットには、入力 1 (最初のテンソル入力として使用される隠し状態のインデックス)、入力 2 (2 番目のテンソル入力の隠し状態のインデックス)、定数 (MAX などの関数で使用できる実数値の定数)、およびディメンション サイズ (出力次元のサイズを表すために使用される整数) など、すべての潜在的な操作パラメータの値が含まれます。特定の操作で使用されないパラメータは単に省略されます。

研究者らはまた、探索空間で最も効果的なモデルアーキテクチャを見つけることを目的とした進化的探索も提案した。主なアプローチは、固定のトレーニング バジェット (TPUv2 を使用して 24 時間に制限) を設計し、その適合性メトリックを Tensor2Tensor の 10 億語ベンチマーク (LM1B) のパープレキシティとして定義することです。

これらのアーキテクチャ検索の取り組みの明確な目標は、効率を最適化しながらトレーニングまたは推論のステップ時間を短縮することです。検索プロセス中に、ステップ時間を 2 倍にしてサンプリング効率を 3 倍にすることが、最終的にモデル アーキテクチャの計算効率を高めるため、適切な変更であることがわかりました。 ReLU を二乗し、深さ方向の畳み込みを注意に追加して、トレーニング ステップ時間を増やすこともできます。

これらの操作により、モデルのサンプル効率が大幅に向上し、目標品質に到達するために必要なトレーニング ステップの数が大幅に削減されるため、目標パフォーマンスに到達するために必要な計算全体が削減されます。

この検索プログラムを通じて発見されたモデルは、研究者によってPrimerと名付けられました。これは、Primitives searched transformERの略です。

Primer の改良点は主に、2乗 ReLU アクティベーションと、自己注意の各 Q、K、V 投影の後に深い畳み込み層を追加することです。

最も効果的な変更は、Transformer フィードフォワード ブロックの ReLU アクティベーションを 2 乗 ReLU アクティベーション関数に改善することです。これは、この修正多項式アクティベーションが Transformer で有用であることが示された初めてのケースです。また、高次多項式の有効性は、GLU、ReGLU、近似 GELU などのさまざまなバリアントなど、他の Transformer 非線形活性化関数でも確認できます。ただし、2乗 ReLU は、最も一般的に使用される活性化関数である ReLU、GELU、および Swish とは非常に異なる漸近特性を持ちます。

2乗 ReLU は ReGLU とかなり重複しており、実際には ReGLU の U および V 重み行列が同じ場合は ReLU と同等になります。また、2乗 ReLU は GLU バリアントの利点を享受しながら、よりシンプルで、追加のパラメータを必要とせず、より優れた品質を提供します。

研究者は、Primer と比較するために 3 つの Transformer バリアントを使用しました。

1. Vanilla Transformer: ReLU アクティベーションとレイヤー正規化を使用するオリジナルの Transformer。

2. Transformer+GELU: Transformerの一般的な変種で、GELUを使用して活性化関数を近似します。

3. Transformer++: フィードフォワード後方ボトルネック (SwiGLU) で RMS 正規化、Swish 活性化、および GLU 乗算ブランチを使用します。これらの変更は T5 でベンチマークされ、効果があることが示されました。

実験では、Primer の Transformer に対するゲインは計算規模とともに増加し、最適なモデル サイズでの品質に関してべき乗法則に従うことが示されています。

研究者らはまた、Primer をさまざまなコード ベースに導入できることを実験的に検証し、追加の調整なしでトレーニングを大幅に高速化できることを発見しました。たとえば、パラメータ サイズが 500M の場合、Primer は C4 自己回帰言語モデリングの元の T5 アーキテクチャを改善し、トレーニング コストを 4 分の 1 に削減します。

さらに、トレーニング コストが削減されるということは、Primer が目標のワンショット パフォーマンスを達成するために必要な計算量が少なくなることを意味します。たとえば、GPT-3 XL と同様の 19 億のパラメータ構成では、Primer は 1/3 のトレーニング コンピューティングを使用して、Transformer と同じワンショット パフォーマンスを実現します。

研究者らは論文の再現性を高めるためにモデルをオープンソース化した。