大きくすると何が起きるか ── スケーリングと事前学習

教科書 › 第III部深層学習：何の前提を足すか

CHAPTER 8

大きくすると何が起きるか ── スケーリングと事前学習

★ 会員12分

第7章までで、1つのモデルを作って学習させる道具はそろいました。この章の問いは別の層にあります。使える計算が増えたとき、それを何に使うか。

選べるのは2つです。モデルを大きくする（パラメータ $N$ ）か、データを増やす（トークン $D$ ）か。どちらにも使えますが、両方に同じだけ使うことはできません。予算が決まっているからです。

この章では、まず予算の式を自分で数えて作ります。そのうえで配分の問題を解きます。

1. 計算量を自分で数える

Transformer 1ブロックの重みは、第7章で数えたとおりです（ $d_{ff} = 4d$ として）。

	重みの数
注意（ $W_Q,W_K,W_V,W_O$ ）	$4d^2$
FFN（2層）	$8d^2$
1ブロック合計	$12d^2$

ここから1トークンあたりの計算量が出ます。

前向き：重み1個につき掛け算1回と足し算1回で 2 FLOPs。全部で $2N$ 。
逆向き：入力への勾配と重みへの勾配の2本を作るので、前向きのおよそ2倍で $4N$ （第3章の式(3)(4)）。

合わせて1トークンあたり $6N$ 。 $D$ トークン流せば

C \;\approx\; 6ND

(1)

この式は本書の範囲の算数だけでできています。使ってみます。

$N$	$D$	$C = 6ND$
$10^{8}$	$10^{9}$	$6.0\times10^{17}$
$10^{9}$	$2\times10^{10}$	$1.2\times10^{20}$
$7\times10^{10}$	$1.4\times10^{12}$	$5.9\times10^{23}$

2. $6ND$ が崩れるところ

式(1)は重みを通る計算しか数えていません。第6章で見た $T \times T$ の注意の表は、重みではないので入っていません。どこで効き始めるかは割り算で出ます。

1層・1トークンあたり、重みの項は $2 \times 12d^2 = 24d^2$ 。注意の表の項（ $QK^{\top}$ と $PV$ ）は $2 \times (2Td) = 4Td$ 。比を取ると

\frac{\text{表の項}}{\text{重みの項}} \;=\; \frac{4Td}{24d^2} \;=\; \frac{T}{6d}

(2)

文脈長 $T$ が $6d$ を超えると、表の項のほうが重くなります。

$d$	$6d$	$T=1$ k	$T=4$ k	$T=16$ k	$T=64$ k
512	3,072	0.33	1.33	5.33	21.33
1,024	6,144	0.17	0.67	2.67	10.67
4,096	24,576	0.04	0.17	0.67	2.67

$d = 512$ で $T = 4$ k なら、表の項はもう重みの項の1.33倍です。 $C \approx 6ND$ は短い文脈での近似で、長文脈のモデルでは第6章の $T^2$ がそのまま効いてきます。

3. 予算をどう割るか

$C$ が決まれば、 $N$ を決めた時点で $D = C/(6N)$ も決まります。あとは損失が最小になる $N$ を選ぶだけ——ですが、損失が $N$ と $D$ にどう依存するかを知らないと解けません。

ここでは形を仮定します。

L(N, D) \;=\; \frac{A}{N^{a}} \;+\; \frac{B}{D^{b}} \;+\; E

(3)

第1項はモデルが小さいことによる損、第2項はデータが少ないことによる損、 $E$ はどちらを増やしても消えない床（第0部の第4章の $H(p)$ と同じ役）。

$A, B, E, a, b$ の値は、この本では測っていません。 以下の計算は $A=B=400$、$E=1.0$、$a=b=0.34$ と置いた場合の形を見るためのもので、実際のモデルで成り立つ値ではありません。実測値は論文にあり、この章では引用しません（第5節で触れる論文の主張だけを出典つきで引きます）。

FIG 1横軸が N（右へ行くほど大きいモデル・少ないデータ）、縦が仮定した形の損失。予算 C を動かすと最適（金の線）が動く。指数 a と b もつまみにしてある——「a と b が等しいときだけ最適が D/N = 1 になる」。「等しく伸ばす」という結論が、どの仮定に乗っているかが見える

この仮定で $C = 10^{21}$ のときを解くと、最適は $N = 1.29\times10^{10}$ 、 $D = 1.29\times10^{10}$ 、 $D/N = 1.00$ でした。予算を変えると、最適はこう動きます。