지금 까지의 LM Scaling Law에는 문제점이 있다?!?! 😶‍🌫️ Chinchilla: Training Compute-Optimal Large Language Models 논문 리뷰

The overview of this paper

이 논문에서는 주어진 compute budget에서 transformer LM을 학습시키기 위한 최적의 모델 사이즈 & 토큰의 수를 조사하였다. 그 결과, 현재의 모델들은 상당히 under-train 되어 있다는 사실을 밝혔다. 현재 많은 연구들은 모델의 사이즈는 scaling 하지만, training data의 양은 크게 변화시키지 않고 있다. 논문에서는 여러 모델 사이즈와 토큰의 수에 대해 실험을 진행하고 compute-optimal training을 위해 모델 사이즈와 training 토큰의 수는 똑같이 scale 되어야 한다는 사실을 밝혀냈다.

 

 논문에서는 이를 검증하기 위해 Gopher와 똑같은 compute budget을 사용하지만 70B 파라미터와 4배 더 많은 데이터에서 학습된 논문의 scaling law에 의한 compute-optimal model인 Chinchilla를 학습시켰다. 그 결과 Chinchilla는 Gopher(280B), GPT-3(175B), Jurassic(178B), Megatron-Turing NLG(530B)를 여러 downstream 평가 task에서 능가하는 모습을 보여줬다!! 😲

 

 이 논문에서는 최적의 모델 파라미터 수와 training token의 수를 찾기 위해 수식적 정보를 사용하는데 이 부분을 말로 조금 쉽게 풀어내보았다!!

 

 

Table of Contents

1. Introduction

2. Estimating the optimal parameter/training tokens allocation

3. Chinchilla

4. Discussion & Conclusion

 

 

1. Introduction

 최근에 소개된 LLM들은 엄청 크고 task에서 상당히 인상적인 결과들을 보여주고 있다. 하지만 LLM을 학습시키기 위한 계산량과 에너지 비용은 상당하고, 심지어 모델의 사이즈가 커짐에 따라 훨씬 더 커진다. compute & energy cost는 사전에 할당되는데, 이는 large model을 한 번 돌릴 정도의 양 밖에 되지 않아서 주어진 compute budget에 대한 최고의 모델 하이퍼 파라미터를 정확하게 측정하는 것은 중요하다.

 

표 1. 현재 LLM들의 상태. 파라미터는 상당히 증가하고 있지만, 그에 반해 training token의 수는 정체되어있는 것을 알 수 있음

 

 이전의 논문에서는 AR LM에서 파라미터의 수와 성능 간에는 power law가 존재한다는 것을 밝혀냈다. 이 논문의 결과에 의하면 더욱 큰 모델에서 학습할수록 모델의 성능은 향상된다. Chinchilla 논문에서도 똑같은 결론에 도달하게 되는데, 다만 이전 논문의 저자와 달리 여기에서는 더 많은 training token에서 학습되어야 한다는 것을 측정하였다. 이전의 논문에서는 computation budget이 10배 상승하면 모델의 사이즈는 5.5배 증가해야 하고, training token의 수는 1.8배 증가해야 한다고 제안하였다. 하지만, Chinchilla 논문의 실험 결과에 따르면 모델 사이즈와 training token은 똑같은 비율로 scale 되어야 한다! ⚠️

 

 논문에서는 다음의 질문을 다시 생각해보았다.

 

'고정된 FLOPs budget이 주어지면, 모델 사이즈와 training token의 수는 어떻게 trade-off 해야 할까?'

 

 이 질문에 대답하기 위해 논문에서는 최종 pre-training loss $L(N, D)$를 모델링하였다. 여기서 $N$은 모델 파라미터의 수이고, $D$는 training token의 수이다. 그리고 $C$는 compute budget인데 $FLOPs (N, D)$이다. 논문에서는 이 $FLOPs(N, D) = C$ 하에서 $L$을 최소화하기 위해 학습된다.

 

 

 이 함수에서 $N_{opt}(C)$와 $D_{opt}(C)$는 computation budget $C$의 최적 할당을 의미한다. 논문에서는 70M to 16B 파라미터 & 5B to 400B tokens의 400개 이상의 모델에서 loss에 기반해서 이 함수를 측정하였다. 논문의 방식은 이전의 논문과는 상당히 다른 결과를 이끌어낸다. 다음의 그림 1이 논문의 결과를 강조하고 있다.

 

그림 1. 모델들에 대한 논문의 서로 다른 3개의 approach 예측을 겹쳐둠. 3개의 approach 모두 현재 모델들이 상당히 작아져야 하고 더욱 길게 학습되어야 한다는 것을 보여줌.

 

 논문에서 측정한 compute-optimal frontier에 기반해서, 논문에서는 Gopher를 train 하기 위해 사용된 compute budget에서  optimal model은 4배 더 적고, 대신에 4배 더 많은 token에서 학습시켰다. 논문에서는 이를 더욱 compute-optimal 한 70B 모델인 Chinchilla로 1.4T token에서 학습시킴으로써 입증시켰다. 그 결과 Chinchilla는 더욱 큰 모델인 Gopher를 능가할 뿐만 아니라 감소된 모델 사이즈는 추론 비용을 상당히 감소시켰고, smaller 하드웨어에서 downstream 사용을 용이하게 하였다.

 

 

2. Estimating the optimal parameter/training tokens allocation

 논문에서는 '고정된 FLOPs budget이 주어지면 어떻게 모델 사이즈와 training token의 수 간에 trade-off를 해야할까?'라는 질문에 대한 답을 내놓기 위해 3개의 서로 다른 approach를 사용하였다. 이 3개의 경우에서 모델 사이즈와 training token의 수를 다양하게 해서 여러 모델을 학습시켰고, 어떻게 scale 해야 하는지 알기 위해 estimator를 결과로 나온 training curve에 사용하였다. 3개의 approach에 대해서 나온 결과는 모두 유사하였고, 파라미터 수와 training token의 수는 더 많은 compute와 함께 똑같이 증가되어야 한다는 사실을 제안하였다. 이는 확실히 이전 논문의 결과와는 확실히 다르다!

 

2-1. Approach 1: 모델 사이즈 고정 & training token의 수 다양하게

 

 고정된 familt model에서 4개의 서로 다른 training sequence의 수를 달리하여 실험을 진행하였다. 필요한 training token의 수와 함께 모든 모델 중 가장 낮은 loss를 달성하는 모델의 사이즈를 찾았다. 그 결과는 그림 2와 같다.

 

그림 2. 왼쪽은 모든 실험에 대한 결과를 보여준다. 70M to 10B 모델 사이즈에서 실험을 진행하고, 각각은 4개의 서로 다른 cosine cycle 길이에 대한 값이다. 이 커브로부터, 논문에서는 FLOP 당 최소 loss를 추출해내고, 주어진 compute budget과 최적 training token의 수(오른쪽)에 대한 최적의 모델 사이즈를 측정(중간).

 

 마지막으로, 주어진 compute budget에 대한 최적의 모델 사이즈와 training token 수를 추정하기 위해 거듭제곱 법칙을 적용하여 관계 $N_{opt} \propto C^{a}$와 $D_{opt} \propto C^{b}$를 얻었다(그림 2의 가운데 및 오른쪽 참조). 그리고 여기서 $a=0.5$와 $b=0.5$가 최적이었다.

 

2-2. Approach 2: IsoFLOP profiles

 

 9개의 서로 다른 training FLOP 카운트의 고정된 세트에 대한 모델 사이즈를 달리하였고, 각 포인트에 대한 최종 training loss를 고려하였다. 전체 training run을 따라 point $(N, D, L)$를 고려하는 approach 1과 반대이다. 이는 '주어진 FLOP budget에 대해, 무엇이 최적의 파라미터 카운트인가?'라는 질문에 대답할 수 있게 해 준다. 그림 3을 살펴보자.

 

그림 3. IsoFLOP curves.  다양한 모델 크기에 대해 최종 FLOP가 일정하도록 training token 수를 선택한다. 타깃 FLOP count를 맞추기 위해 cosine cycle 길이가 맞춰진다. 그리고 loss에서 분명한 골짜기를 볼 수 있는데, 이는 주어진 FLOP budget에 대해 학습시키기 위한 최적의 모델은 있다(왼쪽). 이 골짜기의 위치를 사용해서 최적의 모델 사이즈와 larger model에 대한 토큰의 수를 project한다(중앙 & 오른쪽). 초록색 선은 최적의 모델에 대한 측정된 파라미터와 토큰의 수를 보여준다.

 

 

 각 FLOP budget에 대해, 논문에서는 그림 3(왼쪽)에서 파라미터 카운트에 대한 최종 loss를 보여준다. 모든 경우에서 loss의 분명한 최소를 보기 이해 충분히 다양한 모델 세트를 학습시켰다. 논문에서는 각 IsoFLOPs curve에 포물선을 적용하여 minimum loss가 얻어지는 모델 사이즈를 측정하였다(그림 3(왼쪽)). 이전 appraoch처럼 FLOPs와 loss-optimal model size & training token의 수 간에 power law를 적용하였다(그림 3(중앙, 오른쪽)). 또다시 $N_{opt} \propto C^{a}$와 $D_{opt} \propto C^{b}$를 적용해 본 결과 $a=0.49$와 $b=0.51$을 얻을 수 있었다.

 

2-3. Approach 3: parametric loss 함수 적용

 

 approach 3에 대한 설명은 그림 4를 보고 설명하도록 하겠다. 

 

그림 4. parametric modeling을 적용해서 윤곽(왼쪽)과 isoFLOP slice(오른쪽)을 보여준다. 각각의 isoFLOP slice에 대해, 왼쪽 그래프에 해당 점선을 추가한다. 왼쪽 그래프에서는 효율적인 frontier를 파란색으로 나타내고, 이는 log-log 공간에서의 선이다. 구체적으로 곡선은 FLOP가 가장 적은 지점에서 각 등손실 윤곽선을 통과한다. 논문에서는 주어진 Gopher FLOP budget에 대해 최적의 모델 사이즈를 40B 파라미터로 정의하였다.

 

2-4. Optimal model scaling

 

 3개의 approach에 대해 증가된 training compute에 따른 측정된 파라미터와 데이터 scaling을 다음과 같다.

 

표 2. 증가된 training compute에 따른 측정된 파라미터와 데이터 scaling. 3 approach 모두 비슷한 결과를 보여주고 이 결과는 이전 논문의 결과와는 상당히 다른 모습을 보여준다.

 

 논문에서 3개의 approach는 서로 다른 fitting methodology와 trained model을 사용했음에도 불구하고 파라미터와 토큰의 최적 scaling에 대해 비슷한 예측을 내놓는다. 모든 approach는 compute budget이 증가함에 따라 모델 사이즈와 training data의 양은 거의 똑같은 비율로 증가해야 한다고 보여준다.

 

표 3. 다양한 모델에 대한 최적의 training FLOPs & training tokens 추정.

 

 논문의 결과에 의하면 현재 LLM들은 상당히 over-size 하다. 그리고 모델 사이즈와 training token의 수는 동등한 비율로 scale 되어야 한다는 것을 보여준다.

 

 

3. Chinchilla

 섹션 2의 분석에 기반해서, Gopher compute budget에 대한 최적의 모델 사이즈는 40B에서 70B 사이의 파라미터이다. 논문에서는 이 가설을 가장 큰 모델인 70B 파라미터 모델에 1.4T 개의 토큰을 사용해서 평가하였다. 이 섹션에서는 Chinchilla라 부르는 이 모델을 Gopher와 다른 LLM들과 비교하였다. Chinchilla와 Gopher는 똑같은 수의 FLOPs에 대해 학습되지만 모델의 사이즈와 training token의 수는 다르다.

 

 LLM을 pre-training 하는 것은 상당한 계산 비용이 들지만, downstream fine-tuning과 inference 또한 상당한 계산량을 필요로 한다. Chinchilla는 Gopher보다 4배 작기 때문에 memory footprint & 추론 비용이 더 작다.

 

 Chinchilla를 학습시키기 위해 사용된 모든 하이퍼파라미터는 표 4에 나타나 있다. Chinchilla는 아래 표의 예외를 제외하고 Gopher와 똑같은 모델 아키텍처와 training setup을 사용한다.

 

표 4. Chinchilla architecture detail

 

3-1. Results

 

 논문에서는 Chinchilla를 다양한 LLM과 비교하면서 광범위한 평가를 수행하였다. 평가를 진행한 task subset은 다음의 표 5와 같다.

 

표 5. 모든 평가 tasks

 

Language modeling.  Chinchilla는 그림 5에서 보이는 것처럼 The Pile의 모든 평가 서브셋에서 Gopher를 상당히 능가하였다.

 

그림 5. The Pile 평가 결과. 모든 서브셋에서 Gopher를 능가하는 성능을 보여줌

 

MMLU.  표 6에 MMLU에서 Chinchilla의 평균 5-shot 성능이 기록되어 있다. 이 벤치마크에서 Chinchilla는 Gopher보다 훨씬 작음에도 불구하고 상당히 능가하는 모습을 보여줬다.

 

표 6. MMLU 결과. 57개의 task에서 평균 5-shot 결과를 측정

 

 그림 6에서는 세분화된 task에서 Chinchilla와 Gopher 간의 비교를 보여준다. 전반적으로 Chinchilla는 다양한 주요 task에서 성능을 개선시켰다.

 

그림 6. Gopher와 비교한 MMLU 결과

 

Rading comprehension.  최종 단어 예측 데이터셋 LAMBADA에서 Chinchilla는 다른 모델들보다 더 좋은 결과를 얻었다.

 

표 7. Reading comprehension 결과

 

BIG-bench.  똑같은 BIG-bench task에서 Chinchilla를 분석하였다. MMLU와 유사하게 Chinchilla는 Gopher를 대부분의 주요 task에서 능가하였다(그림 7).

 

그림 7. Gopher와 비교한 BIG-bench 결과

 

 

4. Discussion & Conclusion

 LLM training에 대한 현재 추세는 모델 사이즈는 엄청 증가하였지만, 토큰의 수는 정체되어 있다. 논문에서는 더욱 큰 모델을 학습시키기 위한 경쟁은 똑같은 compute budget에서 상당히 underperform 하는 모습을 보여준다는 가설을 세웠다.

 

 논문에서는 이 가설을 증명하기 위해 총 3 가지의 approach를 제안해서 총 400개 이상의 모델을 실행 결과에 기반해서 모델 사이즈와 training token 수 간의 최적의 세팅을 찾았다. 그 결과 모든 approach는 Gopher가 상당히 over-size 하고 똑같은 compute budger에 대해 더 많은 데이터에서 학습된 smaller model이 더 잘 수행할 것이라고 예측하였다. 

 

 논문에서는 높은 compute budget에서 몇몇 $log(N_{opt}$의 오목함을 관찰하였다. 이는 우리가 large model의 최적 사이즈를 overestimate 하였다는 것을 제안한다. Chinchilla와 Gopher 간의 비교는 성능 예측을 입증하게 만들었고, 똑같은 compute budget에서 더 나은 training을 가능하게 만들었다.

 

 결과적으로 논문에서 말하고자 하는 것은 기존의 scaling law가 잘못되었다는 것이다. 기존에는 모델 사이즈를 키우는 것이 더 중요하다고 주장했는데, 이 논문의 실험에 의하면 모델 사이즈와 데이터의 양은 동등한 비율로 증가되어야 모델의 성능을 더욱 향상시킨다는 점을 발견하였다. 

 

 

 

 

출처

https://arxiv.org/abs/2203.15556

Training Compute-Optimal Large Language Models