Mixture of Experts — sparse 컴퓨트, dense 지식

Mixture of Experts(MoE)는 FFN을 expert FFN 집합 + 학습된 router로 대체. 각 토큰에 대해 router가 top-K expert(보통 K=1 또는 2)를 골라, 그 expert들만 계산. 전체 파라미터는 큼, 토큰당 active 파라미터는 훨씬 작음.

고품질이면서 서빙 싸야 하는 모델의 아키텍처적 선택. 100B 파라미터 모델의 지식을 20B 파라미터 모델의 토큰당 컴퓨트로 얻어.

모델	총 파라미터	active 파라미터	expert	routing
Mixtral 8×7B	47B	13B	8	top-2
Mixtral 8×22B	141B	39B	8	top-2
Llama 4 Scout	109B	17B	16 (routed)	top-1
Llama 4 Maverick	400B	17B	128 + 1 shared	top-1
DeepSeek-V3	671B	37B	128 + 1 shared	top-K
Mistral Small 4	119B	6B	128	top-4

왜 MoE가 통하나

토큰마다 필요한 "사고" 종류가 달라. 코드 토큰은 어떤 expert가 좋고, 한국어 토큰은 다른 expert, 수치 추론 단계는 또 다른 expert. router가 expert를 특화하게 하면 모델이 파라미터를 효율적으로 써 — 각 토큰에 맞는 expert만 발화. 비용: 학습이 더 어려움(load balancing, routing 붕괴), 서빙에 커스텀 추론 프레임워크 필요, 배치당 tail latency 변동.

Code

Top-2 MoE block (sketch)·python

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, n_experts, top_k=2):
        super().__init__()
        self.experts = nn.ModuleList(
            [SwiGLU(d_model, d_ff) for _ in range(n_experts)]
        )
        self.router = nn.Linear(d_model, n_experts, bias=False)
        self.top_k = top_k
    def forward(self, x):
        # x: (B, L, d_model)
        logits = self.router(x)                            # (B, L, n_experts)
        topk_vals, topk_idx = logits.topk(self.top_k, dim=-1)
        gates = F.softmax(topk_vals, dim=-1)              # normalize over chosen
        out = torch.zeros_like(x)
        for k in range(self.top_k):
            for e in range(len(self.experts)):
                mask = (topk_idx[..., k] == e)
                if mask.any():
                    out[mask] += gates[..., k][mask].unsqueeze(-1) * self.experts[e](x[mask])
        return out
# Real implementations dispatch tokens to experts in batched form
# (custom kernel territory) for performance.

Mixture of Experts — sparse 컴퓨트, dense 지식

왜 MoE가 통하나

Code

External links

Exercise

Progress

댓글 0