Skip to content

Split View: Mixture of Experts(MoE) 아키텍처 논문 리뷰와 실전 스케일링 2026

✨ Learn with Quiz
|

Mixture of Experts(MoE) 아키텍처 논문 리뷰와 실전 스케일링 2026

Mixture of Experts(MoE) 아키텍처 논문 리뷰와 실전 스케일링 2026

개요

대규모 언어 모델의 파라미터 수가 수조 단위로 확장되면서, 모든 파라미터를 매 토큰마다 활성화하는 Dense 모델의 한계가 분명해졌다. 학습과 추론 비용이 파라미터 수에 비례해 증가하기 때문이다. Mixture of Experts(MoE)는 이 문제에 대한 가장 유력한 해법으로 자리 잡았다. MoE는 전체 파라미터 중 일부 전문가(Expert)만 선택적으로 활성화하여, 모델 용량은 크게 유지하면서 실제 연산량은 Dense 모델 대비 극적으로 줄인다.

2022년 Switch Transformer가 단일 전문가 라우팅으로 조 단위 파라미터 확장의 가능성을 열었고, GShard는 600B 파라미터 모델의 분산 학습 파이프라인을 제시했다. ST-MoE는 학습 안정성과 전이 학습 품질을 동시에 끌어올렸으며, Mixtral 8x7B는 오픈소스 MoE 모델의 실용성을 증명했다. DeepSeek-MoE와 DeepSeek-V3는 세분화된 전문가 분할과 보조 손실 없는 로드 밸런싱이라는 새로운 방향을 열었다. 이 글에서는 각 논문의 핵심 기여를 분석하고, 라우팅 메커니즘, 학습 안정성, 추론 최적화까지 MoE를 실전에서 운영하기 위한 전략을 정리한다.

MoE 핵심 개념

Sparse Activation의 원리

MoE의 핵심 아이디어는 조건부 연산(Conditional Computation)이다. 입력 토큰마다 전체 전문가 중 일부만 활성화하므로, 모델의 총 파라미터 수와 실제 연산에 사용되는 파라미터 수가 분리된다. 예를 들어 Mixtral 8x7B는 총 47B 파라미터를 보유하지만, 각 토큰은 13B 파라미터만 사용한다. Dense 47B 모델 대비 연산량이 약 1/3.6로 줄어들면서도 품질은 유사하거나 더 높다.

Dense vs Sparse MoE 비교

항목Dense 모델Sparse MoE
활성 파라미터전체 파라미터 = 활성 파라미터전체의 10-30%만 활성화
FLOPs 효율파라미터 수에 비례활성 파라미터 수에 비례
메모리 요구량파라미터 크기만큼전체 파라미터 저장 필요 (메모리는 더 큼)
학습 안정성상대적으로 안정라우팅 불안정, 전문가 붕괴 위험
확장성선형 비용 증가전문가 추가로 저비용 확장
추론 지연예측 가능라우팅 오버헤드, 전문가 로딩 지연
대표 모델LLaMA, GPT-4(추정)Mixtral, DeepSeek-V3, Switch Transformer

기본 MoE 레이어 구현

MoE 레이어의 기본 구조를 PyTorch로 구현하면 다음과 같다. 게이팅 네트워크가 입력 토큰에 대해 각 전문가의 가중치를 계산하고, Top-K 전문가를 선택하여 출력을 합산한다.

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

class Expert(nn.Module):
    """단일 전문가 FFN 모듈"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(self.dropout(F.gelu(self.w1(x))))

class MoELayer(nn.Module):
    """Mixture of Experts 레이어"""
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int,
        top_k: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout) for _ in range(num_experts)
        ])
        # 게이팅 네트워크: 입력 -> 전문가 점수
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # (B*S, d_model)

        # 게이팅 점수 계산
        gate_logits = self.gate(x_flat)  # (B*S, num_experts)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-K 전문가 선택
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )  # (B*S, top_k)
        # 선택된 전문가 가중치 재정규화
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 전문가별 출력 합산
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, k]  # (B*S,)
            weight = top_k_probs[:, k].unsqueeze(-1)  # (B*S, 1)
            for i in range(self.num_experts):
                mask = (expert_idx == i)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += weight[mask] * expert_output

        return output.view(batch_size, seq_len, d_model)

이 구현은 교육 목적이며, 실제 프로덕션에서는 Megablocks나 Tutel 같은 최적화 라이브러리를 사용하여 전문가별 배치 처리와 All-to-All 통신을 효율적으로 수행한다.

주요 논문 분석

Switch Transformer (Fedus et al., 2022)

Switch Transformer의 핵심 기여는 MoE 라우팅을 극단적으로 단순화한 것이다. 기존 MoE가 Top-2 이상의 전문가를 조합했다면, Switch Transformer는 각 토큰에 단 하나의 전문가만 배정하는 Switch Routing을 도입했다. 이 결정은 직관에 반하지만, 단일 전문가 라우팅이 통신 비용을 절반으로 줄이고 라우팅 연산 자체의 오버헤드도 감소시킨다는 실험 결과로 뒷받침되었다.

T5-Base 아키텍처 기반으로 동일 연산량 대비 최대 7배 사전학습 속도 향상을 달성했고, 조 단위 파라미터 모델까지 확장했다. bfloat16 혼합 정밀도 학습을 통해 메모리 효율도 개선했으며, 전문가 드롭아웃(Expert Dropout)과 보조 로드 밸런싱 손실을 함께 적용하여 학습 불안정성을 완화했다.

Switch Transformer가 도입한 전문가 용량 인자(Capacity Factor)는 각 전문가가 처리할 수 있는 최대 토큰 수를 결정한다. 용량을 초과하는 토큰은 드롭되며, 이는 학습 중 정보 손실을 유발하지만 연산 부하의 균형을 잡는 데 필수적이다.

GShard (Lepikhin et al., 2021)

GShard는 MoE를 대규모 분산 환경에서 실용적으로 학습시키기 위한 시스템 설계에 초점을 맞춘 논문이다. 600B 파라미터 다국어 번역 모델을 2048개 TPU v3 가속기에서 4일 만에 학습시켰다. 핵심 기술 기여는 세 가지다.

첫째, Expert Parallelism이다. 전문가를 여러 디바이스에 분산 배치하고, 토큰을 해당 전문가가 위치한 디바이스로 라우팅하는 All-to-All 통신 패턴을 정립했다. 둘째, Random Routing이다. Top-1 전문가를 결정론적으로 선택한 뒤, 두 번째 전문가는 게이트 가중치에 비례하는 확률로 확률적으로 선택한다. 이 방식은 학습 시 탐색(exploration)을 촉진한다. 셋째, XLA 컴파일러 확장을 통해 최소한의 코드 변경으로 다양한 병렬화 패턴을 표현할 수 있는 경량 API를 제공했다.

GShard의 그룹별 용량 계산 공식 C = 2N / (G * E)는 그룹 크기 G, 전문가 수 E, 토큰 수 N을 고려하여 각 전문가의 처리 용량을 동적으로 결정한다.

ST-MoE (Zoph et al., 2022)

ST-MoE의 이름이 의미하듯, 이 논문의 핵심 질문은 "MoE 모델을 안정적으로(Stable) 학습시키고 효과적으로 전이(Transferable)할 수 있는가"이다. 269B 파라미터 Sparse 모델(ST-MoE-32B)을 32B Dense 모델 수준의 연산 비용으로 학습시키면서, SuperGLUE, ARC, XSum 등 다양한 벤치마크에서 Sparse 모델 최초로 최고 성능을 달성했다.

학습 안정성을 위해 Router Z-Loss를 도입했다. 이 손실 함수는 라우터 로짓의 크기를 제어하여 학습 중 발산(divergence)을 방지한다. 기존 보조 손실(auxiliary loss)과 함께 사용하면 학습 안정성이 크게 개선된다. 또한 Top-2 라우팅, 학습 시 용량 인자 1.25, 코어당 최대 1개 전문가 배치 등 실전 설계 권장사항을 체계적으로 정리했다.

파인튜닝 시에는 전문가 드롭아웃 비율을 높이고, 배치 크기를 줄이며, 학습률을 보수적으로 설정하는 것이 전이 학습 품질 향상에 효과적이라는 실험 결과를 제시했다.

Mixtral 8x7B (Mistral AI, 2024)

Mixtral 8x7B는 오픈소스 MoE 모델의 실용성을 결정적으로 증명한 모델이다. Mistral 7B와 동일한 Transformer 아키텍처에서 각 레이어의 FFN을 8개 전문가로 교체하고, 토큰마다 Top-2 전문가를 선택한다. 총 47B 파라미터 중 13B만 활성화되므로, 추론 속도는 13B Dense 모델과 유사하면서 품질은 LLaMA 2 70B에 필적한다.

항목Mixtral 8x7BLLaMA 2 70BMistral 7B
총 파라미터46.7B70B7.3B
활성 파라미터12.9B70B7.3B
추론 FLOPs~26B~140B~15B
MMLU70.669.860.1
ARC-Challenge66.464.655.5
전문가 수8--
활성 전문가 수2--

Mixtral의 전문가 특화 분석 결과, 전문가들은 도메인보다는 구문론적 패턴에 따라 특화되는 경향을 보였다. 특정 전문가가 코드나 수학에만 반응하는 것이 아니라, 문장 구조나 토큰 위치에 따라 전문가 선택이 달라진다는 것이다.

DeepSeek-MoE와 DeepSeek-V3

DeepSeek-MoE는 "궁극의 전문가 특화(Ultimate Expert Specialization)"를 목표로 두 가지 핵심 전략을 제시했다.

첫째, 세분화된 전문가 분할(Fine-Grained Expert Segmentation)이다. 기존 N개 전문가를 mN개로 세분화하되, 각 전문가의 FFN 중간 차원을 1/m로 줄여 총 파라미터 수를 동일하게 유지한다. 대신 mK개의 전문가를 활성화하여 더 유연한 전문가 조합이 가능하다. DeepSeek-MoE 16B는 LLaMA 2 7B와 유사한 성능을 연산량 40%만으로 달성했다.

둘째, 공유 전문가(Shared Expert) 격리다. 모든 토큰이 공통으로 사용하는 전문가를 별도로 분리하여, 라우팅 전문가 간 중복 지식 저장을 방지하고 전문가별 특화도를 높인다.

DeepSeek-V3는 이를 671B 파라미터로 확장하면서 혁신적인 보조 손실 없는(Auxiliary-Loss-Free) 로드 밸런싱을 도입했다. 256개 라우팅 전문가와 1개 공유 전문가를 사용하며, 토큰당 8개 전문가를 활성화한다. 기존 보조 손실 방식은 로드 밸런싱과 모델 성능 사이에 트레이드오프가 존재했지만, DeepSeek-V3는 각 전문가에 바이어스 항을 추가하고 학습 중 전문가 부하를 모니터링하여 바이어스를 동적으로 조정한다. 이 방식으로 학습 전체 과정에서 토큰 드롭 없이 균형 잡힌 로드를 유지한다.

MoE 모델 비교 종합

모델총 파라미터활성 파라미터전문가 수활성 전문가라우팅 방식로드 밸런싱
Switch Transformer1.6T~1/E1281Top-1보조 손실
GShard 600B600B~2/E20482Top-1 + Random용량 제한 + 보조 손실
ST-MoE-32B269B32B 상당642Top-2Router Z-Loss + 보조 손실
Mixtral 8x7B46.7B12.9B82Top-2보조 손실
DeepSeek-V3671B37B256+18+1Top-8Auxiliary-Loss-Free

라우팅 메커니즘 심층 분석

Top-K 라우팅 구현

라우팅 메커니즘은 MoE의 핵심이다. 입력 토큰의 히든 상태를 받아 어떤 전문가를 활성화할지 결정하며, 이 결정의 품질이 모델 전체 성능을 좌우한다. 다양한 라우팅 전략의 구현을 살펴보자.

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class TopKRouter(nn.Module):
    """Top-K 라우팅 with 노이즈 주입"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 2,
        noise_std: float = 1.0,
    ):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.noise_std = noise_std
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(
        self, x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch * seq_len, d_model)
        Returns:
            top_k_indices: (batch * seq_len, top_k)
            top_k_weights: (batch * seq_len, top_k)
            gate_logits: (batch * seq_len, num_experts)
        """
        gate_logits = self.gate(x)

        # 학습 시 노이즈 주입으로 탐색 촉진
        if self.training:
            noise = torch.randn_like(gate_logits) * self.noise_std
            gate_logits = gate_logits + noise

        gate_probs = F.softmax(gate_logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )
        # 재정규화
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_indices, top_k_weights, gate_logits

라우팅 전략 비교

전략논문활성 전문가장점단점
Top-1Switch Transformer1최소 통신 비용, 단순전문가 붕괴 위험 높음
Top-2ST-MoE, Mixtral2안정적, 성능 우수Top-1 대비 2배 통신
Top-1 + RandomGShard2탐색 촉진비결정적 추론
Top-K (K >= 4)DeepSeek-V38세분화된 전문가 활용전문가 수 많아야 유효
Expert ChoiceEC Routing가변완전 균형 보장토큰별 전문가 수 불균일

Expert Choice Routing은 전통적 라우팅과 관점을 뒤집는다. 토큰이 전문가를 선택하는 대신, 각 전문가가 자신이 처리할 토큰을 선택한다. 이 방식은 로드 밸런싱을 완벽하게 보장하지만, 특정 토큰이 어떤 전문가에도 선택되지 않을 수 있다는 문제가 있다.

학습 안정성과 로드 밸런싱

로드 밸런싱 손실 구현

MoE 학습에서 가장 흔한 문제는 전문가 붕괴(Expert Collapse)다. 특정 전문가에 토큰이 집중되면 해당 전문가만 학습되고, 나머지 전문가는 활용되지 않아 사실상 Dead Expert가 된다. 이를 방지하기 위해 보조 로드 밸런싱 손실을 사용한다.

def load_balancing_loss(
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    top_k: int = 2,
) -> torch.Tensor:
    """
    Switch Transformer 스타일 로드 밸런싱 손실 계산

    각 전문가에 할당된 토큰의 비율(f_i)과
    게이트 확률의 평균(p_i)의 내적을 최소화한다.
    균등 분포일 때 최솟값을 가진다.

    Args:
        gate_logits: (num_tokens, num_experts)
        top_k_indices: (num_tokens, top_k)
        num_experts: 전문가 수
        top_k: 활성 전문가 수
    Returns:
        스칼라 로드 밸런싱 손실
    """
    num_tokens = gate_logits.shape[0]
    gate_probs = F.softmax(gate_logits, dim=-1)

    # f_i: 각 전문가에 할당된 토큰 비율
    # one-hot 인코딩으로 할당 카운트
    expert_mask = F.one_hot(
        top_k_indices, num_classes=num_experts
    ).float()  # (num_tokens, top_k, num_experts)
    expert_mask = expert_mask.sum(dim=1)  # (num_tokens, num_experts)
    tokens_per_expert = expert_mask.sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / (num_tokens * top_k)

    # p_i: 각 전문가에 대한 평균 게이트 확률
    p = gate_probs.mean(dim=0)  # (num_experts,)

    # 밸런싱 손실: num_experts * sum(f_i * p_i)
    balance_loss = num_experts * (f * p).sum()

    return balance_loss

Router Z-Loss

ST-MoE에서 도입한 Router Z-Loss는 라우터 로짓의 크기를 제어하여 학습 발산을 방지한다. 로짓 값이 지나치게 커지면 softmax의 기울기가 사라지면서 학습이 불안정해지는데, Z-Loss가 이를 억제한다.

def router_z_loss(gate_logits: torch.Tensor) -> torch.Tensor:
    """
    ST-MoE Router Z-Loss 구현

    라우터 로짓의 log-sum-exp 값의 제곱 평균을 최소화.
    로짓이 과도하게 커지는 것을 방지하여 학습 안정성 확보.

    Args:
        gate_logits: (num_tokens, num_experts)
    Returns:
        스칼라 z-loss 값
    """
    # log(sum(exp(x)))의 제곱
    log_z = torch.logsumexp(gate_logits, dim=-1)  # (num_tokens,)
    z_loss = (log_z ** 2).mean()
    return z_loss


# 학습 루프에서의 통합
def compute_moe_loss(
    task_loss: torch.Tensor,
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    alpha_balance: float = 0.01,
    alpha_z: float = 0.001,
) -> torch.Tensor:
    """
    전체 MoE 손실 = 태스크 손실 + 밸런싱 손실 + Z-Loss
    """
    bal_loss = load_balancing_loss(
        gate_logits, top_k_indices, num_experts
    )
    z_loss = router_z_loss(gate_logits)

    total_loss = task_loss + alpha_balance * bal_loss + alpha_z * z_loss
    return total_loss

alpha_balance와 alpha_z의 설정이 중요하다. 보조 손실 계수가 너무 크면 모델 성능이 하락하고, 너무 작으면 로드 밸런싱 효과가 없어 전문가 붕괴가 발생한다. ST-MoE는 alpha_balance = 0.01, alpha_z = 0.001을 권장했다.

Auxiliary-Loss-Free 접근법

DeepSeek-V3가 도입한 보조 손실 없는 로드 밸런싱은 이 트레이드오프를 근본적으로 해소한다. 각 전문가에 학습 가능한 바이어스 항을 추가하고, 학습 중 전문가별 부하를 실시간 모니터링하여 바이어스를 조정한다.

class AuxLossFreeRouter(nn.Module):
    """DeepSeek-V3 스타일 Auxiliary-Loss-Free 라우터"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 8,
        bias_update_speed: float = 0.001,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gamma = bias_update_speed
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        # 바이어스는 그래디언트 업데이트가 아닌 규칙 기반 조정
        self.register_buffer(
            "expert_bias", torch.zeros(num_experts)
        )
        self.register_buffer(
            "target_load", torch.ones(num_experts) / num_experts
        )

    def forward(self, x: torch.Tensor):
        gate_logits = self.gate(x)
        # 바이어스를 더해 라우팅 결정에 반영
        biased_logits = gate_logits + self.expert_bias.unsqueeze(0)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-K 선택은 바이어스 포함 로짓으로 수행
        _, top_k_indices = torch.topk(
            biased_logits, self.top_k, dim=-1
        )
        # 가중치는 원본 게이트 확률에서 추출
        top_k_probs = torch.gather(gate_probs, 1, top_k_indices)
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 학습 시 바이어스 업데이트 (그래디언트 불필요)
        if self.training:
            with torch.no_grad():
                expert_load = F.one_hot(
                    top_k_indices.view(-1), self.num_experts
                ).float().sum(0)
                expert_load = expert_load / expert_load.sum()
                # 과부하 전문가는 바이어스 감소, 저부하 전문가는 증가
                self.expert_bias += self.gamma * (
                    self.target_load - expert_load
                )

        return top_k_indices, top_k_weights, gate_logits

이 방식의 핵심은 라우팅 결정(Top-K 선택)에는 바이어스를 반영하지만, 전문가 출력의 가중치 계산에는 원본 게이트 확률을 사용한다는 점이다. 바이어스가 모델의 그래디언트 흐름에 간섭하지 않으므로 성능 저하 없이 로드 밸런싱을 달성한다.

추론 최적화

Expert Offloading과 Capacity Budgeting

MoE 모델의 추론에서 가장 큰 도전은 메모리다. 모든 전문가를 GPU 메모리에 상주시키기 어려운 경우, 전문가 오프로딩(Expert Offloading)이 필수적이다. 활성화되지 않는 전문가를 CPU 메모리나 디스크에 내려놓고, 필요할 때만 GPU로 로딩하는 전략이다.

2025년부터 등장한 연구들은 투기적 디코딩(Speculative Decoding)과 전문가 오프로딩을 결합하여 성능을 크게 개선했다. SpecMoEOff는 투기적 디코딩으로 전문가 작업량을 확대하여 오프로딩 지연을 숨기고, 최대 2.5배 디코드 처리량 향상을 달성했다. MoE-Spec은 학습 없이 각 레이어의 전문가 용량을 고정하여 자주 사용되지 않는 전문가를 드롭함으로써 10-30% 처리량 향상을 보였다.

from dataclasses import dataclass
from typing import Dict, Optional
import torch

@dataclass
class ExpertCacheConfig:
    """전문가 캐시 설정"""
    gpu_cache_size: int = 4   # GPU에 상주할 전문가 수
    prefetch_count: int = 2    # 선제 로딩할 전문가 수
    device: str = "cuda:0"

class ExpertOffloadManager:
    """
    전문가 오프로딩 관리자.
    LRU 기반으로 GPU 캐시를 관리하고,
    다음 레이어 전문가를 선제 로딩한다.
    """
    def __init__(
        self,
        experts: Dict[int, torch.nn.Module],
        config: ExpertCacheConfig,
    ):
        self.experts = experts
        self.config = config
        self.gpu_cache: Dict[int, torch.nn.Module] = {}
        self.access_order: list = []

    def get_expert(self, expert_id: int) -> torch.nn.Module:
        """전문가를 GPU로 로딩 (LRU 캐시)"""
        if expert_id in self.gpu_cache:
            self.access_order.remove(expert_id)
            self.access_order.append(expert_id)
            return self.gpu_cache[expert_id]

        # GPU 캐시가 가득 찬 경우 LRU 전문가 제거
        if len(self.gpu_cache) >= self.config.gpu_cache_size:
            evict_id = self.access_order.pop(0)
            self.gpu_cache[evict_id].cpu()
            del self.gpu_cache[evict_id]

        # 전문가를 GPU로 이동
        expert = self.experts[expert_id].to(self.config.device)
        self.gpu_cache[expert_id] = expert
        self.access_order.append(expert_id)
        return expert

    def prefetch(self, predicted_expert_ids: list):
        """다음 레이어에서 사용될 전문가를 비동기 선제 로딩"""
        for eid in predicted_expert_ids[:self.config.prefetch_count]:
            if eid not in self.gpu_cache:
                self.get_expert(eid)

양자화와 Expert Pruning

MoE 모델의 추론 효율을 높이는 또 다른 축은 양자화와 전문가 프루닝이다. 모든 전문가를 동일한 비트로 양자화하는 대신, 자주 활성화되는 전문가는 고정밀도로 유지하고 드물게 사용되는 전문가는 공격적으로 양자화하는 적응형 양자화 전략이 효과적이다. Expert Pruning은 학습 후 활성화 빈도가 낮은 전문가를 완전히 제거하여 메모리 사용량을 줄인다. Mixtral 8x7B에서 2개 전문가를 프루닝해도 성능 하락이 미미하다는 실험 결과가 보고된 바 있다.

트러블슈팅

MoE 모델 학습과 추론에서 빈번하게 발생하는 문제와 해결책을 정리한다.

전문가 붕괴 (Expert Collapse)

증상: 학습 중 소수의 전문가만 활성화되고 나머지 전문가의 게이트 확률이 0에 수렴한다. 텐서보드에서 전문가별 토큰 할당 비율을 모니터링하면 특정 전문가에 90% 이상 토큰이 집중되는 것을 확인할 수 있다.

원인은 대부분 보조 손실 계수가 너무 작거나, 학습 초기 라우터 가중치 초기화가 불균형한 경우다. 해결책으로 alpha_balance를 0.01에서 시작하여 점진적으로 조정하고, 라우터 가중치를 작은 값으로 균일 초기화하며, 학습 초기 일정 스텝 동안 라우팅을 랜덤으로 수행하는 워밍업을 적용한다.

학습 발산 (Training Divergence)

증상: 학습 중 loss가 갑자기 급증하거나 NaN이 발생한다. MoE 모델에서 특히 bfloat16 학습 시 빈번하다.

Router Z-Loss를 적용하고, 라우터 로짓에 대해 별도의 float32 연산을 유지하는 것이 효과적이다. 학습률을 Dense 모델 대비 2-5배 낮추는 것도 권장된다.

All-to-All 통신 병목

증상: 분산 학습 시 전문가 병렬화의 All-to-All 통신이 전체 학습 시간의 30% 이상을 차지한다.

전문가 수를 디바이스 수의 배수로 설정하고, 용량 인자를 줄여 전송 데이터량을 감소시킨다. 통신과 연산을 오버랩하는 파이프라인 스케줄링도 효과적이다. Megablocks 라이브러리는 블록 희소 연산을 통해 All-to-All 통신을 회피하는 접근법을 제공한다.

추론 시 메모리 부족

증상: 전문가가 모두 GPU 메모리에 적재되지 않아 OOM이 발생한다.

Expert Offloading을 적용하거나, 전문가 프루닝으로 메모리 요구량 자체를 줄인다. 텐서 병렬화(Tensor Parallelism)와 전문가 병렬화를 조합하여 여러 GPU에 분산시키는 것도 방법이다.

실패 사례

보조 손실 계수 과다 설정

한 팀이 Mixtral 아키텍처 기반 커스텀 MoE 모델을 학습하면서 전문가 붕괴를 방지하기 위해 alpha_balance를 0.1로 높게 설정했다. 로드 밸런싱은 완벽했지만, 모델이 토큰의 의미와 무관하게 전문가를 균등하게 분배하는 방향으로 학습되면서 실제 태스크 성능이 Dense 모델보다 낮아졌다. 보조 손실이 주 손실을 압도하여 라우터가 의미 있는 전문가 특화를 학습하지 못한 것이다. alpha_balance를 0.005로 낮추고 Router Z-Loss를 함께 적용한 후에야 성능이 개선되었다.

추론 시 Top-1 전환의 함정

학습 시 Top-2 라우팅을 사용한 모델의 추론 비용을 줄이기 위해 Top-1으로 전환한 사례가 있다. 연산량은 절반으로 줄었지만, 벤치마크 점수가 8-12% 하락했다. 두 전문가의 출력이 상호 보완적으로 학습되었기 때문에, 하나만 사용하면 표현력이 크게 손실된다. 추론 비용 절감이 목적이라면 학습 단계에서부터 Top-1으로 설계하거나, 양자화와 프루닝으로 접근해야 한다.

공유 전문가 없는 세분화 실패

DeepSeek-MoE의 세분화된 전문가 분할을 적용하면서 공유 전문가(Shared Expert)를 생략한 경우가 있다. 모든 전문가가 기본적인 언어 패턴(관사, 전치사, 구두점 등)을 중복 학습하면서 전문가당 고유 지식의 양이 줄었고, 세분화의 이점이 상쇄되었다. 공유 전문가가 공통 지식을 흡수해야 라우팅 전문가가 고유한 패턴에 집중할 수 있다.

운영 체크리스트

MoE 모델을 학습부터 서빙까지 운영할 때 확인해야 할 항목을 정리한다.

학습 설정

  • 라우터 가중치는 작은 표준편차(0.01 이하)로 초기화했는가
  • 보조 손실 계수(alpha_balance)를 0.005-0.01 범위에서 시작했는가
  • Router Z-Loss를 활성화했는가
  • 학습률이 동일 크기 Dense 모델 대비 보수적(2-5배 낮음)인가
  • bfloat16 학습 시 라우터 로짓 연산은 float32로 유지하는가
  • 전문가 수가 디바이스 수의 배수인가
  • 용량 인자(Capacity Factor)를 1.0-1.5 범위로 설정했는가

모니터링

  • 전문가별 토큰 할당 비율을 실시간 추적하고 있는가
  • Dead Expert(할당 비율 1% 미만) 수를 경보로 설정했는가
  • 라우터 로짓의 평균 크기와 분산을 추적하고 있는가
  • All-to-All 통신 시간 대비 연산 시간 비율을 측정하고 있는가
  • 학습 loss의 급격한 변동(spike)에 대한 자동 감지가 있는가

추론 서빙

  • 전문가 오프로딩 전략을 결정했는가 (전량 GPU / 부분 오프로딩 / 디스크)
  • 양자화 수준을 전문가별로 차등 적용했는가 (자주 쓰이는 전문가는 고정밀도)
  • 전문가 프루닝 후 성능 저하를 벤치마크로 확인했는가
  • 배치 크기에 따른 라우팅 오버헤드를 측정했는가
  • 전문가 캐시 적중률(Hit Rate)을 모니터링하고 있는가
  • 투기적 디코딩 적용 시 Draft 모델과 전문가 예측 정확도를 확인했는가

파인튜닝

  • 전문가 드롭아웃 비율을 사전학습 대비 높였는가 (ST-MoE 권장)
  • 배치 크기를 사전학습 대비 줄였는가
  • 라우터를 고정(freeze)하고 전문가만 파인튜닝하는 옵션을 실험했는가
  • 파인튜닝 후 전문가 특화 분포가 태스크에 맞게 변화했는가 확인했는가

참고자료

Mixture of Experts (MoE) Architecture Paper Review and Production Scaling 2026

Mixture of Experts (MoE) Architecture Paper Review and Production Scaling 2026

Overview

As the parameter counts of large language models scale to trillions, the limitations of Dense models that activate all parameters for every token have become clear. Training and inference costs increase proportionally with parameter count. Mixture of Experts (MoE) has established itself as the most promising solution to this problem. MoE selectively activates only a subset of experts from the total parameters, dramatically reducing actual computation compared to Dense models while maintaining large model capacity.

In 2022, Switch Transformer opened the possibility of trillion-parameter scaling with single-expert routing, and GShard presented a distributed training pipeline for a 600B parameter model. ST-MoE improved both training stability and transfer learning quality simultaneously, while Mixtral 8x7B proved the practicality of open-source MoE models. DeepSeek-MoE and DeepSeek-V3 opened new directions with fine-grained expert segmentation and auxiliary-loss-free load balancing. In this article, we analyze the key contributions of each paper and organize strategies for operating MoE in production, covering routing mechanisms, training stability, and inference optimization.

MoE Core Concepts

The Principle of Sparse Activation

The core idea of MoE is conditional computation. Since only a subset of experts is activated for each input token, the model's total parameter count and the parameters actually used for computation are decoupled. For example, Mixtral 8x7B has 47B total parameters, but each token uses only 13B parameters. Computation is reduced to approximately 1/3.6 compared to a Dense 47B model, while quality remains similar or even better.

Dense vs Sparse MoE Comparison

ItemDense ModelSparse MoE
Active parametersTotal parameters = active parametersOnly 10-30% of total activated
FLOPs efficiencyProportional to parameter countProportional to active parameter count
Memory requirementEqual to parameter sizeFull parameters must be stored (more memory)
Training stabilityRelatively stableRouting instability, expert collapse risk
ScalabilityLinear cost increaseLow-cost scaling by adding experts
Inference latencyPredictableRouting overhead, expert loading delays
RepresentativeLLaMA, GPT-4 (estimated)Mixtral, DeepSeek-V3, Switch Transformer

Basic MoE Layer Implementation

The basic structure of an MoE layer implemented in PyTorch is as follows. The gating network computes weights for each expert given an input token, selects the Top-K experts, and sums their outputs.

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

class Expert(nn.Module):
    """Single expert FFN module"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(self.dropout(F.gelu(self.w1(x))))

class MoELayer(nn.Module):
    """Mixture of Experts layer"""
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int,
        top_k: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout) for _ in range(num_experts)
        ])
        # Gating network: input -> expert scores
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # (B*S, d_model)

        # Compute gating scores
        gate_logits = self.gate(x_flat)  # (B*S, num_experts)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Select Top-K experts
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )  # (B*S, top_k)
        # Renormalize selected expert weights
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # Sum expert outputs
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, k]  # (B*S,)
            weight = top_k_probs[:, k].unsqueeze(-1)  # (B*S, 1)
            for i in range(self.num_experts):
                mask = (expert_idx == i)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += weight[mask] * expert_output

        return output.view(batch_size, seq_len, d_model)

This implementation is for educational purposes. In actual production, optimized libraries like Megablocks or Tutel are used for efficient per-expert batch processing and All-to-All communication.

Key Paper Analysis

Switch Transformer (Fedus et al., 2022)

The key contribution of Switch Transformer was radically simplifying MoE routing. While previous MoE used Top-2 or more expert combinations, Switch Transformer introduced Switch Routing, which assigns only a single expert to each token. This decision is counterintuitive, but it was supported by experimental results showing that single-expert routing halves communication costs and reduces the overhead of the routing computation itself.

Based on the T5-Base architecture, it achieved up to 7x pre-training speed improvement at the same computational cost and scaled to trillion-parameter models. It also improved memory efficiency through bfloat16 mixed-precision training, and mitigated training instability by applying Expert Dropout together with auxiliary load balancing loss.

The expert Capacity Factor introduced by Switch Transformer determines the maximum number of tokens each expert can process. Tokens exceeding capacity are dropped, which causes information loss during training but is essential for balancing computational load.

GShard (Lepikhin et al., 2021)

GShard is a paper focused on system design for practically training MoE at massive distributed scale. It trained a 600B parameter multilingual translation model on 2,048 TPU v3 accelerators in just 4 days. There are three key technical contributions.

First, Expert Parallelism. It established the All-to-All communication pattern for distributing experts across multiple devices and routing tokens to the device where the target expert resides. Second, Random Routing. After deterministically selecting the Top-1 expert, the second expert is probabilistically selected with probability proportional to the gate weight. This approach promotes exploration during training. Third, it provided a lightweight API through XLA compiler extensions that can express various parallelism patterns with minimal code changes.

GShard's per-group capacity calculation formula C = 2N / (G * E) dynamically determines each expert's processing capacity considering group size G, number of experts E, and token count N.

ST-MoE (Zoph et al., 2022)

As the name ST-MoE implies, the core question of this paper is "Can MoE models be trained stably (Stable) and transferred effectively (Transferable)?" It trained a 269B parameter Sparse model (ST-MoE-32B) at the computational cost of a 32B Dense model, and achieved state-of-the-art performance for Sparse models for the first time across diverse benchmarks including SuperGLUE, ARC, and XSum.

For training stability, it introduced Router Z-Loss. This loss function controls the magnitude of router logits to prevent divergence during training. When used together with the existing auxiliary loss, training stability is significantly improved. It also systematically organized practical design recommendations: Top-2 routing, capacity factor 1.25 during training, and a maximum of 1 expert per core.

For fine-tuning, it presented experimental results showing that increasing the expert dropout rate, reducing batch size, and setting the learning rate conservatively are effective for improving transfer learning quality.

Mixtral 8x7B (Mistral AI, 2024)

Mixtral 8x7B is the model that decisively proved the practicality of open-source MoE models. Using the same Transformer architecture as Mistral 7B, it replaces each layer's FFN with 8 experts and selects Top-2 experts per token. Since only 13B of the total 47B parameters are activated, inference speed is similar to a 13B Dense model while quality rivals LLaMA 2 70B.

ItemMixtral 8x7BLLaMA 2 70BMistral 7B
Total parameters46.7B70B7.3B
Active parameters12.9B70B7.3B
Inference FLOPs~26B~140B~15B
MMLU70.669.860.1
ARC-Challenge66.464.655.5
Number of experts8--
Active experts2--

Analysis of Mixtral's expert specialization revealed that experts tend to specialize based on syntactic patterns rather than domains. Rather than specific experts responding only to code or math, expert selection changes based on sentence structure or token position.

DeepSeek-MoE and DeepSeek-V3

DeepSeek-MoE presented two key strategies aimed at "Ultimate Expert Specialization."

First, Fine-Grained Expert Segmentation. It subdivides the existing N experts into mN experts, reducing each expert's FFN intermediate dimension to 1/m to maintain the same total parameter count. Instead, mK experts are activated, enabling more flexible expert combinations. DeepSeek-MoE 16B achieved performance similar to LLaMA 2 7B with only 40% of the computation.

Second, Shared Expert isolation. Experts commonly used by all tokens are separated, preventing redundant knowledge storage among routed experts and increasing per-expert specialization.

DeepSeek-V3 scaled this to 671B parameters while introducing innovative Auxiliary-Loss-Free load balancing. It uses 256 routed experts and 1 shared expert, activating 8 experts per token. While traditional auxiliary loss approaches had a trade-off between load balancing and model performance, DeepSeek-V3 adds a bias term to each expert and dynamically adjusts the bias by monitoring expert load during training. This approach maintains balanced load throughout training without any token dropping.

MoE Model Comparison Summary

ModelTotal ParamsActive ParamsExpertsActive ExpertsRoutingLoad Balancing
Switch Transformer1.6T~1/E1281Top-1Auxiliary loss
GShard 600B600B~2/E20482Top-1 + RandomCapacity limit + aux loss
ST-MoE-32B269B32B equiv.642Top-2Router Z-Loss + aux loss
Mixtral 8x7B46.7B12.9B82Top-2Auxiliary loss
DeepSeek-V3671B37B256+18+1Top-8Auxiliary-Loss-Free

Routing Mechanism Deep Dive

Top-K Routing Implementation

The routing mechanism is the heart of MoE. It receives the hidden state of an input token and decides which experts to activate, and the quality of this decision determines overall model performance. Let us examine implementations of various routing strategies.

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class TopKRouter(nn.Module):
    """Top-K routing with noise injection"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 2,
        noise_std: float = 1.0,
    ):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.noise_std = noise_std
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(
        self, x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch * seq_len, d_model)
        Returns:
            top_k_indices: (batch * seq_len, top_k)
            top_k_weights: (batch * seq_len, top_k)
            gate_logits: (batch * seq_len, num_experts)
        """
        gate_logits = self.gate(x)

        # Inject noise during training to promote exploration
        if self.training:
            noise = torch.randn_like(gate_logits) * self.noise_std
            gate_logits = gate_logits + noise

        gate_probs = F.softmax(gate_logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )
        # Renormalization
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_indices, top_k_weights, gate_logits

Routing Strategy Comparison

StrategyPaperActive ExpertsProsCons
Top-1Switch Transformer1Minimum comm cost, simpleHigh expert collapse risk
Top-2ST-MoE, Mixtral2Stable, high performance2x communication vs Top-1
Top-1 + RandomGShard2Promotes explorationNon-deterministic inference
Top-K (K >= 4)DeepSeek-V38Fine-grained expert usageOnly effective with many experts
Expert ChoiceEC RoutingVariableGuarantees perfect balanceUneven expert count per token

Expert Choice Routing flips the perspective from traditional routing. Instead of tokens choosing experts, each expert chooses which tokens to process. This approach guarantees perfect load balancing, but has the issue that certain tokens may not be selected by any expert.

Training Stability and Load Balancing

Load Balancing Loss Implementation

The most common problem in MoE training is Expert Collapse. When tokens concentrate on specific experts, only those experts are trained, and the remaining experts go unused, effectively becoming Dead Experts. Auxiliary load balancing loss is used to prevent this.

def load_balancing_loss(
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    top_k: int = 2,
) -> torch.Tensor:
    """
    Switch Transformer style load balancing loss calculation

    Minimizes the dot product of the token assignment ratio per expert (f_i)
    and the mean gate probability (p_i).
    Achieves minimum value at uniform distribution.

    Args:
        gate_logits: (num_tokens, num_experts)
        top_k_indices: (num_tokens, top_k)
        num_experts: number of experts
        top_k: number of active experts
    Returns:
        Scalar load balancing loss
    """
    num_tokens = gate_logits.shape[0]
    gate_probs = F.softmax(gate_logits, dim=-1)

    # f_i: token assignment ratio per expert
    # Count assignments via one-hot encoding
    expert_mask = F.one_hot(
        top_k_indices, num_classes=num_experts
    ).float()  # (num_tokens, top_k, num_experts)
    expert_mask = expert_mask.sum(dim=1)  # (num_tokens, num_experts)
    tokens_per_expert = expert_mask.sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / (num_tokens * top_k)

    # p_i: mean gate probability per expert
    p = gate_probs.mean(dim=0)  # (num_experts,)

    # Balancing loss: num_experts * sum(f_i * p_i)
    balance_loss = num_experts * (f * p).sum()

    return balance_loss

Router Z-Loss

Router Z-Loss, introduced by ST-MoE, controls the magnitude of router logits to prevent training divergence. When logit values become excessively large, the softmax gradient vanishes and training becomes unstable -- Z-Loss suppresses this.

def router_z_loss(gate_logits: torch.Tensor) -> torch.Tensor:
    """
    ST-MoE Router Z-Loss implementation

    Minimizes the squared mean of log-sum-exp of router logits.
    Prevents logits from growing excessively to ensure training stability.

    Args:
        gate_logits: (num_tokens, num_experts)
    Returns:
        Scalar z-loss value
    """
    # Square of log(sum(exp(x)))
    log_z = torch.logsumexp(gate_logits, dim=-1)  # (num_tokens,)
    z_loss = (log_z ** 2).mean()
    return z_loss


# Integration in the training loop
def compute_moe_loss(
    task_loss: torch.Tensor,
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    alpha_balance: float = 0.01,
    alpha_z: float = 0.001,
) -> torch.Tensor:
    """
    Total MoE loss = task loss + balancing loss + Z-Loss
    """
    bal_loss = load_balancing_loss(
        gate_logits, top_k_indices, num_experts
    )
    z_loss = router_z_loss(gate_logits)

    total_loss = task_loss + alpha_balance * bal_loss + alpha_z * z_loss
    return total_loss

The settings of alpha_balance and alpha_z are critical. If the auxiliary loss coefficient is too large, model performance degrades; if too small, load balancing has no effect and expert collapse occurs. ST-MoE recommended alpha_balance = 0.01 and alpha_z = 0.001.

Auxiliary-Loss-Free Approach

The auxiliary-loss-free load balancing introduced by DeepSeek-V3 fundamentally resolves this trade-off. It adds a learnable bias term to each expert and monitors per-expert load in real-time during training to adjust the bias.

class AuxLossFreeRouter(nn.Module):
    """DeepSeek-V3 style Auxiliary-Loss-Free router"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 8,
        bias_update_speed: float = 0.001,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gamma = bias_update_speed
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        # Bias is adjusted rule-based, not via gradient updates
        self.register_buffer(
            "expert_bias", torch.zeros(num_experts)
        )
        self.register_buffer(
            "target_load", torch.ones(num_experts) / num_experts
        )

    def forward(self, x: torch.Tensor):
        gate_logits = self.gate(x)
        # Add bias to influence routing decisions
        biased_logits = gate_logits + self.expert_bias.unsqueeze(0)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-K selection uses biased logits
        _, top_k_indices = torch.topk(
            biased_logits, self.top_k, dim=-1
        )
        # Weights are extracted from original gate probabilities
        top_k_probs = torch.gather(gate_probs, 1, top_k_indices)
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # Update bias during training (no gradient needed)
        if self.training:
            with torch.no_grad():
                expert_load = F.one_hot(
                    top_k_indices.view(-1), self.num_experts
                ).float().sum(0)
                expert_load = expert_load / expert_load.sum()
                # Decrease bias for overloaded experts, increase for underloaded
                self.expert_bias += self.gamma * (
                    self.target_load - expert_load
                )

        return top_k_indices, top_k_weights, gate_logits

The key to this approach is that while the bias influences routing decisions (Top-K selection), the weight computation for expert outputs uses the original gate probabilities. Since the bias does not interfere with the model's gradient flow, load balancing is achieved without performance degradation.

Inference Optimization

Expert Offloading and Capacity Budgeting

The biggest challenge in MoE model inference is memory. When it is not feasible to keep all experts resident in GPU memory, Expert Offloading becomes essential. The strategy involves offloading inactive experts to CPU memory or disk and loading them to GPU only when needed.

Research emerging since 2025 has significantly improved performance by combining speculative decoding with expert offloading. SpecMoEOff hides offloading latency by expanding expert workload through speculative decoding, achieving up to 2.5x decode throughput improvement. MoE-Spec achieves 10-30% throughput improvement without training by fixing each layer's expert capacity and dropping infrequently used experts.

from dataclasses import dataclass
from typing import Dict, Optional
import torch

@dataclass
class ExpertCacheConfig:
    """Expert cache configuration"""
    gpu_cache_size: int = 4   # Number of experts to keep on GPU
    prefetch_count: int = 2    # Number of experts to prefetch
    device: str = "cuda:0"

class ExpertOffloadManager:
    """
    Expert offloading manager.
    Manages GPU cache based on LRU and
    prefetches next-layer experts.
    """
    def __init__(
        self,
        experts: Dict[int, torch.nn.Module],
        config: ExpertCacheConfig,
    ):
        self.experts = experts
        self.config = config
        self.gpu_cache: Dict[int, torch.nn.Module] = {}
        self.access_order: list = []

    def get_expert(self, expert_id: int) -> torch.nn.Module:
        """Load expert to GPU (LRU cache)"""
        if expert_id in self.gpu_cache:
            self.access_order.remove(expert_id)
            self.access_order.append(expert_id)
            return self.gpu_cache[expert_id]

        # Evict LRU expert if GPU cache is full
        if len(self.gpu_cache) >= self.config.gpu_cache_size:
            evict_id = self.access_order.pop(0)
            self.gpu_cache[evict_id].cpu()
            del self.gpu_cache[evict_id]

        # Move expert to GPU
        expert = self.experts[expert_id].to(self.config.device)
        self.gpu_cache[expert_id] = expert
        self.access_order.append(expert_id)
        return expert

    def prefetch(self, predicted_expert_ids: list):
        """Asynchronously prefetch experts expected in the next layer"""
        for eid in predicted_expert_ids[:self.config.prefetch_count]:
            if eid not in self.gpu_cache:
                self.get_expert(eid)

Quantization and Expert Pruning

Another axis for improving MoE model inference efficiency is quantization and expert pruning. Instead of quantizing all experts at the same bit level, an adaptive quantization strategy that maintains frequently activated experts at high precision while aggressively quantizing rarely used experts is effective. Expert Pruning completely removes experts with low activation frequency after training to reduce memory usage. It has been reported that pruning 2 experts from Mixtral 8x7B results in minimal performance degradation.

Troubleshooting

The following section organizes commonly occurring problems and solutions in MoE model training and inference.

Expert Collapse

Symptom: Only a small number of experts are activated during training, and the gate probabilities of the remaining experts converge to 0. Monitoring per-expert token assignment ratios in TensorBoard shows that over 90% of tokens are concentrated on specific experts.

The cause is usually that the auxiliary loss coefficient is too small, or the initial router weight initialization is imbalanced. Solutions include starting alpha_balance at 0.01 and gradually adjusting, initializing router weights uniformly with small values, and applying a warmup that performs random routing for a certain number of steps at the beginning of training.

Training Divergence

Symptom: Loss suddenly spikes or NaN occurs during training. This is especially common during bfloat16 training with MoE models.

Applying Router Z-Loss and maintaining separate float32 computation for router logits is effective. Lowering the learning rate by 2-5x compared to Dense models is also recommended.

All-to-All Communication Bottleneck

Symptom: During distributed training, All-to-All communication for expert parallelism accounts for over 30% of total training time.

Set the number of experts as a multiple of the number of devices, and reduce the capacity factor to decrease data transfer volume. Pipeline scheduling that overlaps communication and computation is also effective. The Megablocks library provides an approach that avoids All-to-All communication through block-sparse operations.

Inference Memory Shortage

Symptom: OOM occurs because all experts cannot fit in GPU memory.

Apply Expert Offloading, or reduce memory requirements through expert pruning. Combining Tensor Parallelism with Expert Parallelism to distribute across multiple GPUs is also an option.

Failure Cases

Excessive Auxiliary Loss Coefficient

A team training a custom MoE model based on the Mixtral architecture set alpha_balance to 0.1 (high) to prevent expert collapse. Load balancing was perfect, but the model learned to distribute tokens equally among experts regardless of token semantics, resulting in actual task performance lower than a Dense model. The auxiliary loss dominated the main loss, preventing the router from learning meaningful expert specialization. Performance only improved after lowering alpha_balance to 0.005 and applying Router Z-Loss together.

The Pitfall of Switching to Top-1 at Inference

There was a case where a model trained with Top-2 routing was switched to Top-1 to reduce inference costs. Computation was halved, but benchmark scores dropped by 8-12%. Because the outputs of two experts were trained to be complementary, using only one results in significant loss of expressiveness. If reducing inference costs is the goal, either design with Top-1 from the training stage, or approach through quantization and pruning.

Fine-Grained Segmentation Failure Without Shared Experts

There was a case where DeepSeek-MoE's fine-grained expert segmentation was applied without the Shared Expert. All experts redundantly learned basic language patterns (articles, prepositions, punctuation, etc.), reducing the amount of unique knowledge per expert and negating the benefits of segmentation. The shared expert must absorb common knowledge so that routed experts can focus on unique patterns.

Operations Checklist

The following organizes items to verify when operating MoE models from training to serving.

Training Setup

  • Router weights initialized with small standard deviation (0.01 or less)
  • Auxiliary loss coefficient (alpha_balance) started in the 0.005-0.01 range
  • Router Z-Loss activated
  • Learning rate conservative compared to same-size Dense model (2-5x lower)
  • Router logit computation maintained in float32 during bfloat16 training
  • Number of experts is a multiple of the number of devices
  • Capacity Factor set in the 1.0-1.5 range

Monitoring

  • Real-time tracking of per-expert token assignment ratios
  • Alert set for Dead Expert count (assignment ratio under 1%)
  • Tracking average magnitude and variance of router logits
  • Measuring All-to-All communication time vs computation time ratio
  • Automatic detection for sudden training loss spikes

Inference Serving

  • Expert offloading strategy decided (full GPU / partial offloading / disk)
  • Quantization level applied differentially per expert (high precision for frequently used experts)
  • Performance degradation after expert pruning verified via benchmarks
  • Routing overhead measured for different batch sizes
  • Expert cache hit rate monitored
  • Draft model and expert prediction accuracy verified when applying speculative decoding

Fine-tuning

  • Expert dropout rate increased compared to pre-training (ST-MoE recommendation)
  • Batch size reduced compared to pre-training
  • Experimented with freezing the router and fine-tuning only experts
  • Verified that expert specialization distribution changed appropriately for the task after fine-tuning

References

Quiz

Q1: What is the main topic covered in "Mixture of Experts (MoE) Architecture Paper Review and Production Scaling 2026"?

Core paper analysis of Mixture of Experts (MoE) architecture. Covering Switch Transformer, GShard, ST-MoE, and Mixtral -- routing mechanisms, expert parallelism, and load balancing strategies from a production perspective.

Q2: What is MoE Core Concepts? The Principle of Sparse Activation The core idea of MoE is conditional computation. Since only a subset of experts is activated for each input token, the model's total parameter count and the parameters actually used for computation are decoupled.

Q3: Explain the core concept of Key Paper Analysis. Switch Transformer (Fedus et al., 2022) The key contribution of Switch Transformer was radically simplifying MoE routing. While previous MoE used Top-2 or more expert combinations, Switch Transformer introduced Switch Routing, which assigns only a single expert to each token.

Q4: What are the key aspects of Routing Mechanism Deep Dive? Top-K Routing Implementation The routing mechanism is the heart of MoE. It receives the hidden state of an input token and decides which experts to activate, and the quality of this decision determines overall model performance.

Q5: How does Training Stability and Load Balancing work? Load Balancing Loss Implementation The most common problem in MoE training is Expert Collapse. When tokens concentrate on specific experts, only those experts are trained, and the remaining experts go unused, effectively becoming Dead Experts.