Computing Library › Neural Architectures
Neural Architectures

Mixture-of-Experts Routing

A mixture of experts routes each input to a few specialized subnetworks out of many, growing total capacity while keeping per-input computation low.

Conditional computation

A mixture of experts (MoE) layer replaces a single dense subnetwork with many parallel expert subnetworks and a small router that decides which experts handle each input. Rather than every input passing through every parameter, each token is sent to only a few experts. This is conditional computation: the model can hold a very large number of parameters, and therefore a large capacity, while the cost of processing any single input stays roughly constant because only a fraction of the experts activate.

The router

Kronos motion — many body

For each token the router computes a score for every expert, usually with a small linear layer followed by a softmax, and selects the top few, often one or two. The token is processed by just those experts, and their outputs are combined weighted by the router's scores. Because only the selected experts run, the added parameters do not add proportional compute. The router is trained jointly with the experts, learning which inputs each expert should specialize in.

python
scores = softmax(router(x))              # (num_experts,)
topk = scores.topk(k)                    # choose k experts
y = sum(w * experts[i](x) for i, w in zip(topk.indices, topk.values))

Load balancing

Left alone, the router tends to favor a few experts, leaving others untrained and wasting capacity. To prevent this collapse, MoE layers add an auxiliary load-balancing loss that encourages tokens to spread evenly across experts, and they cap how many tokens each expert accepts per batch, dropping or rerouting overflow. Getting this balance right is central to training MoE models well.

Where it fits

MoE layers commonly replace the feedforward blocks of a transformer, letting a model scale parameter count far beyond a dense model of the same per-token cost. The trade-offs are systems complexity, since experts must be distributed across devices, and memory, since all experts must be stored even though few run at a time. Routing to specialized modules is a recurring theme, related in spirit to the conditional weight generation of hypernetworks.