LabHub

Blog

Mixture of Experts (MoE) Architecture Deep Analysis: Evolution from Switch Transformer to Mixtral and Efficient Scaling Strategies

2026-03-11 · 17 min read · #ai-papers#moe#switch-transformer#mixtral#model-architecture#2026-03

한국어English日本語中文

Mixture of Experts Architecture Analysis

Introduction

The most intuitive way to raise the performance of a large language model (LLM) is to add parameters. In a dense model, however, doubling the parameter count nearly doubles the compute needed for training and inference as well. The Mixture of Experts (MoE) architecture offers an elegant answer to this. It raises the model's total parameter count sharply while activating only a fraction of those parameters for each input token, keeping the compute cost constant.

The idea goes back to the original MoE paper proposed by Jacobs et al. in 1991, but it has advanced sharply over the last few years. Google's Switch Transformer (2021) cut the complexity of MoE dramatically with single-expert routing, Mistral AI's Mixtral 8x7B (2024) set a new bar for open-source LLMs with Top-2 routing, and DeepSeek-MoE (2024) maximized expert specialization with fine-grained expert splitting. DeepSeek-V3 (2024) then reached top-tier efficiency with a structure that activates only 37B of its 671B parameters.

This article covers the whole picture: the basic principles of the MoE architecture, the design differences between the major models, a comparison of routing mechanisms, training and inference optimization strategies, operational cautions, failure cases, and a production checklist.

Core Principles of the MoE Architecture

The Gating Mechanism

The heart of an MoE layer is the gating network. Given an input token x, the gating network outputs a probability distribution over the experts.

G(x)=softmax(Wg⋅x)G(x) = \text{softmax}(W_g \cdot x)

Here W_g is the learnable gating weight matrix. Each element of the output vector G(x) is the selection probability of the corresponding expert.

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

class GatingNetwork(nn.Module):
    """Basic MoE gating network"""
    def __init__(self, d_model, num_experts, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        logits = self.gate(x)  # (batch_size, seq_len, num_experts)

        # Select the Top-k experts
        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

Sparse Activation

In a dense model every parameter is activated on every input, but in Sparse MoE only the Top-k experts are. Since only k of the N experts join the computation, capacity can grow by a factor of N while compute stays at roughly a factor of k.

The output of an MoE layer is computed as the weighted sum of the selected experts.

y=∑i∈Top-kG(x)i⋅Ei(x)y = \sum_{i \in \text{Top-k}} G(x)_i \cdot E_i(x)

Here E_i(x) is the output of the i-th expert network.

class MoELayer(nn.Module):
    """Basic Sparse MoE layer"""
    def __init__(self, d_model, d_ff, num_experts, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = GatingNetwork(d_model, num_experts, top_k)

        # Each expert is an independent FFN
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model)
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        gates, indices = self.gate(x)  # gates: (B,T,k), indices: (B,T,k)

        # Initialize the output
        output = torch.zeros_like(x)

        # Process the tokens expert by expert
        for i in range(self.num_experts):
            # Mask of the tokens assigned to the i-th expert
            expert_mask = (indices == i).any(dim=-1)  # (B, T)
            if not expert_mask.any():
                continue

            # Extract those tokens and run the expert
            expert_input = x[expert_mask]
            expert_output = self.experts[i](expert_input)

            # Apply the gate weights
            gate_values = gates[indices == i]
            output[expert_mask] += gate_values.unsqueeze(-1) * expert_output

        return output

Capacity Factor

The upper bound on how many tokens each expert can process is controlled by the Capacity Factor (CF). Under an ideal even split each expert handles T / N tokens, and multiplying by CF gives the actual capacity.

Expert Capacity=CF×TN\text{Expert Capacity} = \text{CF} \times \frac{T}{N}

A CF of 1.0 assumes a perfectly even split; in practice 1.25~1.5 is used. Too low a CF drops tokens (overflow), too high a CF wastes memory.

Switch Transformer Analysis

The Innovation of Single-Expert Routing

Google Brain's Switch Transformer (Fedus et al., 2021) broke the conventional wisdom of MoE. Activating Top-2 or more experts had been considered necessary for stable training, but Switch Transformer showed that even Top-1 single-expert routing can reach excellent performance.

Single-expert routing has three advantages.

  1. Less router computation: selecting only one expert simplifies the gating computation
  2. More efficient expert capacity: since each token is assigned to exactly one expert, the batch size can be doubled at the same Capacity Factor
  3. Lower communication cost: in distributed training a token only has to be sent to one expert device
class SwitchRouter(nn.Module):
    """Switch Transformer router (Top-1)"""
    def __init__(self, d_model, num_experts, capacity_factor=1.25):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        # Gating logits
        logits = self.gate(x)  # (B, T, num_experts)
        probs = F.softmax(logits, dim=-1)

        # Top-1 selection
        gate_values, expert_indices = probs.max(dim=-1)  # (B, T)

        # Compute the expert capacity
        capacity = int(self.capacity_factor * T / self.num_experts)

        # Drop the tokens that exceed capacity
        dispatch_mask = torch.zeros(B, T, self.num_experts, dtype=torch.bool)
        for i in range(self.num_experts):
            expert_mask = (expert_indices == i)
            # Drop the overflow beyond capacity
            positions = expert_mask.nonzero(as_tuple=True)
            if len(positions[1]) > capacity:
                drop_indices = positions[1][capacity:]
                expert_mask[positions[0][capacity:], drop_indices] = False
            dispatch_mask[:, :, i] = expert_mask

        return gate_values, expert_indices, dispatch_mask

Load Balancing Loss

Switch Transformer introduces an auxiliary loss to prevent load imbalance across experts. It is defined as the dot product of the fraction of tokens assigned to each expert (f_i) and that expert's mean routing probability (P_i).

Lbalance=α⋅N∑i=1Nfi⋅PiL_{\text{balance}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i 
def load_balancing_loss(gates, expert_indices, num_experts, alpha=0.01):
    """Switch Transformer load balancing auxiliary loss"""
    B, T = expert_indices.shape

    # f_i: fraction of tokens assigned to each expert
    f = torch.zeros(num_experts, device=gates.device)
    for i in range(num_experts):
        f[i] = (expert_indices == i).float().sum() / (B * T)

    # P_i: mean routing probability of each expert
    probs = F.softmax(gates, dim=-1)  # probabilities over all experts
    P = probs.mean(dim=[0, 1])  # (num_experts,)

    # Balancing loss
    loss = alpha * num_experts * (f * P).sum()
    return loss

Performance Results

Switch Transformer reached a 7x faster pretraining speed at the same compute as T5-Base. It kept training stable while scaling to 1.6 trillion (1.6T) parameters, thanks to bfloat16 training and the introduction of the router z-loss.

Mixtral 8x7B Architecture

Structural Design

Mistral AI's Mixtral 8x7B (2024) is the representative case of an open-source MoE model. Built on the same Transformer architecture as Mistral 7B, it replaces the FFN block of every layer with 8 experts and applies Top-2 routing.

The key specifications are as follows.

class MixtralMoELayer(nn.Module):
    """Mixtral 8x7B-style MoE layer"""
    def __init__(self, d_model=4096, d_ff=14336, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Router: linear gating
        self.gate = nn.Linear(d_model, num_experts, bias=False)

        # 8 independent SwiGLU FFN experts
        self.experts = nn.ModuleList([
            SwiGLUExpert(d_model, d_ff) for _ in range(num_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (B*T, D)

        # Compute the routing probabilities
        logits = self.gate(x_flat)  # (B*T, 8)
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)  # Normalize within the Top-k

        # Weighted sum of the expert outputs
        output = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = indices[:, i]  # (B*T,)
            gate_weight = weights[:, i]  # (B*T,)

            for j in range(self.num_experts):
                mask = (expert_idx == j)
                if mask.any():
                    expert_out = self.experts[j](x_flat[mask])
                    output[mask] += gate_weight[mask].unsqueeze(-1) * expert_out

        return output.view(B, T, D)


class SwiGLUExpert(nn.Module):
    """SwiGLU-based FFN expert used in Mixtral"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

Benefits of Top-2 Routing

Mixtral adopted Top-2 for the following reasons.

DeepSeek-MoE: Fine-Grained Expert Segmentation

Core Strategies

DeepSeek-MoE (2024) maximized expert specialization with two core strategies.

Strategy 1: Fine-grained Expert Segmentation

Where a conventional MoE picks Top-2 out of 16 large experts, DeepSeek-MoE splits the same total parameters into 64 small experts and picks Top-8. Cutting the FFN intermediate dimension of each expert to 1/4 and keeping 4x as many experts increases the number of possible expert combinations dramatically.

class DeepSeekMoELayer(nn.Module):
    """DeepSeek-MoE-style layer: fine-grained experts + shared experts"""
    def __init__(
        self,
        d_model=2048,
        d_ff=10944,
        num_routed_experts=64,
        num_shared_experts=2,
        top_k=6,
    ):
        super().__init__()
        self.top_k = top_k

        # Fine-grained routed experts (small FFN)
        expert_d_ff = d_ff // 4  # 1/4 the size of the conventional one
        self.routed_experts = nn.ModuleList([
            SwiGLUExpert(d_model, expert_d_ff)
            for _ in range(num_routed_experts)
        ])

        # Shared experts (always active)
        self.shared_experts = nn.ModuleList([
            SwiGLUExpert(d_model, d_ff)
            for _ in range(num_shared_experts)
        ])

        # Router
        self.gate = nn.Linear(d_model, num_routed_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)

        # Shared expert output (always active)
        shared_out = sum(expert(x_flat) for expert in self.shared_experts)

        # Routed expert output
        logits = self.gate(x_flat)
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)

        routed_out = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = indices[:, i]
            gate_weight = weights[:, i]
            for j in range(len(self.routed_experts)):
                mask = (expert_idx == j)
                if mask.any():
                    out = self.routed_experts[j](x_flat[mask])
                    routed_out[mask] += gate_weight[mask].unsqueeze(-1) * out

        return (shared_out + routed_out).view(B, T, D)

Strategy 2: Shared Expert Isolation

Some experts are designated as shared experts and stay active for every token. The shared experts take on the common knowledge (general language patterns, syntactic structure and so on) so that the routed experts can concentrate on specialized knowledge. This reduces knowledge duplication among the routed experts and raises the level of specialization.

Performance

DeepSeek-MoE 16B matched the performance of the dense DeepSeek 7B with roughly 40% of the compute. DeepSeek-MoE 2B matched GShard 2.9B, which uses 1.5x more expert parameters and compute.

Comparing Routing Mechanisms

Top-k Routing

The most common approach, which selects the top k experts from the gating network.

Expert Choice Routing

In the Expert Choice approach proposed by Google Research (2022), instead of tokens choosing experts, each expert chooses the tokens it will process. Because every expert picks its own Top-k tokens, perfect load balancing is guaranteed.

class ExpertChoiceRouter(nn.Module):
    """Expert Choice routing: the expert picks the tokens"""
    def __init__(self, d_model, num_experts, capacity_factor=1.0):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (N, D) where N = B*T
        N = x_flat.shape[0]

        logits = self.gate(x_flat)  # (N, num_experts)
        scores = F.softmax(logits, dim=0)  # softmax over the token dimension

        # Number of tokens each expert will handle
        k = int(self.capacity_factor * N / self.num_experts)

        # Select the Top-k tokens for each expert
        top_k_scores, top_k_indices = torch.topk(
            scores.t(), k, dim=-1
        )  # (num_experts, k)

        return top_k_scores, top_k_indices

Hash Routing

This assigns tokens to experts with a hash function instead of learnable gating. There is no routing computation, so the overhead is close to zero and routing instability disappears at the root. The trade-off is that an assignment optimized for the characteristics of the input is impossible, so performance is somewhat lower.

MoE Model Comparison Table

ItemSwitch TransformerMixtral 8x7BDeepSeek-MoE 16BDeepSeek-V3
Release date2021.012024.012024.012024.12
Total parameters1.6T (max)46.7B16.4B671B
Active parametersVariable12.9B2.8B37B
Expert count128864 routed + 2 shared256 routed + 1 shared
Active experts1 (Top-1)2 (Top-2)6 (Top-6)8 (Top-8)
Routing methodLearned Top-1Learned Top-2Learned Top-kLearned Top-k
Load balancingAuxiliary LossNot disclosedAuxiliary LossAuxiliary-Loss-Free
Shared expertsNoneNoneYes (2)Yes (1)
Base architectureT5 (Encoder-Decoder)Mistral 7B (Decoder)Decoder-onlyDecoder-only
Headline result7x faster training6x faster inference than Llama 2 70BMatches dense 7B with 40% of the computeBest open-source performance

Training and Inference Optimization

Expert Parallelism

Expert Parallelism (EP) is central to distributed training of MoE models. Each expert is placed on a different GPU, and tokens are sent to their expert over All-to-All communication.

# DeepSpeed MoE configuration example
deepspeed_config = {
    "train_batch_size": 256,
    "fp16": {"enabled": True},
    "zero_optimization": {"stage": 2},
    "moe": {
        "enabled": True,
        "ep_size": 8,            # Expert Parallelism: experts spread over 8 GPUs
        "num_experts": 64,
        "top_k": 2,
        "capacity_factor": 1.25,
        "min_capacity": 4,
        "use_residual": True,    # Residual MoE
        "moe_param_group": True  # Separate parameter group for the experts
    }
}

EP is generally combined with Data Parallelism (DP). Configuring EP=8 and DP=8 across 64 GPUs, for example, puts one expert group on 8 GPUs while 8 replicas process data in parallel.

Load Balancing Strategies

Here are the load balancing techniques used for training stability.

1. Auxiliary Loss (Switch Transformer)

The basic auxiliary loss, which minimizes the product of the expert assignment ratio and the routing probability. The coefficient alpha is tuned in the 0.01~0.1 range.

2. Router z-loss (ST-MoE)

This penalizes the magnitude of the gating logits to improve training stability. Suppressing large logit values that enter the router reduces floating-point rounding error.

Lz=1BT∑b,t(log⁡∑i=1Negb,t,i)2L_z = \frac{1}{BT} \sum_{b,t} \left(\log \sum_{i=1}^{N} e^{g_{b,t,i}}\right)^2

3. Auxiliary-Loss-Free (DeepSeek-V3)

This removes the auxiliary loss entirely and adds a learnable bias term to each expert to adjust the routing probability dynamically. It eliminates at the root the negative effect an auxiliary loss has on model performance.

Inference Optimization

The biggest challenge in MoE inference is memory usage. Every expert has to be resident in GPU memory, so even the experts that never take part in the computation occupy memory.

The main optimization techniques are as follows.

Operational Considerations

Memory Management

An MoE model's total parameter count is far larger than its active parameter count, so memory planning needs care. For Mixtral 8x7B the active parameters come to 12.9B, but the full 46.7B has to be loaded onto the GPU. At FP16 that needs roughly 93GB of GPU memory, which a single A100 80GB cannot cover.

Expert Collapse

Expert collapse is the most common problem in MoE training. Most tokens concentrate on a handful of experts while the rest go effectively unused. The cause and the countermeasures are as follows.

Routing Instability

Routing instability, where routing decisions change abruptly early in training, can occur. The following techniques prevent it.

Failure Cases and Fixes

Case 1: One Expert OOMs at Inference

Symptom: tokens concentrate excessively on a single expert for a particular batch, causing OOM on that GPU

Cause: when the distribution of the input data differs greatly from the training data, the router makes decisions biased toward one expert

Fix:

Case 2: Loss Diverges During Training

Symptom: after several thousand steps early in training, the loss suddenly diverges

Cause: the router's logit values grow too large, the softmax output approaches 0 or 1, and the gradient explodes

Fix:

Case 3: Expert Parallelism Communication Bottleneck

Symptom: GPU utilization is low and most of the time goes to All-to-All communication

Cause: there are many GPUs relative to the number of experts, so the communication overhead exceeds the computation time

Fix:

Production Checklist

These are the items to check when deploying an MoE model to production.

Memory and infrastructure

Training stability

Inference optimization

Monitoring

References

Conclusion

The MoE architecture is the most practical answer to the central challenge of LLM scaling: a bigger model for less cost. Switch Transformer opened up the possibility of single-expert routing, Mixtral proved MoE practical in the open-source ecosystem, and DeepSeek-MoE/V3 set a new standard for efficiency with fine-grained expert strategies.

Challenges remain all the same. Expert collapse and routing instability need continuous management during training, and the high memory requirement relative to the active parameters raises deployment cost. Recent techniques such as Expert Choice routing and Auxiliary-Loss-Free balancing are gradually solving these problems, but a complete answer has not arrived yet.

The expected direction for MoE from here is as follows. Dynamic creation and removal of experts, multimodal expert specialization, adaptive control of the number of experts at inference time, and MoE architectures co-designed closely with the hardware will appear. As models like DeepSeek-V3 have shown, MoE has already become a core component of frontier AI models, and the trend will only accelerate.

Comments

No comments yet.

Sign in to leave a comment