arthrop0d

Brief articles on varied topics.

Mixture-of-Experts Routing and Load Balancing

artificial intelligenceperformance

A mixture-of-experts (MoE) model contains many alternative expert networks but activates only a small number of them for each token. A router chooses the experts, so the model can have many more parameters than the hardware would execute for every token.

Why use experts?

A conventional transformer is dense: every token passes through the same feed-forward network in each layer. Increasing that network's size generally increases both the model's parameter count and the computation required for every token.

An MoE layer replaces one dense feed-forward network with a collection of smaller, independently parameterized feed-forward networks called experts. The model might store dozens or hundreds of experts, while routing each token to only one or two. This increases the model's capacity without multiplying computation by the total number of experts.

The trade-off is that routing becomes a scheduling problem. If too many tokens choose the same expert, that expert becomes a bottleneck even though other experts are idle. The layer must also move token data to the machines that host the selected experts, making uneven routing a performance problem as well as a modeling problem.

How a router selects experts

For each token representation, the router calculates one score per expert. A router is usually a small learned linear layer: it reads the token's current vector and produces these scores, called logits. A softmax converts the logits into a probability-like distribution whose values sum to one.

The router then keeps the highest-scoring experts. This is called top-k routing. In top-1 routing, each token goes to one expert. In top-2 routing, it goes to the two experts with the largest router scores. The selected scores can be used to weight the experts' outputs before combining them.

Conceptually, top-2 routing looks like this:

scores = softmax(router(token))
experts = top_k(scores, k=2)
outputs = [expert(token) for expert in experts]
result = sum(weight * output for weight, output in outputs)

This is illustrative pseudocode: production implementations normally batch tokens by destination expert rather than call an expert once per token.

The actual execution is usually split into four stages:

  1. Route: compute scores and select the expert IDs for each token.
  2. Dispatch: group selected token vectors by expert and send those groups to the devices holding the experts. A distributed transfer of this kind is often called an all-to-all, because every device may send data to and receive data from many others.
  3. Compute: run each expert on its assigned batch of tokens.
  4. Combine: return the results to the original token positions and apply the router weights if more than one expert was selected.

The rest of the transformer layer can then continue as if the MoE layer had produced one output vector per input token.

Capacity limits keep one expert from blocking the batch

A router can assign an unexpectedly large fraction of a batch to one expert. Without a limit, that expert would receive a large temporary batch, take longer to run, and delay the whole layer. On accelerators, experts are often processed in parallel, so the step finishes only when the slowest expert finishes.

Each expert therefore receives a capacity, the maximum number of token assignments it may process in that layer. A common capacity calculation for top-1 routing is approximately:

capacity = capacity factor × number of tokens / number of experts

The capacity factor is a tunable allowance above the perfectly even average. A factor of 1.0 gives each expert roughly its average share; a larger factor leaves room for normal variation but uses more memory and computation. For top-k routing, the expected number of assignments is multiplied by k, so the capacity calculation must account for those additional assignments.

If an expert reaches capacity, later assignments to it are overflow tokens. Implementations may drop the expert computation for those tokens, pass them through a residual or shared path, or use another fallback policy. Dropping does not necessarily mean dropping the token from the whole model; it usually means skipping that MoE expert's contribution for that layer.

Capacity limits put a hard ceiling on the damage caused by a hot expert, but they do not make routing balanced. A low capacity factor can cause many overflow tokens and reduce quality. A high factor reduces overflow but consumes resources and weakens the performance benefit of sparse execution.

The auxiliary loss encourages balanced choices

The router is trained along with the rest of the model. Its main objective—usually the model's language-modeling loss—may find it useful to send many tokens to the same expert. Nothing in that objective inherently values using all experts.

Training therefore adds an auxiliary load-balancing loss, an extra term whose job is to make expert usage more even. A representative formulation tracks two quantities for each expert:

A commonly used loss is proportional to:

number of experts × sum over experts of f_i × P_i

The exact coefficient and normalization vary between architectures. The important behavior is that the loss is smallest when both actual assignments and router probabilities are spread across experts. If one expert receives most assignments and the router strongly prefers it, the balancing penalty increases. During training, the main task loss and this auxiliary loss pull the router in different directions: the first asks for useful expert choices, while the second discourages concentration.

Some systems also use a router-logit regularizer, often called a router z-loss, to keep router scores from growing excessively. That can improve numerical stability, but it is separate from the load-balancing loss: one controls score magnitude, while the other controls expert usage.

What to monitor in production

Capacity and balancing solve related but different problems:

The loss is not a guarantee. Routing is decided for discrete tokens in each batch, so a batch can still be uneven even when average usage is healthy. Systems commonly monitor the number of tokens assigned to each expert, overflow or dropped-token rates, per-expert processing time, and communication time during dispatch.

When an incident report mentions a high expert load, capacity-factor changes, dropped tokens, or an all-to-all slowdown, it is usually describing this routing system. MoE performance depends not just on how many experts exist, but on whether the router keeps their work—and the data movement needed to reach them—close to evenly distributed.

← All articles