ニューラル言語モデルのスケーリング則

論文解説

問題意識と動機

Transformerアーキテクチャの登場により言語モデリングは急速に進展しましたが、大規模モデルの学習には莫大な計算量とデータが必要です。本研究以前は、性能（交差エントロピー損失で測定）がアーキテクチャの形状（深さ vs 幅）、パラメータ数、データセットサイズ、学習計算量にどのように依存するのかが明確ではありませんでした。定量的なスケーリング関係が存在しなかったため、研究者は固定された計算予算を最適に割り振る方法を理解できず、通常は比較的小さなモデルを収束近くまで学習させていました。

主要なアイデア

• アーキテクチャよりもスケールが支配的: 広い範囲内において、Transformerの性能は非埋め込みパラメータ数 $N$、データセットサイズ $D$、計算量 $C$ に強く依存し、深さ、幅、アテンションヘッド数などのアーキテクチャのハイパーパラメータは、$N$ を固定した場合ほとんど影響を与えません。
• 滑らかなべき乗則スケーリング: テストの交差エントロピー損失は、$N$、$D$、および計算量 $C_{\text{min}}$ に対して6〜8桁以上の範囲にわたりべき乗則に従ってスケールします。
• 過学習を避けるための劣線形なデータ増加: 過学習を防ぐには、データセットサイズをパラメータ数に対して劣線形にスケールさせる必要があります: $D\propto N^{0.74}$。
• 計算効率化とは早期停止する大規模モデルを意味する: 固定された計算予算 $C_{\text{min}}$ のもとで性能を最大化するには、非常に大きなモデルを学習させ、収束より大幅に手前で停止させるべきです。計算リソースは $N\propto C_{\text{min}}^{0.73}$、バッチサイズ $B\propto C_{\text{min}}^{0.24}$、ステップ数 $S_{\text{min}}\propto C_{\text{min}}^{0.03}$ として配分します。

手法（主要な数式を含む）

著者らはWebText2上でデコーダオンリーのTransformerモデルを評価しています（BPEトークナイザーを使用、コンテキスト長 $n_{\text{ctx}}=1024$）。モデルサイズ $N$ には埋め込みパラメータを意図的に除外しており、$N\approx 12n_{\text{layer}}{d}_{\text{model}}^2$ と計算されます。

フォワード・バックワードパス1回あたりの非埋め込み計算量は $C\approx 6NBS$ として近似されます。

1. 個別のべき乗則:

   • パラメータ制約下: $$L(N)={\left(\frac{N_c}{N}\right)}^{{\alpha }_N},{\alpha }_N\approx 0.076,N_c\approx 8.8\times 10^{13}$$
   • データ制約下: $$L(D)={\left(\frac{D_c}{D}\right)}^{{\alpha }_D},{\alpha }_D\approx 0.095,D_c\approx 5.4\times 10^{13}$$
   • 計算量制約下（バッチサイズ $B\ll B_{\text{crit}}$ において）: $$L(C_{\text{min}})={\left(\frac{C_c^{\text{min}}}{C_{\text{min}}}\right)}^{{\alpha }_C^{\text{min}}},{\alpha }_C^{\text{min}}\approx 0.050,C_c^{\text{min}}\approx 3.1\times 10^8\text{ PF-days}$$

2. パラメータとデータの同時スケーリング: $$L(N,D)={\left[{\left(\frac{N_c}{N}\right)}^{\frac{{\alpha }_N}{{\alpha }_D}}+\frac{D_c}{D}\right]}^{{\alpha }_D}$$

3. 学習ダイナミクスとクリティカルバッチサイズ: $$B_{\text{crit}}(L)=\frac{B_{\ast }}{L^{1/{\alpha }_B}},B_{\ast }\approx 2\times 10^8\text{ tokens},{\alpha }_B\approx 0.21$$ 学習曲線はステップ数 $S_{\text{min}}$（$B\gg B_{\text{crit}}$ におけるステップ数）にフィットします: $$L(N,S_{\text{min}})={\left(\frac{N_c}{N}\right)}^{{\alpha }_N}+{\left(\frac{S_c}{S_{\text{min}}}\right)}^{{\alpha }_S},{\alpha }_S\approx 0.76,S_c\approx 2.1\times 10^3$$

4. 最適配分: $${\alpha }_C^{\text{min}}=\frac{1}{1/{\alpha }_S+1/{\alpha }_B+1/{\alpha }_N}\approx 0.052$$ 効率的な学習では、収束時の損失より約 $10\%$ 高い損失を目標とします（$\frac{{\alpha }_N}{{\alpha }_S}\approx 0.10$）。

結果（最も重要な数値を含む）

• これらの傾向は、計算量で8桁、パラメータ数で6桁、データ量で2桁の範囲にわたり成り立ちます。
• $N$ を2倍にすると、損失は係数 $2^{-{\alpha }_N}\approx 0.95$ で低下します。
• 過学習回避の閾値: $D\gtrsim (5\times 10^3)N^{0.74}$。モデルサイズを8倍にする際、必要なデータ量は約5倍の増加にとどまります。
• 計算予算 $C_{\text{min}}$ を10倍に増やす場合、計算効率的な学習ではモデルサイズ $N$ を約5倍（$N\propto C_{\text{min}}^{0.73}$）、バッチサイズを約1.7倍（$B\propto C_{\text{min}}^{0.24}$）に増やし、学習ステップ数の増加はわずか約1.07倍（$S_{\text{min}}\propto C_{\text{min}}^{0.03}$）に抑えます。
• ほぼ収束するまで学習させるアプローチ（収束値からの乖離2%）と比較して、計算最適（compute-optimal）な学習では計算量を65%削減でき、パラメータ数は2.7倍、パラメータ更新ステップ数は7.7倍少なくなります。
• 分布外のコーパス（Books Corpus、Common Crawl、英語版Wikipedia）への転移損失は、モデルサイズに伴って滑らかに改善し、WebText2検証損失に対してほぼ一定のペナルティオフセットを伴うのみです。

限界と未解決の課題

• 正確な指数を説明する厳密な理論的導出（例えば統計物理学的な基礎づけなど）が欠けています。
• 計算量のスケーリング $C\approx 6NBS$ は、$n_{\text{ctx}}^2$ に比例するアテンション特有のコストや、$n_{\text{ctx}}\gtrsim 12d_{\text{model}}$ となる場合のコンテキスト依存の計算量を無視しています。
• 正則化、ドロップアウトの調整、極端なスケールにおける学習率スケジュールなどのハイパーパラメータは、ほぼ固定されているか、データセットサイズごとにチューニングされていません。
• 計算最適な傾向を外挿すると、$C^{\ast }\sim 10^4\text{ PF-days}$、$N^{\ast }\sim 10^{12}$ パラメータ、損失 $L^{\ast }\sim 1.7\text{ nats/token}$ において $L(C_{\text{min}})$ が $L(D)$ と交差するという矛盾が生じるため、実証的な傾向は最終的に破綻するか、自然言語のエントロピーの下限付近で頭打ちになるはずです。