[논문 리뷰] Scaling Laws for Neural Language Models

About this offer

Summary Language model의 test loss는 model 크기 $N$, dataset 크기 $D$, 학습 compute $C$ 각각에 대해, 나머지 둘이 병목이 아닐 때 power law를 따른다. 반면 depth, width, head 수 같은 architecture shape는 non-embedding parameter 수 $N$이 같으면 loss에 거의 영향을 주지 않는다. $N$과 $D$를 함께 키울 때 overfitting 정도는 $N^{0.74}/D$라는 비율 하나로 결정된다. 즉 model을 8배 키우면 data는 약 5배만 늘리면 된다. 고정된 compute budget에서 최적의 전략은 model을 크게 키우고($N \propto C^{0.73}$), 수렴하기 훨씬 전에 학습을 멈추는 것이다. 1. Introduction Language modeling의 성능은 model architecture, parameter의 수, compute의 수, 주어진 data 양 등에 따라 달라진다. 해당 논문은 Transformer language modeling을 대상으로 위 요인들 중 실제 성능을 결정하는 주요 요인들이 무엇일지를 실험적으로 측정한다. 조건을 바꾸어가며 진행한 여러 실험 결과에 따르면 성능은 학습 시간, context 길이, dataset 크기, model 크기, compute budget에 대해 power law 형태로 scaling 되는 것을 확인할 수 있다. 여러 조건에 따른 결과를 비교하는만큼 해당 논문에서는 몇 가지 notation들을 정의하여 사용한다. $L$: Cross-entropy loss $N$: Number of model parameters $D$: Dataset size $C$: Total non-embedding training compute $C_{min}$: Estimate of minimum amount of non-embedding compute to reach a given loss 2. Background & Methods 2.1. Parameter and Compute Scaling of Transformers Transformer architecture에 대한 hyperparameter를 다음과 같이 설정하였다. $n_{layer}$: layer 수 $d_{model}$: residual stream의 차원 $d_{ff}$: feed-forward layer의 중간 차원 $d_{attn}$: attention output의 차원 $n_{head}$: layer당 attention head 수 1 Model 크기 $N$ layer 하나는 다음과 같은 parameter들로 이루어져 있다. Attention의 Q, K, V projection과 output projection: $4 \times d_{model} \times d_{attn}$ FFN의 두 linear layer: $2 \times d_{model} \times d_{ff}$ 따라서 embedding과 bias를 제외한 전체 parameter 수는 $$ N \approx 2, d_{model}, n_{layer} \left(2 d_{attn} + d_{ff}\right) = 12, n_{layer}, d_{model}^2 \qquad (d_{attn} = d_{model} = d_{ff}/4) $$ 2 Compute token 하나의 forward pass에 드는 연산량(FLOPs)은 $$ C_{forward} \approx 2N + 2, n_{layer}, n_{ctx}, d_{attn} $$ 이때, 두 번째 항은 attention score 계산에서 나오는 context 길이 의존 항이므로 $d_{model} \gg n_{ctx}/12$인 경우에 $N$값에 비해 훨씬 작아지기 때문에 이 항을 무시할 수 있다. 또한, backward pass의 compute는 forward의 약 2배이므로 token 당 전체 학습 compute는 다음과 같이 근사할 수 있다. $$ C \approx 6N $$ 3. Empirical Results and Basic Power Laws 3.1. Transformer Shape and Hyperparameter Independence Transformer의 shape가 성능에 끼치는 영향을 확인하기 위하여 Non-embedding parameter 수 $N$을 고정한 상태에서 $n_{layer}$, $n_{head}$, $d_{ff}$ 중 하나씩 바꿔 가며 loss를 측정하였다. 그 결과 loss는 다음과 같다. Feed-forward 비율, aspect ratio($d_{model}/n_{layer}$), head 차원을 넓은 범위에서 바꿔도 loss 변화는 수 % 이내로 aspect ratio가 40배 차이 나는 구조들도 비슷한 성능을 낸다. 즉 $N$이 고정되면 Transformer의 성능은 shape에 크게 의존하지 않는다. 3.2. Performance with Non-Embedding Parameter Count $N$ Model의 크기를 나타내기 위하여 사용하는 값은 model 속 parameter의 개수이다. 이때, parameter는 실제 objective function에 맞게 학습하는 parameter만 셀 것인지, 혹은 embedding parameter까지 포함할 것인지에 따라 크게 두 가지 관점으로 볼 수 있다. 그렇기에 parameter 수에 따른 model 성능 관계를 더욱 정확하게 나타내기 위해서 두 parameter 중 어떠한 값을 parameter 값으로 사용할 것인지를 정해야 하며 그 검증 과정으로 다음과 같은 실험을 진행하였다. Left: embedding parameter까지 포함해서 세면, loss가 parameter 수뿐만 아니라 $n_{layer}$에도 의존하는 것처럼 보인다. Right: embedding parameter를 빼고 세면, depth가 다른 model들이 하나의 직선(power law)으로 모인다. 이때, layer가 1개뿐이거나 depth/width 비율이 극단적인 경우만 예외적이었다. 왼쪽처럼 보이는 이유는, 작은 model일수록 전체 parameter에서 embedding이 차지하는 비중이 크기 때문이다. parameter 수는 많아 보이지만 실제로 계산에 쓰이는 부분은 적어서 loss가 나쁘게 나온다. 그래서 이 논문은 이후 모든 분석에서 non-embedding parameter 수를 $N$ 으로 사용한다. 이렇게 정의하면 loss를 $N$에 대한 식으로 쓸 수 있다. $$ L(N) \approx \left(\frac{N_c}{N}\right)^{\alpha_N}, \qquad \alpha_N \approx 0.076,\quad N_c \approx 8.8 \times 10^{13} $$ 3.3. Performance with Dataset Size and Compute Dataset 크기 $D$ Data의 수가 model의 bottleneck이 되도록 큰 model($n_{layer}=36$, $d_{model}=1280$)을 WebText2의 부분집합으로 학습하여 특정 data 수로 도달할 수 있는 최저 loss를 측정한다. 이때, 최저 loss를 비교하기 위해서 test loss가 더 이상 줄어들지 않으면 학습을 멈춘다. 그 결과 loss와 Dataset 크기의 관계는 다음과 같다. $$ L(D) \approx \left(\frac{D_c}{D}\right)^{\alpha_D}, \qquad \alpha_D \approx 0.095,\quad D_c \approx 5.4 \times 10^{13} $$ Compute $C$ 학습 compute는 $C = 6NBS$이다. $C$를 고정하고 $N$을 바꿔 가며, 그 compute 안에서 가장 낮은 loss를 내는 model을 찾는다. 이 최적점들을 이으면 다음 관계가 나온다. $$ L(C) \approx \left(\frac{C_c}{C}\right)^{\alpha_C} $$ 각각의 식에서 $X_c$로 나타낸 값들은 log-log 그래프 상에서 관계를 나타내기 위한 fitting으로 구해낸 값이다. 이러한 parameter, Dataset, Compute에 대한 loss의 관계를 정리하면 다음과 같다. 세 그래프 모두 log-log 축에서 linear하다. 왼쪽 그래프의 하늘색 곡선들은 크기가 다른 model 각각의 learning curve이고, 이들의 아래쪽 envelope(검은 선)이 compute에 대한 power law $L = (C_{\min}/2.3 \cdot 10^8)^{-0.050}$ 을 이룬다. 단, 이 power law는 나머지 두 요인이 bottleneck이 아닐 때만 성립한다. 4. Charting the Infinite Data Limit and Overfitting 이전까지는 하나의 변수만 bottleneck이 되도록 하여 각 변수별로 독립적인 영향을 확인하였다. 이번에는 $N$과 $D$ 두 값을 동시에 바꿀 때 loss가 어떻게 되는지, 특히 data가 부족할 때 생기는 overfitting을 다룬다. 4.1. Proposed $L(N, D)$ Equation $$ L(N, D) = \left[\left(\frac{N_c}{N}\right)^{\frac{\alpha_N}{\alpha_D}} + \frac{D_c}{D}\right]^{\alpha_D} $$ 이러한 $N$, $D$에 따른 loss는 다음의 원칙들을 만족하도록 정하였다. Rescaling : Vocabulary 크기나 tokenization이 바뀌면 loss 전체가 일정 배율로 달라질 것이므로, 식이 이런 rescaling을 자연스럽게 허용해야 한다. 한 변수를 무한히 키워도 다른 변수에 의한 한계 : $D$를 고정하고 $N \to \infty$로 보내면 loss는 $L(D)$로, $N$을 고정하고 $D \to \infty$로 보내면 $L(N)$으로 수렴해야 한다. $L(N, D)$는 $D = \infty$에서 analytic : $1/D$에 대한 정수 거듭제곱 급수로 전개할 수 있어야 한다. 4.2. Results $N$과 $D$를 바꿔 가며 학습하고, test loss가 더 이상 줄지 않으면 멈춘 뒤, 위 식의 parameter를 fitting했다. Parameter $\alpha_N$ $\alpha_D$ $N_c$ $D_c$ Value 0.076 0.103 $6.4 \times 10^{13}$ $1.8 \times 10^{13}$ ** Overfitting** Data가 무한할 때의 loss $L(N, \infty)$에 비해 얼마나 손해를 보는지를 $\delta L$로 정의한다. $$ \delta L \equiv \frac{L(N, D)}{L(N, \infty)} - 1 \approx \left(1 + \left(\frac{N}{N_c}\right)^{\frac{\alpha_N}{\alpha_D}} \frac{D_c}{D}\right)^{\alpha_D} - 1 $$ $\delta L$이 클수록 overfitting이 심하다. 식을 보면 $\delta L$은 $N^{\alpha_N/\alpha_D}/D \approx N^{0.74}/D$ 하나에만 의존 한다. Random seed에 따른 loss 변동이 약 0.02이므로, $\delta L < 0.02$이면 overfitting이 없다고 본다. 이를 만족하는 조건은 $$ D \gtrsim (5 \times 10^3), N^{0.74} $$ 즉 model 크기를 키울 때 data는 sub-linear하게만 늘려도 overfitting을 피할 수 있다. 이처럼 Data $D$와 parameter $N$으로 loss 식을 fitting 할 수 있는 것은 물론, overfitting을 방지하기 위한 $N$, $D$의 관계 또한 확인할 수 있었다. 5. Scaling Laws with Model Size and Training Time 이제는 학습 step 수 $S$까지 고려하여 $N$과 학습 시간에 따른 loss를 하나의 식으로 묶는다. 그 전에 batch 크기에 따라 달라지는 step 수와 compute를 보정하기 위해 critical batch size 를 도입한다. 5.1 Critical Batch Size $B_{crit}(L)$ Batch 크기 $B$에 따라 학습 시간(step 수)과 compute가 달라진다. $B \ll B_{crit}$: Batch를 키우면 step 수가 거의 비례해서 줄어든다. Compute 효율이 좋다. $B \gg B_{crit}$: Batch를 더 키워도 step 수가 거의 줄지 않는다. Step 수는 최소에 가깝지만 compute를 낭비한다. 목표 loss에 도달하는 데 필요한 step 수 $S$와 처리한 data 수 $E = BS$ 사이에는 다음 관계가 있다. 이때 E는 Dataset 크기인 D와 다른 실제 처리한 token 수 이다. $$ \left(\frac{S}{S_{\min}} - 1\right)\left(\frac{E}{E_{\min}} - 1\right) = 1 $$ $S_{\min}$: 목표 loss에 도달하는 데 필요한 최소 step 수 ($B \to \infty$일 때) $E_{\min}$: 목표 loss에 도달하는 데 필요한 최소 data 수 ($B \to 0$일 때) Critical batch size는 이 둘의 비율로 정의한다. $$ B_{crit}(L) \equiv \frac{E_{\min}}{S_{\min}} $$ 목표 loss에 따라 $E_{\min}$, $S_{\min}$이 정해지므로 $B_{crit}$도 loss의 함수다. $B = B_{crit}$으로 학습하면 $S = 2S_{\min}$, $E = 2E_{\min}$이 되어, 시간과 compute 사이의 균형점이 된다. 실험적으로 $B_{crit}$은 model 크기와 무관하게 loss에만 의존 하며, 다음 power law를 따른다. $$ B_{crit}(L) \approx \frac{B_ }{L^{1/\alpha_B}}, \qquad B_ \approx 2 \times 10^8 \text{ tokens},\quad \alpha_B \approx 0.21 $$ 실제 측정한 loss에 따른 critical batch를 나타내면 다음과 같다. loss가 13% 줄어들 때마다 $B_{crit}$은 약 2배가 된다. 이를 이용하면 임의의 batch 크기 $B$로 학습한 결과($S$, $C$)를, $B \gg B_{crit}$일 때의 최소 step 수와 $B \ll B_{crit}$일 때의 최소 compute로 환산할 수 있다. $$ S_{\min}(S) = \frac{S}{1 + B_{crit}(L)/B}, \qquad C_{\min}(C) = \frac{C}{1 + B/B_{crit}(L)} $$ 5.2. Performance with Model size and Compute $S_{\min}$을 쓰면 loss를 $N$과 학습 step의 함수로 쓸 수 있다. $$ L(N, S_{\min}) = \left(\frac{N_c}{N}\right)^{\alpha_N} + \left(\frac{S_c}{S_{\min}}\right)^{\alpha_S}, \qquad \alpha_S \approx 0.76,\quad S_c \approx 2.1 \times 10^3 $$ 첫 번째 항은 무한히 학습했을 때 model 크기가 정하는 한계, 두 번째 항은 학습이 덜 되어서 생기는 손실이다. Left: compute를 고정하면 loss를 최소로 만드는 최적 model 크기가 존재한다. 너무 작은 model은 용량이 부족하고, 너무 큰 model은 같은 compute로 충분히 학습되지 못한다. Right: step 수를 고정하면 model이 클수록 loss가 낮다. 점선이 $L(N, S)$ 식이고, 학습 초반(작은 $S$)을 제외하면 실험값과 잘 맞는다. 5.3. Lower bound on Early Stopping Step Data가 유한할 때와 무한할 때의 learning curve는 $S_{\min} \approx S_{stop}$ 근처까지 거의 같다. 따라서 overfitting 때문에 학습을 멈춰야 하는 step 수의 하한을 다음과 같이 추정할 수 있다. $$ S_{stop}(N, D) \gtrsim \frac{S_c}{\left[L(N, D) - L(N, \infty)\right]^{1/\alpha_S}} $$ data가 부족해서 생기는 loss 차이 $L(N,D) - L(N,\infty)$가 작을수록, 더 오래 학습한 뒤에야 overfitting이 시작된다. 6. Optimal Allocation of the Compute Budget 6.1. Optimal Performance and Allocations $B_{crit}$으로 보정한 $C_{\min}$을 쓰면, loss를 $C_{\min}$의 함수 $L(C_{\min})$으로 다시 쓸 수 있다. 그리고 주어진 compute에서 loss를 최소로 만드는 최적 model 크기 $N(C_{\min})$을 구하면 $$ N(C_{\min}) \propto C_{\min}^{0.73} $$ $C_{\min} = 6NB_{crit}S$이고 $B_{crit} \propto C_{\min}^{0.24}$이므로, 남는 step 수는 $$ S_{\min} \propto C_{\min}^{0.03} $$ 이다. 즉, 결론은 다음과 같다. Compute를 최적으로 배분해서 language model을 키울 때는 model 크기 $N$을 주로 키우고, batch 크기는 $B_{crit}$에 맞춰 함께 늘리며, step 수는 거의 늘리지 않는다. 6.2 Predictions from $L(N, S_{\min})$ 5.2절의 $L(N, S_{\min})$ 식과 $B_{crit}(L)$ 식을 결합하면, $L(C_{\min})$의 지수를 이론적으로 예측할 수 있다. $$ L(C_{\min}) = \left(\frac{C_c^{\min}}{C_{\min}}\right)^{\alpha_C^{\min}}, \qquad \alpha_C^{\min} = \frac{1}{1/\alpha_S + 1/\alpha_B + 1/\alpha_N} \approx 0.054 $$ 실험에서 직접 측정한 값 0.050과 잘 맞는다. 세 개의 독립적인 power law가 서로 일관된다는 근거다. 6.3 Contradictions and a Conjecture 이 scaling law를 계속 적용하면 모순이 생긴다. Overfitting을 피하려면 data가 $D \propto N^{0.74} \propto C_{\min}^{0.54}$만큼 늘어야 한다. 그런데 compute-efficient하게 학습할 때 실제로 보는 data는 $$ D(C_{\min}) = \frac{2C_{\min}}{6N(C_{\min})} \propto C_{\min}^{0.26} $$ 으로, 필요한 속도보다 훨씬 느리게 늘어난다. 따라서 data를 재사용하지 않더라도 결국 overfitting 영역에 들어가게 된다. $L(C_{\min})$(주황 점선)은 compute를 늘릴수록 빠르게 줄어든다. 교차점 왼쪽에서는 compute가 bottleneck이라, compute를 늘리면 loss가 줄어든다. $L(D(C))$(빨간 선)는 그 compute에서 보는 data 양으로 달성할 수 있는 loss의 하한이다. 이 선은 더 천천히 줄어든다. 교차점 오른쪽에서는 $L(C_{\min})$이 $L(D(C))$보다 낮아지는데, 이는 data가 부족한데도 loss가 그 data로 가능한 한계보다 낮아진다는 뜻이라 불가능하다. 따라서 scaling law는 교차점 이전에 깨져야 한다. 교차점은 대략 $C^* \sim 10^4$ PF-days, $N^* \sim 10^{12}$ parameters, $D^* \sim 10^{12}$ tokens, $L^* \sim 1.7$ nats/token이다. 논문은 이 지점의 loss $L^*$가 자연어의 token당 entropy에 대한 추정치일 수 있다고 추측한다. 다만 교차점의 위치는 power law 지수에 매우 민감하므로 정확한 값으로 볼 수는 없다. 7. Discussion Language model의 loss는