What is a load balancing loss in mixture of experts?

In this article

1. Before you start
2. Why interviewers ask this
3. Load balancing loss from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

Mixture-of-experts layers are trained as part of how large language models are trained, and each one sits inside the transformer.

Why interviewers ask this

A mixture-of-experts model that is left unbalanced wastes most of its experts, so a question on balancing shows whether a candidate understands how sparse models are trained rather than only how they are drawn. The question, in our words: why does a mixture-of-experts model need a load balancing loss, and what does its formula do?

A weak answer says "it spreads tokens evenly." A strong answer writes the formula, says which factor carries the gradient and why, explains the factor N and the weight alpha, and names the cost of setting alpha too high. A Staff-level answer adds that some recent models balance load without a large loss at all, using a bias on the routing scores, and says when each approach is preferable.

Load balancing loss from first principles

Mixture-of-experts layer

A mixture-of-experts layer is a block that replaces one feed-forward block of a transformer with several parallel feed-forward blocks, called experts, plus a router, so each token runs through only a few experts and most of the layer's parameters stay unused for that token.

Load balancing loss

A load balancing loss is an auxiliary term added to the main training loss of a mixture-of-experts model that penalises uneven routing, so that each expert receives about the same share of tokens in every batch, and it is smallest when routing is uniform.

Routing collapse

Routing collapse is a failure in which the router sends almost all tokens to a few experts, so the rest receive little training and stay weak, which makes the router favour the strong few even more.

In plain words

Picture a restaurant with four chefs and one host who seats the guests. The host learns that chef one cooks fastest and starts seating everyone there. The other chefs never practise, so they stay slow, and the queue at chef one grows. A manager fixes this by charging the host a fine whenever one chef's queue is long. The analogy breaks in one place. The fine acts on the host's habits, meaning the router's probabilities, and does not move any guest directly.

In one precise sentence, the load balancing loss multiplies each expert's share of tokens by the router's average probability for that expert, sums the products, and adds the result to the training loss.

Why does a mixture-of-experts model need a load balancing loss?

A mixture-of-experts model needs a load balancing loss because an unbalanced load leads to routing collapse or extra computation cost, as Wang et al. (2024) put it. There are three concrete reasons.

First, collapse: experts that receive few tokens learn little. Second, hardware: when experts are spread over several machines, unbalanced load lowers efficiency under expert parallelism, as the DeepSeek-V3 report notes; the data parallelism versus model parallelism page covers that split. Third, capacity: expert capacity is the fixed number of tokens one expert may take per batch, equal to tokens per batch divided by the number of experts, times a capacity factor (Fedus et al., 2022, Eq. 3). Tokens beyond it are skipped by that expert and passed to the next layer through the residual connection.

Step 1: The router counts how many tokens each expert receives

The first quantity is f, the fraction of tokens in a batch whose top router choice is each expert. Input: the router's softmax probabilities for the T tokens in a batch. Output: one fraction per expert.

For expert i, the count is n_i, the number of tokens whose highest probability lands on expert i, and the fraction is:

Switch Transformer routes each token to a single expert, so this count is simple (Fedus et al., 2022). Mixtral instead selects two experts out of eight for every token at each layer (Jiang et al., 2024); this step follows the single-expert Switch definition, and the trade-offs section shows the top-k version.

The count f comes from a hard argmax and carries no gradient. A hard pick is a jump: changing a router weight slightly either leaves the pick unchanged or flips it, so the count has no smooth slope to follow. You now know what is measured; the next step explains how to measure it in a way training can use.

Step 2: Only the router probability P carries the gradient

The second quantity is P, the router's average probability for each expert across the batch, and it is the part that training can adjust. Input: the same softmax probabilities. Output: one average per expert, written as:

Here p_i(x) is the router's probability for expert i on token x.

The mean router probability P is differentiable. Fedus et al. (2022) state it plainly: "the P-vector is differentiable, but the f-vector is not." Skywork-MoE gives the reason a surrogate is needed: a hard statistic does not allow back-propagation, so a loss built from it alone cannot train the router.

Multiplying the two gives Σ f_i · P_i. The product is large only when an expert is both busy (f large) and favoured (P large). Under uniform routing every f_i and P_i equals 1/N, so the sum equals 1/N (Fedus et al., 2022). That is the minimum, which is why the loss rewards balance.

Step 3: Scaling by N and alpha sets the strength

The full Switch loss is the sum scaled by two constants: loss = α · N · Σ f_i · P_i, where N is the number of experts and α is the weight of the penalty (Fedus et al., 2022, Eq. 4).

The factor N keeps the loss constant as the number of experts varies: under uniform routing the sum is 1/N, so the loss equals α for any N. Fedus et al. (2022), §2.2 used α = 0.01, which they describe as "sufficiently large to ensure load balancing while small enough to not overwhelm the primary cross-entropy objective". They swept α from 0.1 to 0.00001 in powers of ten and found 0.01 balanced load quickly without interfering with the training loss.

The scaled product is added to the cross-entropy loss: Switch adds this auxiliary term to the model's total loss for every Switch layer during training. Input: f and P from the earlier steps. Output: one extra number per layer.

Step 4: The gradient lowers the probability of busy experts

Training treats the counts f as constants and differentiates only through P, so the pressure on each expert's router probability is proportional to how busy that expert already is. Input: the scaled loss. Output: a gradient on each P_i.

Because f_i is a constant during the backward pass, the derivative of the loss with respect to P_i is α · N · f_i. A busy expert has a large f_i, so its router probability is pushed down hardest. Because the probabilities come from a softmax and sum to one, the probability removed from busy experts flows to idle ones, and the next batch sends them more tokens.

In plain terms: the router is not told where to send tokens. It is made slightly less fond of whichever expert is currently overloaded, and its cross-entropy training still decides which expert suits each token.

1 · Count tokens2 · Average probability3 · Scale the product4 · Add to the lossReplay

Step-by-step description

1. The router's top choices give the count f per expert.
2. The mean probability P is the part that has a gradient.
3. Alpha times N times the sum of f times P over experts.
4. The result joins the cross-entropy loss.

• The count f comes from a hard argmax and carries no gradient.
• The mean router probability P is differentiable.
• The scaled product is added to the cross-entropy loss.

Worked example

Take four experts, eight tokens in a batch and α = 0.01; every number here is illustrative and computed in a script, not taken from a paper.¹ In the balanced case each expert gets two tokens, so f = (0.25, 0.25, 0.25, 0.25), and the router's average probabilities are also 0.25 each. Then Σ f·P = 4 × 0.0625 = 0.25, N × 0.25 = 1.0, and the loss equals α, which is 0.01.¹

In the skewed case five tokens go to expert one, two to expert two, one to expert three and none to expert four, so f = (0.625, 0.25, 0.125, 0) with P = (0.55, 0.25, 0.15, 0.05). Then Σ f·P = 0.425, N × 0.425 = 1.7, and the loss is 0.017.¹ In full collapse f = P = (1, 0, 0, 0), so the sum is 1, the scaled value is N = 4, and the loss is 0.04.¹

The gradient row shows Step 4 in numbers: in the skewed case expert one, with five of the eight tokens, receives the largest push downward on its probability, and expert four, with no tokens, receives none.¹ For capacity, 1,024 tokens over 8 experts is 128 per expert, and a factor of 1.25 raises that to 160; the 161st token routed to one expert is skipped and carried to the next layer by the residual connection.¹ Switch reports that dropped tokens were typically under 1% and did not depend on the number of experts (Fedus et al., 2022).

Text description of the diagram

• In the balanced case the loss equals α, which is 0.01.
• In full collapse the loss is 0.04.

Trade-offs and edge cases

The loss buys balance at a price, and the price has three parts: model quality, spare capacity, and the gradient interference that has led some recent systems to drop the loss as their main tool.

Balance against quality

A larger weight alpha buys more balance but costs model quality. Too large an auxiliary loss impairs model performance, according to the DeepSeek-V3 report, and Wang et al. (2024) say a large auxiliary loss introduces non-negligible interference gradients into training. The gradient from the balancing term competes with the gradient from the language-modelling loss, because both act on the same router. Switch settled on 0.01 after a sweep, which is a starting point rather than a constant of nature; other models and expert counts need their own sweep.

Capacity factor and dropped tokens

The capacity factor is a second safety valve, and it does not replace the loss. A factor above 1.0 creates a buffer for tokens that are not perfectly balanced, but high values waste computation and memory (Fedus et al., 2022). The loss reduces the chance of overflow by evening out the load; the capacity factor decides what happens when overflow occurs anyway.

What changes with top-k routing?

With top-K routing each token picks K experts, so the count f must be rescaled. The complementary sequence-wise balance loss of the DeepSeek-V3 report (§2.1.2) defines f_i as N / (K · T) times the number n_i of tokens that select expert i among their top K, where T is the number of tokens in the sequence. P_i is the mean normalised affinity of expert i over those tokens, and the loss keeps the same shape:

Under uniform routing f_i = 1 and P_i = 1/N, so the sum is 1 and the loss is α. The factor N/K plays the role that N plays in Switch: it fixes the uniform value however many experts exist and however many each token picks. The same gradient logic holds, because f_i is a hard count and only P_i carries a gradient.

Loss-free balancing

Loss-free balancing is a method that adds a per-expert bias to the routing scores before the top-K choice and adjusts that bias according to recent load, so no balancing gradient exists. Wang et al. (2024) report better performance and better load balance than auxiliary-loss control on models up to 3 billion parameters trained on up to 200 billion tokens. The DeepSeek-V3 report lowers the bias by γ for an overloaded expert and raises it by γ for an underloaded one at the end of each step. The bias affects expert selection only, and the gate value still comes from the original score.

With γ = 0.001 an overloaded expert's bias goes 0, −0.001, −0.002, −0.003 over three steps. In an illustrative pair of scores, 0.600 for expert A and 0.599 for expert B, with A overloaded and B underloaded, a single step is enough to make B win the selection, while the gate value for either expert stays its original score.¹

DeepSeek-V3 used γ = 0.001 for the first 14.3 trillion tokens and 0.0 for the remaining 500 billion, and still keeps a tiny sequence-wise balance loss with α = 0.0001 to avoid extreme imbalance within one sequence (DeepSeek-AI, 2024). That report also notes that a batch-wise auxiliary loss can match the loss-free method at a similar level of batch-wise balance, so the scope of the balancing matters as well as the mechanism.

Text description of the diagram

• Loss-free balancing adds a per-expert bias to routing scores before the top-K choice.
• The gradient from the balancing term competes with the gradient from the language-modelling loss.

Mixtral 8x7B is a top-2 mixture-of-experts model, a released example of the routing described above.

How to answer it in an interview

The spoken structure: name the failure, give the formula and which factor has the gradient, explain N and alpha, then mention loss-free balancing as the alternative.

Common mistakes

• "The loss makes every expert learn the same thing." (wrong) It only evens out how many tokens each expert receives, and it is a soft pressure weighed against cross-entropy by alpha. (right) "It equalises token counts, not expert behaviour."
• "We back-propagate through f." (wrong) f comes from a hard argmax and has no gradient. (right) "The gradient flows through the mean router probability P, with f acting as a constant weight."
• "Multiplying by N is arbitrary." (wrong) It has a purpose. (right) "N keeps the loss equal to alpha under uniform routing for any number of experts."
• "More balance is always better, so raise alpha." (wrong) A large auxiliary loss adds interference gradients and hurts quality. (right) "I sweep alpha and keep it small, or use a bias-based method."
• "Auxiliary-loss-free means no balancing loss at all." (wrong) DeepSeek-V3 still keeps a tiny sequence-wise balance loss. (right) "The main balancing comes from the bias update; a very small loss guards against extreme imbalance in one sequence."

Follow-up questions

How does mixture of experts load balancing work?

The router scores every expert for each token. A load balancing loss measures how unevenly tokens and probabilities fall across experts and adds that to the training loss, so the router learns to spread tokens. Some models instead adjust a bias on the scores according to recent load.

Why is f not differentiable, and does that break training?

The count f depends on an argmax, a hard choice with no smooth slope. Training still works because the gradient flows through P, the mean router probability, with f acting as a constant weight on each expert.

How do you choose the coefficient alpha?

Sweep it. Switch Transformer (Fedus et al., 2022) tried values from 0.1 to 0.00001 in powers of ten and kept 0.01, because it balanced load quickly without interfering with the training loss. Watch both the load per expert and the main loss, since a larger alpha trades quality for balance.

What is loss-free balancing, and why did DeepSeek-V3 adopt it?

It adds a per-expert bias to routing scores before the top-K choice and nudges the bias down for overloaded experts and up for underloaded ones. DeepSeek-V3 adopted it as its main balancer because a large auxiliary loss can impair model performance, while the bias produces no balancing gradient.

How is f defined when each token picks more than one expert?

The count is rescaled by N / (K · T), where K is the number of experts each token selects, so uniform routing again gives f_i = 1 and a loss of α; the section on top-k routing above gives the DeepSeek-V3 form.

What happens to a token when its expert is full?

The expert skips it, and the token's representation passes straight to the next layer through the residual connection. That is a dropped token; Switch reports typically under 1%. Balancing happens during pretraining; adapting the finished model afterwards is covered in the large language model (LLM) fine-tuning guide, and small language models cover a different route to cheap inference, a smaller dense network.

Sources

1. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (arXiv 2101.03961) — arXiv (Fedus, Zoph, Shazeer), 16 Jun 2022
2. Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts (arXiv 2408.15664) — arXiv (Wang et al.), 28 Aug 2024
3. DeepSeek-V3 Technical Report (arXiv 2412.19437) — arXiv (DeepSeek-AI), 27 Dec 2024
4. Mixtral of Experts (arXiv 2401.04088) — arXiv (Jiang et al.), 8 Jan 2024
5. Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models (arXiv 2406.06563) — arXiv (Wei et al.), 3 Jun 2024