A Routing Testbed for MoEs, Part 1: Measuring Routing Effectiveness

The first post of this series revisited the central promise of sparse MoE models: sparsity decouples model quality from cost and allows training high quality models for less compute. But one aspect that framing quietly assumes is that the total parameters are actually put to good use. In a recent paper that I co-authored with colleagues at Amazon, MoE Routing Testbed, we argue that whether an MoE can fully realize its potential depends heavily on two properties of the routing, load balance and expert specialization, and we propose a testbed to measure that. In this (two-part) post, I will walk you through that paper.

Balance and specialization determine the sparsity benefit

A sparse MoE with $A$ active parameters and $N$ total parameters comes with the promise of enabling better model quality than a dense model with $A$ parameters. In the worst case, which is the router collapsing to route every token to just a single expert (or fixed set of $k$ experts) each time, the quality in fact collapses to just that dense counterpart (or worse), making it the lower bound for the MoE quality, or at least a baseline it should certainly beat. A dense model with $N$ parameters is an upper reference, and presumably an ambitious one that the MoE will never reach given that it limits FLOPs per token to $A$. The interesting question of practical MoE training is where exactly between the dense quality of $A$ and $N$ the MoE will land.

Our paper builds on the idea that there are two key properties of the learned routing that determine where on that model quality spectrum the MoE will land:

• Load Balance: The router should distribute tokens equally across the experts. This is the well-known thing about MoE training, usually discussed in the context of training stability (a collapsed router can throw off the whole training) or efficiency (the more balanced the load, the more efficient the EP setup). But beyond these, balance is also important for model quality more generally. An imbalanced load that over-utilizes some experts implies that other under-utilized experts receive fewer tokens, and hence less training data, leaving their parameters under-trained compared to the other experts. As a result, the full model capacity $N$ is not leveraged as much as in a perfectly balanced case.

• Expert Specialization: The router should route tokens in a way that allows experts to specialize. This is a somewhat overloaded term in the MoE literature, referring to different concepts and even more proposed metrics to capture them. For our purposes here, we define good specialization as avoiding duplication: The additional parameters a sparse model has are maximally useful if they are used to learn distinct things. That happens mostly automatically in a dense model, but in a sparse MoE it depends also on the router. Only if it learns to separate the input space in a meaningful way, each expert can learn a distinct slice of the overall problem, while a less optimal routing can lead to duplicated learning across the experts. Such duplication wastes model capacity and lowers the effective capacity $N$ to some $N’<N$, reducing as a result the quality benefit over $A$ that the sparsity brings.

The first goal is what the well-known balancing techniques for MoE training address. The second is just as important, but much harder to even see clearly. And that difference in visibility is what motivated us to set up a special training setup for MoE models in which expert specialization can be directly quantified – the MoE routing testbed.

Specialization needs a reference

Specialization is a claim about what an expert has learned, but what we can observe easily is where the router sends tokens. The two are connected: if an expert receives tokens from all over the data distribution, it cannot specialize on any coherent subset; if it receives a coherent subset, it at least has the chance to. So routing behavior is the natural place to look for specialization and that’s also how most previous work approached it. The problem is that for real training data, we don’t know what coherent subsets exist at the token level, and therefore we cannot say whether any given routing pattern is good, bad or just different.

The solution is to control the data. If the training mix is composed of clearly separated domains, we have a strong prior for what an ideal router should do: send tokens from the same domain to the same group of experts, so that each group can specialize on its domain. That immediately allows us to measure specialization in a very intuitive way: Given the tokens that are being routed to an expert, we can simply calculate to what extent they come from a single domain – the higher that fraction, the better specialized the expert is.

The domains even provide a reference routing, which is a routing that would route all tokens from each domain cleanly to a dedicated (set of) experts. For that reference to be meaningful, it should in fact be a good routing in the sense of leading to a very well-trained MoE on that training data: ideally known to be optimal, and failing that, at least demonstrably better than what routers learn on their own. How much of that certainty we can get depends entirely on the data, which made the data mix the central design decision of the testbed.

Designing the data mix

We considered various options to build data mixes for this testbed that represent different trade-offs between two core goals: knowing what optimal routing looks like vs. being close to realistic LLM pretraining data. The following figure illustrates that:

Previous work such as Dikkala et al. (2023) sits at one end of the spectrum: With fully synthetic data sampled from mixtures of Gaussians, the optimal routing is known by construction and can be used to assess any learned routing. But this data is far from what MoEs work with in the real LLM pretraining case, namely natural language, and raises questions about whether findings from this setup transfer. Synthetic data such as sentences sampled from context-free grammars as used in the Physics of Language Models line of work get closer to real language, but also lose the property of having a known optimal routing. Text classification datasets bring natural text whose labels give a routing prior, but they are orders of magnitude too small to train language models on.

The other end of the spectrum is the actual target use case, a large and diverse pretraining mix as used to train general-purpose LLMs. That is the real data we care about, but we have no idea about the size, shape and number of slices of that data that different experts should specialize in to obtain the optimal language modeling performance. Taking distinct domains from such a corpus and mixing them in equal proportions introduces a stronger prior for how experts should specialize, but domains are a weak signal for the individual token: many tokens of a news article could equally appear in a novel.

We landed on a data mix that we think offers the most interesting trade-off, given that you can’t have fully realistic data with a known optimal routing anyway: Wikipedia in eight languages with different scripts: German, Korean, Japanese, Chinese, Hebrew, Thai, Hindi and Arabic. Language is a domain boundary that survives down to the token level, and using distinct scripts sharpens it further, since tokens from different scripts mostly don’t even share surface forms. At the same time, every sequence is real human text. This provides strong token-level priors for how tokens should be routed, as I show with experiments later in this post. We believe this data option is also the most interesting addition to the existing literature: While Dikkala et al. occupy the left end of the spectrum, and various other papers sit on the right end by trying to measure specialization in the real setup (see 2.3 in the paper), we are, to my knowledge, the first to build a testbed from the middle of that spectrum.

Two structural decisions complete the mix: Each domain gets exactly 1/8 of the mix, so that all domains are equally important and – crucially – a perfectly balanced expert load and a perfectly specialized routing are compatible with each other by construction. And the model is matched to the mix: we pair the 8-domain data with 32 experts under top-4 routing, so that the experts divide evenly into eight groups of four, one group per domain.

Not every token has a domain

The clean picture above has one flaw: not all tokens are domain-specific, even in this mix. Digits, punctuation, whitespace and many named entities appear across all eight languages. A reference that forces every token into its sequence’s domain group has to make arbitrary calls for these generic tokens – the digit “7” in a Thai article is not meaningfully Thai. Naively enforcing domain routing for all tokens builds that arbitrariness into the reference, and as we’ll see later, that measurably hurts.

We therefore split tokens into a domain-specific set and a generic set. A small feed-forward classifier is trained to predict the domain from the token embedding, reaching 79% test accuracy¹ on that dataset, and its confidence serves as a proxy for domain-specificity: tokens it classifies correctly with high confidence are treated as domain-specific, the rest as generic. The reference routing then prescribes domain-aligned routing only for the domain-specific tokens and leaves the generic ones to be routed freely.

Measuring against the reference

With domains attached to tokens, two metrics fall out naturally that capture the two goals I introduced initially:

• Routing purity captures expert specialization: for each expert, take the fraction of its received tokens that come from its most common domain, and average across experts: $\frac{1}{E} \sum_{i=1}^{E} \max_{d} \frac{f_{id}}{f_i}$ where $E$ is the number of experts, $f_i$ the share of all token-expert assignments that go to expert $i$ and $f_{id}$ the share that goes to expert $i$ and comes from domain $d$. A purity of $1$ means every expert serves exactly one domain; with $D$ domains the lower end is $1/D$, which is where a routing that ignores domains entirely ends up. Note that purity does not check which expert serves a domain – tokens from the same domain should travel together, but which expert hosts them is irrelevant.

• Expert utilization captures the load balancing goal: it is maximal when all experts receive equal token counts and decreases as load concentrates: $\sum_{i=1}^{E} \min\left(f_i, \frac{1}{E}\right)$ reflecting that an over-used expert implies under-utilized ones elsewhere and hence reduces the effective total capacity of the model due to under-training some experts. Since top-$k$ routing sends every token to $k$ distinct experts, the worst case is a collapse onto a fixed set of $k$ experts, which puts the lower end at $k/E$.

Expert utilization can be computed simply from the routing statistics of how many tokens go to each expert. It is also independent of our testbed. Purity, on the other hand, looks at how the tokens are routed, i.e. which specific tokens go where, not just pure statistics, and for that it relies on our testbed setup and the domain-induced reference routing. We compute it only over the domain-specific tokens, ignoring the generic ones. Note that matching the model to the mix, i.e. $D = E/k$, givens both metrics the same range as their lower ends coincide at $1/D = k/E$ so the common range is $[1/D, 1]$. Both metrics are computed over the full global batch, which is an important detail I’ll discuss in Part 2 of the post.

Verifying the reference routing

Having given up the certainty of the synthetic Gaussian world, we need to make the empirical case instead: a reference routing is only a meaningful comparison target if following it is actually better than anything the model learns on its own – otherwise “closeness to the reference” measures closeness to a mediocre policy.

We verify this empirically by training MoEs on the data mix while constraining the routing to follow the domain structure. Router logits are masked so that domain-specific tokens can only select experts from their domain’s assigned group, while generic tokens route freely. Sweeping the classifier-confidence threshold that splits the tokens into domain and generic moves the setup continuously from fully learned routing (no tokens constrained) to naive domain routing (all tokens constrained), changing the achieved validation loss:

Naive domain routing – forcing every token into its sequence’s domain group – lands at the highest validation loss of the sweep, and higher than many fully learned routings, confirming that arbitrary domain assignments for generic tokens damage the model. Learned routing itself spans a wide range depending on the balancing configuration that Part 2 will discuss. The reference with a roughly 50:50 split, constraining the half of tokens the classifier is most confident about, reaches the lowest loss, below every fully learned configuration we trained. That makes it a useful reference and we can meaningfully measure routing purity against that target.

What’s next

We now have a routing testbed: a data mix with clearly defined domains, an empirically validated reference routing and two metrics – routing purity and expert utilization – that make the router’s two competing goals separately visible instead of collapsed into a single loss number. In Part 2, I show how that can clearly visualize the effects of various popular token balancing methods on learned routing and that scope is what matters most, much more than the actual balancing approach.