기술 요약: 강력한 교사 모델이 필요한가? LLM 사전 학습에서의 증류에 관하여

문제 제기

대형 언어 모델 (LLM) 사전 학습에서의 지식 증류 (KD) 는 오랫동안 '강한 모델에서 약한 모델로'의 관계가 필수적이라는 가정 하에 운영되어 왔습니다. 즉, 성능 향상을 얻기 위해서는 교사 모델이 더 큰 아키텍처나 더 많은 학습 토큰을 통해 학생 모델보다 훨씬 더 유능해야 한다는 것입니다. 본 논문은 이러한 가정이 완전하지 않다고 주장하며, 약한 교사 모델이 더 강한 학생 모델을 개선할 수 있는지, 그리고 교사 모델을 지속적으로 강화 (파라미터 또는 토큰 확장) 하는 것이 학생 모델의 성능을 단조롭게 향상시키는지 여부를 조사합니다. 저자들은 지배적인 관점이 교사 - 학생 간의 호환성과 전달되는 지식의 구체적인 본질이라는 결정적인 역할을 간과하고 있다고 논합니다.

방법론

저자들은 다양한 차원에 걸쳐 교사 및 학생 구성을 변화시키며 체계적인 실증 연구를 수행했습니다:

• 아키텍처: Llama3 계열을 기반으로 한 네 가지 모델 크기 (0.7B, 1.7B, 3.8B, 8.0B 파라미터).
• 학습 예산: 교사 모델은 10B 에서 300B 까지의 토큰 예산으로 사전 학습됩니다. 학생 모델은 일관되게 50B 토큰으로 학습됩니다.
• 관계: 이 설정은 다양한 교사 - 학생 관계를 생성합니다:
  • 강한 모델에서 약한 모델로: 교사가 학생보다 크거나 더 많이 학습된 경우.
  • 동일 수준: 교사와 학생이 아키텍처와 토큰 예산을 공유함 (무작위 시드만 다름).
  • 약한 모델에서 강한 모델로: 교사가 학생보다 작거나 더 적게 학습된 경우.
• 학습 목적: 학생 모델은 표준 자기회귀 언어 모델링 (LM) 과 지식 증류 (KD) 를 결합한 혼합 손실 함수를 사용하여 학습됩니다:
  $$\mathcal{L} = (1 - \alpha)\mathcal{L}_{lm} + \alpha\mathcal{L}_{kd}$$
  저자들은 각 구성에 대해 최적의 손실 혼합 계수 $\alpha \in \{0.2, 0.4, 0.5, 0.6, 0.8, 1.0\}$를 체계적으로 탐색합니다.
• 평가: 성능은 세 가지 차원에서 측정됩니다:
  1. 도메인 내 퍼플렉시티 (In-domain Perplexity): 학습 분포 (FineWeb-Edu) 에 대한 적합도.
  2. 도메인 외 (OOD) 퍼플렉시티: 11 가지 다양한 코퍼스 (예: 수학, 코드, 법률, 바이오의학) 에 대한 일반화 능력.
  3. 하류 작업 정확도: 추론, 상식, 사실적 지식을 다루는 15 개 벤치마크에서의 성능.

주요 발견

1. 약한 모델에서 강한 모델로 및 동일 수준 증류의 유효성

강한 모델에서 약한 모델로라는 교조와 달리, 본 연구는 손실 혼합 계수 $\alpha$가 적절하게 조정될 때 (일반적으로 낮은 $\alpha$), 학생보다 약한 교사 모델 (더 작은 아키텍처 또는 더 적은 학습 토큰) 이 학생 모델의 성능을 향상시킬 수 있음을 발견했습니다.

• 30B 토큰으로 학습된 0.7B 교사 모델조차 1.7B 학생 모델의 하류 작업 정확도를 1.3% 향상시켰습니다.
• 동일 수준 증류 (동일한 아키텍처와 토큰 예산, 다른 시드) 는 모든 지표에서 일관되게 긍정적인 향상을 보여주어, 증류가 아키텍처와 데이터 양만으로는 포착되지 않는 지식을 전달함을 시사합니다.

2. 더 강력한 교사 모델이 항상 더 나은 학생 모델을 만드는 것은 아님

교사 모델의 강도와 학생 모델의 성능 간의 관계는 단조롭지 않습니다.

• 포화 및 역전: 높은 토큰 예산에서 더 큰 교사 모델은 더 작은 교사 모델보다 성능이 떨어질 수 있습니다. 예를 들어, 300B 토큰의 경우 1.7B 교사 모델은 +4.3% 의 하류 작업 정확도 향상을 보인 반면, 8.0B 교사 모델은 +2.8% 만 보였습니다.
• 과학습 (Overtraining): 큰 교사 모델 (예: 8.0B) 의 경우 특정 지점 (예: 30B 에서 300B 로) 을 넘어 학습 토큰을 늘리면 증류 이득이 저하됩니다. 이러한 과학습된 교사 모델에 대한 최적의 $\alpha$도 감소하여, 성능 저하를 피하기 위해 학생 모델이 교사 모델의 신호에 덜 의존해야 함을 나타냅니다.
• 호환성: 결과는 절대적인 교사 모델의 규모보다 '호환성' (학생 모델의 용량 및 학습 체계와 일치) 이 더 중요함을 시사합니다.

3. 일반화가 도메인 내 적합보다 더 쉽게 향상됨

증류 사전 학습은 평가 유형별로 뚜렷한 패턴을 보입니다:

• 일반화: 도메인 외 퍼플렉시티와 하류 작업 정확도는 도메인 내 퍼플렉시티보다 더 크고 일관된 향상을 보입니다. 일부 구성은 도메인 내 성능을 저하시키더라도 도메인 외 및 하류 지표는 향상시킵니다.
• 메커니즘: 이는 증류가 단순히 학습 분포에 적합시키는 것보다는 일반화에 유용한 지식을 전달함을 시사합니다.
• 토큰 난이도 분석: 기계적 연구는 증류가 쉬운 토큰보다 '어려운 토큰' (높은 엔트로피) 에 훨씬 더 큰 이점을 준다는 것을 밝혔습니다. 이는 KD 가 단순한 균일 정규화 (라벨 스무딩과 같은) 로 작용한다는 가설을 반박하며, 대신 기준 모델이 가장 불확실한 영역에서 지식을 표적적으로 전달한다는 것을 보여줍니다.

4. 손실 혼합 ($\alpha$) 의 역할

최적의 $\alpha$는 교사 모델의 강도에 크게 의존합니다:

• 약한 교사 모델: 신뢰할 수 없는 교사 신호와 실제 지시 (ground-truth supervision) 를 균형 있게 맞추기 위해 낮은 $\alpha$(예: 0.2–0.4) 가 필요합니다. 약한 교사 모델로 순수 증류 ($\alpha=1.0$) 를 수행하면 성능이 크게 저하됩니다.
• 강한 교사 모델: 더 높은 $\alpha$(최대 1.0) 를 선호하지만, 과학습된 경우에도 강한 교사 모델조차 체감 수익이 감소합니다.
• 비단조성: 막대한 토큰 예산으로 학습된 매우 큰 교사 모델의 경우 최적의 $\alpha$가 감소할 수 있으며, 이는 "더 강력함"이 항상 "더 좋음"은 아니라는 발견을 더욱 뒷받침합니다.

의의 및 주장

본 논문은 증류를 통한 LLM 사전 학습이 강력한 교사 모델을 엄격하게 요구한다는 일반적인 신념에 도전한다고 주장합니다. 그 의의는 다음과 같습니다:

1. 교사 모델 선정의 재정의: 최첨단 모델이 이전 세대 모델 (더 작거나 덜 학습된 모델일 수 있음) 로 훈련되는 시나리오에서 증류는 여전히 유효하고 유익한 전략임을 보여줍니다.
2. 효율성 함의: 점점 더 크거나 더 많이 학습된 교사 모델에 투자하는 것이 체감 수익을 내거나 심지어 부정적인 수익을 낼 수 있음을 시사하며, 대신 교사 - 학생 호환성과 손실 혼합 전략에 초점을 맞추도록 장려합니다.
3. 일반화 중점: 증류가 도메인 내 적합뿐만 아니라 일반화 (도메인 외 및 하류 작업) 를 개선하는 데 특히 효과적임을 강조하여, 사전 학습에서 KD 의 유용성에 대한 새로운 관점을 제공합니다.

저자들은 "강한 모델에서 약한 모델로"라는 프레임이 불완전하며, 효과적인 증류는 교사 및 학생 간의 호환성, 적절한 손실 혼합, 그리고 학습 과정의 구체적인 일반화 목표에 달려 있다고 결론지었습니다.