単一のGPUではGPT-3をトレーニングすることはできませんが、これを使用するとハイパーパラメータを調整できます。

偉大な科学的成果は試行錯誤だけでは達成できません。たとえば、宇宙計画におけるすべての打ち上げは、空気力学、推進力、天体に関する何百年にもわたる基礎研究に基づいています。同様に、大規模な人工知能システムを構築する場合、基礎研究によって試行錯誤の回数が大幅に削減され、明らかなメリットが得られます。

ハイパーパラメータ (HP) の調整はディープラーニングの鍵となりますが、特に数十億のパラメータを持つニューラル ネットワークの場合はコストのかかるプロセスでもあります。 HP が適切に選択されない場合、モデルのパフォーマンスが低下し、トレーニングが不安定になります。これらの問題は、非常に大規模なディープラーニング モデルをトレーニングするときにさらに悪化します。

最近の研究[54]では、異なるニューラルネットワークのパラメータ化が異なる無限幅の制限につながることが示されました。彼らは、制限内で「最大限の」特徴学習を達成できる最大更新パラメータ化（µP）を提案しました。直感的に言えば、幅に関係なく、トレーニング中に各レイヤーが同じ順序で更新されることが保証されます。対照的に、標準的なパラメータ化（SP）は、初期化時に活性化がユニット順序付けされることを保証しますが、実際のトレーニング中[54]、各層の学習率の不均衡により、幅広いモデルで活性化が爆発します。

Microsoft と OpenAI の研究者は、大規模なニューラル ネットワーク (何度もトレーニングするには大きすぎるもの) を調整する方法に関する初の基礎研究を発表しました。彼らは、特定のパラメータ化がさまざまなモデル サイズに対して最適なハイパーパラメータを保持することを示すことによってこれを行います。 µP を使用して、HP を小規模モデルから大規模モデルに移行します。つまり、この研究では、大規模モデルでほぼ最適な HP を達成しました。

論文の著者の一人であるマイクロソフトのグレッグ・ヤン氏は、「GPT-3 を単一の GPU でトレーニングすることはできません。ましてやハイパーパラメータ (HP) を調整することはできません。しかし、新たな理論的進歩のおかげで、単一の GPU で HP を調整できるようになったのです」と述べています。

この論文のアイデアは非常にシンプルです。この論文では、狭いニューラル ネットワークと広いニューラル ネットワークが最適なハイパーパラメータのセットを共有する特別なパラメータ化 µP を紹介しています。幅→∞の場合でも同じです。

具体的には、この研究では、µP では、モデルのサイズが変化しても多くの最適な HP が安定したままであることが実証されています。これにより、HP チューニングの新しいパラダイムである µTransfer が生まれます。µTransfer では、ターゲット モデルが µP でパラメータ化され、HP は小さいモデルで間接的にチューニングされ、フルサイズ モデルをチューニングせずにゼロ ショットをフルサイズ モデルに転送します。この研究では、Transformer と ResNet で µTransfer を検証しました。たとえば、1) 1300 万パラメータのモデルから事前トレーニング済みの HP を移行することで、この研究では BERT-large (3 億 5000 万パラメータ) よりも優れたパフォーマンスを発揮し、総チューニング コストは事前トレーニング済みの BERT-large 1 つ分と同等でした。2) 4000 万パラメータから移行することで、この研究では公開されている 67 億 GPT-3 モデルよりも優れたパフォーマンスを発揮し、チューニング コストは事前トレーニングの総コストのわずか 7% でした。

• 論文アドレス: https://arxiv.org/pdf/2203.03466.pdf
• プロジェクトアドレス: https://github.com/microsoft/mup

この技術は、どのトレーニングハイパーパラメータを使用するかを推測する必要性を大幅に減らすことで、GPT-3 のような巨大なニューラルネットワークの研究を加速し、将来的にはさらに大規模な後継ネットワークの研究を加速する可能性があります。

拡張機能の初期化は簡単だが、拡張機能のトレーニングは難しい

大規模なニューラル ネットワークは、サイズが大きくなるにつれて動作がどのように変化するかを理解していないため、トレーニングが困難です。ディープラーニングの初期の研究では、研究者はヒューリスティックなアルゴリズムを使用していました。一般に、ヒューリスティックは、モデルが初期化されるときに活性化の広がりを一定に保とうとします。ただし、トレーニングが始まると、図 1 の左側に示すように、この一貫性はさまざまなモデル幅で崩れます。

ランダム初期化とは異なり、トレーニング中のモデルの動作を数学的に分析するのはより困難です。この研究は µP を使用して取り組まれ、図 1 の右側に示されています。これは、モデル幅が増加するにつれて、最初の数回のトレーニング ステップにわたってネットワーク活性化スケールの安定性が示されることを示しています。

図 1: PyTorch のデフォルトのパラメータ化 (左) では、1 つのトレーニング ステップ後に活性化拡張の幅が変化します。ただし、右側の µP では、トレーニング ステップの幅に関係なく、活性化の広がりは一貫して変化します。

実際、µP は、トレーニング全体を通じて活性化の広がりを一定に保つことに加えて、十分に広い異なるニューラル ネットワークがトレーニング中に同様に動作し、望ましい限界に収束することを保証します。この研究では、これを特徴学習限界と呼んでいます。

図に示すように、µP は幅全体にわたって最良の学習率を維持する唯一のパラメータ化であり、幅が 213 ～ 8192 のモデルで最高のパフォーマンスを実現します。また、特定の学習率では、幅の広いモデルの方がパフォーマンスが向上します (つまり、曲線は交差しません)。

図 2 の左側では、この研究では、異なる幅 (異なる色とパターンの曲線に対応) の多層パーセプトロン (MLP) を CIFAR10 で異なる学習率 (x 軸に沿って表示) でトレーニングし、トレーニング損失を y 軸に沿ってプロットしました。右側では、1) PyTorch デフォルトと µP (x 軸) 間の初期化の広がり、および 2) PyTorch デフォルトと µP (y 軸) 間の学習率の広がりを補間することによって、パラメーター化された 2D 平面が形成されます。この平面では、PyTorch はデフォルトで (0,0) で表され、µP はデフォルトで (1,1) で表されます。

µTransfer は、Tensor プログラムの理論的基礎に基づいて、Transformer や ResNet などの高度なアーキテクチャに自動的に適用されます。さらに、さまざまなハイパーパラメータを同時に移行することもできます。

Transformer を例にとると、図 3 は主要なハイパーパラメータが幅全体にわたって安定している様子を示しています。ハイパーパラメータには、学習率、学習率スケジュール、初期化、パラメータ乗数などが含まれる場合があり、各パラメータ テンソルに個別に固有のものになる場合もあります。この研究では、最大幅 4096 の Transformer でこれを検証しました。

図 3: µP でパラメータ化され、WikiText-2 でトレーニングされたさまざまな幅のトランスフォーマー。モデル幅が広くなっても、最適な学習率、クロスエントロピー温度、初期化サイズ、学習率スケジュールは安定したままです。ネットワークのハイパーパラメータを調べると、より広範なネットワークに最適なハイパーパラメータを予測するのに役立ちます。右下の図では、この研究では次の学習率スキームを試しました：(a)線形減衰、(b)StepLR @ [5k, 8k]、減衰係数は0.1、(c)StepLR @ [4k, 7k]、減衰係数は0.3、(d)コサインアニーリング、(e)定数、(f)逆平方根減衰。

モデルの深さの実験的拡張

現代のニューラル ネットワークは、幅という 1 つの次元だけではなく、複数の次元で拡張されます。この研究では、幅以外の次元に対する単純なヒューリスティックと組み合わせることで、µP を現実的なトレーニング シナリオに適用する方法についても調査しました。下の図 4 では、同じトランスフォーマー設定を使用して、幅以外の寸法の妥当な範囲にわたって最適な学習率が安定している様子を示しています。

図 4: µP でパラメータ化され、Wikitext-2 でトレーニングされたさまざまなサイズのトランスフォーマー。図 3 に示すように、最適な学習率は幅にわたって転送可能であるだけでなく、テスト範囲内の深さ、バッチ サイズ、シーケンス長などの他のスケーリング次元にわたって実験的に転送可能です。これは、幅にわたる理論的な転送を、他のスケーリング次元にわたる実験的に検証された転送と組み合わせて、小さなモデルのハイパーパラメータを間接的に調整し、大きなモデルに転送できる µTransfer を取得できることを意味します。

学習率に加えて、他のハイパーパラメータは次の図に示されています。

µTransfer のテスト

研究者たちは、個々のハイパーパラメータの転送を検証した後、それらをより現実的なシナリオに組み合わせようと試みました。下の図 5 は、µTransfer が小さなプロキシ モデルから調整されたハイパーパラメータを転送する場合と、大きなターゲット モデルを直接調整する場合の 2 つのケースを比較しています。どちらの場合も、調整はランダム検索によって行われます。

図 5: µTransfer により計算効率が約 1 桁向上します。

プロキシモデルは、たとえ小さい場合でも最適なハイパーパラメータを有意に予測できるため（図 3 および 4 を参照）、研究者は、数十億のパラメータを持つより大きなターゲットモデルをトレーニングするにつれて、パフォーマンスのギャップが拡大すると予想しています。

将来の方向性: µP + GPT-3

この作業以前は、モデルが大きくなるほどチューニングにかかる​​コストが高くなり、チューニング効果も悪くなることが予想されていました。研究者たちは、µTransfer が最大規模のモデルに最大の利益をもたらすと期待しており、OpenAI と協力して GPT-3 で µTransfer を評価しました。

この研究では、µP の相対的注意を使用して GPT-3 のバージョンをパラメーター化した後、4000 万個のパラメーターを持つ小さなプロキシ モデルを調整し、µTransfer のアプローチに従って、最適なハイパーパラメーターの組み合わせを 67 億個のパラメーターを持つ GPT-3 のバリアントにコピーしました。このチューニング フェーズで使用される合計コンピューティングは、67 億のモデルの事前トレーニングに使用されるコンピューティングのわずか 7% です。下の図 6 に示すように、µTransfer を使用したこのモデルは、GPT-3 論文の同じサイズのモデル (Absolute Attention) よりも優れており、そのパフォーマンスは、パラメータ数が 2 倍の GPT-3 論文のモデル (Absolute Attention) に匹敵します。

理論的意義

µP は、トレーニング損失の観点から、さまざまな幅のモデル間で最適なハイパーパラメータの組み合わせを一意に保持する拡張ルールを提供します。対照的に、他のスケーリング ルール (PyTorch のデフォルトの初期化や NTK パラメーター化など) では、ネットワークが広くなるにつれて、ハイパーパラメーター空間内で最適値がどんどん遠ざかります。研究者たちは、実際に特徴学習ニューラル ネットワークを使用する際に適用可能な洞察を得るためには、NTK 制限よりも µP 特徴学習制限の方が自然であると考えています。したがって、過剰パラメータ化されたニューラル ネットワークは、広い幅の設定で µP の特徴学習の制限を再現するはずです。

この進歩は、過去数年にわたって開発されたテンソル プログラム (TP) の理論によって可能になりました。 TP 理論により、研究者は行列次元が大きくなるにつれて、一般的な計算グラフの限界を計算することができます。 TP アプローチは、ニューラル ネットワークとガウス過程の対応に関するアーキテクチャの普遍性や動的二分定理などの基本的な理論的結果をもたらし、µP と特徴学習の制約を導出することで µTransfer につながります。研究者たちは、TP 理論を深さやバッチ サイズなどの拡張次元に拡張することが、幅を超えて大規模なモデルを確実にスケーリングするための鍵であると考えています。

研究者らは、基礎研究は試行錯誤を費用対効果の高い方法で補完するものであり、今後も大規模な機械学習のためのより原理的な方法を導き出していくだろうと述べている。