Training the MoE router — specialization without expert collapse

Yesterday we explored Mixture of Experts (MoE): a Transformer can contain many feed-forward experts while activating only a few for each token. This separates total model capacity from compute per token.

But we left an important question open: who teaches the router to distribute work sensibly?

The intuition is:

A router must learn which experts are useful, while preventing a few popular experts from monopolizing the workload.

That sounds like a scheduling problem. In an MoE model, it is also a training problem, because routing decisions determine which experts receive gradients and therefore which experts improve.

The router makes two decisions at once¶

For a token representation \(h\), a router produces logits and probabilities over \(E\) experts:

\[ z= W_rh,\qquad p_e=\frac{e^{z_e}}{\sum_j e^{z_j}}. \]

With top-\(k\) routing, the model selects the \(k\) experts with the largest scores. For example, if

\[ p=[0.50,\ 0.30,\ 0.15,\ 0.05] \]

and \(k=2\), experts 0 and 1 run. If their selected probabilities are renormalized, the output is

\[ y=\frac{0.50}{0.80}F_0(h)+\frac{0.30}{0.80}F_1(h). \]

This combines two distinct choices: which experts execute (discrete selection) and how much their outputs contribute (continuous weights).

Ordinary backpropagation can adjust the selected experts and the continuous routing weights. But the top-\(k\) selection boundary is discontinuous; unselected experts generally receive no task-loss gradient from this token. Different MoE architectures handle routing gradients differently, but that asymmetry is fundamental.

Why expert collapse can reinforce itself¶

Imagine four experts at the beginning of training. Expert 0 happens to perform slightly better on common tokens, so the router sends it more work.

Expert 0 now receives more gradient updates and may improve faster. The router observes its improving performance and selects it even more often.

Meanwhile, expert 3 receives few tokens, learns slowly, and becomes less attractive to the router.

slightly better expert
        ↓
more tokens routed to it
        ↓
more training signal
        ↓
expert improves faster
        ↓
even more tokens routed to it

This is expert collapse or severe routing imbalance. It wastes parameter capacity and can create hardware bottlenecks even when the overall model loss is decreasing.

Perfectly equal routing is not necessarily desirable, however. Some experts may genuinely be more useful for particular inputs. The goal is productive specialization without pathological underuse or overload.

Measuring balance: assignments, not just probabilities¶

Suppose a layer processes \(T\) tokens using \(E\) experts and top-\(k\) routing. It generates

\[ A=Tk \]

token-to-expert assignments. The ideal average is

\[ \bar A=\frac{Tk}{E}. \]

Define \(f_e\) as the fraction of actual assignments going to expert \(e\), and \(P_e\) as the mean router probability assigned to that expert:

\[ f_e=\frac{\text{assignments to expert }e}{Tk}, \qquad P_e=\frac{1}{T}\sum_{t=1}^{T}p_{t,e}. \]

Both vectors sum to one. But they measure different things:

\(f_e\) measures actual compute demand after top-\(k\) selection.
\(P_e\) measures the router's soft preference, including probabilities for experts that were not selected.

A router can have reasonably spread-out probabilities but still send too many tokens to the same top-ranked expert. That is why looking only at average softmax probabilities can hide the real systems problem.

Worked example: 16 tokens, four experts¶

Take \(T=16\), \(E=4\), and \(k=2\). There are 32 assignments, so balanced routing would give 8 assignments per expert.

Instead, suppose the observed loads are

\[ A_e=[14,\ 10,\ 6,\ 2]. \]

The busiest expert receives 1.75 times the average workload, while the least-used receives only one-quarter of it.

If each assignment costs roughly the same and the experts execute concurrently, the busiest expert can determine the layer's completion time. More GPUs do not help much when most are waiting for one straggler.

Now impose a capacity factor \(c=1.25\). The capacity per expert is approximately

\[ C=\left\lceil c\frac{Tk}{E}\right\rceil =\lceil1.25(8)\rceil=10. \]

Expert 0 has four assignments above capacity. Depending on the implementation, those assignments may be dropped, rerouted, or handled through another execution path.

A dropless implementation can preserve all assignments, but it must accommodate uneven workloads and possibly larger temporary buffers. Capacity limits control resource usage; they do not automatically teach the router better decisions.

Teaching the router to balance itself¶

One approach is to add an auxiliary load-balancing objective to the task loss. A commonly used family of objectives has the form

\[ \mathcal L_{\mathrm{balance}} =E\sum_{e=1}^{E} f_eP_e. \]

For the observed assignments,

\[ f=[0.4375,\ 0.3125,\ 0.1875,\ 0.0625]. \]

Suppose the mean router probabilities are

\[ P=[0.40,\ 0.30,\ 0.20,\ 0.10]. \]

Then

\[ \mathcal L_{\mathrm{balance}} =4(0.4375\cdot0.40+0.3125\cdot0.30 +0.1875\cdot0.20+0.0625\cdot0.10) =1.25. \]

For uniform assignments and uniform probabilities, this expression equals 1.0. The higher value here signals that high-probability experts are also receiving disproportionate traffic.

The training objective might be

\[ \mathcal L_{\mathrm{total}} =\mathcal L_{\mathrm{task}} +\lambda\mathcal L_{\mathrm{balance}}. \]

The coefficient \(\lambda\) sets the tradeoff between predictive quality and balanced utilization.

Important caveat: this is a balancing heuristic, not a mathematically guaranteed measure of imbalance. Typically \(f_e\) is treated as a nondifferentiable assignment statistic, while gradients flow through \(P_e\). The objective does not guarantee uniform loads, and pushing balance too hard can interfere with useful specialization. Implementations also differ in how they normalize top-\(k\) assignments and compute the auxiliary loss.

An alternative: adjust routing biases¶

Some MoE designs discourage overload without adding a balancing term directly to the model's task loss.

Give each expert a routing bias \(b_e\) and select experts using

\[ s_e=z_e+b_e. \]

After observing utilization, update the biases with a rule resembling

\[ b_e\leftarrow b_e+\eta\left(\frac{1}{E}-f_e\right). \]

An overloaded expert gets a lower bias; an underused expert gets a higher bias. The router's learned logits still express token preferences, but the biases act like a slowly adjusted traffic-control signal.

This is a simplified feedback-controller model, not a universal implementation recipe. Some architectures use counts rather than normalized fractions, different update rules, or biases only for expert selection—not for weighting the selected expert outputs.

The distinction matters: auxiliary-loss balancing changes the gradient objective; bias-based balancing changes the selection incentives through a separate feedback mechanism.

Training balance and serving balance are not identical¶

A router trained on large batches may look balanced on average, yet inference requests arrive with very different token counts and subject matter.

During prefill, thousands of tokens can be grouped into reasonably sized expert batches. During decode, the system may have only a handful of new tokens per step. A few unlucky routing choices can produce tiny, inefficient expert matrix multiplications or an overloaded remote GPU.

Balance also depends on where it is measured. Equal load across an entire training batch does not guarantee equal load on each device, within each microbatch, or at every inference step.

This connects directly to Day 75's all-to-all token dispatch: the router controls not only which computation occurs, but also the traffic pattern of the distributed system.

Mental model: an evolving specialist marketplace¶

Imagine four consulting teams and a dispatcher that assigns each project to two teams.

The dispatcher initially favors the teams with the best track records. Those teams get more projects, accumulate experience, and become even more attractive. The neglected teams never develop their potential.

A load-balancing loss is like giving the dispatcher a performance target that includes team utilization, not just client satisfaction.

A routing bias is more like a dynamic surcharge on overloaded teams and a discount for underused ones.

Neither mechanism should force every team to receive exactly the same projects. The objective is to develop distinct expertise while keeping the whole organization productive.

The deeper lesson¶

MoE is often described as a way to reduce active FLOPs. But the router is doing something more consequential: it determines the training data distribution experienced by each expert and the runtime load distribution experienced by each GPU.

\[ \boxed{\text{In an MoE model, routing is simultaneously a learning policy and a resource-allocation policy.}} \]

The practical question is: How do we encourage experts to specialize without letting early routing preferences starve some experts or overload others?

That question links model quality, optimization dynamics, and distributed inference performance in a way that dense Transformers largely avoid.