RegMix: Data Mixture as Regression for Language Model Pre-Training

문제의 배경 — 기존 연구의 한계

대형 언어 모델(LLM)은 인터넷에서 수집한 방대한 텍스트 데이터를 기반으로 학습된다. 이 데이터는 학술 논문(arXiv), 책(Project Gutenberg), 코드(GitHub), 웹페이지(CommonCrawl) 등 다양한 도메인(domain)으로 구성되어 있다. 그런데 문제가 있다: 각 도메인을 얼마나 섞어서 학습해야 하는가?

GPT-3를 만들 때부터 이 질문은 중요했다. OpenAI 연구자들은 Wikipedia가 '고품질'이라는 직관에 따라 Wikipedia를 업샘플링(upsample)했다. 그러나 이런 방식은 확장성이 없다 — 도메인이 수십, 수백 개가 되면 인간의 직관만으로 최적 비율을 결정하기 어렵고, 직관적으로 '좋아 보이는' 데이터가 반드시 더 나은 모델을 만들지 않는다.

기존 접근법과 그 한계

데이터 선택 방법은 세 가지 수준으로 구분된다:

"Training a proxy model for current models, such as Llama-3, would require using up to 15T tokens, which is likely too expensive and too slow to make it worthwhile."

현재 Llama-3 수준의 모델에 기존 프록시 방식을 적용하려면 최대 15T 토큰이 필요해, 실용적이지 않다.

그렇다면 어떻게 해야 할까? 논문은 이 문제를 완전히 다른 방향에서 접근한다: "많은 소형 모델을 병렬로 훈련해서, 데이터 혼합이 성능에 미치는 영향을 회귀 모델로 예측하자."

방법론

RegMix는 4단계 파이프라인으로 구성된다:

① 소형 프록시 모델 훈련 — 다양한 혼합 생성

첫 번째 과제는 회귀 모델 학습을 위한 훈련 데이터 생성이다. 단순히 균일하게 혼합을 샘플링하면 특정 극단적 조합을 놓칠 수 있다. RegMix는 Dirichlet 분포를 사용한다:

② 회귀 모델 피팅

N개 프록시 모델의 훈련 결과 {(혼합 비율 벡터, 검증 손실)} 쌍을 가지고 회귀 모델을 학습한다.

선형 회귀 (Linear Regression)

LightGBM 회귀

LightGBM은 결정 트리 앙상블을 그래디언트 부스팅으로 학습한다. 비선형적인 도메인 상호작용까지 포착할 수 있어 선형 회귀보다 훨씬 높은 예측 정확도를 달성한다. 하이퍼파라미터: 1000 반복, 학습률 1e-2.

③ 시뮬레이션 & 예측 — 가장 효율적인 단계

회귀 모델이 학습되면, 이를 이용해 방대한 혼합 공간을 탐색한다:

"Running prediction for 1,000,000 data mixtures takes less than 10 CPU seconds."

100만 개 데이터 혼합에 대한 예측이 CPU 10초도 안 걸린다.

그 중 가장 좋은 예측값을 가진 상위 100개 혼합을 평균 내어 최종 혼합 비율을 결정한다. 상위 1개만 쓰지 않고 100개를 평균내는 이유는 예측 불확실성에 대한 앙상블(ensemble) 효과로 더 안정적인 결과를 얻기 위해서다.

RegMix 알고리즘 — 스텝별 추적

수식 & 알고리즘 요약 (Algorithm 1)

논문 Appendix H에 수록된 의사코드를 풀어서 설명한다:

-- 입력 --
x₀: n개 도메인의 기본 토큰 분포 벡터
N: 프록시 모델 수 (실험에서 N=512)
f: 회귀 모델 (LightGBM)
y: 타깃 메트릭 (Pile-CC 검증 손실)

-- 1단계: 프록시 모델 훈련 --
for i = 1 to N:
  λ ← Uniform([0.1, 5.0])
  xᵢ ← Dirichlet(α = λ · x₀)      // 다양한 혼합 샘플링
  소형 모델을 xᵢ 혼합으로 T 토큰 훈련   // 병렬 가능
  yᵢ ← 검증 손실 측정

-- 2단계: 회귀 모델 피팅 --
f ← Fit({(x₁,y₁), ..., (xₙ,yₙ)})   // LightGBM 학습

-- 3단계: 시뮬레이션 & 예측 --
{x₁,...,x_M} ← 후보 혼합 생성 (M=1,000,000)
x* ← argmin_{xⱼ} f(xⱼ)             // 상위 100개 평균

-- 4단계: 대규모 모델 훈련 --
대규모 모델을 x* 혼합으로 훈련 (1000× 파라미터, 25× 토큰)
return x*

구현 세부사항

모델 아키텍처는 TinyLlama를 기반으로 하며, 레이어 수·헤드 수·임베딩 차원을 조절해 1M, 60M, 1B, 7B 변형을 만든다.

모델	레이어 수	헤드 수	임베딩 차원	FFN 차원	어휘 크기
1M	2	8	256	512	50,432
60M	10	8	768	1,536	50,432
1B	22	16	2,048	5,632	50,432
7B	32	16	4,096	12,288	50,432

공통 훈련 설정:

항목	값
토크나이저	GPTNeoX (어휘 크기 50,432)
배치 크기	1M 토큰
학습률	4e-4
스케줄러	코사인(cosine) 학습률 감소
1M 프록시 모델	1,000 스텝 × 1M 배치 = 1B 토큰
1B 최종 모델	25,000 스텝 × 1M 배치 = 25B 토큰
데이터셋	The Pile (저작권 미문제 17개 서브셋)
회귀: Ridge λ	5-fold CV, {1e-3, 1e-2, 1e-1, 1, 10, 100, 1000}
회귀: LightGBM	1000 반복, 학습률 1e-2, 나머지는 기본값
코드	github.com/sail-sg/regmix

부록 추가 Appendix E (구현 세부사항): 모델 아키텍처는 TinyLlama(Zhang et al., 2024a) 기반이며, GPTNeoX 토크나이저(어휘 크기 50,432)를 사용한다. 선형 회귀에 5-fold 교차 검증을 적용해 L2 정규화 가중치를 {1e-3, 1e-2, ..., 1e3} 중에서 선택한다. LightGBM은 1000 반복, 학습률 1e-2, 나머지 하이퍼파라미터는 기본값.

부록 추가 Appendix F (안정성 분석): DoReMi는 프록시 모델 크기가 달라지면 결과가 불안정하다 — 280M 모델 사용 시 Pile-CC 가중치 0.67, 1B 모델 사용 시 0.20 미만으로 급변한다. RegMix는 1M과 60M 프록시 모델에서 거의 동일한 분포를 도출해 높은 안정성을 보였다 (Figure 14). 이는 여러 모델의 집합적 통계를 사용하기 때문이다.

평가 논문은 FP16/BF16 혼합 정밀도 사용 여부, GPU 수, Flash Attention 적용 여부를 명시하지 않는다. TinyLlama 아키텍처를 기반으로 한다는 점에서 표준 transformer 훈련 설정을 따를 것으로 추정된다. 파라미터 수는 비임베딩 파라미터만 카운트한다 (소형 모델에서 임베딩이 차지하는 비율이 과도하게 크기 때문).

데이터셋: The Pile

실험에 사용된 The Pile 데이터셋의 17개 이용 가능 도메인 (저작권 문제가 있는 Books3, BookCorpus2 등은 제외):

도메인	크기 (GiB)	도메인	크기 (GiB)
Pile-CC (CommonCrawl)	227.12	OpenSubtitles	19.47
PubMed Central	180.55	Wikipedia (en)	19.13
OpenWebText2	125.54	DM Mathematics	15.49
ArXiv	112.42	Ubuntu IRC	11.03
Github	95.16	EuroParl	9.17
FreeLaw	76.73	HackerNews	7.80
Stack Exchange	64.39	YoutubeSubtitles	7.47
USPTO Backgrounds	45.81	PhilPapers	4.76
PubMed Abstracts	38.53	NIH ExPorter	3.79

한계점 & 트레이드오프

Q&A — 연구자의 고민과 독자의 질문