Abstract

대규모 언어 모델(LLM)을 위한 사전학습 데이터셋은 수조 개의 토큰 규모로 성장했으며, 대량의 CommonCrawl(CC) 웹 스크랩 데이터와 더불어 소규모의 도메인 특화 데이터셋들로 구성된다. 이러한 도메인 특화 데이터셋이 모델 성능에 미치는 영향을 이해하는 것은 비용이 많이 든다. 이는 난이도가 높고 emergent한 벤치마크에서 의미 있는 변화를 드러내기 위해 대규모 FLOP 수준의 학습이 필요하기 때문이다. 사전학습 데이터 실험 비용이 점점 증가하는 상황에서, 일반 웹 스크랩 데이터의 다양성과 도메인 특화 데이터의 정보 밀도 사이의 최적 균형은 어떻게 결정할 수 있을까? 본 연구에서는 학습 마지막 단계에서 도메인 특화 데이터셋을 CC 대비 업샘플링함으로써, 이러한 소규모 도메인 데이터셋을 활용해 어려운 벤치마크 성능을 향상시키는 방법을 제시한다. 이 단순한 기법을 통해, 1조(1T) 토큰으로 학습된 7B 모델에서 기본 데이터 혼합(base data mix) 대비 MMLU는 최대 6.90 percentage point(pp), GSM8K는 8.26 pp, HumanEval은 6.17 pp 향상시킬 수 있었으며, 이는 두 배 더 긴 학습을 수행한 Llama-2 (7B) 모델에 필적하는 성능이다. 또한 우리는 도메인 업샘플링 기간을 전체 학습의 5%에서 30%까지 변화시키는 실험을 수행하였고, 일반 언어 모델링 능력과 특정 벤치마크 성능 간의 trade-off를 조절하는 데에는 10%~20% 구간이 최적임을 발견했다. 더 나아가, 학습 마지막 단계에서 특정 데이터셋을 제거하는 방식의 도메인 업샘플링을 사용하여, 개별 데이터셋이 다양한 벤치마크 성능 향상에 얼마나 기여하는지를 대규모 환경에서 분석하였다. 이 방법은 전체 사전학습을 반복 수행하는 것에 비해 한 자릿수 이상 낮은 비용으로, 서로 다른 사전학습 데이터셋의 영향을 대규모로 실험할 수 있는 가능성을 제공한다.

1. Introduction

대규모 언어 모델(LLM)을 위한 사전학습 데이터셋은 Dolma(Soldaini et al., 2023)와 같이 수조 개의 토큰 규모로 성장해왔다. 이러한 대규모를 수용하기 위해, 데이터셋은 일반적으로 두 가지 유형의 데이터 소스로 구성된다. 첫째, CommonCrawl(CC) 덤프에서 처리된 대규모 웹 스크랩 데이터를 포함한다. 이러한 데이터는 보통 수천억에서 수조 개 토큰 규모이며, 매우 다양한 정보 분포를 담고 있다. 그러나 규모가 매우 크기 때문에 정보 밀도가 낮고 필터링도 상대적으로 덜 되어 있다. 둘째, 특정 도메인을 대상으로 하거나 단일 고품질 출처에서 수집된 데이터셋들도 포함된다. 이러한 데이터셋은 훨씬 작으며(대개 수백억 토큰 이하), 보다 정교하게 처리되어 있고 LLM이 잘 수행하길 원하는 도메인의 정보가 고밀도로 포함되어 있다. 하지만 데이터 출처가 제한적이기 때문에 다양성은 상대적으로 부족한 경우가 많다(Computer, 2023).

관련 연구(Related Works).
LLM 사전학습에서 가장 큰 과제 중 하나는 CC 기반 데이터와 소규모 도메인 특화 데이터 간의 최적 혼합 전략을 결정하는 것이다.

일부 선행 연구는 강하게 정제된 CC 데이터만으로 사전학습을 수행했다(Penedo et al., 2023).
또 다른 연구들은 CC 데이터와 보다 도메인 특화된 데이터셋 사이의 균형을 맞추기 위해 다양한 휴리스틱을 사용하였다(Computer, 2023).

그러나 최근 대규모로 학습된 대부분의 언어 모델들은 사전학습 데이터 구성에 대한 정보를 제한적으로만 공개하고 있다(Touvron et al., 2023; Jiang et al., 2023; 2024; Team et al., 2024). 소규모 실험 환경에서는 데이터 혼합 비율을 알고리즘적으로 최적화하려는 시도들이 있었지만, 이러한 방법들이 현대 언어 모델 학습 규모에서 공개적으로 검증된 적은 거의 없다(Xie et al., 2024). 이처럼 대규모 환경에서 데이터 혼합 전략을 검증하는 비용이 막대하기 때문에, LLM 사전학습 데이터에 대한 공개 연구는 매우 부족한 상황이다.

이상적으로는 더 작은 규모에서 데이터 혼합 실험을 수행하여 좋은 데이터 조합을 찾아야 한다. 그러나 이는 종종 효과적이지 않다. 왜냐하면 난이도가 높고 emergent한 벤치마크에서 의미 있는 변화를 관찰하려면 대규모 FLOP 수준의 학습이 필요하기 때문이다. 실제로 작은 규모로 학습된 대부분의 LLM은 MMLU(Wei et al., 2022)와 같은 중요한 벤치마크에서 무작위 수준의 정확도만 보이는 경우가 많다. 따라서 작은 규모에서의 실험은 오해를 불러일으킬 수 있다. 이 규모에서는 중요한 벤치마크에서 데이터 혼합 간의 성능 차이가 데이터셋 품질이 아니라 노이즈 때문인 경우가 많기 때문이다. 반면, 이러한 벤치마크에서 무작위 수준 이상의 성능을 측정할 수 있는 충분한 규모로 여러 학습 실험을 반복 수행하는 것은 비용적으로나 현실적으로 매우 어렵다.

본 연구의 목표는, 합리적인 규모에서 사전학습 데이터 실험을 수행하기 위한 대안적 접근법의 효용을 분석하는 것이다. 우리의 전략은 어려운 벤치마크에서 의미 있는 신호를 측정할 수 있을 정도로 충분한 FLOPs를 사용해 학습을 진행한 뒤, 학습 마지막 단계에서 데이터 혼합 비율을 수정하는 것이다. 우리는 이러한 전략이 전체 학습을 반복 수행하는 것보다 한 자릿수 이상 저렴한 비용으로도 LLM 사전학습 데이터 혼합을 효과적으로 개선할 수 있음을 보인다.

Contributions

공개적으로 사용 가능한 데이터셋들로 구성된 기본 데이터 혼합(base mix)을 구축하였으며, 이를 통해 1조 토큰으로 학습된 7B 모델에서 Llama-2 모델 계열과 동일한 FLOP 대비 성능 스케일링을 달성하였다.
우리는 domain upsampling을 제안한다. 이는 학습 마지막 단계에서 Common Crawl 대비 도메인 특화 데이터셋의 비중을 증가시키는 데이터 개입 기법이며, 이를 통해 어려운 벤치마크 성능을 향상시킬 수 있음을 보였다. 구체적으로, 우리의 학습 설정에서 기본 데이터 혼합 대비 MMLU는 최대 6.90 pp, GSM8K는 8.26 pp, HumanEval은 6.17 pp 향상되었다. 이로써 전체 학습 FLOPs는 절반 수준이지만 성능은 Llama-2 (7B)에 근접하였다.
전체 학습 중 domain upsampling을 적용하는 비율을 변화시키는 실험을 수행하였으며, 일반 언어 모델링 능력과 특정 벤치마크 성능 간의 trade-off를 조절하는 데에는 10%~20% 구간이 최적임을 보였다.
또한 domain upsampling을 이용해 개별 데이터셋이 모델 성능에 어떤 영향을 미치는지를 FLOP 효율적으로 분석할 수 있음을 보였다. 학습 마지막 단계에서 업샘플링되는 데이터 중 수학 중심 데이터셋 일부를 제거함으로써, 해당 데이터셋들이 특정 벤치마크에 미치는 영향을 정량화하였다.

2. Training Details

우리는 1조(1T) 토큰으로 학습된 70억(7B) 파라미터 모델에서 domain upsampling을 연구하였다. 이러한 FLOP 규모를 선택한 이유는, 모델이 MMLU와 같은 핵심 벤치마크에서 노이즈 수준을 넘어서는 성능을 보이게 하여, 데이터 개입(data intervention)이 모델에 미치는 영향을 관찰할 수 있도록 하기 위함이다.

본 연구에서 사용한 7B 모델은 LLM Foundry(MosaicML et al., 2023)의 MPT 아키텍처를 사용하는 decoder-only Transformer이다. 모델 평가에는 최신 버전인 Eval Gauntlet v0.3(MosaicML et al., 2023)을 사용하였다. 이는 LLM base model 평가를 위해 널리 사용되는 35개의 in-context learning 평가 태스크로 구성된 프레임워크이다. Gauntlet v0.3은 6개의 카테고리에 걸친 벤치마크 점수를 집계하며, 자세한 내용은 Appendix A에 설명되어 있다. 또한 우리는 (Zhai et al., 2022)와 유사한 inverse square root learning schedule을 사용하였다.

3. Results

여기서는 domain upsampling을 통해 얻어진 성능 향상과, 데이터셋이 난이도가 높고 emergent한 지표에 어떤 영향을 미치는지를 분석하는 데 있어 이 기법이 얼마나 유용한지를 보여주는 실험 결과를 제시한다.

3.1 Baseline data mix achieves Llama-2 scaling

기본 데이터 혼합(baseline data mix)을 구성하기 위해, 우리는 공개적으로 이용 가능한 데이터셋들을 다음과 같은 네 가지 범주로 그룹화하였다.

Large-Scale Common Crawl:
규모를 중시하는 Common Crawl 기반 데이터셋들이다. 이러한 데이터셋은 엄격한 품질 필터링보다는, 크고 다양한 토큰 집합을 구성하는 데 초점을 둔다.
Small-Scale Common Crawl:
보다 광범위한 필터링을 거친 Common Crawl 기반 데이터셋이지만, large-scale Common Crawl보다 규모는 작다.
Domain Specific data:
특정 도메인을 대상으로 하거나 개별 출처에서 수집된 소규모 고품질 데이터셋들이다(예: Wikipedia).
Code:
다양한 프로그래밍 언어로 구성된 코드 데이터.

우리는 이러한 데이터셋 그룹들이 1조 토큰 학습 동안 몇 epoch 정도 반복될지를 기준으로 하는 대략적인 휴리스틱에 따라 혼합 비율을 설정하였다. 구체적으로는 Small-Scale Common Crawl과 Domain Specific 데이터는 각각 0.5 epoch, Code 데이터는 1 epoch가 되도록 설정하였다. 이후 남은 토큰은 Large-Scale Common Crawl 데이터로 채웠다. 정확한 비율은 Table 2에 제시되어 있다.

이러한 비율을 선택한 이유는 다음과 같다. Small-Scale Common Crawl과 Domain Specific 데이터는 품질이 높다고 예상되므로, 1조 토큰 예산 내에서 충분히 반영되기를 원했다. 또한 코딩 능력을 강조하고자 했기 때문에 코드 데이터를 높은 비율로 샘플링하였다. 초기 실험에서는 약 20% 수준의 높은 코드 데이터 비율이 언어 능력을 저해하지 않으면서 프로그래밍 및 추론 능력을 향상시킨다는 결과를 얻었다. 이후 Large-Scale CC 데이터는 데이터셋의 다양성을 높이고 전체 토큰 예산을 채우는 용도로 사용하였다.

중요한 점은, 본 실험 설정의 목적이 학습 마지막 단계에서의 domain upsampling(Section 3.2에서 논의)을 보여주는 데 있기 때문에, 초기 데이터 혼합 비율을 지나치게 최적화하기보다는 합리적인 휴리스틱을 사용하는 데 초점을 두었다는 것이다.

Table 3과 Figure 1은 이러한 초기 데이터 혼합으로 각각 0.5T 및 1T 토큰 동안 학습한 두 개의 7B 모델 성능을 보여준다.

여기서 Us가 여기서 제안한 모델임

Error 대 FLOPs 관계를 그려보면, 우리의 모델들은 Gauntlet v0.3 Core Average, MMLU, GSM8K, HumanEval에서 Llama-2 scaling line 위 또는 아래에 위치한다. 이는 이러한 휴리스틱이 실제로 효과적이었음을 보여준다.

저자들이 만든 데이터 믹스로 학습한 7B 모델이, FLOPs 대비 성능 측면에서 Llama-2 계열과 동등하거나 더 좋은 scaling behavior를 보인다.
그림을 보면 Llama-2의 사이즈에 따른 FLOPS을 계산하고 벤치마크 error을 y축으로 해서 그림을 그린다.
여기서 가정은 선형으로 FLOPS와 err가 비례한다는것? (log-log 공간에서)
어쨌든 그 가정보다 제안한 모델(Us)가 같은 FLOPS에 비해 error가 낮을것으로 예측되는 그래프가 그려진다.
따라서 이렇게 데이터 혼합해서 학습하는게 효율적이라는 것을 보여줌 (pretraining stage1 이라 보면 될 듯)

흥미롭게도, 전체 성능 스케일링(Gauntlet v0.3 Core Average 기준)은 매우 유사하지만, 데이터 선택과 혼합 계수의 차이로 인해 서로 다른 trade-off가 나타났다. 1T 토큰으로 학습한 모델은 2T 토큰으로 학습된 Llama-2 7B 모델보다 GSM8K와 HumanEval에서 더 높은 성능을 보였다. 이는 우리의 모델이 절반 수준의 토큰만 사용하고도 더 뛰어난 수학 및 프로그래밍 능력을 갖추었음을 의미한다. 또한 우리는 데이터 혼합 비율 일부를 공개한 7B 모델인 OpenLlama 7Bv2(Geng & Liu, 2023)와의 비교도 제공한다.