Sequence-Level Balance Loss

Section 6.3 introduced DeepSeek's clean fix for load imbalance: a learned bias on each expert's router score, updated by a non-differentiable feedback rule. That mechanism is enough to balance the routing on average across the batch — and on average, it works beautifully. But averages can hide bad behavior at finer time scales. A batch can look balanced while individual sequences inside it are quietly collapsing. The sequence-level balance loss is the tiny safety net DeepSeek adds on top, designed to catch exactly this failure.

The promise of sequence-level balance. A nearly-zero auxiliary loss applied per sequence — not per batch — to penalize within-sequence routing collapse. The coefficient is so small ($\alpha = 10^{-4}$) that it does not corrupt the routing decision. Just large enough that, when one sequence quietly funnels 80% of its tokens to a single expert, the gradient has somewhere to push back.

The Hidden Failure of Batch-Level Balance

Imagine a training batch with 32 sequences. The bias-term feedback rule from the previous section nudges biases up for under-used experts and down for over-used ones, measured over the whole batch. After a few hundred steps it converges to a clean state: across those 32 sequences combined, every expert handles roughly its fair share of tokens. The batch-level histogram looks like a flat line.

Now inspect the routing inside a single sequence. You may discover that sequence #7 sent 80% of its tokens to expert 12, while sequence #19 sent 80% of its tokens to expert 3. Both sequences are using only a sliver of the expert pool locally, but they happen to use different slivers — so when you average across the batch, the imbalance cancels out. The bias-term mechanism is satisfied. The model is not.

Two Time Scales, Two Different Failures

Picture a restaurant kitchen. The head chef (bias-term feedback) watches the whole shift and shouts "more salads, fewer steaks" when the totals drift out of balance. By the end of the night the count is even — across the whole shift, the kitchen handled equal numbers of every dish. But within any given hour the kitchen might have been doing nothing but desserts, then nothing but soups. The shift-level average is balanced; the hour-by-hour rhythm is broken.

The sequence-level loss is the line cook — a much quieter signal that operates inside each hour. It does not try to fix the absolute counts (the head chef already handles that). It just whispers "hey, you've sent the last four orders to the same station; spread the next one out." Two different control loops at two different time scales, doing two different jobs.

This split is fundamental to how DeepSeek thinks about load balancing. Thebias-term feedback handles slow, global imbalance and does so without contributing any gradient. The sequence-level loss handles fast, local imbalance and contributes a microscopic gradient — small enough that it cannot reshape the routing decision, but present enough that the optimizer sees a directional hint when a sequence starts to collapse.

The Math: A Per-Sequence Smoothing Loss

Fix one sequence with $T$ tokens and $E$ experts. Top-$k$ routing produces, for each token $t$, a set of chosen experts$\mathcal{T}_t \subset \{1, \dots, E\}$. The sequence-level balance loss DeepSeek-V3 uses is:

$\mathcal{L}_{\text{seq}} = \alpha \sum_{i=1}^{E} f_i \cdot P_i$

with two ingredients. The hard fraction $f_i$ counts how often expert $i$ was actually routed to, normalized so that perfectly uniform routing gives $f_i = 1$:

$f_i = \frac{E}{k T} \sum_{t=1}^{T} \mathbb{1}\{i \in \mathcal{T}_t\}$

The soft probability $P_i$ is the average router probability that expert $i$ received across the sequence, where $s_{t,i}$ are the un-biased router logits for token $t$:

$P_i = \frac{1}{T} \sum_{t=1}^{T} \frac{e^{s_{t,i}}}{\sum_{j} e^{s_{t,j}}}$

Why two terms — one hard, one soft

The hard term $f_i$ is non-differentiable: an argmax / top-$k$ cannot pass a gradient. The soft term $P_i$ is fully differentiable through the softmax. Their product $f_i P_i$ is the engineering trick: gradient flows only through $P_i$, but it is weighted by the actual frequency $f_i$. So the loss says: "if expert $i$ is being routed to a lot $(f_i \text{ large})$, then push its average probability down $(\nabla_{P_i} < 0)$ for next time."

Why the bias is excluded from the logits

The router logits used in $P_i$ are pre-bias — the raw output of the router's linear layer, before the section-6.3 bias is added. If you let the bias appear inside the softmax here, the balance loss would push gradient into the bias, and the bias-term mechanism would lose its "non-differentiable controller" property. Keeping the two mechanisms cleanly separated is critical.

What is conspicuously absent

Notice this loss is averaged over $T$ (tokens within one sequence) and computed independently for each sequence in the batch. The batch dimension $B$ only appears in the very last reduction — the per-sequence losses are averaged together to give one scalar. There is no cross-sequence mixing inside the formula. That is the entire point: the loss sees only what is happening within one sequence at a time.

Manual Numerical Walkthrough

Let us run the loss on one toy sequence of $T = 6$ tokens, $E = 4$ experts, top-$k = 2$. Every number is calculated by hand so the mechanism is fully exposed.

Visualizing Batch vs. Sequence Scope

The diagram below holds two sequences side by side. Sequence A is routed evenly; Sequence B suffers a within-sequence collapse onto expert 0. Toggle between Batch-level and Sequence-level to see how the same routing pattern produces a single averaged loss versus per-sequence losses. The α slider scales the loss without changing the routing.

Three things to lock in. First, in batch mode the f-bars and P-bars for sequence B's collapse get washed out by sequence A's spread — the merged signal is misleadingly clean. Second, in sequence mode each sequence carries its own loss: sequence B pays more, and that extra penalty is exactly the gradient signal that will pull its next-step routing back toward uniform. Third, the absolute scale of the loss is microscopic even when you push α up to $2 \times 10^{-3}$ — DeepSeek-V3's default of $10^{-4}$ sits well below the noise floor of the main LM loss.

Plain Python: One Sequence, By Hand

Below is the entire per-sequence calculation in NumPy. Six tokens, four experts. Every line maps one-to-one to a line in the math above — and the printed L_seq matches the walkthrough we just did by hand.

1import numpy as np
2
3# Tiny toy: T = 6 tokens, E = 4 experts, top-k = 2.
4T, E, k = 6, 4, 2
5alpha = 1e-4
6
7# (a) Per-token, per-expert router scores (T, E).
8scores = np.array([
9    [2.0, 0.4, 1.6, 0.5],
10    [0.3, 2.1, 0.6, 1.5],
11    [1.8, 1.5, 0.4, 0.5],
12    [0.4, 0.5, 1.9, 1.4],
13    [0.4, 1.7, 1.6, 0.5],
14    [1.7, 0.4, 0.5, 1.5],
15])
16
17# (b) Softmax over the experts for each token.
18def softmax(x):
19    z = x - x.max(axis=-1, keepdims=True)
20    e = np.exp(z)
21    return e / e.sum(axis=-1, keepdims=True)
22
23probs = softmax(scores)                                   # (T, E)
24
25# (c) Top-k indices per token define the actual routing.
26topk = np.argsort(-scores, axis=-1)[:, :k]                # (T, k)
27
28# (d) f_i = fraction of tokens routed to expert i, normalized so uniform = 1.
29counts = np.zeros(E)
30for r in topk:
31    counts[r] += 1
32f = (E / (k * T)) * counts                                # (E,)
33
34# (e) P_i = average router probability for expert i over the sequence.
35P = probs.mean(axis=0)                                    # (E,)
36
37# (f) Per-sequence balance loss.
38L_seq = alpha * (f * P).sum()
39print(f"f = {f}")
40print(f"P = {P}")
41print(f"L_seq = {L_seq:.3e}")

Two structural details deserve a second pass. First, the normalization factor $E/(kT)$ on $f$ is what makes the loss insensitive to sequence length. Without it, a longer sequence would mechanically produce a larger loss, and the per-sequence reduction would just become a roundabout per-token average. With it, $f_i$ is dimensionless and uniform-target = 1, so a 4096-token sequence and a 256-token sequence are compared on equal terms.

Second, this implementation is intentionally non-batched. There is no $B$ dimension; you call this function once per sequence. The PyTorch version below adds the batch dim and a final mean, but otherwise the math is identical.

Sanity check. Set every routing decision to $i = 0$ and every score vector to $[100, 0, 0, 0]$. Then $f = [4, 0, 0, 0]$, $P \approx [1, 0, 0, 0]$, and $f \cdot P = 4$ — the worst possible value (= E). At the other extreme, perfectly uniform routing and a flat softmax give $f \cdot P = 1$. The loss spans a factor of E between the worst and best collapse states.

PyTorch: Vectorized Across the Batch

The production version lives inside the MoE forward pass and runs once per layer. The only new ideas are the batch dimension and the careful reduction order — average within each sequence first, then average across the batch.

1import torch
2import torch.nn.functional as F
3
4def sequence_balance_loss(
5    scores: torch.Tensor,  # (B, T, E) — raw router logits, NO bias added
6    topk_idx: torch.Tensor,  # (B, T, k) — actual routed experts per token
7    alpha: float = 1e-4,
8) -> torch.Tensor:
9    """
10    Per-sequence balance loss from DeepSeek-V3. Reduced as the mean across B.
11
12    Note: 'scores' here are pre-bias router logits. The bias from section 6.3
13    is only added to decide topk_idx; it must NOT leak gradients here.
14    """
15    B, T, E = scores.shape
16    k = topk_idx.shape[-1]
17
18    # (a) Smooth term: per-expert mean probability inside each sequence.
19    probs = F.softmax(scores, dim=-1)              # (B, T, E)
20    P = probs.mean(dim=1)                           # (B, E)
21
22    # (b) Hard term: per-expert routed-fraction inside each sequence.
23    onehot = F.one_hot(topk_idx, num_classes=E)     # (B, T, k, E)
24    counts = onehot.sum(dim=(1, 2)).float()         # (B, E)
25    f = counts * (E / (k * T))                       # (B, E)
26
27    # (c) Per-sequence inner product, then mean across the batch.
28    L_per_seq = (f * P).sum(dim=-1)                 # (B,)
29    return alpha * L_per_seq.mean()

Three subtleties worth marking, all about how this loss interacts with the rest of the MoE block:

1. The probs softmax and the routing softmax are different. The softmax inside this loss is purely for measuring imbalance — it is not the one whose top-k decides actual routing. In DeepSeek-V3 the routing decision uses scores-plus-bias, while this loss uses scores-only. Two softmaxes, two purposes.
2. Gradient flows into the router weights, not the bias. Because the bias was not used in probs, the gradient of L_seq with respect to the bias parameters is exactly zero. The bias is updated only by the non-differentiable feedback rule from section 6.3. This is a clean separation of concerns: the bias controls long-run balance, the loss controls short-run balance.
3. One loss per MoE layer. Modern stacks have dozens of MoE layers, and each has its own router and its own potential for collapse. The balance loss is computed independently per layer and summed; with $\alpha = 10^{-4}$ and 58 MoE layers (DeepSeek-V3) the total contribution to the loss is still only $\sim 0.006$ — well under 1% of the LM loss.

What Changes at Massive Scale

At toy scale the sequence-level loss looks like a curiosity — a tiny add-on with a near-zero coefficient. At production scale the situation is dramatically different because the within-sequence collapse failure mode becomes much more likely as the routed pool grows. The DeepSeek-V3 numbers tell the story:

Two patterns to read here. First, as the routed pool grew from 160 to 256, DeepSeek reduced α from $10^{-3}$ to $10^{-4}$. The bias-term mechanism from section 6.3 got better at handling the bulk of the balancing job, so the safety net could be pulled tighter and quieter. The loss is now there mainly to handle the tail — the small fraction of sequences with unusual content distributions.

Second, Mixtral does not use a sequence-level loss at all. It uses the standard auxiliary load-balancing loss with a much larger coefficient. That works adequately for an 8-expert pool — there are not enough experts for within-sequence collapse to become a serious problem. But it would not scale to 256 experts: the auxiliary loss would have to be loud, and routing quality would suffer.

The interaction with sequence packing

Real training rarely uses one logical document per sequence. Multiple documents are concatenated up to $T = 4096$ with an attention mask that prevents cross-document attention. From the sequence-level loss's point of view, this is a feature: a 4k-token packed sequence contains a varied document mix, so balanced routing is genuinely the right target. If you packed a single homogeneous document into the full 4k window — say, a long block of Chinese poetry — the loss would correctly fire and push the router away from the natural domain-driven collapse. Whether that is what you want is a separate design call; for general pre-training, it is.

The interaction with expert parallelism

Each MoE layer's router runs on every GPU that holds a copy of the model (data-parallel rank). The per-sequence loss is computed locally on the rank that owns the sequence — no all-reduce needed for the balance loss itself, because there is no cross-sequence mixing. Only the gradient accumulation does any sync, and that happens once for the whole step regardless. So the sequence-level loss is essentially free in communication terms, which is part of why DeepSeek can afford to add it at every MoE layer.

Engineering Reality and Gotchas

The sequence-level balance loss looks innocuous and is one of the easiest pieces of DeepSeek's MoE stack to get subtly wrong. Three failure modes are worth flagging:

1. Letting the bias into the softmax. If you reuse the same scores_with_bias tensor for both the routing decision and the balance loss, the gradient of L_seq will leak into the bias parameters. The bias-term feedback rule was explicitly designed to be non-differentiable; letting gradient enter it muddies the control loop and reintroduces all the routing-quality problems section 6.2 documented. Keep the two paths separate.
2. Using α too large. Try $\alpha = 10^{-2}$ and the loss starts dominating the gradient direction for the router. Routing collapses into a forced uniform policy, and the model loses the ability to specialize. This is the same failure mode as the section-6.2 auxiliary loss — just slower and meaner because it strikes at the sequence scale.
3. Averaging in the wrong order. Reducing first across the batch and then across the sequence is mathematically different from the intended "sequence first, batch second" order — and the former is exactly the auxiliary loss from section 6.2. Get the reduction order wrong and you are quietly running the old method while believing you are running the new one. Always check the dim arguments in your .mean() calls.

The one sentence to carry forward: the sequence-level balance loss is the smallest piece of the DeepSeek MoE stack and the easiest to dismiss, yet it is what turns the bias-term mechanism from a batch-average controller into a true per-sequence balancer — and that is what makes the architecture behave at single-stream inference, not just at training time.