파라미터 수는 많지만 계산은 빠르게: Mixture of Experts(MoE)란?

MoE(Mixture of Experts)란 무엇인가?

딥러닝 모델의 성능을 높이는 방법 중 하나는 모델의 크기를 키우는 것입니다. 하지만 모든 파라미터를 매번 사용하는 것은 계산 자원이 비효율적일 수 있습니다. 이럴 때 주목할 만한 기술이 바로 Mixture of Experts (MoE)입니다.

MoE는 아주 간단히 말하면, 전체 모델 중 일부만 골라 사용하는 구조입니다. 마치 회의를 할 때 모든 전문가가 동시에 말하는 것이 아니라, 필요한 전문가 몇 명만 발언하는 것과 비슷합니다.

---

1. MoE의 개념

Mixture of Experts(MoE)는 여러 개의 **전문가 모델(Experts)**을 두고, 입력에 따라 그중 일부만 활성화하여 학습과 추론에 사용하는 방식입니다.

• 각 전문가(Expert)는 동일한 구조를 가지며, 서로 다른 가중치를 학습합니다.
• 게이트(Gating) 네트워크가 입력을 받아 어떤 전문가를 사용할지를 결정합니다.
• 이로 인해 계산량은 줄이면서도 모델 전체는 매우 크고 강력할 수 있습니다.

즉, 모델은 거대하지만 입력 하나당 사용하는 파라미터 수는 적기 때문에 효율적입니다.

---

2. 구조 요약

다음은 MoE의 간단한 구조입니다:

입력 → Gating Network → 전문가 중 일부 선택 → 선택된 전문가들의 출력 결합 → 출력

예를 들어, 전문가가 8명 있다면, Gating Network는 그 중 2명만 활성화하도록 결정할 수 있습니다. 이 선택은 보통 Softmax 확률을 기반으로 하며, 선택된 전문가의 출력을 가중합해 최종 출력을 생성합니다.

---

3. 대표적인 MoE 아키텍처

1) GShard (Google Research)

Google이 2020년에 발표한 GShard는 Transformer에 MoE를 적용한 대표적인 예입니다. 수천억 개의 파라미터를 가진 모델을 효율적으로 분산 학습하기 위해 MoE 구조를 사용했습니다.

• 한 레이어에 수천 개의 전문가를 두고, 입력마다 2개의 전문가만 선택
• 다양한 GPU/TPU에 전문가를 분산시켜 학습 속도 개선

2) Switch Transformer

GShard를 발전시킨 구조로, 하나의 전문가만 활성화하는 방식으로 더욱 효율적입니다.

• Gating의 결과로 가장 높은 확률을 가진 전문가 하나만 사용
• 통신 비용 감소
• 훈련이 안정적이고 빠름

---

4. 장점

1. 연산 효율성 향상
   입력마다 일부 전문가만 사용하므로, 전체 모델의 파라미터는 크지만 실제 연산은 작게 유지됩니다.

2. 확장성(Scalability)
   전문가의 수를 늘리는 것은 쉬우며, 모델을 거대하게 확장할 수 있습니다.

3. 모듈화(Modularity)
   전문가마다 서로 다른 특성을 학습할 수 있어서, 특화된 표현 학습이 가능합니다.

---

5. 단점

1. Gating의 불안정성
   게이트 네트워크가 전문가를 편향적으로 선택할 수 있어 학습이 어려울 수 있습니다.

2. 로드 불균형
   특정 전문가에 트래픽이 집중되면 일부 GPU/TPU는 과부하, 일부는 놀게 되는 문제가 발생합니다.

3. 추론의 복잡도
   선택된 전문가만 사용하더라도, 하드웨어 배치가 복잡해지고 프레임워크 수준의 최적화가 필요합니다.

---

6. MoE 구현 예시 (PyTorch 기반)

다음은 간단한 MoE 레이어를 PyTorch로 구현한 예시입니다:

import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.fc = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return F.relu(self.fc(x))

class MoELayer(nn.Module):
    def __init__(self, input_dim, output_dim, num_experts=4, k=2):
        super().__init__()
        self.num_experts = num_experts
        self.k = k  # top-k experts
        self.experts = nn.ModuleList([Expert(input_dim, output_dim) for _ in range(num_experts)])
        self.gate = nn.Linear(input_dim, num_experts)

    def forward(self, x):
        gate_scores = self.gate(x)  # [batch_size, num_experts]
        topk_scores, topk_indices = torch.topk(gate_scores, self.k, dim=1)  # Top-k gating
        topk_probs = F.softmax(topk_scores, dim=1)

        out = 0
        for i in range(self.k):
            idx = topk_indices[:, i]
            expert_out = torch.stack([self.experts[j](x[b].unsqueeze(0)) for b, j in enumerate(idx)])
            out += expert_out.squeeze(1) * topk_probs[:, i].unsqueeze(1)

        return out

---

7. 실제 적용 사례

• GPT-4, PaLM, Switch Transformer 등 대형 언어 모델에 활용
• Google Translate 시스템 내에서도 MoE가 사용된 바 있음
• 최근 HuggingFace, DeepSpeed 등에서도 MoE 레이어를 공식 지원

---

8. 마무리

MoE는 모델을 더 크고 더 강력하게 만들면서도, 실행 비용은 낮출 수 있는 혁신적인 방법입니다. 특히 대규모 언어 모델이나 멀티태스크 학습 등에 적합하며, 앞으로도 다양한 구조와 변형이 연구될 것으로 보입니다.