AI Research Portfolio Post

주제: 5-3 효율화/서빙 (MoE: Switch Transformer / Mixtral)

분류: llm

타입: concept

난이도: 상

선수지식: 있음 — Transformer, Feed Forward Network, Distributed Training

문제 설정

LLM 성능을 높이기 위해 일반적으로 모델 크기를 증가시킵니다.

하지만 dense Transformer는 다음 문제가 있습니다.

모든 파라미터가 매 토큰마다 계산됨
compute 비용 매우 큼
추론 비용 증가

이를 해결하기 위해 등장한 구조가 Mixture of Experts (MoE)입니다.

직관 비유

Dense 모델 -> 모든 전문가가 동시에 일함
MoE -> 필요한 전문가만 선택

즉 모델 파라미터는 크지만 활성 파라미터(active parameters)는 작습니다.

1. Dense Transformer FFN

일반 Transformer의 FFN:

F F N (x) = W 2 σ (W 1 x)

기호 의미

x : hidden representation
W1,W2 : 파라미터
σ : activation function

왜 문제인가

모든 토큰에 대해 동일 FFN이 계산됩니다.

2. MoE 구조

MoE는 여러 FFN expert를 사용합니다.

Input token
↓
Router
↓
Expert 선택
↓
Expert FFN

수식:

y = Σ g_{i} (x) E x p e r t_{i} (x)

기호 의미

Expert_i : i번째 expert network
g_i(x) : routing weight

왜 필요한가

모델 capacity를 크게 늘릴 수 있습니다.

3. Switch Transformer

Switch Transformer는 가장 단순한 MoE 구조입니다.

각 토큰은 1개의 expert만 선택합니다.

수식:

e x p e r t = a r g m a x g (x)

출력:

y = E x p e r t_{s} e l e c t e d (x)

기호 의미

g(x) : router network

왜 중요한가

계산 비용을 크게 줄입니다.

주의점

expert imbalance 문제가 발생할 수 있습니다.

4. Router 함수

router는 softmax 기반 gating을 사용합니다.

g (x) = s o f t m a x (W_{r} x)

기호 의미

W_r : router parameter

왜 필요한가

어떤 expert가 적합한지 결정합니다.

5. Load Balancing 문제

MoE에서는 일부 expert에 트래픽이 집중될 수 있습니다.

이를 방지하기 위해 auxiliary loss를 사용합니다.

L_{t} o t a l = L_{t} a s k + λ L_{b} a l a n c e

기호 의미

L_balance : load balancing loss

왜 필요한가

expert utilization을 균등하게 유지합니다.

6. Mixtral 구조

Mixtral은 MoE 구조의 개선된 모델입니다.

특징:

Top-2 routing
8 experts

수식:

y = g_{i} E x p e r t_{i} (x) + g_{j} E x p e r t_{j} (x)

기호 의미

i,j : top-2 experts

왜 중요한가

Switch보다 안정적인 학습이 가능합니다.

7. Dense vs MoE

모델	계산 방식
Dense	모든 파라미터 사용
MoE	선택된 expert만 사용

8. MoE 장점

모델 capacity 증가
compute 효율 향상
대형 모델 확장 가능

9. MoE 단점

router 학습 어려움
load imbalance 문제
distributed communication 증가

10. 실제 모델 예

모델	구조
Switch Transformer	Top-1 MoE
Mixtral 8x7B	Top-2 MoE

코드-수식 연결

개념	코드	설명
router	`torch.nn.Linear()`	expert selection
expert	`MLP()`	FFN expert
top-k routing	`torch.topk()`	expert 선택

자주 하는 오해 5개

MoE는 항상 dense 모델보다 빠르다고 생각한다
모든 expert가 항상 사용된다고 생각한다
MoE는 모델 구조를 완전히 변경한다고 생각한다
router는 학습되지 않는다고 생각한다
MoE는 inference에서 항상 효율적이라고 생각한다

체크리스트 (스스로 설명 가능해야 하는 질문)

Mixture of Experts 구조의 핵심 아이디어는 무엇인가?
Switch Transformer의 routing 방식은 무엇인가?
Mixtral의 top-2 routing은 왜 필요한가?
MoE에서 load balancing 문제가 발생하는 이유는 무엇인가?
Dense 모델과 MoE 모델의 compute 차이는 무엇인가?