---
kind: concept
slug: expert-load-balancing
title: Expert load balancing
summary: >-
  Keeps the tokens sent to each expert in a sparse layer roughly equal, so no expert sits idle or overflows. One compiled variant assigns tokens with a balanced matching during training and no balancing loss or capacity limit; test-time routing is greedy. Results are author-reported. Only one card so far, so the comparison with loss-based balancing is not yet covered.
tags: [routing, sparse-experts, load-balancing, distributed-training]
papers: ["2103.16716"]
related: [mixture-of-experts]
status: draft
updated_at: 2026-11-09
---

## Overview

In a mixture-of-experts layer, a router sends each token to one or a few
expert networks. Without a constraint, the router tends to favour a few experts,
so those become the bottleneck while the rest are under-used [arxiv:2103.16716 §2.2].
The usual fixes are an auxiliary loss that penalises uneven routing, and a
capacity limit that drops tokens once an expert is full [arxiv:2103.16716 §2.2].
Both add hyperparameters, and the capacity limit discards tokens. The shared
goal is an assignment in which every expert receives about the same share of
tokens, without that cost.

## Variants

### Balanced assignment with no balancing loss (BASE layers)

Each expert has a learned centroid. During training, the tokens of a batch are
assigned to experts by solving a linear assignment problem in which every expert
must receive the same number of tokens (T/E for T tokens and E experts) and the
total affinity is maximised [arxiv:2103.16716 eq. 2] [arxiv:2103.16716 §3.2.1].
The paper solves this with an auction algorithm, which it prefers to the
Hungarian method because it is easier to run in parallel on GPUs [arxiv:2103.16716 §3.2.1].
Tokens are first shuffled across workers so that the correlation within a
sequence does not steer the assignment [arxiv:2103.16716 §3.2.1]. The expert
output is added back to the token through a sigmoid gate on the affinity
[arxiv:2103.16716 eq. 1] [arxiv:2103.16716 §3.1]. Card: [BASE Layers: Simplifying Training of Large, Sparse Models](/papers/2103.16716).

At inference the assignment is not used, because balancing would depend on
future tokens. Each token goes to its single highest-scoring expert
[arxiv:2103.16716 §3.2.2]. The authors report that the trained routing is still
reasonably balanced [arxiv:2103.16716 §5.1].

### Loss-based balancing and capacity limits (background only)

The auxiliary balancing loss and the per-expert capacity factor used by sparse
MoE baselines are described in the BASE paper's background and experiment
setup [arxiv:2103.16716 §2.2] [arxiv:2103.16716 §4.3]. The sources of those
methods have no card here, so this page does not describe them in detail. They
can be read on arXiv: [Sparsely Gated Mixture-of-Experts](https://arxiv.org/abs/1701.06538)
and [Switch Transformer](https://arxiv.org/abs/2101.03961).

## When to use which

| Situation | Use | Why |
|---|---|---|
| Training from scratch with many experts, and a constraint on hyperparameters or tuning time | Balanced assignment (BASE) | The paper states that it adds no balancing loss and no capacity factor [arxiv:2103.16716 §1] |
| Training on a cluster where one expert per device is the design | Balanced assignment (BASE) | The layer is built with one expert per worker and needs an all-to-all exchange per routing step [arxiv:2103.16716 §3.2.1] |
| Inference-time load must be even, or tokens must never be dropped under test-time routing | Not decided by this card | The test-time routing is greedy and the paper reports no inference-time balance guarantee [arxiv:2103.16716 §3.2.2] [arxiv:2103.16716 §5.1] |
| Communication cost is the bottleneck | Check before choosing | Each balanced layer adds four all-to-all steps during training [arxiv:2103.16716 §2.2] |

The paper reports that its layer matches data-parallel training on 8 GPUs and
gains at larger budgets, and that it is comparable to Sparsely Gated MoE and
better than Switch on validation perplexity [arxiv:2103.16716 §4.2] [arxiv:2103.16716 §4.3].
These are author-reported comparisons of curves, and the text gives no
perplexity values. Throughput is reported in one table: 545k tokens per second
for BASE against 600k for data parallel, 292k for MoE and 469k for Switch on
128 GPUs [arxiv:2103.16716 tab. 2].

## Implementation sketch

The core is one step that can be solved in any way, as long as each expert gets
exactly T/E tokens. The paper's solver is an auction; the sketch leaves it as a
function.

```pseudocode
# Balanced routing for one training batch  [arxiv:2103.16716 §3.2.1]
# Input: tokens (T, D), centroids (E, D). Each expert gets T/E tokens.

def route_train(tokens, centroids):
    # Shuffle tokens across workers so that sequence order does not
    # decide the assignment  [arxiv:2103.16716 §3.2.1]
    order = random_permutation(T)
    tokens = tokens[order]

    # Score every token against every expert centroid  [arxiv:2103.16716 eq. 2]
    scores = tokens @ centroids.T                # (T, E)

    # Solve with equal capacity per expert; the paper uses an auction
    # solver  [arxiv:2103.16716 §3.2.1]
    expert_of = balanced_assignment(scores)      # (T,)
    return expert_of, order

def route_test(tokens, centroids):
    # No balancing at test time: each token takes its best expert
    # [arxiv:2103.16716 §3.2.2]
    return argmax(tokens @ centroids.T, dim=1)   # (T,)
```

Where variants differ: the training assignment, the blending of the expert
output, and whether the shuffle is applied. The paper's exact equations, and the
points where the released code differs, are on the card.

## Contradictions

The paper and its released code differ on several details. Both sides are
listed below; they are not resolved here. The full comparison is on the card
[Paper vs code](/papers/2103.16716).

- **Expert output mixing.** The paper writes the expert output as a gated term
  added to the token through a residual connection [arxiv:2103.16716 eq. 1].
  The released code blends the expert output with the input, with no identity
  term [arxiv:2103.16716 eq. 1] [Paper vs code](/papers/2103.16716).
- **Solver termination.** The paper's pseudocode says the solver switches to a
  greedy rule after a fixed number of iterations [arxiv:2103.16716 fig. 9]. The
  released solver keeps iterating until every token is assigned
  [Paper vs code](/papers/2103.16716).
- **Balancing claims.** The paper says the method adds no new hyperparameters
  [arxiv:2103.16716 §1]. The implementation exposes settings for the layer
  count, sublayer count and shuffle, and the solver has its own epsilon and
  iteration cap [Paper vs code](/papers/2103.16716).

## Papers

- [BASE Layers: Simplifying Training of Large, Sparse Models](/papers/2103.16716): balanced assignment of tokens to one expert per worker, with no balancing loss or capacity factor.