単一の GPU で GPT-3 をトレーニングする方法にまだ苦労していますか? HP チューニングの新しいパラダイムをぜひご覧ください。

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

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

最近、ニューラル ネットワークのパラメータ化が異なると、無限幅の制限も異なることが示されました。彼らは、無限幅の制限によって「最大限の」特徴学習を実現する最大更新パラメータ化 (µP) を提案しました。直感的に言えば、幅に関係なく、トレーニング中に各レイヤーのアクティベーションの更新サイズが一定に保たれることを保証します。対照的に、標準的なパラメータ化 (SP) は初期化時にアクティベーションの一貫性を保証しますが、実際のトレーニング中は、各レイヤーの学習率の不均衡により、幅広いモデルでアクティベーションが爆発的に増加します。

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

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

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

具体的には、この研究では、µ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 の左側に示すように、この一貫性はさまざまなモデル幅で崩れます。

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

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

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

図に示すように、µP は幅全体にわたって最良の学習率を維持する唯一のパラメータ化であり、幅 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 の相対位置エンコーディング バージョンをパラメーター化した後、4,000 万個のパラメーターを持つ小さなプロキシ モデルを調整し、µTransfer に従って最適なハイパーパラメーターの組み合わせを GPT-3 の 67 億個のパラメーターのバリアントにコピーしました。このチューニング フェーズで使用される合計コンピューティングは、67 億のモデルの事前トレーニングに使用されるコンピューティングのわずか 7% です。以下に示すように、µTransfer を使用したこのモデルは、GPT-3 論文 (絶対位置エンコーディング) の同じサイズのモデルよりも優れており、そのパフォーマンスは、パラメーター数が 2 倍の GPT-3 論文のモデル (絶対位置エンコーディング) に匹敵します。

理論的意義

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

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

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