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

하임·2026년 1월 9일

MoE

목록 보기

10/14

https://aclanthology.org/2024.emnlp-main.890.pdf

LLaMA-MoE: LLaMA 기반 Mixture-of-Experts 구축과 지속적 사전훈련

이 논문에서는 Mixture-of-Experts (MoE) 모델을 기존의 LLaMA 모델에서 구축하는 방법을 제안한다. 기존 MoE 모델을 처음부터 훈련하는 방식은 비용이 높고 불안정성이 크므로, 기존의 Dense LLM (LLaMA-2-7B)를 MoE 모델로 변환하는 접근법을 연구한다. 이를 위해 전문가(Experts)를 구성하는 방법과 지속적 사전 훈련(Continual Pre-training) 전략을 제안하고 평가한다.

1. 배경 및 문제점

기존 대형 언어 모델(LLMs)은 모델 크기 증가가 성능 향상의 주요 방법이지만, 이는 막대한 계산 비용을 초래함.
이를 해결하기 위해, MoE (Mixture-of-Experts) 기반 모델이 연구됨. MoE는 일부 전문가만 활성화하는 방식으로 연산 비용을 줄이는 동시에 성능을 유지하는 방식.
하지만 MoE 모델을 처음부터 훈련하는 것은 많은 데이터와 연산 자원이 필요하며, 학습 안정성 문제도 있음.
따라서, 기존 Dense 모델을 활용하여 MoE 모델을 구축하는 방법을 연구하며, 이를 통해 훈련 비용 절감 및 성능 향상을 목표로 함.

2. 방법론

LLaMA-MoE 모델을 구축하는 과정은 두 가지 주요 단계로 이루어진다:

전문가 구성(Expert Construction)
지속적 사전 훈련(Continual Pre-training)

2.1 전문가 구성 (Expert Construction)

LLaMA-2의 Feed-Forward Network (FFN) 레이어를 여러 개의 전문가(Experts)로 분할.
기존 FFN의 SwiGLU 활성화 함수를 유지하면서, 각 전문가가 독립적으로 FFN을 담당하도록 함.
전문가를 구성하는 방식은 크게 두 가지로 나뉨:
1. Neuron-Independent 방식: 각 뉴런을 독립적인 전문가 그룹으로 랜덤하게 분할.
2. Neuron-Sharing 방식: 뉴런을 여러 전문가 간에 공유하며, 특정한 뉴런이 특정 전문가에 더 많이 활성화되도록 설계.

본 연구에서는 IndependentRandom (뉴런을 무작위로 나누는 방법)을 사용하여 전문가를 구성함.

전문가 분할 후 보정 (Rescaling)

기존 LLaMA의 FFN을 여러 전문가로 분할하면, 각 전문가의 활성화된 뉴런 수가 줄어들게 됨.
이를 보완하기 위해 출력 스케일링 (Rescaling) 기법을 적용하여 각 전문가의 출력이 원래 Dense FFN의 출력 크기와 유사하도록 보정.

2.2 지속적 사전 훈련 (Continual Pre-training)

모델 구조를 Dense → MoE로 변경하면서 발생하는 성능 저하를 보완하기 위해 추가 학습을 수행.
LLaMA-2의 훈련 목표와 동일한 언어 모델링 목표를 사용하여 사전 훈련을 계속 진행.
데이터 샘플링 전략을 실험하여 가장 효과적인 지속적 학습 방법을 탐색:
- StaticSheared: 사전 정의된 비율로 데이터를 샘플링.
- StaticLLaMA: 기존 LLaMA-2의 데이터 샘플링 전략을 유지.
- DynamicSheared: 특정 간격마다 손실 변화를 평가하여 샘플링 비율을 조정.
- DynamicLLaMA: 기존 LLaMA 방식에서 동적으로 비율을 조정.

실험 결과, StaticSheared 전략이 가장 안정적이고 좋은 성능을 보였음.

3. 실험 및 결과

3.1 데이터셋 및 실험 환경

사전 훈련 데이터: SlimPajama 데이터셋 (총 627B 토큰)
평가 데이터셋:
- 논리 및 일반 지식 평가: SciQ, PIQA, WinoGrande, ARC-e, ARC-c, HellaSwag
- 세계 지식 평가: LogiQA, BoolQ, LAMBADA, NQ, MMLU
모델 크기:
- LLaMA-2-7B (Baseline)
- LLaMA-MoE-3.0B (2/16)
- LLaMA-MoE-3.5B (4/16)
- LLaMA-MoE-3.5B (2/8)

3.2 주요 결과

LLaMA-MoE-3.5B 모델이 기존 Dense 모델과 유사한 수준의 성능을 유지하면서 연산 비용을 절감함.
LLaMA-MoE-3.5B (4/16)은 Sheared-LLaMA 대비 1.3점 높은 평균 점수를 기록하며, Dense 모델 대비 89.2%의 성능을 유지.
LLaMA-MoE는 기존 Dense 모델 대비 빠르게 수렴하며, 성능 저하 없이 FLOPs를 57.7%까지 절감.

3.3 추가 실험: Ablation Study

Rescaling 적용 여부:
- Rescaling을 적용하면 초기 성능이 크게 향상됨.
전문가 구성 방법 비교:
- IndependentRandom (랜덤 뉴런 분할 방식)이 가장 우수한 성능을 보임.
데이터 샘플링 전략 비교:
- StaticSheared 방식이 가장 안정적이고 효과적인 데이터 샘플링 전략임.
MoE 모델을 처음부터 훈련하는 경우 vs. 기존 Dense 모델에서 변환하는 경우:
- 기존 Dense 모델을 변환하는 방식이 훨씬 빠르고 비용 효율적임.

3.4 지시 학습 (Instruction Tuning)

6,000개 ShareGPT 명령어 데이터로 추가 튜닝을 진행.
LLaMA-MoE-3.5B 모델이 Dense 모델 대비 더 강력한 지시 학습 성능을 보임.

4. 결론

본 연구는 기존 Dense 모델을 활용하여 MoE 모델을 구축하는 효율적인 방법을 제안.
LLaMA-2-7B의 FFN 레이어를 전문가로 분할하고, 지속적 사전 훈련을 수행하여 MoE 모델을 학습.
제안된 LLaMA-MoE 모델은 기존 Dense 모델과 유사한 성능을 유지하면서 연산 비용을 절감.
특히 LLaMA-MoE-3.5B는 Open Source LLM 중에서 매우 경쟁력 있는 성능을 보이며, FLOPs를 크게 줄이는 장점이 있음.
향후 연구로는 더 다양한 전문가 수 및 크기에 대한 실험을 진행하여 최적의 MoE 구조를 탐색할 계획.

요약

이 논문은 Dense LLM (LLaMA-2-7B)을 MoE 모델로 변환하는 방법을 연구하며, 이를 위해 전문가 구성 및 지속적 사전 훈련 기법을 제안한다. 실험 결과, LLaMA-MoE 모델은 기존 Dense 모델과 유사한 성능을 유지하면서도 연산 비용을 절감할 수 있음을 확인하였다. 특히, 전문가 분할 방식(IndependentRandom)과 StaticSheared 데이터 샘플링 전략이 가장 효과적이며, 지시 학습(Instruction Tuning)에서도 뛰어난 성능을 보였다.