Mixture-of-Experts models dominate the frontier right now, and almost every explanation stops at the same level: “a router picks a few experts per token.” That sentence hides the most interesting part of the architecture. The gate is where these models live or die — where routing collapse, load balancing losses, and the clever bias-update trick behind DeepSeek-V3 all play out. If you deploy or fine-tune MoE models, the routing mechanics are not trivia; they determine throughput, cost, and how well a fine-tune actually lands.
This post walks through the routing decision exactly as it happens — affinity scores, top-k selection, weight renormalization — then looks at why the naive version breaks, how auxiliary losses fixed it and what they cost, and how the auxiliary-loss-free balancing strategy changed the trade-off. Working PyTorch snippets for each stage, sized to be readable rather than production-complete.
The routing decision, step by step
A dense transformer runs every token through every feed-forward parameter. An MoE layer replaces that single FFN block with many expert FFNs — anywhere from 8 to several hundred — plus a small gating network that decides which experts process each token. The parameter count explodes; the compute per token does not, because only a few experts fire. DeepSeek-V3 carries 671B parameters but activates roughly 37B per token. Kimi K3 scales the same idea much further: 896 routed experts, 16 active per token.
Concretely, the gate is a linear projection from the token’s hidden vector to one score per expert, and routing is a top-k over those scores. Here is a minimal router:
import torch
import torch.nn as nn
import torch.nn.functional as F
class TopKRouter(nn.Module):
def __init__(self, hidden_dim: int, n_experts: int, top_k: int):
super().__init__()
self.proj = nn.Linear(hidden_dim, n_experts, bias=False)
self.top_k = top_k
def forward(self, x):
# x: (batch, seq_len, hidden_dim)
scores = torch.sigmoid(self.proj(x)) # token-expert affinities
weights, idx = torch.topk(scores, self.top_k, dim=-1)
# renormalize so the selected experts fully define the mixture
weights = weights / weights.sum(dim=-1, keepdim=True)
return idx, weights # (.., top_k) eachEach selected expert then transforms the token, and the outputs are combined as a weighted sum using these normalized weights. Two details matter more than they look. First, the top-k operation is non-differentiable — argmax-style selection has no gradient — so the gate only learns through the gradient flowing into the weights of the experts that happened to be picked. Second, earlier models used softmax over the scores (Mixtral does exactly this); DeepSeek-V3 switched to sigmoid, scoring each expert independently rather than forcing the scores to compete for a probability mass of one. Sigmoid made the gate less sensitive to how many strong candidates a token has.
Why routing collapses without help
Left to its own devices, top-k routing develops a taste for a few experts — and the dynamics are self-reinforcing. The favored experts receive more tokens, so they train more, so their outputs improve, so the gate prefers them even more. The rest of the experts starve: their weights get almost no useful gradient updates, and the model effectively pays for capacity it never uses. Worse, because expert computation is sharded across GPUs, hot experts become throughput bottlenecks for the whole layer.
This is not a hypothetical tail risk. It shows up reliably in from-scratch and continued-training runs. The earliest large MoE language model had to fight it with an explicit load-balancing loss, and every major MoE line since has shipped some mechanism for the same purpose.
Auxiliary losses: the first fix and its tax
The classic approach, most cleanly defined in the Switch Transformer work, adds a term to the training loss that penalizes imbalance. For each MoE layer, measure two things per expert: the fraction of tokens actually routed to it, and its average gate probability. When routing is perfectly balanced, both equal 1/N, and their product summed over experts is minimized.
def switch_aux_loss(gate_probs, expert_idx, n_experts, alpha=1e-2):
# gate_probs: (tokens, n_experts) softmax/sigmoid scores
# expert_idx: (tokens,) hard assignments from top-k (top-1 here)
f = F.one_hot(expert_idx, n_experts).sum(0).float()
f = f / f.sum() # hard load, non-differentiable
p = gate_probs.mean(dim=0) # soft probabilities, differentiable
return alpha * n_experts * torch.sum(f * p)The trick is the pairing: the differentiable factor lets gradient reach the gate, while the hard factor reports what actually happened. A closely related regularizer, the router z-loss, penalizes the magnitude of the raw gate logits to keep training numerically stable. Both losses are added on top of the language-modeling loss, with small scaling coefficients.
The tax shows up late in training. The auxiliary loss pushes the gate toward uniformity, which is a slightly different objective from routing well. With a large coefficient, that interference measurably drags on model quality — you are spending optimization pressure on load balancing that the language-modeling loss would have spent on capability. Teams end up tuning the coefficient carefully: too high hurts quality, too low lets imbalance creep back.
The auxiliary-loss-free trick: bias the selection, not the gradient
DeepSeek-V3 popularized a different balancing mechanism. Instead of adding a loss term, maintain one scalar bias per expert. During top-k selection, add the bias to the scores — overloaded experts look slightly worse, underloaded ones slightly better — but compute the gate weights from the raw, unbiased scores. The bias shifts which experts are chosen while leaving the gradient path untouched, and it is updated by a simple rule between optimizer steps, not by backpropagation:
bias = torch.zeros(n_experts) # NOT registered with the optimizer
gamma = 1e-3
def route_with_bias(x, router, bias):
scores = torch.sigmoid(router.proj(x)) # raw affinities
_, idx = torch.topk(scores + bias, router.top_k, dim=-1) # bias guides selection
weights = scores.gather(-1, idx) # gate value uses raw scores
weights = weights / weights.sum(dim=-1, keepdim=True)
return idx, weights
def update_bias(bias, idx):
load = torch.bincount(idx.flatten(), minlength=bias.numel()).float()
load = load / load.sum()
mean = 1.0 / bias.numel()
bias -= gamma * (load - mean) # overloaded expert -> bias down, starved -> upThree lines of real logic, no interference gradients, and load balance comparable to a tuned auxiliary loss with better downstream quality — which is why the strategy shipped in DeepSeek-V3. Two nuances worth knowing. The paper still uses a small complementary sequence-wise balance loss alongside the bias, so “auxiliary-loss-free” describes the primary mechanism, not the entire training recipe. And the bias interacts cleanly with sigmoid gating precisely because sigmoid scores experts independently: there is no softmax normalization for the bias to distort.
The same “parameter the optimizer never sees, updated by a moving rule” pattern is older than LLMs — running statistics in batch norm and target networks in reinforcement learning are the same shape. The insight is structural: some quantities want to adapt faster and more crudely than gradient descent allows.
Shared experts and fine-grained experts
Two more routing-adjacent choices from the DeepSeekMoE line show up in most modern designs. Fine-grained experts: instead of a few large experts, use many small ones and activate a larger number of them per token. With 256 narrow experts and 8 active, a token can draw on far more routing combinations than with 16 wide experts and 2 active, so specialization gets finer without raising compute. Shared experts: one or two experts are always active for every token, no routing involved. They absorb what every token needs — common knowledge that pure sparse routing forces every expert to relearn — leaving the routed experts free to specialize on the long tail.
Layered on top: with sigmoid scoring, independent per-expert biases, shared experts handling the universal baseline, and hundreds of narrow routed experts for specialization, the modern MoE recipe is a stack of small, individually boring decisions that compound into the parameter-to-compute ratio that makes frontier-scale models economical.
What routing means at inference time
For serving rather than training, routing is a systems problem. Expert weights are sharded across GPUs (expert parallelism), so every layer requires an all-to-all exchange: each GPU dispatches its tokens to whichever devices hold the selected experts, then gathers results. A skewed router — the same collapse problem in miniature — turns the busiest expert’s GPU into the layer’s critical path. This is exactly why inference frameworks report expert-load balance in their dashboards, and why uneven token distribution across a batch can hurt MoE throughput in ways a dense model never exhibits.
Practical consequences worth carrying into deployment planning: expert imbalance in your fine-tuning data can shift routing behavior; a quantization or LoRA job that touches experts unevenly will amplify imbalance; and batch composition affects latency because all-to-all traffic scales with how widely a batch’s tokens spread across experts.
If you take one thing from this: the router is a learned system with its own dynamics, and every serious MoE architecture ships explicit machinery to keep those dynamics healthy. When something looks wrong with an MoE model’s behavior or throughput, the gate is one of the first places worth looking.