Training Compute-Optimal Large Language Models

Jordan Hoffmann; Sebastian Borgeaud; Arthur Mensch; Elena Buchatskaya; Trevor Cai; Eliza Rutherford; Diego de Las Casas; Lisa Anne Hendricks; Johannes Welbl; Aidan Clark; Tom Hennigan; Eric Noland; Katie Millican; George van den Driessche; Bogdan Damoc; Aurelia Guy; Simon Osindero; Karen Simonyan; Erich Elsen; Jack W. Rae; Oriol Vinyals; Laurent Sifre

논문 · arXiv 2203.15556

Training Compute-Optimal Large Language Models (연산 최적화된 거대 언어 모델 학습)

Training Compute-Optimal Large Language Models

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland et al.

게시일 2022-03-29cs.CLcs.LG arXiv PDF

XYZNOT로 읽기 (전문 번역과 AI 해설)

세 줄 요약

🎯 고정된 연산 예산 하에서 모델 파라미터와 학습 토큰 간의 최적 트레이드오프 규명.
🔧 최소 손실 포락선(envelope), IsoFLOP 단면, 매개변수적 손실 모델링을 사용해 400회 이상의 트랜스포머 학습 전반에 걸쳐 스케일링 법칙을 실증적으로 피팅.
📈 파라미터와 토큰은 동일한 비율로 스케일링되어야 함( $a \approx 0.5$ , $b \approx 0.5$ ). 1.4조 토큰으로 학습된 70B Chinchilla 모델은 Gopher와 동일한 연산 예산으로 MMLU 67.6%를 달성.

초록

본 연구에서는 주어진 연산 예산 내에서 트랜스포머 언어 모델을 학습시키기 위한 최적의 모델 크기와 토큰 수를 조사합니다. 우리는 현재의 대형 언어 모델들이 상당히 덜 학습(undertrained)되어 있음을 발견했으며, 이는 최근 연구들이 학습 데이터양을 일정하게 유지하면서 언어 모델의 크기만 확장하는 데 집중해 온 결과입니다. 7,000만 개에서 160억 개 이상의 파라미터를 가진 400개 이상의 언어 모델을 50억~5,000억 개의 토큰으로 학습시킨 결과, 연산 최적(compute-optimal) 학습을 위해서는 모델 크기와 학습 토큰 수를 동일한 비율로 확장해야 한다는 결론을 도출했습니다. 즉, 모델 크기가 두 배가 될 때마다 학습 토큰 수도 두 배로 늘려야 합니다. 우리는 이러한 가설을 검증하기 위해 연산 최적 모델로 예측된 Chinchilla를 학습시켰습니다. 이 모델은 Gopher와 동일한 연산 예산을 사용하지만, 파라미터는 70B이며 4배 더 많은 데이터를 사용합니다. Chinchilla는 광범위한 다운스트림 평가 작업에서 Gopher (280B), GPT-3 (175B), Jurassic-1 (178B), Megatron-Turing NLG (530B)를 일관되게 큰 폭으로 능가합니다. 이는 또한 Chinchilla가 미세 조정 및 추론 과정에서 훨씬 적은 연산량을 사용하여 다운스트림 활용성을 크게 향상시킴을 의미합니다. 특히 주목할 만한 점은, Chinchilla가 MMLU 벤치마크에서 최고 수준인 67.5%의 평균 정확도를 기록하여 Gopher 대비 7% 이상의 성능 향상을 보였다는 것입니다.

영어 원문 초록

We investigate the optimal model size and number of tokens for training a transformer language model under a given compute budget. We find that current large language models are significantly undertrained, a consequence of the recent focus on scaling language models whilst keeping the amount of training data constant. By training over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens, we find that for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled. We test this hypothesis by training a predicted compute-optimal model, Chinchilla, that uses the same compute budget as Gopher but with 70B parameters and 4 $\times$ more more data. Chinchilla uniformly and significantly outperforms Gopher (280B), GPT-3 (175B), Jurassic-1 (178B), and Megatron-Turing NLG (530B) on a large range of downstream evaluation tasks. This also means that Chinchilla uses substantially less compute for fine-tuning and inference, greatly facilitating downstream usage. As a highlight, Chinchilla reaches a state-of-the-art average accuracy of 67.5% on the MMLU benchmark, greater than a 7% improvement over Gopher.

논문 해설

문제점 및 연구 배경

이전 연구(특히 Kaplan et al., 2020)는 연산량이 증가할 때 모델 파라미터는 공격적으로 확장( $N \propto C^{0.73}$ )해야 하는 반면, 토큰 수는 훨씬 느리게 증가( $D \propto C^{0.27}$ )시켜야 한다고 제안했습니다. 그 결과 GPT-3 (175B), Jurassic-1 (178B), Gopher (280B), MT-NLG (530B)와 같은 거대 모델들은 일관되게 약 3,000억 개의 토큰만을 사용하여 학습되었습니다. 그러나 이는 투입된 연산량에 비해 거대 모델들이 현저히 덜 학습(undertrained)되는 결과를 낳았습니다. 본 연구는 이러한 트레이드오프를 재검토하여, 고정된 연산 예산 $C$ 하에서 테스트 손실을 최소화하기 위해 파라미터 수 $N$ 과 토큰 수 $D$ 를 어떻게 배분해야 하는지 규명합니다.

핵심 아이디어

동등 스케일링 (Equal scaling): 모델 파라미터가 두 배가 될 때마다 학습 토큰 수도 두 배로 늘려야 합니다 ( $N_{o pt} \propto C^{0.5}$ , $D_{o pt} \propto C^{0.5}$ ).
코사인 스케줄 정렬 (Cosine schedule alignment): 학습률 스케줄은 실제 학습 기간(horizon)과 반드시 일치해야 합니다. 길게 설정된 특정 기간(예: 1,300억 토큰)에 맞춰 감쇠되는 단일 학습 실행의 중간 지점 데이터를 사용하는 것은 더 적은 토큰 수에 대한 손실을 과대평가하여 실증적 스케일링 추정치를 왜곡합니다.
다운스트림 효율성: 훨씬 더 많은 토큰으로 더 작은 크기의 모델을 학습시키면 사전 학습 연산량이 동일하더라도 같거나 더 뛰어난 정확도를 얻을 수 있으며, 메모리 점유율, 미세 조정 비용, 추론 지연 시간을 대폭 줄일 수 있습니다.

연구 방법론 (핵심 수식 포함)

저자들은 세 가지 상호보완적 전략을 사용하여 400개 이상의 모델(7,000만~~160억 개의 파라미터, 50억~~5,000억 개의 토큰으로 학습)을 평가합니다:

학습 곡선 기반 최소 손실 (Minimum over training curves): 고정된 모델 아키텍처에 대해 4가지 학습 기간을 달리하여 실행합니다. FLOP당 최소 손실의 포락선(envelope)을 추출하고 멱법칙을 피팅합니다: $N_{o pt} \propto C^{a}, D_{o pt} \propto C^{b}$ 이를 통해 $a = 0.50$ , $b = 0.50$ 을 도출합니다.
IsoFLOP 프로파일 (IsoFLOP profiles): 9개의 이산적 예산( $6 \times 1 0^{18}$ ~ $3 \times 1 0^{21}$ )에 걸쳐 총 FLOP 수를 고정하고 코사인 감쇠를 학습 기간에 정렬합니다. 파라미터 수에 따른 손실 최솟값을 찾기 위해 포물선을 피팅하여 $a = 0.49$ , $b = 0.51$ 을 얻습니다.
매개변수적 손실 피팅 (Parametric loss fitting): 평활화된 손실을 다음과 같은 위험 분해(risk decomposition) 모델로 정의합니다: $\hat{L} (N, D) ≜ E + \frac{A}{N ^{α}} + \frac{B}{D ^{β}}$ 여기서 $E$ 는 자연어 텍스트 엔트로피, $A / N^{α}$ 는 표현력 한계로 인한 차최적성, $B / D^{β}$ 는 조기 종료 및 유한한 데이터로 인한 차최적성을 나타냅니다. 로그 잔차에 대해 Huber loss를 최소화하면 $E = 1.69$ , $A = 406.4$ , $B = 410.7$ , $α = 0.34$ , $β = 0.28$ 을 얻습니다. $FLOPs (N, D) \approx 6 N D$ 제약 조건 하에서: $a = \frac{β}{α + β} = 0.46, b = \frac{α}{α + β} = 0.54$

연구 결과 (주요 수치 요약)

Chinchilla 70B: Gopher와 동일한 FLOP 예산( $5.76 \times 1 0^{23}$ FLOPs, 정확한 산정 기준 시 $6.3 \times 1 0^{23}$ )을 투입하여 70B 파라미터 및 1.4조 개의 토큰으로 학습되었습니다.
MMLU 벤치마크: Chinchilla는 67.6%의 평균 5-shot 정확도를 달성하여 Gopher (60.0%), GPT-3 (43.9%)를 능가하고, 2023년 6월 인간 전문가 예측치(63.4%)를 뛰어넘었습니다. 또한 4개의 개별 MMLU 작업에서 $> 90%$ 의 정확도를 기록했습니다.
다운스트림 벤치마크: BIG-bench (65.1% vs 54.4%), LAMBADA (77.4% vs 74.5%), closed-book Natural Questions (31.5% 5-shot vs 24.5%) 전반에서 Gopher를 크게 능가했습니다.
편향 및 독성 (Bias and Toxicity): 프롬프트 없는 상태의 독성은 Gopher와 Chinchilla 간에 큰 차이가 없었으나(중앙값 0.064 vs 0.066), Winogender 정확도는 전체적으로 71.4%에서 78.3%로 향상되었습니다.

한계점 및 열린 질문

학습 비용 문제로 인해 동일 연산량 조건의 대규모 프론티어 비교는 단 두 개 모델(Gopher 280B 대 Chinchilla 70B)에 대해서만 직접 검증되었습니다.
이 분석은 웹 데이터셋에 대한 1회 에포크(single-epoch) 학습을 가정한 것으로, 여러 에포크에 걸쳐 데이터를 반복 학습할 때의 동작은 다루지 않았습니다.
높은 연산 예산 영역에서 $lo g (N_{o pt})$ 에 약간의 음의 곡률(위로 볼록한 형태)이 관찰되는 실증적 증거가 있어, 멱법칙이 여전히 최적의 모델 크기를 다소 과대평가할 가능성이 열려 있습니다.

핵심 용어

Chinchilla: Gopher와 동일한 연산 예산을 사용하여 1.4조 개의 토큰으로 학습된 700억(70B) 파라미터의 자기회귀 언어 모델.
Gopher: 3,000억 개의 토큰으로 학습된 2,800억(280B) 파라미터의 트랜스포머 언어 모델로, 연산량이 동일한 비교 기준 모델 역할을 함.
IsoFLOP profiles: 총 학습 FLOP 수를 일정하게 유지하면서 모델 파라미터 크기를 변경하여 손실이 최소화되는 최적의 파라미터 크기를 찾는 방법론.
Compute-optimal frontier: 가용한 학습 FLOP 수와 최적 파라미터 크기 및 학습 토큰 수 배분 간의 관계를 나타내는 수학적 곡선.
Parametric loss function: 언어 모델의 손실을 줄일 수 없는 베이즈 엔트로피, 파라미터 용량 부족, 토큰 표본 최적화 오차로 분해하는 통계적 모델링 수식.
Huber loss: 소규모 연산 영역에서 잔차 이상치의 영향을 완화하면서 실증적 언어 모델 학습 실행 데이터를 피팅하기 위해 접근법 3에서 사용된 로버스트 손실 함수.
MassiveText: 웹 페이지, 서적, 뉴스, 코드 및 위키피디아를 포함하는 대규모 멀티 도메인 사전 학습 코퍼스로, Gopher와 Chinchilla 학습에 모두 사용됨.
FLOPs calculation: 순전파 및 역전파 전반에 걸쳐 대략 $6 N D$ 로 계산되거나 레이어별 정밀 계산을 통해 산출되는 곱셈-누적(multiply-accumulate) 연산 예산.

자주 묻는 질문

Kaplan et al.은 왜 토큰보다 파라미터를 더 빠르게 확장해야 한다고 결론지었나요?

Kaplan et al.은 고정된 코사인 학습률 스케줄(예: 1,300억 토큰까지 감쇠되도록 설정)을 사용하고, 더 적은 수의 토큰 지점에 대해서는 완전히 수렴하지 않은 중간 학습 평가 결과를 사용했습니다. 코사인 학습률 스케줄은 전체 스텝 수가 주기 길이보다 크게 모자랄 경우 성능이 저하되기 때문에, 이러한 중간 손실값들이 인위적으로 높게 측정되어 더 적은 수의 토큰으로 작은 모델을 학습시키는 것의 효과를 과소평가하게 되었습니다.

Chinchilla 연구에서 제시하는 연산 최적 스케일링의 경험 법칙은 무엇인가요?

모델 파라미터 수와 학습 토큰 수 모두 대략 동일한 비율로 확장되어야 합니다. 즉, 연산량이 $2 \times$ 증가할 때마다 모델 크기는 $\approx 2^{0.5} \approx 1.41 \times$ 증가해야 하고, 학습 토큰 수도 $\approx 2^{0.5} \approx 1.41 \times$ 증가해야 합니다.

연산 최적화 관점에서 175B 모델을 학습시키려면 몇 개의 토큰이 권장되나요?

논문의 예측에 따르면 1,750억(175B) 파라미터 모델이 연산 최적이 되려면 약 3.7조~4.3조 개의 토큰과 $3.85 \times 1 0^{24}$ ~ $4.54 \times 1 0^{24}$ FLOPs가 필요하며, 이는 GPT-3에 사용된 전형적인 3,000억 토큰을 크게 상회합니다.

Gopher와 Chinchilla 간에 어떤 아키텍처적 변화가 있었나요?

레이어 수와 은닉 차원을 줄여 280B에서 70B로 모델 크기를 조정한 것 외에도, Chinchilla는 Adam 옵티마이저를 AdamW로 교체하였고, 분산 옵티마이저 상태에서 float32 마스터 가중치를 유지했으며, SentencePiece 토크나이저에서 NFKC 정규화를 제거했습니다.

이 페이지의 해설은 AI가 논문을 바탕으로 작성했으며 오류가 있을 수 있습니다. 자세한 내용은 원문을 확인하세요.

XYZNOT로 읽기 (전문 번역과 AI 해설)