Mixture of Experts (MoE) in LLMs Explained

[LLM] Mixture of Experts (MoE)

Multiple studies have shown that LLM performance scales with parameter count, prompting companies like Meta, Google, and OpenAI to race toward ever-larger models. However, as models grow, so do the compute costs of training and serving them. Mixture of Experts (MoE) emerged as a technique to train and serve much larger LLMs at lower cost.

One of the most prominent MoE models is Mistral AI's Mixtral-8x7B. OpenAI's GPT is widely suspected to use an MoE architecture as well, though the exact design has never been disclosed. These MoE models run inference efficiently by activating only a subset of their total parameters.

Basic Structure and Operating Principle of MoE

The core concepts behind MoE are:

• Expert Networks (Experts): An MoE model consists of multiple independent neural networks, each called an "expert." A typical MoE model contains anywhere from hundreds to thousands of experts.
• Gate Network (Router): A network that decides which experts are activated for a given input. The router takes the input, selects a subset of experts, and combines their outputs — usually via a weighted sum — to produce the final output. Only a few experts are activated per input; the rest remain inactive.

MoE in the Transformer Architecture

A standard LLM Transformer block follows the pattern: Self-Attention → Add & Norm → FFN → Add & Norm. MoE differs in the FFN stage: instead of a single FFN, it maintains multiple FFNs. After a token passes through Self-Attention → Add & Norm, the router directs it to a specific FFN. The next token may be routed to a different FFN entirely. In this way, each token passes through only a subset of the available FFNs — either one or more, depending on configuration. When the router selects the top-k FFNs, each chosen FFN processes the token independently, and the outputs are combined (e.g., via weighted sum or another aggregation method) to produce the final output. The per-expert output weights are typically derived from the routing probabilities computed by the router.

Types of Routing Strategies

Common MoE routing variants include:

• Standard MoE: Combines the outputs of all experts. Computationally expensive, but leverages the full knowledge of every expert.
• Switch Transformer: Selects a single expert per input. Highly efficient, but risks information loss.
• GShard: Developed by Google; applies MoE to very large-scale models with strong scalability across diverse tasks.
• BASE Layers: Trains while maintaining load balance across experts, ensuring all experts are utilized evenly.
• HashMoE: Uses a hash function to assign experts. Computationally efficient but less flexible.
• Soft MoE: Smoothly combines the outputs of all experts; sits between standard MoE and Switch Transformer in behavior.
• ST-MoE (Sparse Top-k MoE): Selectively activates the top-k experts out of many for each input, improving both efficiency and performance in large-scale language models.

Other approaches include Expert Choice Routing and Mixture of Experts with Expert Choice (MoEC).

Token Choice vs. Expert Choice

MoE routing can be categorized along two dimensions.

Token Choice: The standard approach, where each token selects its own top-K experts. It integrates naturally with autoregressive generation, but tends to produce load imbalance as popular experts attract a disproportionate share of tokens. Common mitigations include overprovisioning with a capacity factor of 2–8× and adding an auxiliary load-balancing loss.

Expert Choice: A reverse approach proposed by Google in 2022, where each expert selects the top-K tokens from the batch. This mathematically guarantees perfect load balancing and allows each token to be assigned a variable number of experts. The downside is that expert selection requires access to the entire sequence, which violates causality and makes it difficult to apply directly to standard autoregressive language models.

Key MoE Models: 2024–2025 Comparison

DeepSeek-V3 introduced a loss-free balancing strategy that adds a per-expert dynamic bias term to routing scores instead of using an auxiliary loss. Overloaded experts have their bias reduced while underutilized experts have it increased — but this bias is not reflected in the final gating weights, so it does not affect output quality. Qwen3-MoE takes a different approach: it omits shared experts entirely and relies solely on a global-batch load balancing loss to encourage expert specialization.

Code Example: Mixtral-8x7B

The following example shows how to use Mixtral-8x7B, one of the most representative MoE models.

from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from mistral_common.protocol.instruct.messages import UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest

# Tokenization

mistral_models_path = "MISTRAL_MODELS_PATH"
 
tokenizer = MistralTokenizer.v1()
 
completion_request = ChatCompletionRequest(messages=[UserMessage(content="Explain Machine Learning to me in a nutshell.")])
 
tokens = tokenizer.encode_chat_completion(completion_request).tokens

# Inference by mistral_inference

from mistral_inference.transformer import Transformer
from mistral_inference.generate import generate

model = Transformer.from_folder(mistral_models_path)
out_tokens, _ = generate([tokens], model, max_tokens=64, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)

result = tokenizer.decode(out_tokens[0])

print(result)

# Inference with huggingface transformers

from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
model.to("cuda")
 
generated_ids = model.generate(tokens, max_new_tokens=1000, do_sample=True)

# decode with mistral tokenizer
result = tokenizer.decode(generated_ids[0].tolist())
print(result)

Mixture of Tokens (MoT)

Where MoE asks "which expert should handle this token?", Mixture of Tokens (MoT) approaches the problem from the opposite direction. Presented at NeurIPS 2024, MoT does not route tokens to experts; instead, it first mixes tokens from multiple examples in the batch using learned weights, then feeds the resulting mixed representations into the expert FFNs.

The three core ideas are:

• At each layer, tokens from different examples in the batch are combined into a mixed representation via a learned weighted sum.
• This mixed representation is what passes through the expert FFN.
• The mixing weights are learned; ablations in the paper show this outperforms a fixed uniform weighting (1/n).

Key differences from MoE:

• MoT is a continuous operation, so it avoids the gradient discontinuity that arises from discrete routing decisions.
• When the number of experts equals the number of tokens, the total compute is equivalent to top-1 Token Choice MoE.
• Experiments in the paper show improved training stability and expert utilization compared to both Token Choice and Expert Choice sparse MoE.

Because MoT mixes tokens from different examples in the batch, it operates as designed in batched inference settings rather than single-request inference.

Fine-Tuning Considerations for MoE

Expert Collapse

Fine-tuning a pretrained MoE model often triggers routing collapse, where the router funnels tokens into only a small subset of experts. The remaining experts receive little to no gradient updates, effectively reducing the model's total usable capacity.

Load Balancing Auxiliary Loss

The standard approach to preventing load imbalance is to add an auxiliary loss — typically penalizing the variance in the number of tokens assigned to each expert — on top of the main task loss. However, if the auxiliary loss coefficient is too large, it can interfere with task-specific gradients and hurt performance. For fine-tuning, a coefficient in the range of 0.001–0.01 relative to the main loss is generally recommended.

Loss-Free Balancing

The loss-free approach proposed by DeepSeek in 2024 adds a per-expert dynamic bias term to routing scores without any auxiliary loss. It reduces the bias for overloaded experts and increases it for underutilized ones to maintain balance, while keeping this bias out of the final gating weights. This is worth considering when you want to maintain expert balance during fine-tuning without the gradient interference introduced by an auxiliary loss.

Serving MoE in Practice: vLLM

The choice of parallelism strategy when serving MoE models with vLLM has a significant impact on memory efficiency and throughput.

Tensor Parallelism vs. Expert Parallelism

• Tensor Parallelism (TP): Configured via --tensor-parallel-size. Shards weight tensors across all GPUs for every layer, including attention.
• Expert Parallelism (EP): Enabled with the --enable-expert-parallel flag. For MoE layers only, experts are distributed across GPUs as whole units rather than tensor-sharded. The degree of parallelism matches the configured tensor_parallel_size.

# Mixtral 8x7B — 4-GPU Expert Parallelism example
vllm serve mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --tensor-parallel-size 4 \
  --enable-expert-parallel

MLA-Based Models (DeepSeek Family)

Applying standard TP to models that use Multi-head Latent Attention (MLA), such as DeepSeek-V3, causes the KV cache to be replicated across every GPU, leading to significant memory waste. According to the vLLM documentation, the recommended approach for these models is a hybrid strategy: Data Parallelism for attention layers and Expert Parallelism for expert layers.

# DeepSeek-V3 — 8-GPU DP+EP hybrid example
vllm serve deepseek-ai/DeepSeek-V3 \
  --tensor-parallel-size 1 \
  --data-parallel-size 8 \
  --enable-expert-parallel