Load Balancing Across Experts

"They keep sending everyone to me. I am flattered, I am overloaded, and seven of my colleagues have not seen a single token all morning."

A Popular Expert Holding Up the Whole Batch

The previous section moved tokens to their chosen experts with an all-to-all exchange, taking the routing decision as given. Now we confront the decision itself. In a dense layer every device does an identical share of the work, so balance is automatic; nobody chooses who computes. A mixture of experts breaks that symmetry on purpose: the gate sends each token to a small number of experts, and the experts are scattered across devices for the memory reasons that motivated expert parallelism in the first place. The gate is now, whether or not its designers think of it this way, a load balancer for the cluster. When it does that job badly, the consequences land on the hardware immediately and on model quality more slowly, and the two failures reinforce each other.

The trouble is that nothing in the basic gating objective asks for balance. The gate is trained to route each token to whichever experts most reduce the loss, and real data is not uniform: some patterns are common, some are rare, and a few experts quickly become specialists for the common patterns. Left alone, the router discovers that sending most tokens to a handful of experts is locally optimal, a self-reinforcing collapse sometimes called routing collapse. Those experts improve because they see the most data, which makes the gate prefer them even more, while the starved experts stay weak and the gate learns to avoid them further. The result is a model that has paid for many experts but trained only a few.

1. Why Imbalance Is a Distributed-Systems Problem Beginner

Expert parallelism places a disjoint set of experts on each device. When a batch arrives, the gate's choices induce a load on each device equal to the number of tokens routed to the experts it holds. The forward pass through the expert layer cannot finish until every device has processed its share, so the wall-clock cost of the layer is set by the busiest device, not the average one. This is exactly the straggler effect introduced in Chapter 2: a synchronous step proceeds at the speed of its slowest participant, and one overloaded device makes everyone else wait at the next all-to-all barrier. Adding more experts and more devices does not help if the gate keeps funnelling tokens to the same few; you have bought parallel hardware and then serialized it through a hot spot.

The mechanism that produces the hot spot is also familiar. A gate that concentrates tokens on popular experts is producing a skewed partition of the batch, the same pathology that wrecks a Spark shuffle when one key dominates, studied in Chapter 7. There, one popular join key sends a disproportionate share of rows to a single reducer; here, one popular expert receives a disproportionate share of tokens on a single device. The cause differs (a learned router rather than a data distribution) but the cure rhymes: detect the skew, and reshape the assignment so no single partition carries far more than its fair share.

2. The Auxiliary Load-Balancing Loss Intermediate

The classic fix, introduced with the sparsely gated mixture of experts and carried into Switch Transformer, is to add a term to the training objective that penalizes imbalance directly. For a batch of $T$ tokens and $E$ experts, define $f_e$ as the fraction of tokens dispatched to expert $e$ under the (hard, top-$k$) routing, and $P_e$ as the mean gate probability assigned to expert $e$ across the batch,

where $g_e(x_t)$ is the softmax gate probability of expert $e$ for token $t$. The auxiliary loss is the scaled dot product of these two vectors,

Two facts make this the right object. First, the dot product $\sum_e f_e P_e$ is minimized, subject to each vector summing to one, when both are uniform at $1/E$, so driving it down pushes usage toward the even split we want; the factor $E$ keeps the target value near a constant as $E$ changes, and $\alpha$ (a small weight such as $10^{-2}$) sets how hard balance competes with the task loss. Second, the gradient is well behaved: the hard counts $f_e$ are not differentiable, but $P_e$ is, so the gradient flows through the gate probabilities and gently lowers the logits of experts that are already popular ($f_e$ large) while raising the rest. The term nudges the router toward fairness without dictating any single token's destination, which is what lets the experts still specialize. Section 17.7 picks up the companion control, capacity factors and token dropping, that bounds the damage when balance is imperfect.

The code below trains a tiny gate over synthetic tokens that genuinely cluster toward a few popular experts, first with no balancing term and then with the auxiliary loss switched on. It reports per-expert usage, the load on the busiest of four devices (two experts each), and the max-to-mean ratio that quantifies the skew.

import numpy as np

rng = np.random.default_rng(7)
N, d, E = 6000, 16, 8          # tokens, feature dim, experts

# Tokens that genuinely cluster toward experts 0 and 1, so an unregularized
# gate piles onto them: a learned, self-reinforcing popularity skew.
pop = rng.standard_normal((E, d)); pop[0] *= 3.0; pop[1] *= 2.2
assign = rng.choice(E, size=N, p=[.40, .22, .10, .08, .07, .06, .04, .03])
X = pop[assign] + 0.6 * rng.standard_normal((N, d))

def softmax(z):
    z = z - z.max(axis=1, keepdims=True); e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)

def usage_fraction(W):                      # top-1 routing, fraction per expert
    return np.bincount(np.argmax(softmax(X @ W), axis=1), minlength=E) / N

def train(aux_weight, steps=800, lr=0.5):
    W = 0.01 * np.random.default_rng(0).standard_normal((d, E))   # fixed init
    for _ in range(steps):
        g = softmax(X @ W)
        # Task pull: a rich-get-richer signal that nudges each token toward the
        # expert already most confident for it, so popular experts grow more
        # popular and an unregularized gate tends to collapse onto a few of them.
        hard = np.zeros((N, E)); hard[np.arange(N), np.argmax(g, axis=1)] = 1.0
        grad_task = X.T @ (g - hard) / N
        # Auxiliary balance loss aux = E * ||P||^2, the smooth surrogate for
        # E * sum_e f_e P_e, minimized at uniform usage P_e = 1/E. Its gradient
        # flows through the softmax mean P_e = mean_t g and pushes mass off hot experts.
        P = g.mean(axis=0)
        v = 2.0 * E * P
        grad_aux = X.T @ (g * (v - (g @ v)[:, None])) / N
        W -= lr * (grad_task + aux_weight * grad_aux)
    return W

for label, aw in [("no aux loss ", 0.0), ("with aux loss", 1.0)]:
    u = usage_fraction(train(aw))
    dev = u.reshape(4, 2).sum(axis=1)                  # 4 devices, 2 experts each
    print(f"{label}: usage% = [{', '.join(f'{100*x:4.1f}' for x in u)}]")
    print(f"               busiest device = {100*dev.max():4.1f}%  (ideal 25.0%),"
          f"  max/mean expert = {u.max()/u.mean():.2f}x")

no aux loss : usage% = [ 0.0, 55.4,  0.0,  0.0, 31.1, 13.5,  0.0,  0.0]
               busiest device = 55.4%  (ideal 25.0%),  max/mean expert = 4.43x
with aux loss: usage% = [13.1, 14.1, 13.5, 16.5, 10.3,  7.5, 10.4, 14.6]
               busiest device = 30.0%  (ideal 25.0%),  max/mean expert = 1.32x

The improvement is decisive here, and the honest lesson of the demo is why: left alone, the rich-get-richer pull lets a few experts swallow the batch and starves the rest into dead weight, exactly the routing collapse this section warns about. The auxiliary loss counters that feedback and brings the idle experts back into play. On real data the gate also carries genuine specialization pressure, so a practical $\alpha$ trades a little of that specialization for balance rather than flattening usage all the way; pushing $\alpha$ higher buys more balance at a rising cost to model quality, which is the central tension of the whole topic.

3. Balanced by Construction, and Balanced Without a Loss Advanced

The auxiliary loss treats imbalance as something to penalize after the fact. Two other strategies attack it more directly. The first is to change who does the choosing. In the token-choice routing of Section 17.3, each token picks its top-$k$ experts, and nothing stops a popular expert from being everyone's pick. Expert-choice routing inverts the selection: each expert picks the top-$T k / E$ tokens it wants from the batch. Because every expert selects exactly the same number of tokens, the load is uniform by construction, with no auxiliary loss required and no device able to become a hot spot. The cost is that a token may be chosen by several experts or by none, so expert-choice trades guaranteed device balance for variable per-token capacity, a different knob on the same trade-off that Section 17.7 formalizes.

The second strategy keeps token choice but removes the auxiliary loss entirely. The auxiliary loss has a known side effect: its gradient perturbs the task objective, and an $\alpha$ large enough to enforce balance can measurably hurt quality, an interference the literature calls the balance-versus-specialization tension. The auxiliary-loss-free approach replaces the loss term with a per-expert bias $b_e$ added to the routing logits only for the top-$k$ selection. After each step the biases are nudged by a simple control rule: lower $b_e$ for experts that were overloaded, raise it for experts that were starved,

where $\bar f = 1/E$ is the target fraction and $\gamma$ is a small update rate. The bias steers the routing toward balance, but because it is added only for selection and not to the gate value used in the weighted combination, it never appears in the task gradient, so it balances load without distorting what the experts learn. This is the mechanism DeepSeek-V3 uses at scale, and it is the current state of the art for keeping a large mixture of experts balanced while letting the experts specialize freely.

import torch, torch.nn.functional as F

def load_balance_loss(gate_logits, expert_index, num_experts):
    # gate_logits: (T, E) router scores; expert_index: (T,) the top-1 pick per token.
    P = F.softmax(gate_logits, dim=-1).mean(dim=0)                 # mean prob per expert
    f = torch.bincount(expert_index, minlength=num_experts).float() / gate_logits.size(0)
    return num_experts * torch.dot(f, P)                          # alpha applied by caller

# In the training step the layer hands this back; you just add it to the task loss:
# loss = task_loss + alpha * aux_loss          # alpha ~ 1e-2 in Switch Transformer

4. The Balance-Versus-Specialization Tension Advanced

Perfect balance and perfect specialization pull in opposite directions, and a good mixture of experts lives at a deliberate compromise between them. Push balance to its limit and you would route tokens to experts almost at random, guaranteeing a flat load but destroying the very specialization that makes a sparse model worth more than a smaller dense one; the experts would all learn the same average function. Push specialization to its limit and the gate collapses onto a few experts, which trains well but utilizes the cluster terribly and starves most of the model. The whole engineering art is to enforce just enough balance that no device becomes a straggler, while leaving the gate free to send genuinely different tokens to genuinely different experts.

This is why the field has drifted from heavy auxiliary losses toward lighter-touch controllers. The auxiliary loss enforces balance by adding a force to the loss landscape that every token feels, which inevitably bends the task objective. The bias-update approach enforces balance by adjusting only the selection threshold, leaving the gradient that shapes the experts untouched, so it buys balance at a lower price in specialization. Either way, balance is never free, and the right operating point depends on how skewed your data is and how tight your hardware budget is. The performance models of Chapter 3 give you the language to put a number on the throughput cost of a given imbalance, and that number is what tells you how hard to push balance for a given cluster.

5. Deriving the Auxiliary Loss Formula Intermediate

The auxiliary loss formula introduced in Section 2 warrants a closer derivation to make its minimum explicit and its gradient tractable. For a batch of $T$ tokens routed to $E$ experts, recall that $f_i$ is the fraction of tokens dispatched to expert $i$ under the hard top-$k$ assignment, and $p_i$ is the router's average soft probability for expert $i$ across the batch. Their definitions side by side are

where $g_i(x_t)$ is the softmax gate probability for expert $i$ given token $x_t$. The auxiliary balance loss is the scaled inner product of these two vectors,

with $\alpha$ typically set to $0.01$. To see why minimizing $\mathcal{L}_{\text{aux}}$ drives usage toward the uniform split, apply the Cauchy-Schwarz inequality: $\sum_i f_i p_i \geq \bigl(\sum_i \sqrt{f_i p_i}\bigr)^2 / E$, with equality when $f_i / p_i$ is the same constant for all $i$. When both $f$ and $p$ are valid probability vectors summing to one, the minimum of $\sum_i f_i p_i$ subject to $\sum_i f_i = \sum_i p_i = 1$ and $f_i, p_i \geq 0$ is $1/E$, achieved when $f_i = p_i = 1/E$ for all $i$. Minimising the loss therefore directly incentivises every expert to receive exactly $1/E$ of the batch, which is the uniform balance we want. The factor $\alpha$ controls how strongly that incentive competes with the task loss.

The gradient calculation is what makes the formula practical. The counts $f_i$ are not differentiable through the argmax routing decision, but $p_i$ are, because they come from the softmax of the gate logits. The gradient of $\mathcal{L}_{\text{aux}}$ with respect to the gate logits therefore flows through $p_i$ only, treating $f_i$ as constants computed from the current routing. This asymmetry means the loss penalizes experts proportionally to how often they are already chosen ($f_i$ large), pushing the gate to reduce its soft probability for popular experts and raise it for underused ones, a corrective signal that fights the rich-get-richer pull of the task gradient. Code 17.6.3 below makes this concrete by computing $\mathcal{L}_{\text{aux}}$ over a batch of routing probabilities and plotting how the loss changes as the distribution of tokens across experts shifts from highly concentrated to nearly uniform.

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
E = 8          # number of experts
alpha = 0.01   # balance coefficient

def softmax(z):
    z = z - z.max(); e = np.exp(z); return e / e.sum()

def l_aux(f, p):
    """Auxiliary load-balancing loss: alpha * sum_i f_i * p_i."""
    return alpha * float(np.dot(f, p))

# Sweep a concentration parameter: low kappa = near-uniform; high kappa = peaked.
kappas = np.linspace(0.1, 5.0, 60)
losses = []

for kappa in kappas:
    # Soft logits: concentrate probability on expert 0 as kappa rises.
    logits = np.zeros(E); logits[0] = kappa
    p = softmax(logits)

    # Hard routing: top-1 picks, so f_0 rises with kappa.
    T = 2000
    assignments = rng.choice(E, size=T, p=p)
    f = np.bincount(assignments, minlength=E) / T

    losses.append(l_aux(f, p))

fig, axes = plt.subplots(1, 2, figsize=(10, 4))

axes[0].plot(kappas, losses, color="#14385c", linewidth=2)
axes[0].axhline(alpha / E, color="#b8324a", linestyle="--", label=f"uniform minimum α/E = {alpha/E:.4f}")
axes[0].set_xlabel("Concentration κ (logit for expert 0)")
axes[0].set_ylabel("L_aux")
axes[0].set_title("Auxiliary loss vs. routing concentration")
axes[0].legend(fontsize=9)

# Right panel: show f and p vectors at three representative kappas.
for idx, (kappa, color) in enumerate([(0.5, "#5aab6d"), (2.0, "#e89b2c"), (4.5, "#b8324a")]):
    logits = np.zeros(E); logits[0] = kappa
    p = softmax(logits)
    assignments = rng.choice(E, size=T, p=p)
    f = np.bincount(assignments, minlength=E) / T
    x = np.arange(E) + idx * 0.25 - 0.25
    axes[1].bar(x, f, width=0.22, color=color, alpha=0.8, label=f"κ={kappa:.1f}")
axes[1].axhline(1 / E, color="black", linestyle="--", linewidth=1.2, label="fair share 1/E")
axes[1].set_xlabel("Expert index")
axes[1].set_ylabel("Token fraction f_i")
axes[1].set_title("Per-expert token fractions at three κ values")
axes[1].legend(fontsize=8)

plt.tight_layout()
plt.savefig("fig_17_6_aux_loss.png", dpi=130)
print(f"{'kappa':>8}  {'L_aux':>10}  {'f_0':>8}  {'p_0':>8}")
print("-" * 42)
for kappa in [0.5, 1.5, 3.0, 5.0]:
    logits = np.zeros(E); logits[0] = kappa
    p = softmax(logits)
    assignments = rng.choice(E, size=T, p=p)
    f = np.bincount(assignments, minlength=E) / T
    print(f"{kappa:8.1f}  {l_aux(f, p):10.6f}  {f[0]:8.3f}  {p[0]:8.3f}")
print(f"\nUniform minimum: alpha/E = {alpha/E:.6f}")

The table printed by the code confirms the theoretical minimum: at low concentration ($\kappa = 0.5$) the loss sits near $\alpha/E = 0.00125$, while at high concentration ($\kappa = 5.0$) it climbs by an order of magnitude. The gradient of $\mathcal{L}_{\text{aux}}$ with respect to the gate logits flows through $p_i$ and is proportional to $f_i$ at each expert, so the loss applies its strongest downward pull on the logit of whatever expert is already most popular, a direct corrective to the rich-get-richer dynamic.

6. Capacity Factor and Token Dropping Intermediate

From the distributed-systems perspective, every expert is a bounded queue: it can process at most a fixed number of tokens per batch before it becomes the bottleneck that stalls the all-to-all exchange and delays the whole cluster. The capacity factor is the parameter that sets that queue length, and token dropping is what happens when the queue overflows. Understanding both precisely is essential because token dropping is not a soft quality penalty; it is a hard silence where part of the model simply does not update, and at large scale it compounds into significant degradation if the router stays imbalanced.

For a batch of $T$ tokens spread over $E$ experts, the capacity of each expert is

where $c_{\text{factor}}$ is a hyperparameter, typically in the range $1.0$ to $1.25$. A capacity factor of $1.0$ means each expert can hold exactly its fair share of tokens; a value of $1.25$ adds a $25\%$ buffer above fair share. Any token routed to an expert whose queue is already full is dropped: it bypasses the expert sublayer and receives a zero update for that expert's contribution. Formally, if $\sigma(t, i)$ is the rank of token $t$ in the dispatch queue for expert $i$ (earlier tokens get lower rank), the effective contribution of expert $i$ to token $t$ is

The drop rate is the fraction of (token, expert) dispatch events that are silenced this way. When the router is balanced, $f_i \approx 1/E$ for all $i$, and a capacity factor of $1.0$ suffices with negligible dropping; as imbalance grows, the drop rate rises sharply on the overloaded experts. Code 17.6.4 below simulates this relationship, sweeping over combinations of capacity factor and load imbalance and reporting the resulting drop rate in a table.

import numpy as np

rng = np.random.default_rng(0)

def simulate_drop_rate(T, E, c_factor, imbalance_alpha):
    """
    Simulate token dispatch with a Dirichlet-imbalanced router.

    imbalance_alpha: Dirichlet concentration; low alpha → high imbalance.
    Returns the fraction of dispatched tokens that are dropped.
    """
    C = int(np.ceil((T / E) * c_factor))        # capacity per expert

    # Draw per-expert routing fractions from a Dirichlet distribution;
    # low concentration alpha produces a strongly peaked distribution.
    fracs = rng.dirichlet(np.full(E, imbalance_alpha))
    # Assign each token to one expert (top-1 routing approximation).
    assignments = rng.choice(E, size=T, p=fracs)

    # Count how many tokens land on each expert; anything above C is dropped.
    counts = np.bincount(assignments, minlength=E)
    dropped = np.maximum(counts - C, 0).sum()
    return dropped / T, fracs.max()              # drop rate, busiest expert fraction

T = 4096
E = 8

c_factors     = [0.90, 1.00, 1.10, 1.25, 1.50]
imb_alphas    = [10.0, 2.0, 0.8, 0.3]           # high → near-uniform; low → skewed
imb_labels    = ["low", "moderate", "high", "severe"]

print(f"{'c_factor':>10}  {'imbalance':>10}  {'drop_rate':>10}  {'busiest_f':>10}")
print("-" * 48)
for imb_alpha, imb_label in zip(imb_alphas, imb_labels):
    for c_factor in c_factors:
        drop, busiest = simulate_drop_rate(T, E, c_factor, imb_alpha)
        print(f"{c_factor:10.2f}  {imb_label:>10}  {drop:10.4f}  {busiest:10.4f}")
    print()

The simulation exposes two interacting controls. First, raising the capacity factor directly reduces the drop rate at every imbalance level, but it does so by increasing the maximum per-device compute, which widens the straggler gap and reduces throughput. Second, reducing imbalance (higher Dirichlet concentration, which corresponds to a better-trained auxiliary loss or a working bias-update controller) lowers the drop rate more cheaply than raising the capacity factor, because it removes the root cause rather than masking it. The practical operating point for most large mixture-of-experts models is $c_{\text{factor}} \in [1.0, 1.25]$ combined with an active balancing mechanism, so the buffer absorbs residual spikes without requiring a capacity factor large enough to dominate the step time.

7. Auxiliary-Loss-Free Routing: The 2025 Alternative Advanced

From the distributed-systems angle, the auxiliary loss is a blunt instrument: it modifies the gradient for every token in every step, coupling the routing objective to the task objective in a way that is hard to tune without side effects. A cleaner control-theoretic approach treats the routing decision as a scheduling problem and applies a lightweight feedback controller to the routing logits, leaving the task gradient untouched. This is the auxiliary-loss-free bias-update method, which DeepSeek-V3 demonstrated at 671B scale in late 2024.

The mechanism is straightforward. Maintain a per-expert bias vector $b \in \mathbb{R}^E$ initialized to zero. At each forward pass, add $b_i$ to the router logit for expert $i$ before computing the top-$k$ selection, so the effective selection score becomes $s_i(x_t) + b_i$. Crucially, $b$ is added only for the selection step; the gate values used in the weighted combination of expert outputs are the original unbiased softmax probabilities. After the batch is processed and the dispatch fractions $f_i$ are known, update the biases with a simple integral controller,

where $\bar{f} = 1/E$ is the target equal-share fraction and $\gamma$ is the update step size. When expert $i$ received too many tokens ($f_i > \bar{f}$), its bias is decremented, making it slightly less attractive in future batches; when it was underloaded ($f_i < \bar{f}$), its bias is incremented, nudging more tokens toward it. Because the bias never enters the softmax computation used in the output combination, the task gradient sees none of this correction, and expert specialization is entirely undisturbed. Code 17.6.5 below runs both the auxiliary-loss approach and the bias-update approach on the same synthetic routing scenario over 100 training steps, comparing a perplexity proxy and per-expert utilisation.

import numpy as np

rng = np.random.default_rng(7)
N, d, E = 4000, 16, 8
T_STEPS = 100
gamma = 0.05        # bias update step size
alpha_aux = 1.0     # auxiliary-loss weight for the aux-loss run

# Tokens with natural clustering toward experts 0 and 1.
pop = rng.standard_normal((E, d)); pop[0] *= 3.0; pop[1] *= 2.2
assign = rng.choice(E, size=N, p=[.38,.22,.12,.09,.07,.05,.04,.03])
X = pop[assign] + 0.6 * rng.standard_normal((N, d))

def softmax2d(z):
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z); return e / e.sum(axis=1, keepdims=True)

def task_grad(W, g):
    """Rich-get-richer task gradient (same as Code 17.6.1)."""
    hard = np.zeros_like(g)
    hard[np.arange(N), np.argmax(g, axis=1)] = 1.0
    return X.T @ (g - hard) / N

def aux_grad(W, g):
    """Gradient of the auxiliary balance loss through soft probabilities."""
    P = g.mean(axis=0)
    v = 2.0 * E * P
    return X.T @ (g * (v - (g @ v)[:, None])) / N

def perplexity_proxy(g):
    """Entropy of the average routing distribution: higher = more uniform."""
    P = g.mean(axis=0) + 1e-9
    return float(np.exp(-np.sum(P * np.log(P))))

lr = 0.4

results = {}
for mode in ("aux_loss", "bias_update"):
    W = 0.01 * np.random.default_rng(0).standard_normal((d, E))
    b = np.zeros(E)          # per-expert bias (used only in bias_update mode)
    history = []

    for step in range(T_STEPS):
        # Routing: add bias to logits before top-1 selection in bias_update mode.
        logits = X @ W
        if mode == "bias_update":
            sel_logits = logits + b[None, :]
        else:
            sel_logits = logits
        g_sel = softmax2d(sel_logits)
        g_task = softmax2d(logits)          # unbiased gate for task gradient

        # Dispatch fractions from the selection gate.
        picks = np.argmax(g_sel, axis=1)
        f = np.bincount(picks, minlength=E) / N

        # Gradient step.
        grad = task_grad(W, g_task)
        if mode == "aux_loss":
            grad = grad + alpha_aux * aux_grad(W, g_task)
        W -= lr * grad

        # Bias update (bias_update mode only).
        if mode == "bias_update":
            b += gamma * np.sign(1.0 / E - f)

        proxy = perplexity_proxy(g_task)
        history.append((f.copy(), proxy))

    results[mode] = history

# Report final-step statistics for both modes.
print(f"{'Mode':<15}  {'Perplexity proxy':>18}  {'Max/mean util':>14}  {'Busiest expert f':>17}")
print("-" * 70)
for mode, hist in results.items():
    f_final, ppl = hist[-1]
    print(f"{mode:<15}  {ppl:18.4f}  {f_final.max()/f_final.mean():14.3f}x"
          f"  {f_final.max():17.4f}")

print("\nFinal per-expert utilisation:")
for mode, hist in results.items():
    f_final, _ = hist[-1]
    row = "  ".join(f"{100*fi:5.1f}%" for fi in f_final)
    print(f"  {mode:<14}: [{row}]  (ideal: {100/E:.1f}% each)")

Running Code 17.6.5 reveals the central trade-off between the two mechanisms. The auxiliary-loss mode reaches a similar level of per-expert balance but does so by bending the task gradient, which lowers the perplexity proxy compared to the bias-update mode on the same scenario. The bias-update mode achieves comparable or better balance while keeping the task gradient uncontaminated, because the bias operates as a selection threshold rather than a loss modifier. The gap widens at larger $\alpha$ values, which is why practitioners who need high balance without quality loss increasingly prefer the bias-update approach or its variants.