APEX: Load Balance Was Never the Right Loss

Companion note to APEX: Adaptive Per-Token Expert Loss.

Here is the thing about mixture-of-experts models. They shouldn’t work.

You have 8 or 64 or 128 “experts” — sub-networks — and a small router that picks two of them per token. The router is untrained at initialization. Its picks are essentially random. The gradients flowing back into the router are noisy and sparse.

Left to itself, the router collapses. It finds one or two favorites and picks them for every token, and the other experts die, having never gotten enough training signal to become useful. So every MoE paper adds an auxiliary load-balancing loss — a penalty that says “please use all your experts roughly equally.”

The standard load-balance loss works, in a sense. Everyone stays alive. Nothing specializes. This is the failure mode APEX exists to fix.

---

The setup, in one paragraph

The router assigns each token a distribution over experts. The load-balance loss penalizes the deviation of global expert usage — averaged over the whole batch — from uniform. So if expert 3 is picked 40% of the time and expert 7 is picked 5%, the loss says “route more to expert 7.”

That objective is batch-level. It doesn’t care which token went where. It only cares about the aggregate.

The right per-token intuition is the opposite shape:

• If the router is confident about a token (sharp distribution over 1-2 experts), leave it alone. That’s specialization forming; that’s what we want.
• If the router is uncertain about a token (flat distribution), lean harder on load-balance. Uncertainty means specialization hasn’t formed yet, so nudge harder.

APEX is that intuition, expressed as a loss.

The demo

Run 400 tokens through an 8-expert router. Toggle APEX on and off. Watch what happens.

With APEX on, three things happen. Fewer dead experts. Each surviving expert leans toward one topic (the colored share on the load bar shows which). Mean specialization goes up.

With APEX off — the standard load-balance loss — every expert sees every topic in roughly equal share. Load is balanced. Nothing else is. The whole point of MoE is that experts specialize; the standard loss actively prevents that.

What the loss actually looks like

The load-balancing loss is a weighted sum where the weights adapt per-token, per-step, based on router confidence.

For each token, the router outputs a probability distribution over experts. Take the top-2 gap — the difference between the top expert’s probability and the second’s. Call that c, for confidence.

• The auxiliary loss’s contribution from this token is scaled by 1 / (1 + c). High-confidence tokens contribute weakly. Low-confidence tokens contribute strongly.
• The main loss (task loss) is unaffected. APEX only touches the auxiliary term.

That’s the whole change. Twelve lines of code. Comes for free at training time — a rounding-error compute cost. And it removes an assumption that was hiding in every MoE paper for years.

What it actually buys you

Not a jaw-dropping perplexity number. A few points at matched compute, real but not staggering. If you were chasing a headline benchmark result, APEX by itself won’t get you a headline.

The gains that matter to me are structural:

• Fewer dead experts as you scale expert count. This is the number that lets you actually use large MoE models — a lot of production MoE deployments quietly waste half their experts.
• Cleaner emergent specialization, visible earlier in training. If you probe individual expert activations, you can name what they’re for. That’s not something you can typically do with standard load-balance.
• Better sample efficiency in the router. Because the loss lets router confidence shape training, the router gets to trust its own decisions and refine them, instead of being constantly overridden.

The point I want you to leave with

The whole reason I care about this paper isn’t the number. It’s the principle underneath.

Auxiliary losses that try to shape a distribution should be conditioned on the distribution. Static, batch-averaged auxiliary losses throw away information they didn’t have to.

Once you notice that principle, you see it everywhere:

• KL regularization in VAEs uses a fixed β. Should be per-token.
• Contrastive losses use a fixed temperature. Should be adaptive.
• Dropout uses a fixed rate. Should be per-activation.

Not every one of those is worth chasing. But the shape of the argument is the same. If a loss is trying to reshape a distribution, don’t wave it around uniformly. Wave it harder where the distribution’s wrong, softer where it isn’t.

What I’d do next

• APEX + expert pruning. If the router is confidently ignoring an expert, kill it. Reallocate the parameters. This is a real trick nobody’s fully pulled off yet.
• Curriculum on the confidence weighting. Early in training, use a flatter weighting (help the router bootstrap). Later, use a sharper weighting (let specialization form). We used a fixed schedule; a learned one would probably do better.
• Actually deploy this in a large MoE at frontier scale. The paper is on smaller models. Someone with more compute than me should try it on a real 100B+ MoE. Please.

Full paper: APEX on SSRN.

Research referenced in this post

• Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer
  Shazeer et al. · ICLR 2017 · the MoE paper
• Switch Transformer: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
  Fedus, Zoph, Shazeer · JMLR 2022 · Switch Transformer, standard load-balance loss
• Mixture-of-Experts with Expert Choice Routing
  Zhou et al. · NeurIPS 2022 · alternate routing scheme
• A Closer Look into Mixture-of-Experts in Large Language Models
  Lo et al. · 2024 · analysis of dead experts and specialization
• APEX: Adaptive Per-Token Expert Loss
  Challagundla, Challagundla · 2025 · this paper