Primer: Searching for Efficient Transformers for Language Modeling

https://arxiv.org/abs/2109.08668
2021

近年の自然言語処理では、大規模なトランスモデルが中心的な役割を果たしています。しかし、これらのモデルの学習・推論コストは急速に増大し、非常に高価なものとなっています。ここでは、より効率的な変種を探索することで、Transformerのコストを削減することを目的としています。以前のアプローチと比較して、我々の検索は、Transformer TensorFlowプログラムを定義するプリミティブに対して、より低いレベルで実行されます。その結果、自動回帰型言語モデリングのためのオリジナルのTransformerやその他の変形よりも学習コストが小さいPrimerと名付けられたアーキテクチャを特定した。 Primerの改善は、ReLU活性化の二乗化と、自己注意におけるQ、K、Vの各投影の後に深さ方向の畳み込み層を追加するという、2つの単純な修正に起因するものである。

実験によると、PrimerがTransformerよりも優れている点は、計算規模が大きくなるにつれて増加し、最適なモデルサイズでは品質に関してべき乗則に従うことがわかった。また、Primerを様々なコードベースに落とし込むことで、追加のチューニングなしにトレーニングを大幅に高速化できることを経験的に検証した。例えば、500Mのパラメータサイズでは、Primerは、C4自動回帰言語モデリングにおけるオリジナルのT5アーキテクチャを改善し、学習コストを4倍に削減しました。さらに、学習コストが削減されたことで、目標とするワンショット性能を達成するために必要な計算量が大幅に減少しました。例えば、GPT-3 XLと同様の1.9Bのパラメータ設定では、PrimerはTransformerと同じワンショット性能を得るために必要な学習計算量が1/3になります。再現性を高めるために、モデルといくつかの比較結果をT5でオープンソース化しています。

1 はじめに

Transformers（トランスフォーマー）[1]は，ここ数年，多くのNLP分野で広く使用されています（例：[2, 3, 4, 5, 6, 7]）．しかし，より大きなモデルを学習するためのコストは，非常に高価なものとなっている．本論文では、Transformer言語モデルの学習コストを削減することを目的としています。この目的のために、TensorFlowの計算グラフを修正することで、Transformerのより効率的な代替品を検索することを提案します[10]。

TensorFlowプログラムの検索空間が与えられた場合、evolution[11, 12, 13, 14, 15, 16, 17]を使用して、一定の学習計算量で可能な限り低い検証損失を達成するモデルを検索します。探索空間としてTensorFlowプログラムを使用する利点は，トランスフォーマーを最適化するための単純な低レベルの改良を見つけやすいことである．ここでは、デコーダのみの自動回帰型言語モデリング（LM）に焦点を当てていますが、これはその一般性と成功の理由によるものです[18, 7, 19, 20, 21]。

実験の結果，Primerは，（1）より小さな学習コストで目標品質を達成できる，（2）固定の学習コストでより高い品質を達成できる，（3）より小さな推論コストで目標品質を達成できる，

という利点があることがわかった．これらの利点は，モデルの大きさ（2,000万から1.9Bのパラメータ）や，計算機の種類を問わず，ロバストに維持される． 9Bパラメータ）、計算規模（10～105アクセラレータ時間）、データセット（LM1B、C4、PG19 [22]）、ハードウェアプラットフォーム（TPUv2、TPUv3、TPUv4、V100）、デフォルト設定を用いた複数のTransformerコードベース（Tensor2Tensor、Lingvo、T5）、複数のモデルファミリー（密なTransformer [1]、疎なMixture-of-experts Switch Transformers [8]、Synthesizer [23]）に渡って、これらの利点は頑強に維持される。結果の再現性を高めるため，これらの比較をオープンソースで公開しています．

この節約効果は、最適なサイズのモデルを使用した場合、品質に関してべき乗則に従います。確立された学習環境におけるPrimerの節約効果を実証するために、500MパラメータのPrimerと、Raffelら[5]が自動回帰型言語モデリングに適用した構成を用いたオリジナルのT5アーキテクチャを比較しました。この設定では、Primerは、同じ学習コストで0.9パープレキシティの改善を達成し、4.2倍少ない計算量でT5ベースラインモデルと同等の品質を達成しました。さらに、GPT-3 XL [7]と同様の設定で、PrimerとTransformerを1.9Bのパラメータで比較し、Primerの節約効果がワンショット評価にもつながることを実証しました。 Primerは、3倍少ない学習計算量で、学習前のperplexityと下流のワンショットタスクの両方においてTransformerと同等の性能を達成しました。分析の結果、PrimerがTransformerよりも優れているのは、主に2つの修正に起因していることがわかりました。 ReLU活性の2乗化と、自己注意におけるQ、K、Vの各投影の後に深さ方向の畳み込み層を追加することです。この2つの修正は単純であり、既存のTransformerのコードベースに落とし込むことで、自動回帰型言語モデリングにおいて大きな利益を得ることができます。

2 探索空間と探索方法

TensorFlowプログラム上での探索。

Transformerの代替品の検索空間を構築するために、TensorFlow(TF)の操作を利用します。この検索空間では、各プログラムは、自動回帰型言語モデルのスタッカブルデコーダブロックを定義している。長さn、埋め込み長dの配列を表す入力テンソルX∈R n×dが与えられると、プログラムは同じ形のテンソルを返します。これらのテンソルを積み重ねると、各配列の位置における次のトークンの予測埋め込みが出力されます。私たちのプログラムは、モデルのアーキテクチャを指定するだけで、他には何もありません。言い換えれば、入力と出力の埋め込み行列自体や、入力の前処理や重みの最適化は、我々のプログラムの範囲ではありません。

図1は、進化的探索空間においてプログラムがどのように構築されるかを示しています。各プログラムは、インデックス付きのサブプログラムの集合体である進化的探索DNAから構築されます。サブプログラム0は、実行エントリポイントであるMAIN()関数であり、その他のサブプログラムは、DNAのサブプログラムバンクの一部である。各サブプログラムは、長さに制約のない命令のインデックス付き配列です。命令は、入力引数のセットを持つ操作です。演算は、命令が実行する機能を表します。各操作は、プリミティブ・ボキャブラリーのTF関数か、DNAサブプログラム・バンクの他のサブプログラムにマッピングされます。プリミティブ・ボキャブラリーは、ADD、LOG、MATMULなどのシンプルなプリミティブTF関数で構成されています（詳細は付録A.1を参照）。自己言及のような高レベルの構成要素は、探索空間における操作ではなく、私たちの低レベルの操作から構築することができることを強調しておきます。 DNAのサブプログラムバンクは、命令による関数として実行可能な追加プログラムで構成されています。各サブプログラムは、サブプログラムバンクの上位のインデックスを持つサブプログラムのみを呼び出すことができ、サイクルの可能性を排除しています。各命令の引数セットには、各命令の操作に対する潜在的な引数値のリストが含まれています。引数フィールドのセットは、すべての操作プリミティブが使用するフィールドの組合わせを表しています。

入力1：最初のテンソル入力として使用される隠れた状態のインデックスです。各隠れた状態のインデックスは、それを生成した命令のインデックスであり、サブプログラムの入力状態はインデックス0と1である。これを使用する演算の例はSINである。
入力2：2番目のテンソル入力のインデックス。これは、テンソル入力に対して二値化する演算でのみ使用される。これを使用する演算の例としては、ADDがあります。
Constant：実数値の定数です。これを使用する演算の例としてはMAXがあり、C = 0の場合のtf.math.maximum(x, C)は、TransformerのReLU活性化を表現する方法です。
寸法サイズ。重み行列を利用する変換の出力次元サイズを表す整数です。これを使用する演算の例は、CONV 1X1であり、Transformerの注意プロジェクションとフィードフォワード部分で使用される密なプロジェクションです。相対次元[13]を用いてモデルのサイズを変更する方法については、付録A.2を参照してください。

我々の検索サブプログラムは、各サブプログラムの命令を対応するTFコードの行に変換することで、TFプログラムに変換されますが、その際、インデックス順に1つずつ変換します。 TFコードを生成するために、命令の操作は、対応するTFプリミティブ関数またはDNAサブプログラムにマッピングされ、関連する引数が挿入されます（引数のマッピングを含むTFプリミティブの完全な語彙については、付録A.1を参照してください）；他の引数は無視されます。最終命令で生成されたTFテンソルがサブプログラムの出力となる。我々はTF Eagerを使用しないので、構築されたプログラムの有用な特性は、プログラムの出力に寄与しない無関係なノードが、TFのオリジナルの遅延実行設計[10]に従って無視されることです。サブプログラムがどのようにTFグラフに変換されるかについては、図2を参照してください。また、因果関係のマスキングをどのように処理するかなど、TFグラフがどのように構築されるかについての詳細は、付録A.2を参照してください。

進化的探索

進化的探索の目的は，探索空間において最も学習効率の良いアーキテクチャを見つけることである．これを実現するために，各モデルに固定の学習予算（TPUv2の24時間）を与え，その適合度をTensor2Tensor[25]の10億語ベンチマーク（LM1B）[24]におけるパープレクシティと定義する．固定されたトレーニング予算による暗黙の効率化目的と呼ぶこのアプローチは、効率化を最適化する際にトレーニングや推論のステップ時間を短縮することを明示的に目的とするこれまでのアーキテクチャ探索作品とは対照的である[26, 27, 28, 29]。我々の目的は、ステップタイムとサンプル効率の間のトレードオフが暗黙的であるという点で異なる。例えば，ステップタイムが2倍になっても，サンプル効率が3倍になるような修正は，最終的にアーキテクチャの計算効率を高めることになるので，我々の検索では良い修正となります．実際、我々が最も有益であると判断した修正、すなわち、ReLUの二乗化とアテンションへの深さ方向の畳み込みの追加は、トレーニングのステップタイムを増加させます。しかし、これらの修正はモデルのサンプル効率を向上させるため、目標品質に到達するために必要なトレーニングステップの数を大幅に減らし、総計算量を減少させます。探索アルゴリズムには，ハードル[13]を用いたRegularized Evolution[30]を使用しています． 50パーセンタイルのパスバーを用いてハードルを設定し，各ハードルバンドに同等の計算量が投入されるようにスペースを確保しました．これにより，フルモデル評価を行った同じ実験と比較して，探索コストを6.25倍に削減しました（詳細は付録A.3を参照）．

さらに、7時間のトレーニングを丸1日のトレーニングの代用としています。これは、バニラ・トランスフォーマーが7時間のトレーニングで24時間のトレーニング煩悩の90％以内に収まっているためです。これにより、検索コストはさらに3.43倍削減され、合計で21.43倍の計算量削減効果が得られます。つまり、我々の目標は24時間のパフォーマンスを向上させることですが、個人を評価するのに平均して約1.1時間しかかかりません（変異の詳細やハイパーパラメータなど、探索の詳細については付録A.4を参照してください）。

約25K個の個体に対して検索を行い、上位100個の個体を検索タスクで再学習させ、最適な個体を選択します。我々の検索空間は、これまでの検索空間（[31]によるアーキテクチャ検索の調査を参照）とは異なり、ランダム検索がうまくいくように大きく偏っていることが多い（[32, 33, 34]による分析を参照）。我々の検索空間にはこのような偏りがないため、Transformerプログラムと同じ長さを持つ我々の空間の78%のランダムプログラムは、数値的な不安定さのために5分以上の学習ができません。このように探索空間に偏りがあり、退化したプログラムが多いため、Transformer [13]（入力埋め込みサイズ dmodel = 512、フィードフォワード上方投影サイズ df f = 2048、層数 L = 6）のコピーで探索集団を初期化する必要があります（図3）。この初期化を探索空間に適用するためには、Transformerのプログラムをどのようにサブプログラムに分割するかを決定しなければなりません。そのためには、Transformerプログラムを構成する機械学習の概念に沿って分割します。例えば、自己注目、ReLU、レイヤーノルムについては、一般的に使用されている実装を用いて、それぞれ1つのサブプログラムを作成します（完全なリストは付録A.5を参照）。この方法は、初期化によって探索にバイアスを導入する一方で、進化のための探索空間と変異のための行動空間をオープンエンドのままにしておくことから、概念的初期化と呼んでいる。これは、探索空間を通してバイアスを導入する多くの先行研究とは対照的です。いくつかの作品は、ミニチュアのタスクで私たちのようなオープンエンドの検索空間を模索していますが[35]、私たちの技術はフルサイズの深層学習レジームにまで拡張できることを実証しています（セクション4参照）。

3 Primer

Primer

発見されたモデルを、PRIMitives searched transformERの頭文字をとってPrimerと名付けました（完全なプログラムは付録図23を参照）。 Primerは、検索タスクで再学習した際に、Transformerの半分以下の計算量で、同じ品質に到達するという大きな改善を示しました（図6）。セクション4では、Primerが他のコードベース、学習レジーム、データセット、下流のワンショットタスクに移行しても同様に大きな成果を上げることを示しています。

Primer-EZ

この研究の主な動機は、言語モデリングの実務者が簡単に採用できるシンプルな技術を開発することです。そのために、2つのコードベース（T5[5]とTensor2Tensor[25]）でアブレーションテストを行い、どのPrimerの修正が一般的に有用であるかを判断した（付録図26）。最もロバストな改善をもたらしたのは、フィードフォワードReLUの2乗と、注目のマルチヘッドプロジェクションへの深さ方向のコンボリューションの追加である（図4）。私たちは、この2つの簡単な修正を加えたTransformerをPrimer-EZと呼んでおり、Primerの使用を検討している言語モデリングの実務者には、このPrimer-EZを出発点として使用することを推奨しています。ここでは、これらの修正について説明し、その効果を実証的に測定します。

2乗ReLU。

最も効果的な変更点は、Transformerのフィードフォワードブロックにおいて、ReLU活性化から2乗ReLU活性化への改善です。様々な次数の整流多項式は、ニューラルネットワークの活性化関数の文脈で研究されていますが[36]、一般的には使用されていません。私たちの知る限り、このような整流多項式の活性化がTransformerで有用であることが実証されたのは今回が初めてです。興味深いことに、高次多項式[37]の有効性は、ReGLU[39]（y = Uxmax(V x, 0)、ここでは要素ごとの積）のようなGLU[38]の変種や、近似GELU[40]（y = 0.5x(1 + tanh(p 2/π(x + 0.04715x 3 )))))）のような点ごとの活性化など、他の有効なTransformerの非線形性においても観察されます。しかし、二乗ReLUは、x -→ ∞として、最も一般的に使用されている活性化関数と比較して、劇的に異なる漸近特性を持っています。 ReLU、GELU、Swishと比べて、x -→∞としての漸近特性が大きく異なります（図5左側）。二乗ReLUはReGLUと重なり合う部分があり、実際、ReGLUのUとVの重み行列が同じで、二乗ReLUの直前に重み行列Uによる線形変換が行われている場合には、同等の効果が得られます。

Multi-DConv-Head Attention（MDHA）。

もう1つの効果的な修正は、自己注目におけるクエリQ、キーK、値Vのマルチヘッドプロジェクションのそれぞれの後に、3x1の深さ方向の畳み込みを加えることです。これらの深さ方向の畳み込みは、各密なプロジェクションの出力の空間的な次元に対して実行されます。興味深いことに、ポイントワイズ・コンボリューションとデプスワイズ・コンボリューションの順序は、典型的なセパレート・コンボリューションの逆であり、付録A.6では効果が低いことがわかりました。また、深さ方向の畳み込みと標準的な畳み込みの幅を広げると、性能が向上しないだけでなく、いくつかのケースでは性能が低下することがわかりました。深さ方向の畳み込みは以前からTransformerに使用されているが[41, 42]、各dense head projectionの後に使用することは我々の知る限りでは行われていない。

MDHAはConvolutional Attention [43]に似ていますが、Convolutional Attentionは深さ方向の畳み込みの代わりに分離可能な畳み込みを使用し、我々のようにアテンションヘッドごとに畳み込み演算を適用しません。

その他の修正

他のPrimerの改造はあまり効果がありません。各修正のグラフは付録A.5に、アブレーションの研究は付録A.7に掲載されています。ここでは、それぞれの改良点とその有用性について簡単に説明します。

QとKの深さ方向の表現の共有。プライマーは、QとKのいくつかの重み行列を共有しています。Kは、以前に説明したMDHAプロジェクションを使用して作成され、学習可能な重み行列W∈R d×dに対してQ＝KWとなります。これは一般的にパフォーマンスを低下させることがわかっています。
事前と事後の正規化。Transformersの標準的な手法は、自己注目型変換とフィードフォワード型変換の両方の前に正規化を行うことになっている[44, 45]。Primerは自己注意の前に正規化を行うが、フィードフォワード変換の後に2回目の正規化を行う。これは、すべてのケースではなく、いくつかのケースで役立つことがわかります。
カスタム正規化。Primerは，(x - µ) 2の代わりにx(x - µ)を使用するレイヤー正規化[46]の修正版を使用するが，これは必ずしも効果的ではないことがわかった．
12X Bottleneck Projection．発見されたモデルでは，ベースラインの512に比べて小さい384のdmodelサイズと，ベースラインの2048に比べて大きい4608のdffサイズを使用しています．この大きな投影は、小さいサイズ（約35Mのパラメータ）では劇的に結果を向上させますが、大きいモデルでは効果が少ないことがわかりました[9]。このような理由から、PrimerやPrimer-EZを参照する際には、この修正を含めていません。
Post-Softmax Spatial Gating: 発見されたモデルは、アテンションソフトマックスの後に、チャンネルごとに学習可能なスカラーのセットを持ち、固定長のシーケンスに対するパープレキシティを改善します。しかし、これらのスカラーは可変長の配列には適用できないため、今回の実験ではPrimerにこの修正を加えていません。
その他の修正発見されたアーキテクチャに意味のある違いをもたらさない追加の修正がいくつかあります。例えば，hidden stateに-1.12をかけるなどです．これらの修正が品質に影響を与えないことを確認し、本文では議論の対象外とし、Primerを使った実験ではこれらを含めません。これらの余計な修正は、付録A.5に記載されています。

e4exp / paper_manager_abstract