What is a mixture of experts (MoE)?

What is a mixture of experts?

A mixture of experts, or MoE, is a neural network layer that replaces one big block of computation with many smaller sub-networks, called experts, plus a router. The router picks which experts handle each input. Instead of every parameter processing every token, only a handful of experts fire for any given token.

The idea goes back to Jacobs and Jordan's work in the early 1990s. The Sparsely-Gated Mixture-of-Experts paper by Shazeer et al. (2017) showed it could scale to language models with up to 137 billion parameters, while adding only a small computational cost. That paper framed MoE as a practical form of conditional computation. Parts of the network switch on or off per example, so capacity grows without a matching rise in compute.

How does sparse routing work?

A trainable gating network scores every expert for a given input and keeps only the top few. Shazeer et al. (2017) used noisy top-k gating: add tunable noise to the scores, keep the k highest, and set the rest to effectively zero. Those zeroed experts need never be computed.

Switch Transformers (Fedus et al., 2021) simplified this further by routing each token to a single expert, making k equal to one. The authors argue this cut router computation and halved the batch each expert needed, while still preserving model quality. Mixtral (Jiang et al., 2024) picks two experts per token out of eight, using a softmax over the top-2 router logits. The two experts' outputs are summed and weighted by their gate values.

Because the gate output is mostly zero, the model skips computing any expert whose weight is zero. That sparsity is what makes the layer cheap despite its size.

Active parameters vs. total parameters

A mixture-of-experts model has a total parameter count: all the weights across every expert. It also has an active parameter count: the weights actually used to process one token. Mixtral 8x7B has 47 billion total parameters but only 13 billion active parameters per token, because its router picks 2 of 8 experts at each layer, as Jiang et al. (2024) report.

Active parameters track inference compute cost. Total parameters track memory cost, since every expert's weights must still be held somewhere even if unused on a given token. Jiang et al. (2024) note Mixtral's memory footprint scales with its 47 billion sparse parameters, still smaller than the 70 billion of Llama 2 70B, the model it competes with.

Why are MoE models cheap to run for their size?

Increasing the number of experts adds parameters without adding much compute, since the router still only activates a fixed number per token. Fedus et al. (2021) note that the router's own cost is small, proportional to the embedding dimension times the number of experts. Scaling expert count is, in their words, the most efficient dimension for scaling the model.

This showed up directly in results. Shazeer et al. (2017) reported a model with 4 billion parameters beating the best published LSTM language-modelling result, while using only 6% of the computation, when controlling for training epochs. Jiang et al. (2024) report Mixtral matching or beating Llama 2 70B and GPT-3.5 across most benchmarks, while using five times fewer active parameters than Llama 2 70B.

Why does load balancing matter?

Left alone, a gating network tends to converge on favouring a few experts. The favoured ones get trained faster and so get picked even more often, a self-reinforcing imbalance Shazeer et al. (2017) describe directly. Unbalanced routing wastes the idle experts' capacity and can overload the devices hosting the popular ones.

Shazeer et al. (2017) added two auxiliary losses: one penalising uneven importance across experts, another penalising uneven load. Switch Transformers (Fedus et al., 2021) simplified this to a single differentiable load-balancing loss, added to the training objective with a small weighting coefficient. They found it balanced load quickly without disrupting the main loss.

Switch Transformers also used a capacity factor: each expert is only allowed a fixed batch of tokens. If too many tokens are routed to it, the excess are dropped and passed through via the residual connection instead. A higher capacity factor reduces dropped tokens but increases computation and communication. So Fedus et al. (2021) found keeping it low, near 1.0, worked better once the load-balancing loss did its job.

Do experts specialise by topic?

Not obviously, at least not in Mixtral. Jiang et al. (2024) examined which experts fired on different domains from The Pile dataset, including arXiv papers, biology abstracts, and philosophy texts. They found no clear topic-based pattern; expert assignment distributions looked similar across most subjects.

What they did find was syntactic structure: tokens such as 'self' in Python code or indentation tokens consistently routed to the same expert. Consecutive tokens showed far higher repeated-expert assignment than chance, especially in later layers. Jiang et al. (2024) suggest this positional locality could be exploited for caching, though it also risks overloading particular experts under expert parallelism.