Mixture of Experts Explained: How Mixtral, DeepSeek & Grok Scale to Trillions of Parameters

The pursuit of larger and more capable language models has led researchers to a fundamental tension: model performance tends to improve with scale, but computational costs grow prohibitively. Mixture of Experts (MoE) architectures offer an elegant solution to this dilemma, enabling models with trillions of parameters while keeping inference costs manageable through sparse computation.

In this article, we explore the mathematical foundations of MoE architectures, examine how expert routing works, and discuss the engineering challenges that must be overcome for practical deployment.

1import numpy as np
2import matplotlib.pyplot as plt
3 
4# Simulate 8 experts and their activation patterns for different input types
5np.random.seed(42)
6input_types = ['Math', 'Code', 'Poetry', 'Science', 'History', 'Legal', 'Medical', 'Casual']
7num_experts = 8
8 
9# Each input type activates different experts (sparse activation - top-2)
10expert_activations = np.zeros((len(input_types), num_experts))
11 
12# Simulate learned routing (each input type prefers 2 experts)
13routing = {
14    'Math': [0, 2], 'Code': [1, 2], 'Poetry': [3, 5],
15    'Science': [0, 4], 'History': [4, 6], 'Legal': [6, 7],
16    'Medical': [4, 7], 'Casual': [3, 5]
17}
18 
19for i, input_type in enumerate(input_types):
20    for expert in routing[input_type]:
21        expert_activations[i, expert] = np.random.uniform(0.4, 1.0)
22 
23fig, ax = plt.subplots(figsize=(10, 6))
24im = ax.imshow(expert_activations, aspect='auto')
25 
26ax.set_xticks(np.arange(num_experts))
27ax.set_yticks(np.arange(len(input_types)))
28ax.set_xticklabels([f'Expert {i}' for i in range(num_experts)])
29ax.set_yticklabels(input_types)
30 
31plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
32 
33ax.set_title('Expert Activation Patterns (Sparse Top-2 Routing)')
34ax.set_xlabel('Experts')
35ax.set_ylabel('Input Type')
36 
37cbar = ax.figure.colorbar(im, ax=ax)
38cbar.ax.set_ylabel('Activation Strength', rotation=-90, va="bottom")
39 
40plt.tight_layout()
41plt.show()
42 
43print("Notice how each input type activates only 2 experts (sparse routing)")
44print("This enables efficient scaling while maintaining specialization!")
45 

The Core Insight: Conditional Computation

Traditional dense neural networks activate every parameter for every input. A 70-billion parameter model performs 70 billion operations per forward pass, regardless of input complexity. MoE architectures challenge this assumption with a simple observation: not all parameters need to be active for every input.

Consider how human experts operate. When you have a medical question, you consult a physician - not a lawyer, accountant, and physicist simultaneously. MoE models apply this principle computationally: they maintain multiple "expert" subnetworks and route each input to only a subset of relevant experts.

One MoE transformer layer

1. Input token
2. Attention
3. RouterA small gating network scores every expert
4. routing weights
   Expert 1

   Expert 2selected

   Expert 3selected

   Expert N
   Only the top-k experts run; the rest stay idle
5. Weighted sum of expert outputsEach output scaled by its routing weight
6. Output token

The Gating Mechanism: Mathematical Foundations

The router (or gating network) is the brain of an MoE layer. Given an input representation $x$, it must decide which experts should process this input and with what weights.

Router Architecture

The simplest and most common router uses a learned linear transformation followed by a softmax:

$$p_i = \frac{e^{h(x)_i}}{\sum_j e^{h(x)_j}}$$

Where $h(x) = x \cdot W_g$ is the router's hidden representation, $W_g \in \mathbb{R}^{d \times n}$ is a learned weight matrix, $d$ is the hidden dimension, and $n$ is the number of experts.

Top-K Selection

Computing outputs from all experts would defeat the purpose of sparse computation. Instead, MoE models select only the top-$K$ experts (typically $K=1$ or $K=2$):

$$G(x) = \text{Softmax}(\text{TopK}(x \cdot W_g))$$

The TopK operation zeros out all but the $K$ highest-scoring experts. The final layer output becomes:

$$y = \sum_{i=1}^{n} G(x)_i \cdot E_i(x)$$

Where $E_i(x)$ is the output of expert $i$. Due to the sparsity induced by TopK, this sum only has $K$ non-zero terms, dramatically reducing computation.

A Concrete Example

Consider a layer with 8 experts and $K=2$. For an input token:

Routing one token with 8 experts and K = 2

E1E2E3E4E5E6E7E8

Router scores0.150.420.080.030.120.050.090.06

After top-k00.42000.12000

After softmax00.57000.43000

Only experts 2 and 5 are activated, with weights 0.57 and 0.43 respectively.

Expert Specialization: Emergent Division of Labor

One of the most fascinating aspects of MoE models is that experts develop specializations without explicit supervision. Through gradient descent alone, different experts learn to handle different types of inputs.

Observed Specialization Patterns

Research has revealed consistent specialization patterns:

Syntactic experts tend to activate for specific grammatical constructions - one expert might specialize in relative clauses while another handles conditional statements.

Domain experts emerge for topics like code, mathematics, or specific languages in multilingual models.

Positional experts sometimes develop, with certain experts preferentially activating for tokens at specific positions in the sequence.

Expert specialization, one sentence

E1E2E3E4E5E6E7E8

The

quantum

computer

calculated

the

eigenvalues

Dark: primary expert. Light: secondary expert. E1 takes function words, E2 verbs, E3 technical terms, and E6 and E7 the mathematical context.

The Load Balancing Challenge

A critical challenge in MoE training is expert collapse: without intervention, the router often converges to using only a small subset of experts, wasting capacity. This occurs because popular experts receive more gradient signal, improving further and attracting even more traffic - a rich-get-richer dynamic.

Load Balancing Loss

To encourage uniform expert utilization, MoE models add an auxiliary loss term:

$$\mathcal{L}_{balance} = \alpha \cdot n \cdot \sum_{i=1}^{n} f_i \cdot P_i$$

Where:

• $n$ is the number of experts
• $f_i$ is the fraction of tokens routed to expert $i$
• $P_i$ is the average router probability for expert $i$
• $\alpha$ is a hyperparameter (typically 0.01-0.1)

This loss is minimized when both $f_i$ and $P_i$ are uniform across experts (equal to $1/n$), encouraging balanced routing.

Capacity Factor

Another mechanism involves setting a capacity factor $C$ that limits how many tokens each expert can process:

$$\text{Expert Capacity} = \frac{C \cdot \text{batch\_tokens}}{n}$$

Tokens that would exceed an expert's capacity are either dropped or routed to their second-choice expert. Typical values are $C \in [1.0, 2.0]$.

Efficient Deployment: Pruning and Optimization

While MoE models offer computational savings during training and inference, their large memory footprint (all experts must be stored) poses deployment challenges. Recent research has addressed this through expert pruning and dynamic skipping.

Xie et al. (2024) introduced methods for expert-level pruning and dynamic skipping in their paper "Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models." Their key insight is that expert importance varies significantly across layers and tasks - some experts can be removed entirely with minimal performance degradation, while others can be skipped dynamically based on input characteristics.

Expert pruning by importance

E1E2E3E4E5E6E7E8

Layer 1, before8 experts

Layer 2, before8 experts

Layer 1, after6 experts kept

Layer 2, after5 experts kept

Edge Deployment Considerations

For resource-constrained environments like edge devices, Chen et al. (2024) explored inference offloading strategies in "Efficient Inference Offloading for Mixture-of-Experts Large Language Models in Internet of Medical Things." Their work demonstrates that MoE's inherent sparsity can be leveraged for efficient distributed inference, where different experts reside on different devices and are activated on-demand.

Architectural Variations

Expert Granularity

The definition of "expert" varies across architectures:

Coarse-grained experts use full feed-forward networks (FFN) as experts. This is the classic approach used in models like Switch Transformer.

Fine-grained experts split the FFN into smaller chunks. For instance, some architectures treat each neuron group as a separate expert, enabling more precise routing.

Coarse-grained and fine-grained experts

Coarse-grained (classic MoE)

1. Expert FFNOne full feed-forward network, about 100M parameters

Fine-grained MoE

1. E1

   E2

   E3

   E4

   E5

   E6

   E7

   E8
   The same capacity split into small experts, about 12.5M parameters each

Router Variations

Beyond simple linear routers, researchers have explored:

• Hash-based routing: Deterministic routing based on input hashing (no learned parameters)
• Expert choice routing: Experts select tokens rather than tokens selecting experts
• Hierarchical routing: Two-stage routing for models with many experts

Training Considerations

Batch Size Requirements

MoE models typically require larger batch sizes than dense models. With sparse routing, each expert sees only a fraction of the batch. If the batch is too small, individual experts receive insufficient gradient signal.

Communication Overhead

In distributed training, tokens must be routed to experts that may reside on different devices. This all-to-all communication can become a bottleneck:

Distributed MoE: all-to-all communication

1. Input batch tokensEach token is routed to its selected experts
2. all-to-all
   Device 0E0, E1

   Device 1E2, E3

   Device 2E4, E5

   Device 3E6, E7
   Tokens cross devices to reach their experts, then return

The Future of Sparse Models

MoE architectures represent a fundamental shift in how we think about neural network scaling. Rather than asking "how can we train larger dense models?", the question becomes "how can we allocate computation more intelligently?"

Current research directions include:

1. Adaptive computation: Dynamically adjusting the number of experts based on input complexity
2. Multimodal MoE: Separate expert pools for different modalities (text, image, audio)
3. Retrieval-augmented MoE: Treating retrieved documents as additional "experts"
4. Hardware-aware routing: Optimizing routing decisions based on device topology

As language models continue to scale, sparse architectures like MoE will likely play an increasingly central role - enabling the next generation of AI systems that are both more capable and more efficient than their dense predecessors.

References

1. Xie, X., et al. (2024). "Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models." Semantic Scholar

2. Chen, W., et al. (2024). "Efficient Inference Offloading for Mixture-of-Experts Large Language Models in Internet of Medical Things." Semantic Scholar