[논문리뷰] ChinChilla 논문 리뷰

ChinChilla 논문 리뷰입니다.

 

논문 : 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 and Laurent Sifre (DeepMind)

 

[1] Introduction

최근 몇 년간 LLM 개발이 급속도로 이루어지면서 굉장히 빠른 속도로 많은 모델들이 등장하고 있습니다.

논문에서는 대표적으로 GPT-3, Gopher, LaMDA 등의 모델을 언급하고 있습니다.

 

이러한 모델들은 Autoregressive Transformer 구조를 기반으로 하며Zero-shot, Few-shot, fine-tuning 등 다양한 방식에서 뛰어난 성능을 입증하였습니다.

 

하지만, 본 논문에서 제기하는 문제점은 연산량과 자원입니다.

실제로 LLM을 학습시키는 데는 수백 개의 GPU 혹은 TPU가 동원되며, 이는 연구 및 개발 예산에도 큰 부담으로 작용합니다.

 

따라서, 본 논문은 이러한 "자원의 제약 환경에서 얼마나 효과적으로 모델을 설계하고 학습시킬 수 있는가" 에 집중하였습니다.

 

ChinChilla 연구의 핵심 질문은 다음과 같습니다.

정해진 연산량 (FLOPs) 안에서 모델 파라미터 수와 학습 토큰 수 사이의 균형을 어떻게 설정해야 가장 효율적인 학습이 이루어질까?

 

본 연구팀은 위 문제를 풀기 위해 모델 파라미터 수 $N$과 학습에 사용된 토큰 수 $D$를 변수로 하는 손실 함수를 정의하고 손실이 가장 낮은 $N, D$의 조합을 찾는 최적화 문제를 풀었습니다.

 

쉽게 말해서 연산량이 고정되어 있을 경우, 모델을 너무 크게 만들면 학습할 수 있는 데이터가 부족하고 모델이 너무 작으면 데이터가 많더라도 학습 효율이 낮아지기 때문에 문제가 됩니다.

 

그래서 이러한 $N, D$ 균형점을 실험을 통해 정량적으로 계산해냈고

앞으로의 LLM 연구에서 효율적으로 모델을 설계하는 것의 중요성을 강조합니다.

---

[2] Related Work

본 논문에서 설명하는 LLM은 크게 두 가지 종류가 있습니다.

• Dense Transformer 기반의 전통적인 모델
  • 예를 들어, GPT-3, Gopher, LaMDA
• 전문가 혼합 모델 (Mixture-of-Experts, MoE)
  • 예를 들어, GLaM, Switch Transformer

이러한 모델들도 최근에는 파라미터 수가 500B를 넘는 거대 모델도 등장할 만큼 규모가 계속 성장하고 있는데

모델 크기를 키우는 만큼 연산량과 데이터가 대량으로 필요하기 때문에 막대한 비용이 부담됩니다.

따라서, 효율적인 학습 전략의 필요성이 강조되고 있는 상황입니다.

 

 

여기서 "Kaplan" 이라는 스케일링 법칙을 소개하는 기존 연구가 있었는데

이는 정해진 연산량 내에서 최적의 모델 크기를 찾는 문제를 풀었습니다.

그러나 이 연구는 크기가 다양한 모든 모델에 동일한 학습 토큰 수, 동일한 학습률 스케줄을 적용했다는 한계가 존재합니다.

이 연구를 보완하기 위해 ChinChilla는 모델 크기별로 학습 토큰 수와 학습률 스케줄을 유연하게 적용하였고 그 결과, 모델 크기와 학습 데이터양을 비슷한 비율로 늘리는 것이 훨씬 효율적이라는 결론을 얻었습니다.

 

추가로, 2022년 당시에 MoE 모델이 주목받고 있는 상황이었습니다.

MoE는 전체 파라미터 중 일부 전문가(Expert)만 선택적으로 활성화함으로써 더 적은 연산으로 큰 효과를 낼 수 있도록 설계되었습니다.

 

*여기서 전문가란, 여러 개의 파라미터 혹은 서브 네트워크로 구성된 작은 모델입니다. 게이트 네트워크가 입력을 보고 가장 적절한 내부 네트워크를 선택해 그 일부 네트워크만 연산을 하도록해 효율적으로 연산하는 구조를 가지고 있습니다. 이렇게 되면 연산이 효율적이라는 장점이 있지만 서브 네트워크를 잘 고르지 못하면 성능이 불안해질 수 있다는 단점도 함께 존재합니다.

 

이러한 MoE 방식은 모델이 커질수록 효율이 점차 감소한다는 한계가 확인되었고 무조건적인 구조 변경보다는 데이터와 연산량의 효율적인 배분이 필요해진 시점이라고 시사합니다.

 

---

[3] 접근법

1) 최적의 학습 데이터 양 찾기

 

이 접근법에서는 모델 크기를 고정한 상태에서 학습 토큰 수를 바꿔가며 실험을 진행했습니다.

실험에 사용된 모델의 크기를 70M ~ 10B 까지 다양하게 구성하였다고 합니다.

 

실험 결과로 얻은 loss 그래프를 다음과 같이 처리합니다:

1. 모든 곡선을 부드럽게 보간해서 연속적인 형태로 만든다.

2. 특정 연산량마다 가장 낮은 손실을 달성한 실험을 찾는다.

3. 이렇게 얻은 데이터를 바탕으로 연산량에 따라 최적의 모델 크기($N$)와 학습 토큰 수($D$)가 어떻게 변하는지 추정한다.

 

위의 과정을 통해 연산량이 4배 늘어나면 모델 크기와 학습 토큰 수도 각각 2배씩 늘리는 것이 최적이라는 지수 관계를 발견했습니다.

 

2) IsoFLOP 분석으로 최적의 모델 크기 찾기

 

이번 접근법에서는 고정된 연산량 하에서 모델 크기와 학습 토큰 수를 어떻게 배분하는 것이 효율적인가를 알아봅니다.

 

*IsoFLOP이란, Iso(같은) FLOPs(연산량)을 의미하므로 같은 연산량에서 모델 크기별로 어떤 손실값을 가지는지 비교하는 방법이라고 할 수 있습니다.

 

실험에 사용된 FLOPs 범위는 약 $6 * 10^18 ~ 3 * 10^21$ 이고

각 FLOPs 마다 여러 모델 크기 ($N$)를 실험하고 각 경우의 최종 loss를 비교했습니다.

 

실험 결과는 위와 마찬가지로 모델 크기와 학습 토큰 수를 같은 비율로 늘려야 한다는 결과가 도출되었고

연산량이 늘어나면 모델 크기와 학습 데이터를 함께 늘려야 한다는 점을 시사합니다.

 

3) 손실 함수 모델링 후 최적점 찾기

본 논문의 저자들은 이전 3.1, 3.2 실험에서 얻은 데이터를 바탕으로 loss를 다음과 같은 수식으로 표현하였습니다.

이 식은 세 가지 요소로 구성되는데

1. E - 데이터 자체의 불확실성을 나타냄

2. $A/N^{\alpha}$ - 모델이 충분히 크지 않아 생기는 손실

3. $B/D^{\beta}$ - 학습 데이터가 부족해서 발생하는 손실

 

이 방식은 머신러닝의 리스크 분해 (Risk Decomposition) 개념에서 출발한다고 합니다.

 

학습시킬 때에는 다음의 수식을 따랐습니다.

- log loss 값에 대한 Huber loss를 최소화하는 방식으로 학습

- 최적화는 L-BFGS 알고리즘을 사용

- 다양한 초기값으로부터 시작해 지역 최소값 문제를 회피

 

위와 같이 학습된 손실 함수를 바탕으로 정해진 연산량 내에서 가장 낮은 손실을 달성할 수 있는 모델 크기와 토큰 수를 계산했습니다.

$FLOPs ≈ 6ND$ 

이 제약 조건 하에서 손실함수를 최소화하면 최적의 모델 크기와 학습 데이터 양을 구할 수 있다는 결과가 도출됩니다.

 

*Huber Loss : MSE와 MAE를 결합한 손실 함수로 이상치에 강건하다는 장점이 있음

*L-BFGS : 고차원 최적화에 자주 쓰이는 알고리즘

*Efficient Frontier : 주어진 FLOPs 내에서 손실을 최소화할 수 있는 최적점들의 경계

 

4) 최적의 모델 스케일링 전략

위의 3가지 실험을 통해 모델 크기와 학습 데이터를 비슷한 비율로 동시에 늘려야 한다는 결과를 도출하였고

 

아래 표는 세 가지 접근법에서 도출된 결과입니다.

 

특히, 3-3의 손실함수 모델링 부분에서 큰 연산량일수록 더 작은 모델이 효율적일 수 있다는 결론을 내었는데

이는 비현실적으로 큰 모델보다 적당한 크기의 모델을 오래 훈련하는 쪽이 효율적이라는 점을 시사합니다.

 

 

위 표와 같이 다양한 모델 크기에서 최적의 FLOPs과 학습 데이터 양을 추정하였습니다.

 

위 수치는 현재 상용화된 LLM보다 훨씬 많은 데이터가 필요함을 의미하고 있고

예를 들어, 1조 파라미터 모델을 최적으로 훈련하려면 Gopher 대비 250배 이상의 연산량이 필요하며 21조 이상의 학습 데이터 토큰이 있어야 한다고 얘기하고 있습니다.

 

---

[4] ChinChilla

앞선 분석을 통해 ChinChilla 연구진은 Gopher 수준의 연산량으로 가장 효율적인 모델은 40 ~ 70B 파라미터 수준이라고 결론지었습니다.

이를 바탕으로 약 70B 파라미터의 모델을 1.4조 토큰으로 학습해 ChinChilla 모델을 만들었습니다.

 

기본적으로 Gopher와 동일한 트랜스포머 구조를 사용하되 아래의 몇 가지 개선 사항을 적용했습니다.

 

개선사항 1) 학습 데이터

• 사용한 데이터셋 : Gopher와 동일한 MassiveText
• 다른 점 : 더 많은 토큰 수를 처리할 수 있도록 Subset 구성 비율을 조정

개선사항 2) Optimizer

• Gopher는 Adam, ChinChilla는 AdamW
• 언어모델 손실 감소 및 파인튜닝 성능 향상에 기여

개선사항 3) Tokenizer

• Gopher : NFKC 정규화를 적용한 SentencePiece Tokenizer
• ChinChilla : NFKC 정규화 제거 (수학, 화학 기호 표현력이 개선되고 전체 토큰 중 94.15%는 동일)

개선사항 4) 연산 정밀도 및 가중치 저장 방식

• 순전파 / 역전파 : bfloat16
• 옵티마이저 상태 : float32로 별도 저장 (정밀도 유지하면서 속도와 메모리 균형 확보)

 

ChinChilla와 Gopher의 하이퍼파라미터는 다음과 같이 다릅니다:

 

1) The Pile Evaluation

첫 번째 결과는 The Pile 벤치마크입니다.

 

*The Pile은 15개 이상의 다양한 도메인 (논문, 법률, 소설, 수학, IRC 등)을 포함한 데이터셋입니다.

*성능 평가는 bits-per-byte (bpb) 지표를 기준으로 했으며, 수치가 낮을수록 더 좋은 언어 모델을 의미합니다.

 

결과적으로

1. ChinChilla는 The Pile의 모든 서브셋에서 Gopher보다 낮은 bpb, 즉 더 뛰어난 성능을 기록하였습니다.

2. 특히 Gutenberg PG-19, Europarl 등에서 최대 10% 성능 향상이 관측되었습니다.

 

하지만, ChinChilla는 Gopher 보다 4배 많은 학습 데이터를 사용했기 때문에 일부 평가셋에서는 data leakage가 발생했을 수 있어 논문에서는 MMLU, Big-bench, 질문응답 및 상식 추론처럼 데이터 누수가 적은 Task를 더 중요하게 평가했습니다.

 

*Bits-per-byte (bpb) : 언어 모델이 텍스트를 얼마나 예측할 수 있는지를 나타내는 성능 지표

*data leakage : 학습 데이터 일부가 평가 데이터에 포함되어 모델 성능이 과대평가되는 현상

 

2) MMLU

MMLU는 언어모델이 다양한 학문적 주제에서 얼마나 텍스트를 잘 이해하고 추론할 수 있는지를 평가하는 대표적인 벤치마크입니다.

 

57개 과목의 시험 스타일 문제로 구성되어 있으며 모델의 범용 지식과 응용력을 측정합니다.

 

특히 이번에는 모델에게 5개의 예제를 먼저 보여준 뒤 문제를 해결하게 하는 5-shot 방식으로 평가하였습니다.

 

 

결과적으로

1. 총 57개 과목 중 51개의 과목에서 Gopher보다 성능이 향상되었습니다.

2. 특히 고등학교 정치 및 정부, 국제법, 사회학, 미국 외교 정책의 과목에서는 90% 이상의 정확도를 기록하였습니다.

3. 5-shot 기준으로는 Gopher보다 7.6% 높은 정확도를 기록하였으며 전문가 예측 수치 (2023년도 기준)를 초과한 유일한 모델이라고 합니다.

 

3) Gender bias and toxicity

LLM은 뛰어난 성능만큼이나 잠재적인 위험 요소를 내포합니다.

대표적으로 저희도 잘 알고 있는 사회적 편향이나 차별적인 발화, 개인정보 유출 등이 있을 수 있죠.

ChinChilla 역시 Gopher와 동일한 데이터셋을 기반으로 학습되었기 때문에, 유사한 위험을 가질 수 있다고 가정하고 이에 대한 분석이 수행되었습니다.

 

평가에는 Winogender Dataset을 사용하였고 특정 대명사가 직업명을 정확히 지칭하는지를 판단하는 방식으로 평가했습니다.

 

*공정한 모델이라면 남성/여성 대명사에 관계없이 정확하게 직업 지칭 예측이 가능해야 함

 

 

결과적으로 ChinChilla는 Gopher보다 젠더 편향이 덜하지만, 개선 폭은 성별에 따라 달랐습니다.

 

유해 발화 생성 가능성도 평가했는데

 

25,000개의 무작위 샘플을 생성하고 Perspective API로 유해성 점수를 측정하는 방식으로 평가했습니다.

 

결과적으로 유해 발화 생성 정도는 Gopher와 사실상 유사하며 성능이 높아져도 유해성은 크게 줄어들지 않는 결과를 보였습니다.

 

따라서 모델의 성능 향상과 윤리적 안전성은 별개의 문제이며 지속적인 위험 모니터링과 별도 대응이 필요하다는 점을 시사합니다.

 

---

[5] Conclusions

• 등장 배경
  • 최근 몇 년간 LLM의 발전은 모델 크기를 키우는 방향으로 발전되어 왔음
  • 하지만 연산량과 GPU, 메모리 등 자원의 제약이 많아 효율적으로 학습하는 것의 필요성이 증가되고 있던 추세
  • 그래서 ChinChilla 저자들은 같은 연산량 내에서 가장 효율적으로 좋은 성능을 내기 위한 방법을 탐색함
• 핵심 방법론
  • 모델 크기를 고정한 상태에서 최적의 학습 데이터 양 찾기
  • IsoFLOP으로 같은 연산량 내에서 최적의 모델 크기 찾기
  • 손실함수 모델링 후 최적점 찾기
• 실험 결과
  • Gopher와 동일한 연산량을 사용해 70B 크기의 ChinChilla 모델을 1.4조 토큰의 데이터로 학습
  • ChinChilla는 Gopher보다 대부분의 평가 Task에서 우수한 성능을 달성
  • 특히 MMLU, Big-bench, 질문응답 및 상식 추론 등 실활용도가 높은 영역에서 확실한 성능 차이를 보임
• 한계 및 향후 과제
  • 실험적으로 비교 가능한 대규모 모델이 Gopher와 ChinChilla 뿐이었다는 점
  • 데이터 편향과 유해 발화 문제에서 완전히 자유롭지 않으며 더 다양한 데이터를 학습할수록 개인정보 유출, 편향 증폭 등의 윤리적 문제도 함께 커질 수 있음을 지적