기술 요약: LLM 사전 학습 시 도메인 데이터 반복의 스케일링

문제 정의

대규모 언어 모델(LLM)이 확장됨에 따라, 적절한 토큰 대 파라미터 비율(TPP)을 유지하기 위해서는 총 학습 토큰 예산을 늘려야 합니다. 일반적인 웹 데이터는 비교적 쉽게 확장할 수 있는 반면, 고품질의 도메인 특화 데이터(예: 코드, 수학, 의학 텍티브 등)는 희소하며 동일한 속도로 확장하기 어렵습니다. 결과적으로, 모델 크기와 총 토큰 예산이 커짐에 따라 고품질 도메인 데이터의 혼합 비율이 감소하는 경향이 있으며, 이는 비율이 임계값 아래로 떨어질 경우 목표 지식을 습득하는 데 실패할 가능성을 초래합니다.

이러한 희석 현상을 방지하기 위해 실무자들은 가용 가능한 고품질 데이터를 반복해서 사용하곤 합니다. 그러나 이는 트레이드오프를 동반합니다. 데이터 반복은 지식 습득을 위한 필수적인 데이터 분율을 유지하는 데 도움이 되지만, 과적합(overfitting)의 위험을 높입니다. 기존 문헌은 더 큰 모델이 반복된 데이터에 의한 과적합에 더 취약하다는 점을 시사하며, 이는 반복을 최소화해야 함을 의미합니다. 하지만 이러한 결론은 총 학습 데이터의 크기가 모델 규모에 관계없이 고정된 실험 설정 하에서 도출된 경우가 많습니다. 본 논문은 학습 토큰 예산이 모델 크기에 비례하여 성장하는(고정된 TPP) 보다 실질적인 스케일링 체제 하에서 최적의 반복 횟수를 조사합니다.

방법론

저자들은 네 가지 고품질 도메인(Code, Mathematics, Wikipedia, Medical data)에 대해 체계적인 경험적 연구를 수행했습니다.

• 스케일링 체제: 본 연구는 $D$를 총 학습 토큰, $N$을 모델 크기라고 할 때, 고정된 TPP($D/N$)를 준수합니다. 토큰 예산 $D_N$은 모델 크기에 따라 선형적으로 확장됩니다 ($D_N = \text{TPP} \cdot N$).
• 실험 설정: 각 모델 크기와 도메인에 대해 저자들은 두 가지 파라미터를 변화시켰습니다:
  1. 고유 데이터 분율 ($\alpha$): 총 예산 내에서 고유한 고품질 토큰의 비율 (1/40, 1/20, 1/10로 테스트).
  2. 반복 횟수 ($e$): 고정된 고유 고품질 데이터 서브셋을 재방문하는 횟수 (1에서 7까지 테스트).
• 평가: 모델은 인도메인(in-distribution) 검증 손실(동일 도메인의 홀드아웃 데이터 사용)과 OOD(out-of-distribution) 사전 학습 검증 손실(ArXiv 및 News와 같은 일반 웹 데이터 사용)을 통해 평가되었습니다.
• 이론적 분석: 경험적 발견을 설명하기 위해, 저자들은 멱법칙(power-law) 분포를 따르는 지식 단위(knowledge units)를 가진 $\mathbb{R}^\infty$ 상의 원-핫 선형 회귀 모델을 채택했습니다. 이들은 리스크를 "지식 습득 오차(knowledge-acquisition error)"와 "노이즈 피팅 오차(noise-fitting error)"로 분해하여, 다양한 스케일링 조건 하에서의 정지 시점(최적 반복 횟수)을 도출했습니다.

주요 결과

1. 검증 손실과의 상관관계

최적의 반복 횟수는 도메인의 최종 검증 손실과 강한 음의 상관관계를 보입니다 (피어슨 상관계수 -0.944).

• 낮은 검증 손실을 달성하는 도메인(모델이 해당 도메인을 잘 학습함을 의미)은 더 많은 반복을 견디고 그로부터 이득을 얻을 수 있습니다.
• 높은 검증 손실을 보이는 도메인은 더 빨리 과적합되는 경향이 있어, 더 적은 반복이 필요합니다.
• 이는 검증 손실이 안전한 반복 한계를 결정하는 실용적인 프록시(proxy) 역할을 할 수 있음을 시사합니다.

2. 고정 TPP 하에서의 모델 크기 영향

데이터 예산이 고정된 설정(모델이 커짐에 따라 최적 반복 횟수가 감소함)의 결과와는 대조적으로, 본 연구는 고정된 TPP 하에서 최적 반복 횟수가 모델 크기에 따라 완만하게 증가한다는 것을 발견했습니다.

• 모델 크기에 비례하여 더 큰 토큰 예산으로 학습된 더 큰 모델은 과적합이 발생하기 전까지 약간 더 많은 반복으로부터 이득을 얻습니다.
• 이는 보수적인 전이 전략을 뒷받침합니다: 동일한 TPP로 학습된 더 작은 프록시 모델에서 과적합을 일으키지 않는 반복 횟수는 더 큰 타겟 모델에서도 과적합을 일으킬 가능성이 낮습니다.

3. 고유 데이터 분율에 대한 불감성

최적 반복 횟수는 테스트된 범위 내의 고유 고품질 데이터 분율($\alpha$)에 거의 영향을 받지 않습니다.

• $\alpha$를 높이면 절대적인 검증 손실은 일관되게 낮아지지만, 최적 반복 횟수의 위치(손실을 최소화하는 $e$)는 거의 변하지 않습니다.
• 이는 프록시 모델에서 반복을 튜닝할 때, 고유 데이터의 특정 분율이 타겟 모델과 정확히 일치할 필요는 없으며, 테스트된 범위 내의 대표적인 분율이면 충분함을 의미합니다.

4. 도메인 의존성 및 OOD 효과

• 도메인 특이성: 반복의 효과는 도메인에 따라 크게 다릅니다. 예를 들어, 수학(Mathematics) 데이터는 반복에 강건하지만(4~5회까지), 위키피디아(Wikipedia) 데이터는 2회 반복 후 급격히 저하됩니다.
• OOD 성능: 고품질 토큰의 총 분율이 고정되어 있을 때, 고유 토큰을 반복된 토큰으로 대체하는 것은 OOD 성능에 제한적인 영향만을 미칩니다. 저하는 주로 인도메인 검증 손실에 국한됩니다.

5. 학습률 스케줄의 영향

최적 반복 횟수는 학습률 스케줄에 따라 달라집니다. 조기 감쇠가 일어나는 스케줄(예: 학습 시작 시점 0%에서 감쇠 시작)은 더 적은 반복 후에 성능 저하를 유발합니다. 반대로, 감쇠를 늦추거나 일정한 학습률을 사용하는 것은 모델이 더 많은 반복을 견딜 수 있게 하는데, 이는 아마도 강력한 최적화 노이즈가 암기(memorization)에 대한 암묵적 규제(implicit regularization)로서 작용하기 때문일 것입니다.

이론적 설명

저자들은 지식 습득과 노이즈 피팅 사이의 트레이드오프를 통해 이러한 관찰 결과들을 설명하는 이론적 프레임워크를 제공합니다:

• 노이즈 피팅: 반복된 최적화는 지식 습득 오차를 줄이지만 노이즈 피팅 오차를 증가시킵니다. 최적의 정지 시점은 학습 신호의 한계 이득이 노이즈 피팅의 비용에 의해 압도되는 지점입니다.
• 검증 손실: 낮은 검증 손실은 신호 대비 낮은 노이즈 수준($\sigma^2$)을 의미합니다. 이론적으로 낮은 노이즈는 과적합의 발생을 늦추어 더 많은 반복을 가능하게 합니다 (정리 4.2).
• 스케일링: 고정 TPP 하에서 데이터 예산은 모델 크기에 따라 증가합니다. 이는 모델이 커짐에 따라 지식 습득 오차의 감소가 더 오랜 기간 동안 우세하게 작용하도록 하여, 최적 반복 횟수의 증가를 정당화합니다 (정리 4.4).

의의 및 실무적 시사점

본 논문은 타겟 모델과 동일한 TPP를 공유하는 작은 프록시 모델에서 튜닝된 반복 횟수가 더 큰 모델에 대한 실용적이고 보수적인 추정치로 사용될 수 있다고 결론짓습니다.

고품질 도메인 데이터를 대규모 사전 학습에 구성하기 위해 제안된 워크플로우는 다음과 같습니다:

1. 동일한 TPP를 가진 더 작은 프록시 모델을 학습시킨다.
2. 어떤 대표적인 분율의 고유 데이터를 사용하여 타겟 도메인에 대한 반복 횟수를 스윕(sweep)한다.
3. 검증 손실을 최소화하거나(또는 저하 직전에 멈추는) 반복 횟수를 식별한다.
4. 이 횟수를 더 큰 타겟 모델에 적용한다. 이는 동일한 스케일링 체제 하에서 더 큰 모델이 프록시보다 더 빨리 과적합되지 않을 것이라는 추세를 따릅니다.

이 접근 방식은 종종 가용하기 어려운 방대한 양의 고유 데이터 없이도, LLM 사전 학습에서 고품질 도메인 표현을 유지하기 위한 데이터 효율적인 전략을 제공합니다.