Mixture of Experts

Expert routing & sparse gates

Prerequisites

graph TD
  NN[Neural nets / transformers] --> MoE[Mixture of Experts]
  Prob[Probability / softmax] --> MoE
  TR[Token routing] --> MoE

When to use

Scale model capacity without running every parameter on every token (sparse feed-forward layers).
Specialize sub-networks on domains (code, math, multilingual) while sharing attention stacks.
Inference cost control when Top-K experts keep active FLOPs far below total parameter count.

When not to use

Small models where routing overhead and load-balancing complexity outweigh savings.
Uniform compute budgets per token without auxiliary balancing — routers can collapse to one expert.
Tasks needing full dense mixing on every layer (some fine-grained reasoning paths).

Lab

On the interactive MoE lab, pick a token embedding, tune Top-K, read gate scores on the expert grid, compare sparse vs dense expert activation, try Quantum → physics or Top-1 bio presets, predict how many experts stay idle, then reveal the active fraction.

Simple Fundamental Explanation

Imagine you are building the ultimate consulting firm. You want to be able to answer questions about medicine, law, physics, and history.

The Generalist Approach: You hire one incredibly smart person and force them to study everything. They are good at everything but a master of none, and answering a question takes all their brainpower.
The Mixture-of-Experts (MoE) Approach: You hire 8 different experts (a doctor, a lawyer, a physicist, etc.). You also hire a receptionist. When a client asks a question, the receptionist quickly looks at the question, decides which 2 experts are the best fit, and routes the question only to them. The other 6 experts sleep.

A Mixture-of-Experts (MoE) architecture in AI works exactly like this. Instead of passing every piece of data through every single neuron in a massive neural network, the network is divided into sub-networks (“Experts”). A probabilistic “Router” decides which few experts are best suited to process each specific piece of data (each token).

This allows you to train a massive, highly capable model without requiring impossible amounts of computational power to run it.

Deep Dive: How It Works Under the Hood

Modern LLMs require massive amounts of parameters (weights) to store world knowledge. However, running a 1-Trillion parameter model for every single word generated requires too much GPU compute (FLOPs).

MoE breaks the dense feed-forward layers of the Transformer into multiple smaller, independent feed-forward networks (Experts).

The Router (Gating Network)

For every token processed, a small neural network called the Router looks at the token’s mathematical representation. The router outputs a probability distribution over all available experts (e.g., 8 experts).

The system then selects the top-K experts (usually $K = 1$ or $K = 2$ ). Only those specific experts are activated. The token passes through them, their outputs are combined based on the routing probabilities, and the data continues to the next layer.

Visual Diagram

Input Token: "Quantum"

      [ Router Network ]
           |
   (Probabilities: E1:0.8, E2:0.1, E3:0.0, E4:0.1)
           |
   -------------------------
   |           |           |
[Expert 1]  [Expert 2]  [Expert 3] ... [Expert 8]
(Physics)   (Math)      (Law)
  |            |
(Active)    (Active)    (Inactive)     (Inactive)
  |            |
   ---(Combine)---
           |
    Output Token State

Practical Example: GPT-4 and Mixtral 8x7B

The open-weight model Mixtral 8x7B is a famous example of MoE. It has 8 distinct experts. Each expert has 7 billion parameters. If it were a dense model, it would have ~47 billion parameters (some layers are shared).

However, because the router only selects the Top-2 experts for any given token, the active compute required per token is only roughly 13 billion parameters.

You get the intelligence and knowledge capacity of a 47B parameter model, but it runs as fast and cheap as a 13B parameter model. Rumors strongly suggest OpenAI’s GPT-4 relies heavily on a similar (but much larger) MoE architecture to achieve its massive scale without bankrupting the company on inference costs.

The Mathematics: Equations and In-Depth Analysis

1. The Gating Equation

Let $x$ be the input representation of a token. Let $E_{i} (x)$ be the output of the $i$ -th expert network. Let $G (x)$ be the output of the Gating network (router), which produces an $N$ -dimensional vector of probabilities.

The final output of the MoE layer $y$ is the probabilistically weighted sum of the experts:

y = i = 1 \sum N G (x)_{i} \cdot E_{i} (x)

2. Sparse Routing (Top-K)

If we evaluate all $N$ experts, we haven’t saved any compute! We must make $G (x)$ sparse (mostly zeros). The standard approach is to apply a Softmax function, keep the top $K$ values, and set the rest to $- \infty$ before computing the final probabilities.

TopK (x, K)_{i} = {x_{i} - \infty if x_{i} is in top K elements otherwise

The sparse gating function becomes:

G (x) = Softmax (TopK (W_{g} \cdot x + Noise, K))

3. The Load Balancing Problem

If the router accidentally learns that Expert 1 is slightly better at everything, it will route 100% of tokens to Expert 1. The other 7 experts will die (never receive gradients during training).

To prevent this, MoE training mathematically enforces load balancing using an auxiliary loss function. Let $f_{i}$ be the fraction of tokens routed to expert $i$ in a batch. Let $P_{i}$ be the average probability mass assigned to expert $i$ .

The auxiliary loss $L_{a ux}$ penalizes the model if the experts are not used equally:

L_{a ux} = α \cdot N i = 1 \sum N f_{i} \cdot P_{i}

By multiplying this penalty by a factor $α$ and adding it to the main training loss, the optimizer is forced to probabilistically distribute tokens across all experts evenly, ensuring the entire network is utilized.

Lab

Toy gating router: token embedding dot-products against expert weight rows, Top-K selection, and gate probabilities — same sparse pattern as Mixtral-style MoE in the README.

Top-K expertstop 2

Sparse Top-K routeE0: gate=0.90 · P=41.7% · E2: gate=0.20 · P=20.7%

Dense baselineDense: all 4 experts · Sparse: 2 (50% active)

With Top-K=2 on this token, how many of the 4 experts stay idle?

Token routed to experts 0, 2; gate scores E0=0.90, E1=0.10, E2=0.20, E3=0.10.