Mixture of Experts

# Mixture of Experts 개념 구현 (PyTorch)
import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    """개별 Expert 네트워크 (FFN)"""
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = F.gelu(self.fc1(x))
        return self.fc2(x)

class TopKGating(nn.Module):
    """Top-K Gating Network (라우터)"""
    def __init__(self, input_dim, num_experts, top_k=2):
        super().__init__()
        self.gate = nn.Linear(input_dim, num_experts)
        self.top_k = top_k

    def forward(self, x):
        # 각 Expert에 대한 점수 계산
        logits = self.gate(x)  # (batch, seq, num_experts)

        # Top-K Expert 선택
        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
        top_k_gates = F.softmax(top_k_logits, dim=-1)

        return top_k_gates, top_k_indices

class MixtureOfExperts(nn.Module):
    """Sparse MoE Layer"""
    def __init__(self, input_dim, hidden_dim, output_dim,
                 num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Expert 네트워크들
        self.experts = nn.ModuleList([
            Expert(input_dim, hidden_dim, output_dim)
            for _ in range(num_experts)
        ])

        # Gating Network
        self.gating = TopKGating(input_dim, num_experts, top_k)

    def forward(self, x):
        batch_size, seq_len, dim = x.shape

        # Gating: 어떤 Expert를 활성화할지 결정
        gates, indices = self.gating(x)
        # gates: (batch, seq, top_k) - 각 Expert의 가중치
        # indices: (batch, seq, top_k) - 선택된 Expert 인덱스

        # Expert 출력 계산 및 가중 합산
        output = torch.zeros_like(x)

        for i in range(self.top_k):
            expert_idx = indices[:, :, i]  # (batch, seq)
            expert_gate = gates[:, :, i:i+1]  # (batch, seq, 1)

            for e in range(self.num_experts):
                mask = (expert_idx == e)
                if mask.any():
                    expert_input = x[mask]
                    expert_output = self.experts[e](expert_input)
                    output[mask] += expert_gate[mask].squeeze(-1).unsqueeze(-1) * expert_output

        return output

# 사용 예시: Mixtral 스타일 MoE
moe_layer = MixtureOfExperts(
    input_dim=4096,    # 임베딩 차원
    hidden_dim=14336,  # FFN 히든 차원
    output_dim=4096,
    num_experts=8,     # 8개 Expert
    top_k=2            # 토큰당 2개 활성화
)

# 입력: (batch=2, seq_len=10, dim=4096)
x = torch.randn(2, 10, 4096)
output = moe_layer(x)

print(f"입력 shape: {x.shape}")
print(f"출력 shape: {output.shape}")
print(f"총 Expert 수: 8, 활성화 Expert: 2")
print(f"총 파라미터 대비 활성 파라미터: 2/8 = 25%")