Load-balancing losses: write the denominator before tuning the coefficient

For a nonempty top-one routing batch, define f_e as the fraction of tokens whose requested winning expert is e, before capacity dropping. Define P_e as the mean full-softmax probability for that expert over all eligible tokens. Both vectors sum to one under these assumptions, but they need not be equal. The stated auxiliary loss multiplies their dot product by expert count and a coefficient. A balanced reference point with both vectors uniform gives coefficient alpha, independent of expert count.

During ordinary backpropagation through this hard-routing construction, counts are observations rather than a smooth argmax derivative. Gradients flow through the soft probabilities. A busy expert's larger observed fraction increases the pressure attached to its probability mass, but this is a surrogate optimization signal rather than a proof of equal future workloads. Do not claim every possible pair of nonuniform hard and soft statistics has a loss above the uniform reference or that minimizing one minibatch guarantees an even next batch.

With k requested experts per token, raw counts divided by T sum to k. Dividing by Tk instead makes assignment fractions sum to one. Either can appear in an implementation if the remaining scale compensates, so record the factor explicitly. DeepSpeed 0.19.7's general top-k path includes a division by k in its count/probability auxiliary expression, while its top-two path uses first-choice counts in the corresponding loss. A label such as top-two balancing therefore does not uniquely identify which assignments drive its gradient.

Averaging losses computed separately on routing groups is generally different from multiplying globally averaged f and P. The loss contains a product, and correlations between group demand and soft mass do not disappear by renaming the reduction global. Sequence-wise balancing can intentionally pressure every sequence rather than only an aggregate batch. Excluding padding tokens, counting only accepted assignments, or averaging across experts instead of summing also changes scale or meaning. Inspect masks and reduction groups alongside the equation, especially when data or expert parallelism changes.

A low auxiliary value alone does not establish useful specialization, accepted workload balance, or good model quality. Track hard requested counts, soft mass, accepted counts, overflow, and per-rank work with clearly named denominators. Sweep the coefficient while keeping main-loss scaling and token budget fixed. Compare quality and resource use, because a router can improve count symmetry while harming the task or moving work onto slower communication paths. Report whether the logged value includes its coefficient and how layer losses are combined.

The function lab illustrates a concentration statistic, not this router loss. For a two-way fraction vector [(1+x)/2,(1−x)/2] with x between zero and one, squared concentration is (1+x²)/2. Reset, set a=1 and x=0, predict y=0 and concentration one half, then Play or Step. Separately Reset, set a=1 and x=1, predict y=1 and concentration one, then Play or Step. The fixed graph y=ax² has no hard assignments, P vector, autograd, or distributed groups. Its curve cannot substitute for measuring the specified f–P product.