A dense transformer pays for every parameter on every token. Double the model, double the compute, no matter how easy the token is. Mixture of Experts (MoE) breaks that bargain. It keeps a huge pool of parameters available, but activates only a small subset for each token. The result is a model that can be far larger in total size while costing roughly the same per token as a much smaller dense model. That decoupling of capacity from compute is the core idea, and it explains both the excitement around MoE and most of its failure modes.
This article explains how routing works, why the load balancing loss matters, what happens when routing collapses, and when MoE is the wrong choice.
The intuition: specialists, not generalists
Think of a hospital. You would not ask every doctor to examine every patient. You route a cardiac case to cardiology and a fracture to orthopaedics. The hospital is large, but each patient sees only a few specialists. MoE applies the same logic inside a feedforward layer.
In a standard transformer block, the feedforward network (FFN) transforms each token representation with the same weights. In an MoE layer, that FFN is replaced by a set of expert networks, usually also FFNs, plus a small router. For each token, the router scores the experts, picks the top-k (often 1 or 2), and only those experts run. Their outputs are combined with weights from the router.
The total parameter count grows with the number of experts. The active parameter count, and therefore the floating point cost per token, grows only with k. A model with 8 experts and top-2 routing has roughly 8 times the FFN parameters of its dense counterpart, but only about twice the FFN compute. In practice the router, attention, and embeddings are shared, so the compute increase is smaller still.
That is the promise. The catch is that the router must learn a useful division of labour, and nothing in ordinary training guarantees that it will.
Routing, formally
Let x be a token representation, and let there be N experts E_1, ..., E_N. A router is usually a linear layer followed by a softmax:
where is the gating weight for expert i. With top-k routing, we keep only the k largest gating values, renormalise them, and set the rest to zero. The layer output is:
Two details matter here. First, the gating function is discrete: a tiny change in router scores can flip which experts are selected, which makes gradients through the selection step tricky. In practice, gradients flow through the gating weights of the selected experts, and the router learns from that signal alone. Second, the experts that are never selected receive no gradient at all. That creates a rich-get-richer dynamic, which is the root of routing collapse.
Why naive routing collapses
Suppose one expert starts slightly better than the others, perhaps by chance. The router sends it more tokens. It gets more training signal, improves faster, and the router sends it even more tokens. Meanwhile, other experts starve. The model ends up using a fraction of its capacity, and the unused experts are dead weight.
The standard fix is an auxiliary load balancing loss. Let f_i be the fraction of tokens routed to expert i, and let P_i be the mean router probability assigned to expert i across a batch. The load balancing loss is:
This loss is minimised when routing is uniform, that is when for all i. It is added to the main task loss with a small coefficient, typically around 0.01. Too small, and collapse still happens. Too large, and the router is forced to spread tokens evenly even when specialisation would help, which defeats the purpose of having experts at all.
Note what this loss does not do: it does not guarantee meaningful specialisation. Uniform routing is not the same as useful routing. A router can balance load perfectly by splitting tokens arbitrarily, for example by token position or by surface features, without learning any semantic division of labour. Balanced routing is a necessary condition for using capacity, not a sufficient condition for good experts.
A related technique is expert capacity. Each expert is given a maximum number of tokens it may process per batch, often expressed as a capacity factor times the average load. Tokens beyond capacity are dropped, or passed through unchanged via a residual connection. Capacity limits prevent a single expert from being overwhelmed, but dropped tokens lose information, so the factor must be tuned carefully. A capacity factor of 1.0 means no slack at all, and some tokens will always be dropped. A factor of 1.25 to 2.0 is common during training, with a higher factor at inference to avoid dropping tokens when it matters most.
What experts actually learn
In language models, experts do not usually correspond to clean human categories like "science" and "history". Studies of trained MoE models find a mix: some experts specialise in syntactic roles, punctuation, or formatting, others in particular languages or domains, and many remain diffuse. That is fine. Specialisation does not need to be interpretable to be useful. What matters is that different experts handle different parts of the input distribution, so that the model as a whole can represent more functions than a dense model of the same active size.
There is a subtlety worth stating plainly. MoE increases capacity, not necessarily sample efficiency. A model with more parameters can memorise more, and routing adds optimisation difficulty. MoE shines when training compute is the bottleneck and data is plentiful, which is exactly the regime of large-scale pretraining. On a small dataset, a dense model is usually easier to train and harder to break.
A minimal example
The following PyTorch-style sketch shows the core mechanics: a router, top-k selection, and a weighted combination of expert outputs. It omits load balancing and capacity handling for clarity, but the structure is the same as in production implementations.
import torch
import torch.nn as nn
import torch.nn.functional as F
class MoELayer(nn.Module):
def __init__(self, d_model, d_ffn, num_experts, top_k):
super().__init__()
self.experts = nn.ModuleList([
nn.Sequential(
nn.Linear(d_model, d_ffn),
nn.GELU(),
nn.Linear(d_ffn, d_model),
)
for _ in range(num_experts)
])
self.router = nn.Linear(d_model, num_experts)
self.top_k = top_k
def forward(self, x):
# x: (batch, seq, d_model)
logits = self.router(x)
probs = F.softmax(logits, dim=-1)
top_vals, top_idx = torch.topk(probs, self.top_k, dim=-1)
top_vals = top_vals / top_vals.sum(dim=-1, keepdim=True)
out = torch.zeros_like(x)
for k in range(self.top_k):
idx = top_idx[..., k]
weight = top_vals[..., k].unsqueeze(-1)
# Route each token to its k-th chosen expert.
for expert_id, expert in enumerate(self.experts):
mask = (idx == expert_id)
if mask.any():
out[mask] += weight[mask] * expert(x[mask])
return out
Two practical points follow from this code. First, the loop over experts is a simplification: real implementations batch tokens by expert to avoid the per-expert masking overhead. Second, the implementation already shows why MoE complicates distributed training. Experts live on different devices in expert-parallel setups, and tokens must be sent across devices to their experts and back, which adds communication cost that dense models do not pay.
Common pitfalls
Routing collapse. Watch the fraction of tokens each expert receives. If a few experts dominate and the rest are near zero, increase the load balancing coefficient or check whether the router learning rate is too low relative to the experts.
Token dropping at inference. If you train with a capacity factor, test with a higher one, or with no dropping at all. A model that drops tokens silently at inference will show puzzling quality drops on long or unusual inputs.
Expert parallelism overhead. MoE saves compute but adds communication. On a single device, or with slow interconnects, the all-to-all communication between tokens and experts can erase the savings. Measure wall-clock time, not just floating point operations.
Fine-tuning a pretrained MoE. Fine-tuning data is often much narrower than pretraining data. The router may collapse onto a few experts for the new distribution, and other experts may be effectively lost. Lower learning rates for the router, or freezing the router during early fine-tuning, often help.
Assuming more experts is always better. More experts mean more capacity, but also a harder routing problem and more communication. Many strong results use 8 to 64 experts with top-1 or top-2 routing. Scaling the expert count without scaling data and tuning usually disappoints.
When to use MoE, and when not to
Use MoE when you need a very large model, you have abundant training data, and per-token inference cost is a hard constraint. That is why MoE is common in frontier language models: it lets a provider serve a model with hundreds of billions of total parameters at the cost of a much smaller dense model.
Do not use MoE when your dataset is small, your deployment is a single device, or your team cannot afford the extra tuning and debugging. A dense model is simpler, more predictable, and often better in those regimes. MoE is a scaling tool, not a magic improvement. It trades simplicity for capacity at fixed compute, and that trade is only worth it when capacity is actually your bottleneck.
The takeaway
Mixture of Experts separates two quantities that dense models fuse together: how big a model is, and how much it costs to run on one token. A learned router sends each token to a few specialists, a load balancing loss keeps those specialists alive, and capacity limits keep the system honest. When those pieces work together, you get a much larger model for roughly the same per-token cost. When they do not, you get a large model that behaves like a small one, plus extra engineering pain. The difference is almost always in the routing.
Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial
Top comments (0)