How Mixture of Experts Works: Sparse Gating, Expert Routing, and the Architecture Behind Efficient Large Language Models

Dense transformers scale predictably: double the parameters and you roughly double the compute per token. That relationship is useful for benchmarking but expensive in practice. Mixture of Experts (MoE) breaks that coupling. A 140B parameter MoE model can consume the same compute per forward pass as a 14B dense model, because only a fraction of its parameters are active for any given token. Understanding exactly how that works, and where it breaks, is what this article covers.

The Core Idea: Conditional Computation

In a standard transformer, every parameter in every layer participates in every forward pass. The MLP sublayer after attention, which typically expands the hidden dimension by a factor of 4 and contracts it back, processes every token identically. This is compute-efficient for the hardware but parameter-inefficient for learning: the same weights must handle every kind of token, every linguistic context, every domain.

MoE replaces the single MLP sublayer with a set of N independent MLPs called experts, and adds a router (also called the gating network) that decides which experts each token should be sent to. Only the selected experts compute their outputs for that token. The rest are skipped entirely.

The result is a model where total parameters scale with N (number of experts) but FLOPs per token scale with the number of active experts per token, which is typically 1 or 2 regardless of N.

The Gating Mechanism

The router is a small linear layer followed by a softmax. It takes the token’s hidden state as input and produces a probability distribution over all N experts.

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

class MoERouter(nn.Module):
    def __init__(self, hidden_dim: int, num_experts: int, top_k: int):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        # The router is just a linear projection: no bias, no activation
        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)

    def forward(self, hidden_states: torch.Tensor):
        # hidden_states: [batch_size * seq_len, hidden_dim]
        router_logits = self.gate(hidden_states)  # [tokens, num_experts]
        router_probs = F.softmax(router_logits, dim=-1)

        # Select top-k experts per token
        top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)

        # Renormalize so selected weights sum to 1
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_probs, top_k_indices, router_logits

The top-k selection is the key operation. With top_k=2, each token is routed to exactly 2 experts. The token’s output from this layer is a weighted sum of those two experts’ outputs, weighted by the normalized router probabilities.

class MoELayer(nn.Module):
    def __init__(self, hidden_dim: int, intermediate_dim: int, num_experts: int, top_k: int):
        super().__init__()
        self.router = MoERouter(hidden_dim, num_experts, top_k)
        self.experts = nn.ModuleList([
            FeedForwardExpert(hidden_dim, intermediate_dim)
            for _ in range(num_experts)
        ])
        self.top_k = top_k

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_tokens, hidden_dim = hidden_states.shape
        top_k_probs, top_k_indices, router_logits = self.router(hidden_states)

        # Flatten for dispatch: [batch_tokens * top_k]
        flat_expert_ids = top_k_indices.view(-1)
        flat_probs = top_k_probs.view(-1)

        output = torch.zeros_like(hidden_states)

        for expert_id in range(len(self.experts)):
            # Find which (token, k) pairs go to this expert
            token_mask = (flat_expert_ids == expert_id)
            if not token_mask.any():
                continue

            # Map flat indices back to token indices
            token_indices = torch.where(token_mask)[0] // self.top_k
            expert_weights = flat_probs[token_mask]

            expert_input = hidden_states[token_indices]
            expert_output = self.experts[expert_id](expert_input)

            # Weighted accumulation
            output.index_add_(
                0,
                token_indices,
                expert_output * expert_weights.unsqueeze(-1)
            )

        return output, router_logits

This loop is illustrative. Production implementations use grouped GEMM operations to process all experts in parallel, dispatching tokens to experts in batched matrix multiplications rather than iterating.

The Load Balancing Problem

If you train an MoE model with only the task loss (cross-entropy on predicted tokens), the router learns to route all tokens to a small number of experts. The model discovers that a few generalist experts are easier to optimize and the router collapses onto them. The remaining experts receive almost no gradient signal and never specialize. This is called expert collapse, and it renders MoE no better than a much smaller dense model.

The fix is an auxiliary load balancing loss that penalizes uneven routing.

def load_balancing_loss(router_logits: torch.Tensor, num_experts: int, top_k: int) -> torch.Tensor:
    """
    Compute the auxiliary load balancing loss from the Switch Transformer paper.
    Encourages uniform expert utilization.

    router_logits: [num_tokens, num_experts]
    """
    num_tokens = router_logits.shape[0]

    # Fraction of tokens routed to each expert
    router_probs = F.softmax(router_logits, dim=-1)
    _, top_k_indices = torch.topk(router_probs, top_k, dim=-1)

    # One-hot encode the routing decisions
    expert_mask = F.one_hot(top_k_indices, num_classes=num_experts).float()
    # expert_mask: [num_tokens, top_k, num_experts]
    # Collapse top_k dimension
    expert_mask = expert_mask.sum(dim=1)  # [num_tokens, num_experts]

    # f_i: fraction of tokens dispatched to expert i
    tokens_per_expert = expert_mask.sum(dim=0)  # [num_experts]
    f = tokens_per_expert / num_tokens

    # P_i: average router probability assigned to expert i
    P = router_probs.mean(dim=0)  # [num_experts]

    # Loss is the dot product of f and P, scaled by num_experts
    # Minimizing this encourages both f and P to be uniform (1/num_experts each)
    loss = num_experts * (f * P).sum()
    return loss

The total training loss is task_loss + alpha * load_balancing_loss, where alpha is typically in the range 0.001 to 0.01. Too small and collapse still happens. Too large and the router ignores task-relevant routing and distributes tokens randomly, which hurts model quality.

Mixtral 8x7B uses alpha=0.01. The Switch Transformer paper found that values above 0.01 caused quality degradation. GShard and later Mixtral papers confirm this range is fairly stable across model sizes.

The Capacity Factor

Even with load balancing, token distribution is not perfectly uniform at runtime. Some experts receive more tokens than others in a given batch. This creates a throughput problem for parallel hardware: the slowest expert determines the step time.

The capacity factor (CF) sets a hard limit on how many tokens each expert can process per batch. If an expert would receive more tokens than its capacity, the excess tokens are dropped (their output is set to zero or they fall back to the residual stream).

def compute_expert_capacity(num_tokens: int, num_experts: int, top_k: int, capacity_factor: float) -> int:
    """
    Capacity factor of 1.0 means each expert gets exactly its fair share.
    CF > 1.0 adds buffer for imbalance. CF < 1.0 drops tokens aggressively.
    """
    tokens_per_expert_ideal = (num_tokens * top_k) / num_experts
    capacity = int(tokens_per_expert_ideal * capacity_factor)
    return max(capacity, 1)  # At least 1 token per expert

With capacity_factor=1.25, each expert can handle 25% more than its ideal share, absorbing normal routing variation without dropping tokens. With capacity_factor=1.0, any imbalance causes drops. Values below 1.0 aggressively drop tokens to keep GPU steps uniform but hurt output quality for dropped tokens.

Token dropping is not catastrophic in practice because dropped tokens still pass through the residual stream (their value from the previous layer is carried forward unchanged). The MoE layer contributes nothing for those tokens, but the model still produces output. Quality degrades gracefully rather than failing.

The tradeoff is clear: higher capacity factors reduce drops but require larger buffers in GPU memory and potentially slower expert compute kernels (because they must handle variable-length inputs). Lower capacity factors maximize hardware efficiency at the cost of some quality.

Dense Transformer vs. Sparse MoE: What Actually Changes

A sparse MoE transformer looks like a standard transformer with one substitution: every other MLP sublayer (or every MLP sublayer, depending on the design) is replaced by an MoE layer. The attention layers remain dense and identical. Token positions and sequence structure are unchanged.

The differences that matter in practice:

Parameter count vs. active parameter count. Mixtral 8x7B has 46.7B parameters but activates roughly 13B per token (2 of 8 experts, plus shared attention weights). This is why it is compared to a 13B dense model in compute terms while achieving quality closer to a 70B dense model.

Memory requirement vs. compute requirement. An 8-expert MoE model requires all 8 expert weight matrices to be in memory, because you do not know in advance which experts will be needed for a batch. Memory requirement scales with total parameters. Compute scales with active parameters per token. This means MoE models are memory-bound even when they are compute-efficient. Serving a 46.7B MoE model requires the full 46.7B worth of memory, not 13B.

Training stability. MoE models are harder to train than dense models of equivalent parameter count. The auxiliary loss, capacity factor, and top-k selection all introduce discontinuities or hyperparameters that require tuning. Switch Transformer (top-1 routing) was famously unstable in early experiments and required specific initialization and precision choices (float32 for router, bfloat16 elsewhere).

Expert Parallelism for Serving

The standard parallelism strategies (tensor parallelism, pipeline parallelism) still apply to MoE models, but expert parallelism adds a third dimension.

Expert parallelism assigns different experts to different devices. With 8 experts across 8 GPUs, each GPU holds one expert’s weights and only activates compute for tokens routed to its expert.

Token dispatch flow with expert parallelism:

GPU 0: [Router] -> dispatch tokens to expert GPUs via all-to-all collective
GPU 0: Expert 0 -> compute for tokens assigned here
GPU 1: Expert 1 -> compute for tokens assigned here
...
GPU 7: Expert 7 -> compute for tokens assigned here
         -> all-to-all collective to gather results back
GPU 0: [Combine] -> weighted sum of expert outputs

The all-to-all collective is the bottleneck. Each device must send and receive a variable number of tokens from every other device. With fast NVLink (600 GB/s on H100), this is acceptable. Over PCIe or InfiniBand it becomes a significant fraction of step time.

For serving rather than training, expert parallelism reduces per-GPU memory requirements (each GPU holds only its expert’s weights plus the shared attention weights) at the cost of requiring tight interconnect and coordination overhead per forward pass.

Token Dropping and Overflow Handling

When routing decisions overflow expert capacity, there are three common strategies:

Zero output (skip). Dropped tokens contribute nothing from the MoE layer. Their hidden state from the previous layer passes through the residual connection unchanged. This is the most common default. Quality impact is proportional to drop rate.

Random reassignment. Overflow tokens are randomly reassigned to expert slots with remaining capacity. This avoids the zero output but may route tokens to suboptimal experts. Router probabilities for the reassigned expert are used for weighting.

Auxiliary routing network. A secondary router handles overflow tokens. More complex to implement but used in some production models to avoid quality degradation from drops.

In practice, a well-tuned capacity factor with the auxiliary load balancing loss keeps drop rates under 1-2% of tokens during normal inference, which has negligible quality impact.

Production Implications for Inference

MoE models are memory-bound, not compute-bound. A dense 13B model and a sparse MoE 46.7B model (with 13B active parameters per token) have similar compute per token. But the MoE model needs 3-4x more GPU memory. This shifts the inference bottleneck from compute to memory bandwidth, especially during decode where each step reads all expert weights that could potentially be active for any token in the batch.

Expert caching. Because the same experts tend to be activated repeatedly for similar inputs, you can cache recently-used expert outputs or even pre-load frequently-used experts into faster memory (L2 cache, or HBM vs. system RAM). DeepSpeed-MoE’s expert offloading moves inactive expert weights to CPU RAM and streams them back on demand, which allows serving models whose total weight exceeds GPU memory at the cost of latency spikes during cache misses.

Batching changes the expert activation pattern. With a large enough batch, all experts become active because different tokens route to different experts. With a small batch (single request, few tokens), only 2-3 experts may activate. Small batch serving is less efficient for MoE than for dense models because the per-token savings do not apply at the routing overhead level: you still pay for the all-to-all dispatch even if most experts go unused.

vLLM MoE support. vLLM supports Mixtral and similar MoE models out of the box with fused MoE kernels (using Triton or CUDA) that handle the dispatch and combine operations efficiently without the Python loop shown earlier. The PagedAttention memory management applies to the attention KV cache as usual; the expert weights are loaded into GPU memory in full.

Expert routing cache for repeated prompts. Empirically, routing decisions are relatively consistent across similar prompts. This means prefix caching (caching KV tensors for shared prefixes) is especially valuable for MoE models: the cached activations encode which expert outputs were produced, not just the attention states. Some experimental inference systems cache the routing decisions themselves for known input prefixes to skip the router computation.

Tradeoffs Comparison

Dimension	Dense Transformer	Sparse MoE (top-2)	Switch Transformer (top-1)	Soft MoE
Active params per token	100% of total	2/N of total	1/N of total	100% (soft assignment)
Total memory required	Proportional to active params	Much larger than active params	Largest per-quality unit	Same as dense at equivalent quality
Compute per token	High (all params)	Low (top-k experts)	Lowest (one expert)	High (all experts, soft weighted)
Training stability	High	Moderate (aux loss needed)	Low (known instability issues)	High (no discrete routing)
Expert collapse risk	N/A	Real, needs aux loss	Severe without careful tuning	None (gradients flow to all experts)
Token dropping	None	Low with CF>1.0	High with CF=1.0	None
Inference hardware req	Memory proportional to quality	Memory >> compute requirement	Same as sparse MoE	Same as dense
Serving complexity	Low	High (expert dispatch)	Moderate	Low
Quality per FLOP	Baseline	Better (validated at scale)	Good but instability cost	Competitive, less studied
Examples	GPT-4 (assumed), Llama	Mixtral 8x7B, Mixtral 8x22B	Google Switch-C	Google Soft MoE

Soft MoE is worth a note. Instead of discrete top-k selection, Soft MoE computes a soft assignment where every expert receives a weighted combination of all tokens and produces a weighted contribution to all token outputs. This eliminates token dropping, load imbalance, and routing instability entirely. The cost: all experts compute on every forward pass, so you lose the compute savings of sparse routing. Soft MoE is better thought of as a structured form of parameter sharing than a compute-efficient architecture.

Production Considerations

Memory planning is the primary constraint. Before deploying any MoE model, calculate the full weight memory requirement, not the active parameter count. Mixtral 8x7B at bfloat16 requires approximately 93 GB for weights alone. Two H100 80GB GPUs with tensor parallelism across the shared layers and expert parallelism across experts is the minimum practical configuration.

Monitor expert utilization in production. Load balancing during training does not guarantee uniform expert load on your specific inference workload. Log the per-expert token counts during serving. Consistently overloaded experts under production traffic indicate your capacity factor is too low for the actual input distribution.

Quantization for MoE. GPTQ and AWQ quantization apply to MoE models the same as dense models. Quantizing expert weights to 4-bit (AWQ) reduces the Mixtral 8x7B memory footprint from 93 GB to roughly 26 GB, making it feasible on a single A100 80GB. Quality degradation is similar to dense model quantization at the same bit width. This is often the right tradeoff for serving: 4-bit Mixtral 8x7B on one GPU rather than bfloat16 across two.

Latency vs. throughput characteristics. MoE models have higher latency at low batch sizes (routing overhead, expert dispatch) compared to dense models with equivalent active parameters. Throughput at high batch sizes is competitive or better because the GPU compute is used efficiently once all experts are active. If you are optimizing for p50 latency on low-concurrency requests, a dense model of equivalent quality is simpler and faster. If you are optimizing for throughput on a high-concurrency API serving many requests, MoE is worth the operational complexity.

The capacity factor is a knob you will need to tune. The published values from research papers are tuned for training. For inference, the optimal capacity factor depends on your batch size and the entropy of routing decisions for your workload. Start with capacity_factor=1.25 and measure drop rates on representative traffic before reducing it.

Closing

Mixture of Experts solves a specific problem: how to scale model capacity without scaling compute proportionally per token. The mechanism is elegant but the implementation surface is larger than it looks from the outside. The gating network, auxiliary loss, capacity factor, and token dropping all interact, and each has production implications that do not appear in the benchmark numbers. The memory requirement in particular is the thing that surprises teams most often: the compute savings are real, but you still have to fit all those expert weights somewhere.