Mixture of Experts — Core Concepts

The Original Mixture of Experts

Jacobs et al. (1991) introduced Mixture of Experts as a general machine learning approach: multiple “expert” networks specialize in different regions of the input space, with a “gating” network learning to route inputs to the appropriate experts.

This idea lay relatively dormant until Shazeer et al. (2017) “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” applied MoE to language models at Google Brain, scaling to 137 billion parameters on 4 GPUs — remarkable for 2017.

The key insight: replace the dense feed-forward (FFN) layer in transformers with a collection of expert FFNs, activating only a few for each input token.

Architecture: Sparse MoE in Transformers

In a standard transformer, each layer has:

• Multi-head attention (attends to all positions)
• Feed-forward network (FFN): two linear layers with activation

In a Sparse MoE transformer, the FFN is replaced:

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

Where:

• $N$ = total number of experts (e.g., 64)
• $E_i(x)$ = expert $i$‘s FFN applied to input $x$
• $G(x)_i$ = gating weight for expert $i$

The gating function uses top-k selection:

$$G(x) = \text{softmax}(\text{top-k}(H(x) \cdot W_g))$$

Where $H(x)$ is the input hidden state, $W_g$ is the gate weight matrix, and top-k sets all but the k highest values to $-\infty$ before softmax (making them zero). Typically $k=2$.

For each token, only 2 experts’ FFNs are computed and their outputs are weighted and summed.

Load Balancing: A Critical Training Challenge

Without explicit constraints, the gating network quickly learns to favor a small number of experts — a few get all the traffic while others are never trained. This is called expert collapse or token dropping.

Collapsed routing ruins MoE’s efficiency: you have 64 experts but only 2 actually learn; the others waste memory.

Auxiliary loss for load balancing (Switch Transformer, Fedus et al., 2022):

$$\mathcal{L}{aux} = \alpha \cdot N \cdot \sum{i=1}^N f_i \cdot P_i$$

Where:

• $f_i = \frac{1}{T}\sum_x \mathbf{1}[\text{token } x \text{ routed to expert } i]$ (fraction of tokens to expert $i$)
• $P_i = \frac{1}{T}\sum_x G(x)_i$ (mean gating score for expert $i$)
• $\alpha$ is a hyperparameter (typically 0.01)

This penalizes routing where $f_i$ and $P_i$ are large for the same expert (expert is both frequently chosen and highly scored). The product $f_i P_i$ being large means that expert dominates — the loss pushes toward more uniform routing.

Mixtral 8x7B: A Concrete Example

Mistral AI (2023) released Mixtral 8x7B as an open-source sparse MoE:

Architecture:

• 32 transformer layers
• Each layer: attention block + MoE FFN block
• MoE block: 8 expert FFNs per layer, 2 active per token
• Expert hidden dimension: 14,336 (each expert is a standard FFN with this width)
• Total parameters: 46.7B (8 × FFN params per layer × 32 layers)
• Active parameters: ~12.9B per token (2/8 of the FFN params active)

Performance vs. compute:

• Mixtral 8x7B vs. Llama 2 70B: Mixtral wins on most benchmarks with 5-6x less compute per inference
• Mixtral 8x7B vs. Llama 2 13B: Mixtral wins significantly (with comparable inference speed)

The sweet spot: experts’ combined capacity provides knowledge breadth, while sparse activation maintains inference efficiency.

Emergent Expert Specialization

Do experts actually specialize? Mistral AI released an analysis showing that expert selection in Mixtral does correlate with input domain:

• Certain experts are more frequently selected for code
• Others for mathematical expressions
• Others for natural language in different styles

The specialization isn’t as clean as “expert 1 handles math, expert 2 handles code” — it’s more gradient than binary. But the correlation is statistically significant and increases in deeper layers.

This emergent specialization is important: it’s not designed, it emerges from the load balancing + optimization pressure. Experts that handle similar inputs together develop similar capabilities.

Expert Parallelism for Training

MoE introduces a new dimension for distributed training: expert parallelism. Different experts reside on different GPUs:

• GPU 0 hosts experts 1, 2, 3, …
• GPU 1 hosts experts 9, 10, 11, …

When a token is routed to expert $i$, its activations must be sent to whichever GPU holds expert $i$ — an “all-to-all” communication step. This is the primary communication overhead of MoE training.

Google’s Pathways / MoE at scale: Switch Transformer used expert parallelism across 2048 TPU cores, with experts spread across all cores. The routing overhead was significant but outweighed by the training efficiency gains at very large scale.

One thing to remember: MoE’s efficiency gain is real but comes with training complexity — load balancing auxiliary losses, expert parallelism communication overhead, and routing instability require careful engineering, which is why sparse MoE models took years to successfully deploy despite the clear theoretical benefits.