Expert load balancing

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 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 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 2103.16716 eq. 2 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 2103.16716 §3.2.1. Tokens are first shuffled across workers so that the correlation within a sequence does not steer the assignment 2103.16716 §3.2.1. The expert output is added back to the token through a sigmoid gate on the affinity 2103.16716 eq. 1 2103.16716 §3.1. Card: BASE Layers: Simplifying Training of Large, Sparse Models.

At inference the assignment is not used, because balancing would depend on future tokens. Each token goes to its single highest-scoring expert 2103.16716 §3.2.2. The authors report that the trained routing is still reasonably balanced 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 2103.16716 §2.2 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 and Switch Transformer.

When to use which

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 2103.16716 §4.2 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 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.

# 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.

• Expert output mixing. The paper writes the expert output as a gated term added to the token through a residual connection 2103.16716 eq. 1. The released code blends the expert output with the input, with no identity term 2103.16716 eq. 1 Paper vs code.
• Solver termination. The paper's pseudocode says the solver switches to a greedy rule after a fixed number of iterations 2103.16716 fig. 9. The released solver keeps iterating until every token is assigned Paper vs code.
• Balancing claims. The paper says the method adds no new hyperparameters 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

• BASE Layers: Simplifying Training of Large, Sparse Models: balanced assignment of tokens to one expert per worker, with no balancing loss or capacity factor.