Mixture of Experts Routing: Top-k Gating, Load Balancing and Aux-Loss-Free Balancing

A sparse model is only as good as its router. Mixture-of-experts (MoE) layers promise a huge parameter count at a fixed per-token compute bill, but that promise is kept or broken by one tiny component: a single linear layer and a top-k selection that decides which experts see which token. Get mixture of experts routing wrong and you get expert collapse, where a handful of experts absorb the traffic and the rest never learn. You also get dropped tokens, loss spikes, and inference clusters where one hot GPU holds back every other device.

This matters more in 2026 than it did in 2021. Open-weight frontier models such as Mixtral and DeepSeek-V3 are MoE models, and the routing recipe has moved from a hand-tuned auxiliary penalty to bias-based, gradient-free control. This post walks through the router math, the capacity mechanics, the loss terms that fight collapse, and the aux-loss-free alternative, then connects all of it to inference.

What this covers: the gating equations from Shazeer 2017 to DeepSeek-V3, how expert capacity and token dropping work with worked numbers, why balance losses hurt and how bias-based balancing avoids that, and a runnable PyTorch router you can adapt.

Context and Background

The sparsely-gated MoE layer came from Shazeer et al. (2017), who put up to 131,072 experts in a recurrent language model and reported “up to 137 billion parameters in the MoE layer” (arXiv:1701.06538). Two ideas from that paper survive in every successor: a learned gate that keeps only the top-k scores, and an explicit pressure to keep experts evenly used. GShard (Lepikhin et al., 2020) moved the idea into Transformers, using top-2 gating and a fixed per-expert capacity to scale a translation model to 600B parameters across 2,048 experts on 2,048 TPU v3 cores (arXiv:2006.16668).

Switch Transformer (Fedus et al., 2021) simplified routing to a single expert per token, reporting up to 7x faster pre-training than the T5 baselines at matched compute and a trillion-parameter model (arXiv:2101.03961). ST-MoE (Zoph et al., 2022) added the router z-loss to fix training instability. Expert-choice routing (Zhou et al., 2022) flipped the direction of selection and reported more than 2x faster convergence versus Switch and GShard gating (arXiv:2202.09368).

The modern era of mixture of experts routing is defined by two open models. Mixtral 8x7B uses 8 experts per layer with top-2 routing, 47B total parameters and roughly 13B active per token, released under Apache 2.0 (arXiv:2401.04088). DeepSeek-V3 uses 671B total parameters with 37B active per token, 256 routed experts plus 1 shared expert per MoE layer, and top-8 routing (arXiv:2412.19437).

Why does routing deserve its own deep dive? Because the MoE layer is not the hard part; the experts are ordinary feed-forward networks. The difficulty is a discrete, non-differentiable choice sitting inside a differentiable network, coupled to a hardware constraint that every expert should receive about the same amount of work. For the serving side of that constraint, see our companion post on tensor, expert and data parallelism for MoE inference. Everything below is about the training-time half of mixture of experts routing.

Mixture of Experts Routing: Gating Math and Reference Architecture

Direct answer: An MoE router is a learned linear map from a token’s hidden state to one score per expert. The layer keeps the k highest scores, converts them into gate weights, sends the token to those k experts, and returns the gate-weighted sum of their outputs. Everything else in mixture of experts routing research is a fix for the failure modes of this simple mechanism.

Figure 1: The mixture of experts routing pipeline for a single MoE layer, including the always-on shared expert used by DeepSeek-style models.

Figure 1 shows one MoE layer. The token’s hidden state enters the router, scores are computed, top-k indices are selected, the token is shipped to the devices that host those experts, and the expert outputs are combined and added back to the residual stream. In expert-parallel training the dispatch and combine steps are all-to-all collectives, which is why routing decisions have a direct network cost.

Softmax top-k gating

The classic formulation is a softmax over expert logits followed by top-k selection. For a token with hidden state u and router matrix W_g, the logits are l = u W_g, one per expert. Shazeer’s noisy variant computes H(x)_i = (x W_g)_i + StandardNormal() * Softplus((x W_noise)_i), then keeps the top k and applies a softmax over the kept values: G(x) = Softmax(KeepTopK(H(x), k)). The learned noise matrix W_noise lets the network control exploration per expert.

Two normalisation orders exist, and they behave differently. Shazeer and Mixtral take the softmax over only the k selected logits, so the gates always sum to exactly 1. Switch-style routers compute the softmax over all N experts first and then pick the top entries, so the gate on the chosen expert carries information about how confident the router was relative to the whole pool. With top-1 and softmax-over-all, the gate is a probability below 1, and that probability is what carries the gradient back to the router.

That last point is subtle and worth stating plainly. The top-k index selection has no gradient. The only path by which the router learns is the gate value multiplied into the expert output. If you renormalise over the selected set and k equals 1, the gate is identically 1 and the router receives no learning signal from the task loss, which is why Switch kept the unnormalised probability.

Sigmoid gating and normalised affinities

DeepSeek-V3 replaces the softmax with a sigmoid. The paper defines affinity s_i,t = Sigmoid(u_t^T e_i), where e_i is the centroid vector of routed expert i, selects the top-K_r affinities, and normalises across the selected experts: g_i,t = g’_i,t / sum_j g’_j,t. The DeepSeek-V3 report notes this differs from its predecessor, DeepSeek-V2, which used softmax affinities.

Why prefer sigmoid at 256 experts? With softmax, scores compete: raising one expert’s probability lowers all others, and the distribution gets very flat or very peaked as N grows. Sigmoid scores are independent per expert, so an expert’s affinity does not shrink just because many other experts also match. Normalising only over the selected top-8 then restores a clean convex combination. This is a design rationale, not a proven ablation result, so treat it as the likely reason rather than a measured one.

Why top-k at all: compute versus capacity

Active parameters per token scale with k, while total parameters scale with N. Mixtral’s 47B total versus ~13B active shows the ratio: two of eight experts fire, but attention and embeddings are shared, so active parameters are more than a quarter of the total. DeepSeek-V3’s ratio is more extreme at 37B of 671B, roughly 5.5%.

Larger k means each token blends more experts, which smooths the function the layer represents but raises both FLOPs and all-to-all traffic linearly. Smaller k saves communication but makes each routing decision higher stakes. DeepSeekMoE (Dai et al., 2024) argues for the other direction at fixed cost: split each expert into finer pieces and route to proportionally more of them, which gives a combinatorially larger set of possible expert mixes (arXiv:2401.06066).

Shared experts and fine-grained experts

The DeepSeekMoE design adds two structural ideas. Fine-grained segmentation divides each conventional expert into several smaller ones and activates more of them. A shared expert, always active for every token, absorbs common knowledge such as syntax so the routed experts do not each waste capacity relearning it.

The count of possible expert combinations explains the appeal. Choosing 2 of 16 experts gives 120 combinations. Splitting each expert into four and choosing 8 of 64 gives about 4.4 billion. The compute is the same, since each of the 8 small experts is a quarter the width of an original, but the router has vastly more freedom to specialise. The abstract of the DeepSeekMoE paper reports that its 2B model matches GShard 2.9B, which has 1.5 times the expert parameters and compute; these are the authors’ own claims.

Expert Capacity, Token Dropping and Load Balancing

Direct answer: Expert capacity is the maximum number of tokens one expert may process per batch, set as (tokens per batch / number of experts) x capacity factor. Tokens that exceed capacity are dropped from that expert and pass through on the residual connection. Capacity keeps tensor shapes static for accelerators, at the price of silently skipping computation.

The capacity factor, with numbers

Hardware wants static shapes. XLA on TPUs, and CUDA graphs on GPUs, are far more efficient when every expert processes a fixed-size buffer. So Switch Transformer defines expert capacity as (tokens per batch / number of experts) x capacity factor. A factor of 1.0 allocates exactly the average load. A factor above 1.0 adds buffer for imbalance, at the cost of padding compute and communication.

Take a worked example. A batch has 4,096 tokens and the layer has 64 experts. With top-1 routing the average load is 4,096 / 64 = 64 tokens per expert. At capacity factor 1.25 each expert’s buffer holds 80 tokens. With top-2 routing, as in GShard and Mixtral, each token creates two assignments, so there are 8,192 assignments, an average of 128 per expert, and a buffer of 160 at factor 1.25.

If the router were perfectly uniform, no token would be dropped even at factor 1.0. Real routers are not uniform. If one expert wants 130 tokens in the top-1 case, 50 tokens beyond its 80-slot buffer are dropped, a 1.2% loss of the whole batch from one hot expert alone. Several hot experts compound this.

Figure 2: How expert capacity works. Each expert has a fixed-size queue, tokens beyond capacity skip that expert, and the residual connection carries them to the next layer.

Figure 2 traces the mechanism. Tokens are assigned to per-expert queues in a deterministic order, commonly sequence position, so tokens late in a sequence or batch are the ones dropped. A dropped token still flows through the layer via the residual connection; it simply receives no update from that expert. The long-description version: capacity turns an unbounded, data-dependent load into a bounded one, and the overflow is the price.

What the papers measured about dropping

Switch Transformer reports that capacity factors of 1.0 and 1.25 worked better for the Switch model than larger values, and that dropped tokens were typically under 1% when the balance loss was set properly. It also tried “No-Token-Left-Behind,” rerouting overflow tokens to their second-choice expert, and found no empirical benefit.

ST-MoE quantified how tolerant fine-tuning is. With a train capacity factor of 0.75, about 10.6% of tokens were dropped and SuperGLUE scored 86.5, versus 86.7 at factor 1.25 with 0.3% dropped. Raising the train factor from 1.25 to 2.0 improved perplexity by only about 0.009 but increased step time by roughly 14% for the 269B-parameter ST-MoE-32B (4.244 s to 4.819 s), so the authors kept 1.25. The lesson is that moderate dropping is cheaper than moderate padding.

GShard added a refinement: random routing of the second expert. The second expert is used only with probability proportional to its gate weight, so tokens whose second choice has a tiny gate skip it and free capacity for others. GShard also dispatches tokens in local groups, each with its own fractional capacity of 2N/(G*E), which keeps the capacity check local to a group instead of requiring a global sort.

Why balance is needed at all: expert collapse

Capacity limits the damage of imbalance but does not cure its cause. Mixture of experts routing has positive feedback. An expert that receives more tokens gets more gradient updates, becomes better, earns higher router scores, and receives still more tokens. Left alone, the router converges to using a few experts, a failure called expert collapse or routing collapse. The unused experts are dead parameters, and the dense compute you thought you were spreading concentrates on a subset.

Collapse also has a hardware face. In expert-parallel training, the slowest device sets the step time. A single overloaded expert makes every other GPU wait at the next collective. Balance is therefore both a modelling problem and a throughput problem, which is why it is worth spending loss budget on.

The auxiliary load-balancing loss

Shazeer used two losses: an importance loss, w_importance x CV(Importance)^2, with Importance being the batchwise sum of gate values per expert, and a load loss that uses a smooth estimate of how many examples each expert receives, w_load x CV(Load)^2. Here CV is the coefficient of variation. The importance loss alone is insufficient because an expert can have equal total gate weight yet receive very different token counts. The paper used weights of 0.1 on language modelling and 0.01 on translation.

Switch Transformer collapsed these into one term, loss = alpha x N x sum_i f_i x P_i. Here f_i is the fraction of tokens dispatched to expert i, a hard, non-differentiable count, and P_i is the mean router probability assigned to expert i, which is differentiable. The product is minimised by uniform f and P, and the gradient flows through P. Switch used alpha = 10^-2 and swept values from 10^-1 down to 10^-5. GShard’s version is structurally similar, using mean gate m_e as a differentiable surrogate for the token count c_e.

A quick sanity check on the scale. With perfectly uniform routing, f_i = P_i = 1/N, so the sum is N x (1/N^2) = 1/N and the loss equals alpha. With total collapse onto one expert, f = P = 1 for that expert, and the loss equals alpha x N. For N = 256 the collapsed loss is 256 times the balanced loss. The N factor exists so the balanced optimum stays at a constant alpha independent of expert count.

Router z-loss

ST-MoE introduced a second regulariser for a different problem: logits drifting to large magnitudes, which makes the softmax saturate and amplifies bf16 rounding error. The z-loss is L_z = (1/B) x sum over tokens of (log sum_j exp(x_j))^2, with coefficient c_z = 0.001. In the paper’s ablation, the baseline was stable in 4 of 6 runs, while the z-loss runs were stable in 3 of 3 with quality -1.741 versus -1.755. Noise-based fixes such as input jitter and dropout also stabilised training but lowered quality, which motivated the z-loss.

Switch’s own stability recipe is complementary: run the router softmax in float32 locally and cast back to bfloat16 afterwards (selective precision), and scale down initialisation by 10x. In its Table 2 pure bfloat16 diverged, while selective precision matched float32 quality (-1.716 versus -1.718) at nearly bfloat16 speed.

Aux-Loss-Free Balancing: The DeepSeek-V3 Router

Direct answer: Auxiliary-loss-free balancing keeps experts evenly loaded without adding a balance term to the training loss. A per-expert bias is added to the routing scores only when selecting the top-k, and after each step the bias of overloaded experts is nudged down and that of underloaded experts nudged up. No gradient flows through the correction.

The problem with the auxiliary loss

An auxiliary balance loss is a trade. The gradient that pushes the router toward uniform load is unrelated to next-token prediction, so it competes with the language-modelling gradient. A small alpha leaves collapse risk; a large alpha forces uniformity at the expense of specialisation. Wang et al. (2024) state this directly: a large auxiliary loss introduces interference gradients that hurt performance (arXiv:2408.15664).

The tuning surface is also awkward. The right alpha depends on expert count, batch size, sequence length and learning rate, and the loss is computed per micro-batch, so it enforces balance over a few thousand tokens when what hardware cares about is balance across a global batch. A router can satisfy a global-batch balance while being locally lumpy, and a micro-batch loss penalises exactly the domain-specific skew that specialisation produces. A code sample routed to a few “code” experts is correct behaviour yet looks like imbalance to a sequence-level penalty.

Loss-Free Balancing: the bias mechanism

Loss-Free Balancing, from Wang, Gao, Zhao, Sun and Dai, adds an expert-specific bias b_i to each expert’s routing score before the top-k decision, and updates b_i dynamically from the expert’s recent load. The authors tested it on MoE models of up to 3B parameters trained on up to 200B tokens and reported both better performance and better balance than auxiliary-loss strategies. Those are the authors’ results at modest scale, and the claim at frontier scale rests on the DeepSeek-V3 report.

Figure 3: Two ways to balance experts. The auxiliary loss injects a second gradient into the router; the bias method uses a feedback controller outside backpropagation.

Figure 3 contrasts the two. In the bias approach the top-k selection uses s_i + b_i, but the gate weight multiplied into the expert output is computed from the original s_i. That separation is the key design point. The bias steers who gets chosen, while the gate decides how much each chosen expert contributes. Because b_i never enters the forward value of the layer’s output, it never needs a gradient and never distorts what the expert’s contribution is worth.

It is, in effect, a proportional-style feedback controller sitting beside the optimiser. If expert i received more tokens than the batch average, lower b_i; if fewer, raise it. The only knob is the step size.

What DeepSeek-V3 actually does

The DeepSeek-V3 report states that it pioneers an auxiliary-loss-free strategy for load balancing, and specifies the recipe in enough detail to reproduce. Per MoE layer there is 1 shared expert and 256 routed experts, with 8 routed experts active per token and an expert intermediate hidden dimension of 2048. The first three layers are dense feed-forward networks; every later layer is MoE.

At each training step the bias of each expert is decreased by gamma if the expert is overloaded in the batch and increased by gamma if underloaded. The report sets this bias update speed to 0.001 for the first 14.3T tokens and to 0.0 for the remaining 500B, freezing the bias at the end. The 14.3T plus 0.5T split matches the 14.8T total training tokens.

Pure bias balancing is not the only mechanism. DeepSeek-V3 also keeps a very small complementary sequence-wise balance loss, with coefficient alpha = 0.0001, to stop extreme imbalance inside any single sequence. This is worth emphasising because the model is often described as “no auxiliary loss”. Precisely, the primary balancing mechanism is loss-free and the auxiliary term is a tiny safeguard that is two orders of magnitude below Switch’s alpha of 0.01.

Two further decisions support balance and efficiency. Node-limited routing sends each token to at most M = 4 nodes, which caps the cross-node all-to-all fan-out, because InfiniBand is much slower than NVLink inside a node. And DeepSeek-V3 drops no tokens during training or inference, because its balancing is effective enough that capacity buffers are unnecessary. That removes the dropped-token failure mode from the previous section entirely.

A worked feedback example

Suppose a layer has 8 experts, top-2 routing and a batch of 1,000 tokens, so there are 2,000 assignments and the target load is 250 per expert. Expert 3 receives 400 and expert 5 receives 100. With a sign-style update and gamma = 0.001, b_3 falls by 0.001 and b_5 rises by 0.001 each step.

Is that enough to move traffic? It depends on the gap between scores. If sigmoid affinities of competing experts typically differ by 0.05 near the selection boundary, it takes about 50 consecutive steps of the same sign to flip marginal decisions. That timescale is the point: the bias moves slowly relative to individual batches, so it reacts to persistent skew rather than noise, yet fast relative to a training run of hundreds of thousands of steps. (The 0.05 gap is an illustrative assumption for the arithmetic, not a measured number.)

The same dynamic explains why freezing the bias at the end of training is sensible. Once the model has converged, continuing to chase batch-level load fluctuations would add noise to the selected experts for no benefit.

Expert-choice routing as the other answer

Token-choice routing, everything above, lets each token pick its experts, so load is a statistical outcome. Expert-choice routing (Zhou et al., 2022) inverts this: each expert picks its top-C tokens from the batch. Load is balanced by construction, every expert processes exactly C tokens, and no auxiliary loss or token dropping from capacity is needed. The paper reports more than 2x faster training convergence than Switch top-1 and GShard top-2 gating at the same compute, and better results on 11 GLUE and SuperGLUE tasks.

The cost is variable per-token compute. Some tokens are picked by many experts, some by none, so a token may receive zero expert computation. Expert choice also has a causality problem for autoregressive decoding: choosing the top-C tokens across a sequence uses information about later tokens, so it does not transfer cleanly to generation with a growing prefix. That is one reason the open frontier models above remain token-choice with better balancing, rather than expert-choice. Treat this as an analysis of the mechanism rather than a statement any of those papers make about their own design choices.

Deeper Analysis: A Runnable Router and Its Failure Signatures

Direct answer: A production-quality router needs five pieces: a linear gate in float32, top-k selection, normalised gate weights, a balance loss or bias controller, and a z-loss. The snippet below implements a softmax top-k router with the Switch-style balance loss and ST-MoE z-loss, and exposes the per-expert load you need for monitoring.

Reference code for mixture of experts routing

The snippet is a minimal, single-device reference. It is written to be read and adapted, not as a fast kernel; real systems fuse dispatch with permutation kernels and expert-parallel collectives. I wrote it against the PyTorch 2.x API but did not benchmark it for this post.

import torch
import torch.nn as nn
import torch.nn.functional as F


class TopKRouter(nn.Module):
    """Softmax top-k router with Switch-style balance loss and ST-MoE z-loss."""

    def __init__(self, d_model, n_experts, k=2, capacity_factor=1.25,
                 balance_coef=1e-2, z_coef=1e-3, normalize_topk=True):
        super().__init__()
        self.n_experts, self.k = n_experts, k
        self.capacity_factor = capacity_factor
        self.balance_coef, self.z_coef = balance_coef, z_coef
        self.normalize_topk = normalize_topk
        self.gate = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        # x: [T, d_model] tokens flattened across batch and sequence
        T = x.size(0)
        logits = self.gate(x).float()              # route in fp32 (selective precision)
        probs = F.softmax(logits, dim=-1)          # [T, N]
        topv, topi = probs.topk(self.k, dim=-1)    # [T, k]
        if self.normalize_topk:                    # Mixtral-style: gates sum to 1
            topv = topv / topv.sum(dim=-1, keepdim=True)

        # --- load statistics ---
        one_hot = F.one_hot(topi, self.n_experts).sum(dim=1).float()  # [T, N] 0/1
        f = one_hot.mean(dim=0) / self.k           # fraction of assignments per expert
        P = probs.mean(dim=0)                      # mean router prob per expert
        balance = self.balance_coef * self.n_experts * (f * P).sum()
        z = self.z_coef * torch.logsumexp(logits, dim=-1).pow(2).mean()

        # --- capacity: keep first C assignments per expert, in token order ---
        capacity = int(self.capacity_factor * self.k * T / self.n_experts)
        position = one_hot.cumsum(dim=0) * one_hot  # 1-based slot of each token
        keep = (position <= capacity).float() * one_hot
        kept_mask = keep.gather(1, topi)           # [T, k] 1 if assignment kept
        gates = topv * kept_mask                   # dropped assignments get weight 0

        stats = {"load": f.detach(), "dropped_frac": 1 - kept_mask.mean().item()}
        return topi, gates, balance + z, stats


if __name__ == "__main__":
    torch.manual_seed(0)
    router = TopKRouter(d_model=64, n_experts=8, k=2)
    x = torch.randn(1024, 64)
    idx, gates, aux, stats = router(x)
    print(idx.shape, gates.shape, float(aux), stats["dropped_frac"])
    print("max/mean load:", (stats["load"].max() / stats["load"].mean()).item())

Three details deserve comment. The f term uses hard one-hot counts and so carries no gradient; the balance loss gradient flows entirely through P, exactly as in Switch. The router runs in float32 even if the model is in bf16, following Switch’s selective-precision finding. And the capacity check uses a cumulative sum to give each expert’s first C assignments priority in token order, matching the deterministic drop policy described earlier. In a distributed implementation the same logic runs per group or per device.

To switch this to bias-based balancing, add a non-trainable buffer bias of shape [N]. Select with (probs + self.bias).topk(...) but gather the gate values from probs, and update after each step with self.bias += gamma * torch.sign(mean_load - load). Because the bias is a buffer and not a Parameter, no optimiser state or gradient exists for it. If you make this change, set the balance coefficient to a very small value or zero.

Monitoring: the failure signatures

Mixture of experts routing bugs rarely crash a run; they degrade it quietly. Four signals are worth logging per layer.

The first is max-to-mean load ratio. A healthy layer is near 1.0 to 1.3; a ratio above 2 means one expert does double its share and, in expert-parallel training, sets the step time. The second is dropped-token fraction under capacity-based routing. Sustained values above a couple of percent suggest the balance loss is too weak or the capacity factor too tight, and ST-MoE’s results suggest that 10% is survivable for fine-tuning but not something you want from pre-training.

The third is router entropy: the mean entropy of the softmax distribution. Entropy collapsing toward zero early in training means the router has become near-deterministic, a precursor of collapse. The fourth is the maximum router logit magnitude. Climbing logits are the instability the z-loss was designed to catch, and a spike here often precedes a loss spike by many steps.

Figure 4: Expert-parallel execution of a routed layer. Tokens are dispatched to expert-hosting GPUs and combined after the experts finish; the layer cannot complete until the slowest expert does.

Inference implications: all-to-all and expert imbalance

Figure 4 shows why mixture of experts routing matters at serving time. In expert-parallel inference, each MoE layer performs a dispatch all-to-all and a combine all-to-all, and the layer finishes only when the most loaded expert finishes. Training-time balance does not guarantee inference-time balance. Production traffic is not the training distribution: a burst of code or one language can concentrate on a subset of experts.

A rough per-token traffic estimate illustrates the scale. Take DeepSeek-V3’s hidden size of 7,168 and top-8 routing. Dispatching one token’s activation to 8 experts in 8-bit precision moves about 8 x 7,168 = 57 KB, and the combine step in 16-bit precision moves about 115 KB, roughly 172 KB per token per MoE layer. This is my arithmetic on stated dimensions, ignoring any deduplication from node-limited routing, and the precision choices follow the report’s statement that dispatch uses lower precision than combine. Multiply by 58 MoE layers and thousands of tokens per batch and the interconnect, not the arithmetic units, becomes the binding constraint.

DeepSeek-V3’s serving design addresses the imbalance directly. During prefill the minimum deployment unit is 4 nodes (32 GPUs) with expert parallelism of 32; during decoding the minimum unit is 40 nodes (320 GPUs) with expert parallelism of 320, each GPU hosting one expert, with 64 GPUs hosting redundant and shared experts. Hot experts are duplicated: the prefill stage uses 32 redundant experts and re-determines the set periodically, about every 10 minutes, based on observed load. That is routing-aware load balancing applied at the cluster level, and the same logic reappears in every MoE serving stack. Our guide to sparse attention for long-context LLM training and inference covers the attention half of the same efficiency story, where the long-context KV traffic competes with expert traffic for bandwidth.

Comparison matrix

Method	Who decides	Balance mechanism	Token dropping	Main weakness
Noisy top-k (Shazeer 2017)	Token	Importance + load loss	No capacity in original	Two tuned loss weights
Top-2 + capacity (GShard)	Token	Aux loss, random 2nd expert	Yes, via capacity	Drops tokens; group-local balance
Top-1 (Switch)	Token	Aux loss alpha 0.01	Yes, via capacity factor	Router gradient needs unnormalised gate
Top-2 softmax (Mixtral)	Token	See Mixtral paper	Not described in the paper	Few large experts; coarse specialisation
Sigmoid top-8 + shared (DeepSeek-V3)	Token	Bias controller plus tiny seq loss	None	Bias lag; needs tuning of gamma
Expert choice	Expert	Perfect by construction	Tokens can get zero experts	Causality in autoregressive decoding

For Mixtral, the paper describes the top-2 gating and the model’s size but, in what I could verify, does not publish a detailed balancing recipe; I have therefore not asserted one.

Trade-offs, Gotchas, and What Goes Wrong

In mixture of experts routing, balance and quality pull in opposite directions. Perfect balance means every expert sees an identical slice of the data, which is the opposite of specialisation. Every method above is a different point on that curve, and no setting is free. The aux-loss-free approach moves the point toward quality, but it adds a controller that can oscillate if gamma is too large or lag if it is too small.

Capacity-based dropping fails silently. A dropped token receives no expert computation for that layer, and the loss curve may look fine while specific content, usually the tail of long sequences because of the token-order drop policy, is systematically under-processed. Evaluating with a different batch composition than training can change the drop rate. This is a reason to prefer dropless designs for inference and to log the dropped fraction as a first-class metric.

Micro-batch versus global-batch balance is another trap. A balance loss computed per micro-batch penalises legitimate domain skew. If your data pipeline serves homogeneous micro-batches, for example one source per batch, the loss will fight specialisation harder than you intend.

Router precision bites at scale. Pure bf16 routing diverged in Switch’s experiments, and router logits that grow unchecked will saturate the softmax. Use float32 for the gate, add a z-loss, and consider a smaller initialisation for the router weights.

Fine-grained experts move cost into communication. More, smaller experts with larger k means more messages per token. DeepSeek-V3 needs node-limited routing precisely because top-8 over 256 experts would otherwise fan out to nearly every node. The cheaper-per-FLOP design is only cheaper if the network keeps up.

Finally, claims about balancing methods are often scale-dependent. Loss-Free Balancing was published on models up to 3B parameters; its adoption at 671B in DeepSeek-V3 is strong evidence, but there is no public, controlled ablation at that scale that I could verify. Do not assume a gamma of 0.001 transfers to your expert count and batch size.

Practical Recommendations

If you are training an MoE from scratch in 2026 and designing its mixture of experts routing, start from the mechanism that has the most evidence at your scale. Below a few billion parameters or with a handful of experts, a Switch-style balance loss with alpha around 0.01 and a z-loss of 0.001 is well understood, and the ST-MoE and Switch ablations tell you what to expect. With many fine-grained experts, try a bias controller with a small complementary sequence-level loss, following the DeepSeek-V3 recipe, and sweep gamma on a short proxy run.

Prefer dropless routing if your stack supports it. If you must use capacity, begin at a factor of 1.25 and watch the dropped fraction; the ST-MoE step-time numbers show the cost of padding up to 2.0 is real.

For serving, plan for imbalance rather than assuming it away. Measure per-expert load on production traffic, not on the training mix, and budget redundant experts for the hottest ones. Our post on parallelism strategies for MoE inference covers how to place them.

A short checklist:

Run the router in float32 and add a z-loss around 1e-3.
Log max-to-mean load, dropped fraction, router entropy and maximum logit per layer.
Compute balance over the global batch where possible, not each micro-batch.
Keep any auxiliary balance coefficient small when using a bias controller.
Freeze or decay the bias update late in training once load is stable.
Re-measure expert load on real traffic before sizing expert-parallel groups.

Frequently Asked Questions

What is mixture of experts routing?

Mixture of experts routing is the mechanism that decides which expert sub-networks process each token. A small learned router scores every expert for a token, selects the top-k, converts the scores to gate weights, and sums the selected experts’ outputs weighted by those gates. Only the chosen experts run, so total parameters can grow far faster than per-token compute. The router is trained jointly with the rest of the network.

What is top-k gating?

Top-k gating keeps only the k largest router scores for a token and discards the rest. In Shazeer’s formulation the kept scores pass through a softmax and the others are set to negative infinity, giving them zero weight. Switch uses k = 1, GShard and Mixtral use k = 2, and DeepSeek-V3 uses k = 8 routed experts. Larger k blends more experts per token but raises compute and all-to-all traffic proportionally.

What is expert capacity and the capacity factor?

Expert capacity is the maximum number of tokens an expert can process per batch. Switch Transformer defines it as tokens per batch divided by the number of experts, multiplied by a capacity factor. A factor of 1.0 allocates exactly the average load; values like 1.25 add headroom. Tokens beyond capacity are dropped from that expert and carried forward by the residual connection, so capacity trades compute waste against skipped computation.

How does auxiliary-loss-free load balancing work?

Each expert gets a bias that is added to its routing score only when choosing the top-k. After each step, the bias of overloaded experts decreases and that of underloaded experts increases, by a small step size. The gate weights still come from the unbiased scores, and the bias receives no gradient, so no interfering loss term is added. DeepSeek-V3 used a step of 0.001 plus a tiny sequence-wise loss of 0.0001.

Why do MoE models suffer expert collapse?

Routing has a positive feedback loop: experts that receive more tokens train more, become better, and earn higher scores, so they receive still more tokens. Without a counter-pressure the router concentrates traffic on a few experts and the rest stay undertrained. Auxiliary balance losses, bias controllers, noisy gating and expert-choice routing are all different ways of breaking that loop.

What is the router z-loss?

The router z-loss, introduced in ST-MoE, penalises the squared log-sum-exp of the router logits, discouraging them from growing large. Large logits make the softmax saturate and amplify rounding error in bfloat16, which can destabilise training. ST-MoE used a coefficient of 0.001 and reported stable runs in 3 of 3 trials, against 4 of 6 for the baseline, without a quality loss.

Mixture-of-Experts Routing: Top-k Gating, Load Balancing and Aux-Loss-Free Balancing (2026)