기술 요약: 어디서 시작할 것인가: 서브 네트워크 선택 및 증류를 통한 효율적인 사전 학습

문제 정의

대규모 언어 모델(LLM)은 거대한 규모를 통해 최첨단 성능을 달고 있지만, 이들의 학습 및 배포 비용은 많은 연구 그룹과 자원이 제한된 환경에 있어 매우 부담스럽습니다. 소규모 언 language 모델(SLM)은 더 접근 가능한 대안을 제공하지만, 여전히 사전 학습을 위해 상당한 계산 자원을 필요로 합니다. 지식 증류(knowledge distillation)나 프루닝(pruning)과 같은 기존 전략들은 오픈 웨이트 LLM(예: Llama 3.2 변형, Minitron)에 의존하거나, 최적의 서브 네트워크 초기화를 발견하기 위한 체계적이고 재현 가능한 프레임워크가 부족한 경우가 많습니다. 또한, 교사 모델(teacher models)이 SLM 사전 학습 효율성을 개선하는 메커니즘에 대해서는 아직 명확히 이해되지 않았습니다.

방법론

저자들은 SLM 사전 학습을 실질적으로 더 효율적으로 만들기 위해 세 가지 구성 요소로 이루어진 프레임워크를 제안합니다:

1. 서브 네트워크 초기화: 무작위 초기화 대신, 이 프레임워크는 거대한 사전 학습된 교사 모델로부터 구조적으로 희소한(structurally sparse) 서브 네트워크를 추출합니다. 저자들은 가중치가 어떻게 선택되는지 탐색하기 위해 네 가지 뚜렷한 탐색 공간을 정의합니다:

   • Coarse(조립식) vs. Fine-grained(세밀한): Coarse 선택은 컴포넌트(예: 첫 $n$개의 어텐션 헤드)의 처음 $n$개 항목을 가져오는 반과, Fine-grained 선택은 임의의 인덱스를 샘플링합니다.
   • Uniform(균일) vs. Layer-wise(층별): Uniform은 모든 레이어에 동일한 구성을 적용하는 반면, Layer-wise는 각 레이어가 독립적인 구성을 가질 수 있도록 허용합니다.
   • 이는 Coarse Uniform, Coarse Layer-wise, Fine-grained Uniform, Fine-grained Layer-wise라는 네 가지 탐색 공간을 생성합니다.

2. 제약 조건이 있는 진화적 탐색: 특정 파라미터 예산 내에서 고품질의 서브 네트워크를 식별하기 위해, 저자들은 진화 알고리즘을 채택합니다.

   • 파라미터 공간은 균형 잡힌 모델 크기 커버리지를 보장하기 위해 서로 겹치지 않는 "빈(bins)"(예: 385M–426M, 961M–1.06B)으로 분할됩니다.
   • 각 빈 내에서, 후보 아키텍처들은 변이(mutation) 및 교차(crossover) 연산자를 사용하여 진화됩니다.
   • 제약 조건 강제: 후보 아키텍처가 대상 빈의 파라미터 수를 엄격히 준수하도록 하기 위해 거부 샘플링(rejection sampling)이 사용됩니다.
   • 적합도 함수(Fitness Function): 최적의 초기화를 유도하기 위해, 후보들은 작은 토큰 예산(2B 토큰)에서의 검증 퍼플렉시티(validation perplexity)를 기준으로 평가됩니다.

3. 지식 증류: 선택된 서브 네트워크는 추출된 원천인 동일한 교사 모델으로부터 지식 증류를 통해 사전 학습됩니다.

   • 손실 함수는 표준 교차 엔트로피와 교사 및 학생의 로짓 분포 간의 전방 KL 발산(forward KL divergence)을 결합합니다.
   • 계산 효율성을 높이기 위해, 저자들은 교사의 분포를 가장 돌출된 $k$개의 출력(예: $k=1024$)으로 절단하는 Top-k 로짓 증류를 활용합니다.

주요 기여

• 체계적인 오픈 소스 프레임워크: 저자들은 Hugging Face 모델로부터 서브 네트워크를 추출하고, 유연한 탐색 공간 설계를 지원하며, 탐색부터 사전 학습 및 증류에 이르는 전체 파이프라인을 용이하게 하는 완전한 오픈 소스 라이브러리인 Whittle을 소개합니다.
• 서브 네트워크 초기화 전략: 교사 모델에서 추출한 서브 네트워크가 무작위 초기화보다 훨씬 더 나은 시작점을 제공함을 입증했습니다. 예를 들어, 더 작은 변형 모델(~410M 파라미터)은 동일한 검증 퍼플렉시티를 달성하기 위해 Pythia 모델보다 1.71배 적은 FLOPs가 필요합니다.
• 포괄적 분석: 본 논문은 지식 증류 하에서의 서브 네트워크 초기화와 표준 교차 엔트로피 학습을 체계적으로 비교한 첫 사례를 제공하며, 교사 크기, 탐색 공간의 세밀도, 증류 하이퍼파라미터의 영향을 분석합니다.
• 실증적 검증: Pythia 제품군(6.9B 및 12B 교사)에 대한 실험 결과, 진화적 탐색을 통해 발견되고 교사 가중치로 초기화된 최적의 모델들이 유사한 크기의 Pythia SLM의 성능을 달로하면서도 훨씬 적은 부동 소수점 연산(FLOPs)을 요구함을 보여줍니다.

결과

• 효율성 이득:
  • 410M 파라미터 모델(Bin 1)의 경우, 이 방법은 표준 Pythia-410M 사전 학습과 동일한 검증 퍼플렉시티를 달성하는 데 1.71배 적은 FLOPs를 사용합니다.
  • 1B 파라미터 모델(Bin 2)의 경우, 절감 효과는 1.75배로 증가합니다.
  • 2.8B 파라미터 모델(Bin 3)의 경우, 10B 토큰 예산으로 학습된 유사한 Pythia-2.8B 모델과 검증 퍼플렉시티를 맞추기 위해 5.16배 적은 FLOPs가 필요합니다.
  • 100B 토큰 예산에서, 2.8B 변형 모델은 베이스라인보다 낮은 최종 검증 퍼플렉시티를 달성하면서도 1.26배의 FLOP 감소를 유지합니다.
• 탐색 공간 발견 사항:
  • 세밀도(Granularity): 최적의 탐색 공간은 모델 크기에 따라 달라집니다. 더 큰 변형(Bin 3)의 경우, 가장 덜 세밀한 공간인 Coarse Uniform이 최상의 결과를 냅니다. 더 작은 변형(Bin 1)의 경우, Fine-grained Uniform과 같이 더 세밀한 공간이 더 효과적입니다.
  • 증류: 증류는 표준 사전 학습보다 일관되게 더 나은 퍼플렉시티를 제공합니다. 전체 로짓 분포를 사용하는 것이 Top-k 절단보다 일반적으로 더 낮은 퍼플렉시티를 제공하지만, 주요 실험에서는 효율성을 위해 Top-k가 사용되었습니다.
  • 탐색 지표: 탐색 단계에서 퍼플렉시티를 직접 최적화하는 것이 가중치 크기나 활성화 기반 중요도 점수와 같은 대리 지표(proxy metrics)보다 우수한 성능을 보입니다.
• 다운스트림 성능: 슈퍼넷(supernet)에서 초기화된 서브 네트워크는 동일한 아키텍처의 무작위 초기화 버전과 유사한 크기의 원본 Pythia 모델 모두를 일관되게 능가합니다(COPA, MMLU, HellaSwag 등 일반 상식 및 질의응답 벤치마크 기준). 특히, 작은 6.9B 교사로부터 추출된 서브 네트워크가 12B 교사로부터 추출된 것보다 더 나은 성능을 보이기도 했습니다.

의의 및 주장

본 논문은 더 큰 교사 모델로부터 학생 모델을 웜 스타트(warm-start)하는 것에 대한 최초의 체계적인 오픈 소스 연구를 제시한다고 주장합니다. 이 연구의 의의는 비용 효율적인 SLM 개발을 위한 실용적이고 재현 가능한 경로를 제공한다는 데 있습니다. 진화적 탐색을 통한 서브 네트워크 발견과 지식 증류를 결합함으로써, 저자들은 유사한 크기의 베이스라인 SLM과 비교했을 때(탐색 비용 제외 시 최대 9.2배, 탐색 비용 포함 시 5.16배) SLM 사전 학습을 가속화할 수 있음을 입증했습니다.

저자들은 자신들의 접근 방식이 전통적인 신경망 구조 탐색(NAS) 방법에서 요구되는 값비싼 슈퍼넷 사전 학습을 피할 수 있다는 점을 강조합니다. 그들은 이 방법이 매우 효과적이지만, 이러한 초기화 전략에 대한 스케일링 법칙(scaling laws)을 도출하고 교사 모델의 선택(예: 다국어 vs 단일 언어)이 학생 모델의 성능에 미치는 영향을 조사하기 위한 후속 연구가 필요하다고 명시했습니다. 저자들은 자신들의 발견이 다양한 규모에 걸쳐 계산 최적화된 SLM 사전 학습을 위한 가이드라인을 제공한다고 결론지었습니다.