LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-Training

Mixture-of-Experts (MoE) has gained increasing popularity as a promising framework for scaling up large language models (LLMs). However, training MoE from scratch in a large-scale setting still suffers from data-hungry and instability problems. Motivated by this limit, we investigate building MoE models from existing dense large language models. Specifically, based on the well-known LLaMA-2 7B model, we obtain an MoE model by: (1) Expert Construction, which partitions the parameters of original Feed-Forward Networks (FFNs) into multiple experts; (2) Continual pre-training, which further trains the transformed MoE model and additional gate networks. In this paper, we comprehensively explore different methods for expert construction and various data sampling strategies for continual pre-training. After these stages, our LLaMA-MoE models could maintain language abilities and route the input tokens to specific experts with part of the parameters activated. Empirically, by training 200B tokens, LLaMA-MoE-3.5B models significantly outperform dense models that contain similar activation parameters.

한 줄 요약

LLaMA-MoE는 기존 LLaMA-2 7B 모델을 기반으로 전문가 구성과 지속적 사전 학습을 통해 MoE 모델을 구축한 연구로, 200B 토큰 학습 후 3.5B 파라미터 모델이 유사 활성 파라미터를 가진 밀집 모델을 상회한다.

핵심 기여도

LLaMA-2 7B 모델의 FFN을 여러 전문가로 분할하는 **Expert Construction** 방법론을 제안.
**Continual Pre-training**을 통해 MoE 모델과 게이트 네트워크를 추가 학습.
**200B 토큰** 학습 후, LLaMA-MoE-3.5B가 유사 활성 파라미터를 가진 Sheared-LLaMA-2.7B, Open-LLaMA-3.0B를 상회.
**Static Sheared 데이터 샘플링**과 **fluency-filtered SlimPajama 데이터셋**을 사용해 빠른 수렴과 성능 향상 달성.

핵심 아이디어

기존 밀집형 대형 언어 모델(LLaMA-2 7B)을 기반으로 **Mixture-of-Experts (MoE)** 구조를 구축함으로써, 모델 크기와 계산 비용을 분리하는 새로운 접근법을 제시한다.
LLaMA-2의 FFN 파라미터를 여러 전문가로 분할하는 **Expert Construction** 과정에서, **non-overlapping randomly splitting** 방법이 가장 우수한 성능을 보였다.
이후, **Continual Pre-training**을 통해 MoE 구조와 게이트 네트워크를 학습하며, **Static Sheared 데이터 샘플링** 전략이 빠른 수렴과 성능 향상을 유도함을 발견했다.
이 연구는 **decoder-only 모델**에 기반한 MoE 구축이 가능하다는 점에서 기존 연구와 차별화된다.

기술적 접근법

**모델 기반**: LLaMA-2 7B
**모듈**: FFN 분할 → 여러 전문가 생성 (Expert Construction)
**게이트 네트워크**: 각 토큰을 top-k 전문가로 라우팅
**학습 전략**: 200B 토큰, 112개 A100 GPU, 글로벌 배치 크기 15M, 최대 학습률 2e-4, 코사인 스케줄링
**데이터셋**: fluency-filtered SlimPajama, Static Sheared 데이터 샘플링
**구현 기술**: transformers, ZeRO-1, FlashAttention v2

주요 결과

**LLaMA-MoE-3.5B**는 200B 토큰 학습 후, 유사 활성 파라미터를 가진 Sheared-LLaMA-2.7B와 Open-LLaMA-3.0B를 상회.
**LLaMA-MoE-3.0B**는 Open-LLaMA-3B와 유사한 성능을 보이지만, 더 적은 활성 파라미터만 사용.
**HellaSwag**, **ARC-c**, **MMLU**, **WinoGrande** 등 다수의 평가셋에서 성능 개선.
**Static Sheared 데이터 샘플링**이 빠른 수렴과 성능 향상에 기여함.

의의 및 한계

LLaMA-MoE는 기존 밀집 모델을 기반으로 MoE 구조를 구축함으로써, **대규모 학습 없이도 성능을 유지**할 수 있다는 점에서 학술적 의의가 있다.
또한, **decoder-only 모델에 적용 가능한 MoE 구조**를 제시하며, **전문가 특화 현상**을 관찰함으로써 모델 내부 동작에 대한 통찰을 제공한다.
그러나, **LLaMA-2 7B 기반 모델만 실험**했으며, **더 큰 밀집 모델에서의 확장성**은 검증되지 않았다.
또한, **게이트 네트워크의 정확한 라우팅 메커니즘**에 대한 심층 분석은 부족한 것으로 보인다.

실용적 활용

LLaMA-MoE는 **저비용 고성능 언어 모델 개발**에 적합하며, 특히 **대규모 학습 데이터가 제한된 상황**에서 유용할 수 있다.
**모델 크기와 계산 비용을 분리**할 수 있으므로, **클라우드 인프라 최적화**, **모바일/엣지 기기에서의 추론** 등에 활용 가능하다.
또한, **전문가 특화 현상**을 활용한 **다양한 도메인별 모델 최적화**에도 적용 가능하다.