Bias-Term Load Balancing: The DeepSeek Solution

The Tension We Inherit

Two sections ago we watched a mixture-of-experts layer collapse. A handful of experts captured almost all the routing decisions, every other expert went idle, and the model effectively shrank from N experts down to four or five. One section ago we tried the standard fix: an auxiliary load-balancing loss. It worked — the load curves flattened — but the gradient of that loss flowed straight into the same parameters that produce the token logits, and the model paid for the balance with a measurable hit on language-modeling quality. Bigger model, lower quality. That is the worst possible trade in large-scale training.

So we arrive at this section with a very specific question: can we get the discipline of the auxiliary loss without paying its gradient bill? DeepSeek-V3's answer — published in late 2024 and refined in the V3 / R1 technical reports — is a strikingly small modification to the router. It introduces a single per-expert scalar, updates it with no gradient at all, and produces flatter load curves than the auxiliary loss while leaving model quality strictly better. This section is about that scalar.

The one-sentence version. Add a per-expert bias $b_i$ to the score $s_{i,t}$ used to select the top-K experts — but not to the score used to weight them. Then nudge $b_i$ down whenever expert $i$ is over-served and up whenever it's under-served. The model gradient never sees $b_i$ at all.

The Trick: Decouple Selection from Weighting

Standard top-K routing fuses two distinct decisions into one function. First, the router decides which experts a token activates. Second, it decides how much each chosen expert contributes to the output. In every MoE layer before DeepSeek-V3, both of those decisions are read off the same per-(token, expert) affinity score $s_{i,t}$.

That fusion is what made the auxiliary-loss approach painful. To change which experts get picked, the auxiliary loss had to change how much each gets weighted, because they came from the same logits. The two effects rode the same gradient.

The bias-term trick unfuses them — physically. Let $s_{i,t}$ be the raw affinity (the model's own opinion of how relevant expert $i$ is to token $t$). Now build the routing decision from a different score:

$\tilde{s}_{i,t} = s_{i,t} + b_i$ where $b_i$ is a per-expert scalar that depends only on recent load — not on the token, not on the model weights.

The router selects top-K from $\tilde{s}_{i,t}$. The gating weights it emits — what mathematically combines the experts' outputs — use the original $s_{i,t}$ with no bias. Selection and weighting are now driven by two different scores. They can be tuned independently. That's the entire structural move.

The Mathematical Formulation

Let the MoE layer have $N$ experts, route top-$K$ per token, and process $T$ tokens per batch. The router holds a learned centroid $e_i \in \mathbb{R}^d$ per expert and, separately, a non-trainable per-expert bias $b_i \in \mathbb{R}$.

Affinity (the model's opinion)

For each token hidden state $u_t \in \mathbb{R}^d$, the per-expert affinity is $s_{i,t} = \sigma(u_t^\top e_i)$, where $\sigma$ is the elementwise sigmoid. Each $s_{i,t}$ lives in $(0, 1)$ independently of the others — independent, not normalized into a softmax. This choice matters: it means adding a bias to expert $i$'s score does not steal probability mass from any other expert. The bookkeeping stays local.

Routing decision (selection only)

Define the bias-adjusted score $\tilde{s}_{i,t} = s_{i,t} + b_i$. The router picks the top-$K$ indices:

$\mathcal{T}_t = \operatorname{Top\text{-}K}_i\,\tilde{s}_{i,t}$ — the set of $K$ expert indices a token $t$ activates.

Gating weights (weighting only)

For each $i \in \mathcal{T}_t$, the gating weight is the renormalized raw affinity:

$g_{i,t} = \frac{s_{i,t}}{\sum_{j \in \mathcal{T}_t} s_{j,t}}$

Notice $b_i$ does not appear here. The model's opinion alone decides how much weight each chosen expert contributes. The token's output is then $y_t = \sum_{i \in \mathcal{T}_t} g_{i,t} \, \mathrm{Expert}_i(u_t)$.

The Bias Update as a Control Loop

With selection and weighting decoupled, we're free to update $b_i$ however we want. DeepSeek picks the simplest possible rule: a sign-based proportional controller.

After each step, let $\ell_i = |\{ (t, k) : i \in \mathcal{T}_t \}|$ — the number of token-slots routed to expert $i$ this batch. The uniform setpoint is $\ell^* = T K / N$. Update:

$b_i \leftarrow b_i - \gamma \cdot \operatorname{sign}(\ell_i - \ell^*)$

The hyperparameter $\gamma > 0$ (the DeepSeek paper writes it as $u$) controls the loop's speed. Read the rule sign by sign:

• Expert $i$ got more tokens than expected — $\ell_i > \ell^*$, so $b_i$ drops by $\gamma$. Its bias-adjusted scores fall, and on the next batch it tends to lose top-K races it would otherwise have won by small margins.
• Expert $i$ got fewer tokens — $b_i$ rises by $\gamma$. Its bias-adjusted scores climb; on the next batch it wins more close races.
• Expert exactly at the setpoint — $b_i$ unchanged. (In practice this never happens; the sign function breaks every tie.)

The fixed point is the place where the per-expert biases have settled into whatever offsets the data distribution actually needs — bigger negative biases for the intrinsically "popular" experts, bigger positive biases for the intrinsically "niche" ones — and load is uniform. This is a standard bang-bang controller. It is the same idea as a thermostat clicking on and off.

Why It Has Zero Gradient Interference

The cleanest way to see the argument is to trace the computation graph. Start at the loss $L$ and walk backward to find which parameters it depends on.

The bias $b_i$ influences $L$ only through the discrete top-K choice, which has no usable gradient. The differentiable path — through $g_{i,t}$ — bypasses $b_i$ entirely. So when the optimizer asks autograd "what should I do with $b_i$?", the honest answer is nothing. We update it ourselves, with the control loop, and the gradient signal that improves token-level prediction flows undisturbed through the centroids $e_i$ and the expert networks.

Interactive: Watch the Bias Loop Stabilize

Before any code, run the loop yourself. The simulator below uses 8 experts, top-K = 2, and a deliberately skewed affinity distribution — experts 0 and 1 have systematically higher mean affinity than the rest. With bias correction off, that skew is enough to send the routing into permanent collapse. With it on, watch the biases drift into negative territory for the popular experts and positive for the niche ones until load is essentially uniform.

Things worth doing in the simulator, in order:

1. Press Skip 30 steps with the bias correction off. The per-step load on experts 0 and 1 sits around 20-25 tokens (out of 32 slot-allocations), and the bottom four experts get almost nothing. The imbalance curve refuses to fall.
2. Reset. Turn the bias correction on with $\gamma = 0.05$. Step a few times — you can see the biases for the popular experts already turning negative after step 1. Press Skip 30 steps. Imbalance falls to near $1.0\times$.
3. Crank $\gamma$ to its max ($0.2$) and step a few times. The biases jitter visibly — you're seeing the cost of an over-aggressive controller. With $\gamma = 0.005$, the convergence is smoother but slower. This trade-off is exactly the one DeepSeek tunes with a schedule in real training.
4. Drop $\gamma$ to zero. The control loop freezes. Routing decisions still get the current biases applied, but they stop adapting — this is what DeepSeek does in the last ~5% of training to lock in the routing pattern.

Manual Numerical Walkthrough

Let's execute the full loop by hand for one step, with $N = 4$ experts, $K = 2$, and $T = 6$ tokens. Small enough to carry every intermediate.

Plain Python Implementation

Here is the entire mechanism — affinity, selection, gating, control loop — in pure NumPy. Nothing is hidden behind an MoE library; every operation is a single named line you can trace. Read it once top-to-bottom, then open the explanation panel.

expected_load = 64 * 2 / 8 = 16 tokens / expert / step

s.shape == (64, 8)

s[i] = 1.4   b[i] = -0.6   decision[i] = 0.8   (gate still uses 1.4)

overload = [+30, +28, -6, -7, -8, -9, -10, -11]  -> bias -= 0.05 * [+1,+1,-1,-1,-1,-1,-1,-1]

1import numpy as np
2
3np.random.seed(0)
4
5# ----- Toy MoE layer ---------------------------------------------------------
6N_EXPERTS   = 8        # number of experts in the layer
7TOP_K       = 2        # each token activates 2 experts
8TOKENS_PER_STEP = 64   # mini-batch size for routing
9GAMMA       = 0.05     # bias update rate (DeepSeek calls this u in the paper)
10N_STEPS     = 400
11
12# Per-expert "centroid" e_i — a real router learns these; we hard-code so
13# expert 0 and 1 are systematically more attractive than the rest.
14# This is the imbalance source.
15popular   = np.array([1.3, 1.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
16
17# ----- State that the control loop updates ----------------------------------
18bias = np.zeros(N_EXPERTS)            # b_i — what we update by hand
19load_history = []                     # to plot later
20
21def affinity(num_tokens, rng):
22    """Per-token affinity s_{i,t} — popular experts get higher mean."""
23    noise = rng.standard_normal((num_tokens, N_EXPERTS)) * 0.7
24    return popular[None, :] + noise
25
26def topk_route(scores, k):
27    """Return the indices of the top-k experts per token. (T, k)"""
28    return np.argpartition(-scores, k, axis=1)[:, :k]
29
30rng = np.random.default_rng(0)
31expected_load = TOKENS_PER_STEP * TOP_K / N_EXPERTS    # uniform target
32
33for step in range(N_STEPS):
34    s = affinity(TOKENS_PER_STEP, rng)            # raw affinity, (T, N)
35    decision = s + bias[None, :]                  # ONLY for selection
36    picks    = topk_route(decision, TOP_K)        # (T, k)
37
38    # Count how many tokens each expert received this step
39    load = np.bincount(picks.ravel(), minlength=N_EXPERTS).astype(float)
40    load_history.append(load.copy())
41
42    # ----- The control loop update --------------------------------------
43    # overloaded  -> push bias DOWN  -> fewer tokens next time
44    # underloaded -> push bias UP    -> more  tokens next time
45    overload = load - expected_load
46    bias = bias - GAMMA * np.sign(overload)
47
48# ----- Final state ----------------------------------------------------------
49print("final bias  :", np.round(bias,   3))
50print("final load  :", load.astype(int))
51print("uniform aim :", expected_load)

Two structural things to take away from this file. First, the bias vector is just a NumPy array we mutate by hand — there is no gradient anywhere in the file. Second, the line decision = s + bias[None, :] is the entire algorithm. Everything before it sets up the data; everything after it counts load and updates the controller.

What to try with this code

1. Run it as-is and print the load every 50 steps. You'll see the popular experts start at ~25 tokens/step and walk down to ~16; the unpopular ones start near zero and climb up to ~16.
2. Set GAMMA = 0. The biases never change, the decision is just $s$, and the load stays permanently lopsided — the no-balancing baseline.
3. Replace np.sign(overload) with overload / expected_load. This is the magnitude-aware variant; convergence is faster, but a single unusual batch can over-correct.
4. Add a slow drift to popular — e.g. rotate the high-affinity slots from (0, 1) to (2, 3) at step 200. Watch the biases re-balance on the fly. This is the actual situation during long pretraining: data distribution drifts, and the loop tracks it.

PyTorch Implementation

The PyTorch version preserves the same control structure but uses the framework's tools properly — centroids as Parameters (so the optimizer learns them) and bias as a registered buffer (so it's saved with the model but invisible to the optimizer and to autograd).

h: (B, T, d_model)  ->  s: (B, T, n_experts)

s = sigmoid(h @ e^T)   range (0, 1) per cell

load tensor shape: (n_experts,)

self.bias[i] -= gamma  if expert i is overloaded

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5
6class BiasBalancedRouter(nn.Module):
7    """
8    Auxiliary-loss-free router from DeepSeek-V3.
9
10      score      s_{i,t} = sigmoid(u_t . e_i)
11      selection  top_k( s + b )                    <-- bias INSIDE top-k
12      gate       g_{i,t} = s_{i,t} / sum_{j in topk} s_{j,t}    <-- bias NOT here
13      update     b_i  <-  b_i - gamma * sign(load_i - expected)
14    """
15
16    def __init__(
17        self,
18        d_model: int,
19        n_experts: int,
20        top_k: int,
21        bias_update_rate: float = 1e-3,
22    ):
23        super().__init__()
24        self.n_experts = n_experts
25        self.top_k     = top_k
26        self.gamma     = bias_update_rate
27
28        # Expert centroids — learned by gradient descent (these DO get grads)
29        self.centroids = nn.Linear(d_model, n_experts, bias=False)
30
31        # Per-expert balancing bias — NOT a parameter; updated by the
32        # control loop. Buffer = saved with the model but excluded from
33        # the optimizer.
34        self.register_buffer("bias", torch.zeros(n_experts))
35
36    def forward(self, h: torch.Tensor):
37        """
38        h : (B, T, d_model)
39        returns:
40          gates    (B, T, k)    weighting for each chosen expert
41          indices  (B, T, k)    which expert each slot picked
42          load     (n_experts,) tokens per expert this batch (for the update)
43        """
44        # ---- 1. Per-token affinity --------------------------------------
45        s = torch.sigmoid(self.centroids(h))          # (B, T, n_experts)
46
47        # ---- 2. Select top-k using BIAS-CORRECTED scores ---------------
48        decision = s + self.bias                      # broadcast (n_experts,)
49        _, idx   = decision.topk(self.top_k, dim=-1)  # (B, T, k)
50
51        # ---- 3. Gate weights use RAW affinity (no bias!) ---------------
52        chosen_s = s.gather(-1, idx)                  # (B, T, k)
53        gates    = chosen_s / chosen_s.sum(-1, keepdim=True).clamp_min(1e-9)
54
55        # ---- 4. Count load for the control-loop update ------------------
56        load = torch.zeros(self.n_experts, device=h.device, dtype=h.dtype)
57        load.scatter_add_(
58            0,
59            idx.reshape(-1),
60            torch.ones_like(idx.reshape(-1), dtype=h.dtype),
61        )
62        return gates, idx, load
63
64    @torch.no_grad()
65    def update_bias(self, load: torch.Tensor, total_tokens: int):
66        """
67        Call AFTER the optimizer step. No autograd anywhere in here.
68        """
69        expected = total_tokens * self.top_k / self.n_experts
70        overload = load - expected
71        self.bias.sub_(self.gamma * overload.sign())

How this module fits into a training loop

The usage pattern is intentionally boring:

• Forward pass: call router(h), dispatch tokens to the chosen experts using idx and gates, combine outputs. Standard MoE.
• Backward pass: loss.backward() flows through gates, which depend on self.centroids but not on self.bias. The optimizer updates centroids and expert weights as usual.
• After optimizer.step(): call router.update_bias(load, total_tokens). This is the control-loop update, fully outside autograd.

Because self.bias is a buffer and not a parameter, model.state_dict() captures it. Resuming a training run starts the controller from its warmed-up state — you don't lose thousands of steps of load history every time you checkpoint.

The γ Schedule in DeepSeek-V3

The simulator above hints at the trade-off: a large $\gamma$ converges quickly but oscillates; a small $\gamma$ is stable but slow to recover from distribution shift. DeepSeek-V3 resolves this the same way every other large-scale training trick does: schedule it.

Two details about the published numbers. First, the reported $\gamma = 0.001$ is for $T \approx 4096$ tokens per device with hundreds of devices in parallel — that is, an all-reduced load of tens of thousands of tokens per step. The effective controller strength is $\gamma$ times the per-step load count, so absolute numbers don't transfer to smaller setups. If you reproduce this with batch size 64, you almost certainly want $\gamma \approx 0.01$ to $0.05$.

Second, the cool-down is not optional. If you train with $\gamma > 0$ right up to the last step, the biases that the deployed model uses are the ones from the noisiest single batch you ever saw. The decay phase averages the biases toward their long-run mean. Models with the decay disabled show measurably worse inference-time load balance.

What Changes at Massive Scale

Everything above runs in 50 lines of NumPy on a laptop. Putting it on a 671B-parameter DeepSeek-V3 cluster surfaces three concerns that don't exist in the toy.

1. Load must be aggregated across ranks before the update

A real MoE layer is sharded with expert parallelism: different GPUs hold different experts. A token routed to expert 42 physically travels (via an all-to-all collective) to the GPU that owns expert 42. The load counter $\ell_i$ for expert 42 is naturally local to that GPU. But the data-parallel dimension is replicated — each data-parallel rank sees its own micro-batch of tokens, and they each accumulate their own partial $\ell_i$. Before applying the update we all_reduce the load vector across data-parallel ranks so every replica of the layer sees the same load and applies the same update to its (replicated) bias. Skip this and the biases drift apart, the routing decisions diverge across ranks, and training corrupts within a few hundred steps.

2. Bias bookkeeping is microscopic; expert dispatch is not

The bias update itself is N additions and N sign checks — for N = 256 experts, that is 256 floating-point ops. The MoE layer's compute is dominated by the expert MLPs (billions of FLOPs) and by the all-to-all token dispatch (tens of GB of communication). The balancing logic is essentially free. This is part of why the bias-term method is preferred over more elaborate controllers: whatever you spend on the controller, you spend per layer per step, and there are 61 MoE layers in DeepSeek-V3 firing thousands of times per second.

3. Dropped tokens are an emergent property of the loop

Real MoE deployments cap each expert at a per-step capacity — tokens beyond the cap are dropped or rerouted. With the auxiliary loss, you tune the loss coefficient and hope the load stays inside the cap. With the bias-term method, if the cap is set near the uniform load $\ell^* = TK/N$, the controller actively steers toward that target, and the drop rate falls naturally as training progresses. DeepSeek reports drop rates below 0.1% by the middle of training — orders of magnitude below auxiliary-loss baselines.

Production Pitfalls and What to Monitor

What to log every step

• Per-expert load $\ell_i$ (a length-N vector). Plot as a heatmap over training steps. Should go from skewed to flat over the first few thousand steps.
• Max-over-min load ratio $\max_i \ell_i / \max(1, \min_i \ell_i)$. Single scalar; should fall from >10× to near 1.5× and stay there.
• Bias magnitude $\|b\|_\infty$. Should grow during warmup then plateau. Unbounded growth means γ is too large or the all-reduce is missing.
• Drop rate (tokens exceeding per-expert capacity). Should fall monotonically. A persistent floor means your capacity cap is below the uniform setpoint and you should raise it.

Summary

DeepSeek-V3's auxiliary-loss-free load balancing is one of the most elegant pieces of mechanical design in modern MoE training. Three structural decisions carry the entire result:

1. Decouple selection from weighting. Use $\tilde{s}_{i,t} = s_{i,t} + b_i$ for top-K selection and the raw $s_{i,t}$ for gate weights.
2. Update the bias by a control loop, not by gradient descent. After each step, nudge $b_i$ down for over-served experts and up for under-served ones.
3. Store the bias as a buffer, not a parameter. Autograd can't reach it, the optimizer can't see it, and the loss has no path through it.

The combined effect is the discipline of an auxiliary loss with the gradient quiet of no loss at all. Compared to MoE models that use an auxiliary loss, DeepSeek-V3 trains with flatter load curves, a sharper drop rate, and a small but consistent improvement in language-modeling quality — not in spite of the load balancer, but because the load balancer stopped fighting the model.

The next section closes the chapter by adding one more layer of defense: a per-sequence balance term that catches the rare extreme imbalances even the bias loop can't prevent — and we'll see why it's applied at the sequence level, not the batch level, before moving on to multi-token prediction in chapter 7.