Mixture of Experts: 희소 활성화로 거대 모델 만들기

지난 글에서 GQA가 KV 헤드를 줄여 메모리를 절감하는 방법을 살펴봤다. 이번에는 파라미터 규모와 연산 비용을 분리하는 전혀 다른 전략인 Mixture of Experts(MoE)를 다룬다. MoE는 1991년 Jacobs et al.의 아이디어를 현대 트랜스포머에 접목한 것으로, 2024년 GPT-4, Mixtral 8×7B, DeepSeek-V3 등에서 적용되며 LLM 아키텍처의 새로운 주류가 됐다.

핵심 아이디어: 희소 활성화

밀집(Dense) 트랜스포머에서 FFN 레이어는 모든 토큰에 대해 동일하게 활성화된다. MoE는 이 FFN을 N개의 독립 전문가(Expert) FFN으로 교체한다. 각 토큰은 라우터(Router)를 통해 이 중 Top-K개만 선택해 실행한다. K=2이면 전체 N개 파라미터 중 2/N만 활성화되는 셈이다.

Dense FFN:    토큰 → FFN(전체 파라미터) → 출력
MoE FFN:      토큰 → Router → Expert₁, Expert₃ → 가중 합산 → 출력

결과적으로 파라미터 수는 N배 증가하지만 추론 FLOPs는 거의 동일하다. “같은 연산 비용으로 더 큰 모델”이 가능해진다. Mixtral 8×7B는 총 47B 파라미터지만 추론 시 활성 파라미터는 13B 수준이다.

라우터와 게이팅 메커니즘

라우터는 간단한 선형 레이어다. 입력 토큰 벡터 x를 받아 각 전문가의 적합도 점수를 계산한다:

# 라우터 게이팅
logits = W_g @ x          # (num_experts,)
gates, indices = logits.topk(K)   # Top-K 선택
gates = softmax(gates)     # 정규화 (합 = 1)

# 출력 = 선택된 전문가의 가중 합산
output = sum(gates[i] * Expert_i(x) for i in indices)

PyTorch 구현

위 구현은 개념 설명용이다. 실제 분산 학습에서는 전문가를 서로 다른 GPU에 배치하고, 토큰을 해당 GPU로 올바르게 라우팅하는 Expert Parallelism과 All-to-All 통신이 필요하다.

로드 밸런싱 문제와 보조 손실

MoE의 고질적 문제는 라우터 붕괴(Router Collapse)다. 학습 초기에 특정 전문가가 우연히 좋은 성능을 보이면, 라우터가 그 전문가만 계속 선택하게 되어 나머지 전문가는 학습이 안 된다.

이를 막기 위해 보조 로드 밸런싱 손실(Auxiliary Load Balancing Loss)을 추가한다:

def load_balancing_loss(router_probs, expert_indices, num_experts):
    # 전문가별 할당 비율 f_i
    expert_mask = F.one_hot(expert_indices, num_experts).float()
    tokens_per_expert = expert_mask.mean(dim=0)   # (num_experts,)
    # 전문가별 라우터 점수 평균 p_i
    avg_probs = router_probs.mean(dim=0)          # (num_experts,)
    # 균등 분배 유도: f_i * p_i의 합을 최소화
    loss = num_experts * (tokens_per_expert * avg_probs).sum()
    return loss

# 총 손실 = 언어 모델링 손실 + α * 로드 밸런싱 손실
total_loss = lm_loss + 0.01 * load_balancing_loss(...)

이 손실은 모든 전문가가 균등하게 토큰을 받도록 유도한다. α(보조 손실 계수)는 0.001~0.01 사이가 일반적이다.

Expert Choice vs Token Choice

기존 방식(Token Choice)은 각 토큰이 Top-K 전문가를 선택한다. 문제는 인기 전문가에 토큰이 몰릴 때 처리 용량 제한(Expert Capacity)을 초과하면 일부 토큰이 전문가를 건너뛰어야 한다는 것이다.

Expert Choice(Zhou et al., 2022) 방식은 반대다. 각 전문가가 처리할 Top-K 토큰을 선택한다. 전문가 용량이 보장되고 로드 밸런싱이 자연스럽지만, 일부 토큰이 어떤 전문가에도 선택받지 못할 수 있다.

DeepSeek은 공유 전문가(Shared Expert)를 추가로 두어, 모든 토큰이 반드시 거치는 공통 표현 학습을 담당하게 했다. 나머지 전문가들은 특화된 능력을 학습한다.

주요 MoE 모델 비교

모델	총 파라미터	활성 파라미터	전문가 수	Top-K
Mixtral 8×7B	47B	13B	8	2
Mixtral 8×22B	141B	39B	8	2
DeepSeek-V3	671B	37B	256	8
GPT-4 (추정)	~1.8T	~220B	~16	2

DeepSeek-V3는 256개 전문가 중 8개를 선택하는 극단적 희소성으로 효율을 극대화했다. 총 파라미터가 671B지만 추론 비용은 Dense 37B 수준이다.

추론 최적화

MoE 추론에서는 배치 내 다른 토큰이 서로 다른 전문가를 선택하므로, GPU 활용률이 Dense 모델보다 낮을 수 있다. 이를 해결하는 방법:

• Expert 배치 최적화: 같은 전문가를 선택한 토큰끼리 묶어 행렬 곱 한 번에 처리
• 전문가 캐싱: 최근에 자주 사용된 전문가를 GPU 메모리에 유지
• vLLM, TensorRT-LLM: MoE 최적화 커널 내장

다음 글에서는 이 거대한 LLM들이 텍스트를 처음 받아들이는 방법, 토크나이저를 다룬다.

---

지난 글: MQA와 GQA: KV Cache 경량화 전략

다음 글: 토크나이저와 토큰: LLM이 텍스트를 보는 방법

읽어주셔서 감사합니다. 😊