Mixture of experts: capacity, routing, and the cost of imbalance

An MoE batch performs sixteen expert assignments, yet one device receives nearly half of them. How much does the average hide? Follow token dispatch and inspect the busiest expert before interpreting utilization.

Before you begin: Complete Mistral and Mixtral, and understand feed-forward layers and data parallelism.

Sparse mixture of experts, or MoE, separates the total capacity a model stores from the computation it activates for a token. The router selects a subset of learned networks, usually within feed-forward blocks. Making that selection useful and efficient is the central engineering problem.

Follow one token through dispatch

A token representation enters the router. The router produces scores for experts, a selection rule chooses a subset, and dispatch sends the representation to those experts. Their outputs are weighted and combined into a representation of the original width.

A token can select different experts at different layers. Tokens in one sentence need not travel together. Experts are trained numerical transformations, not separately prompted chatbots, and they do not come with guaranteed subject specialties.

In a top-two design, the output might be 0.7 times expert A's vector plus 0.3 times expert B's vector. Those are weights on computed vectors, not calibrated probabilities that either expert is correct. Shared attention and other dense operations still run.

Count the actual load

Consider eight tokens, each routed to two experts. There are sixteen assignments. With four experts, a perfectly equal assignment count would be four per expert. The following Python 3 program measures a fictional batch using the standard library.

# Runnable: Python 3, standard library.
from collections import Counter

assignments = [
    (0, 1), (0, 1), (0, 1), (0, 2),
    (0, 2), (0, 3), (0, 3), (1, 2),
]
loads = Counter(expert for pair in assignments for expert in pair)
average = sum(loads.values()) / 4
print(dict(loads), "peak/average:", max(loads.values()) / average)
assert sum(loads.values()) == 16
assert loads[0] == 7

Expert zero receives seven assignments while expert three receives two. If they live on separate devices, the busiest expert can determine when the batch finishes. An average across devices can conceal the bottleneck.

Assignment counts are a useful starting point, but also inspect sequence lengths, kernel shapes, communication, and padding. Equal token counts do not guarantee equal wall-clock work in every implementation.

Keep the total work, redistribute the assignments

The local batch has loads 7, 4, 3, and 2. Its average is four and its peak-to-average ratio is 1.75. Construct another set of eight top-two selections that gives four assignments to each expert. Total assignments stay sixteen while the peak falls to four.

For practice, pair experts (0,1) four times and (2,3) four times. Each expert receives four tokens. Then compare with the original assignments and explain why a mean alone cannot distinguish them. This gives the load-balancing discussion a concrete reference point.

Why balancing enters training

A router that sends most tokens to one expert leaves other capacity underused and may overload that expert. Some training methods add a load-balancing objective. Others use different routing or bias-adjustment strategies. The exact mechanism is architecture-specific.

Capacity limits determine how many assignments an expert can process in a batch. Depending on the method, overflow may be dropped, rerouted, or handled without a fixed capacity bound. Never assume overflow behavior from the phrase “MoE” alone.

Balancing also creates a tension: spreading work evenly may compete with the router's preferred selections. Evaluate task quality as well as utilization when changing router settings.

Communication can erase arithmetic savings

With expert parallelism, experts live on different devices. Token representations move to the selected experts and their results move back. This dispatch often involves all-to-all communication.

A model may activate relatively few parameters per token but still require many stored weights across the system. Small batches can use expert kernels poorly. Slow interconnects, uneven routing, and offloading can limit performance. The active-parameter number is neither a memory specification nor a latency measurement.

Practice an investigation

Your MoE server slows down only for code prompts. Overall device utilization looks acceptable. What would you inspect?

To evaluate a change, preserve the same model, hardware, request mix, and quality checks. Record total and active parameter conventions, expert placement, batch conditions, and end-to-end results.

The original Switch Transformers paper develops sparse routing and balancing ideas. Mixtral of Experts supplies a concrete top-two architecture. Their recipes are historical examples, not universal configuration defaults.

Continue: Long-context models, where accepting a large input is only the start of using it well.