Standing
Current standardPart of the default recipe. Present in most frontier models trained today, and the thing a new design departs from rather than argues for.
Now the default at frontier scale. The unsettled parts are routing stability and expert granularity, not whether to use it.
judged as of 2026-09 · what the labels mean
Theory
A router produces logits over experts, keeps the top , renormalises them into weights and sums the selected experts’ outputs. Everything interesting about MoE is a consequence of the routing being discrete: the gate is not differentiable through the selection, only through the weights of what was selected.
Load balancing
Left alone, routers collapse: a few experts win early, receive more gradient, and win harder. The standard remedy is an auxiliary loss on the product of the dispatch fraction and the mean gate probability , which is minimised when both are uniform.
Implementation
import torch
import torch.nn.functional as F
from torch import Tensor, nn
class TopKRouter(nn.Module):
def __init__(self, d_model: int, n_experts: int, k: int = 2):
super().__init__()
self.gate = nn.Linear(d_model, n_experts, bias=False)
self.k = k
def forward(self, x: Tensor):
logits = self.gate(x) # [T, E]
probs = logits.softmax(dim=-1)
weights, idx = probs.topk(self.k, dim=-1)
weights = weights / weights.sum(-1, keepdim=True)
# auxiliary balance loss: dispatch fraction x mean gate prob
n_experts = probs.size(-1)
one_hot = F.one_hot(idx[..., 0], n_experts).float()
f = one_hot.mean(0)
p = probs.mean(0)
aux = n_experts * (f * p).sum()
return weights, idx, auxThe balance loss here is computed over the top-1 assignment only, which is what the Switch formulation specifies; summing over all selected experts changes the fixed point. In a distributed setting must be all-reduced across data-parallel ranks, or each rank balances its own shard and nothing balances globally.