Mixture of Experts vs Mixture of Tokens: Two Sparsity Strategies Compared

Mixture of Experts vs. Mixture of Tokens: How Two Sparsity Strategies Differ

If you've been serving large LLMs lately, MoE (Mixture of Experts) is hard to avoid. Mixtral, DeepSeek-V2, Qwen-MoE — there's a reason models targeting lower inference costs gravitate toward MoE architectures. The logic is straightforward: instead of activating all parameters for every token, activate only a subset of experts and cut the compute.

Mixtral 8x7B is the canonical example. Total parameters: 47B. Active parameters per token: roughly 13B. The attention layers (~5B) are shared, and the FFN uses only 2 of 8 experts (top-2). Compute is on par with a 13B dense model, but the model's knowledge capacity is that of a 47B model. You still need to load all 47B into memory, but GPU compute time is dramatically cheaper.

The hard part is making that "activate only 2" decision. The structural tension that decision creates is what makes MoE serving difficult — and it's the motivation behind the idea called Mixture of Tokens (MoT).

The Structural Tensions in MoE Top-k Routing

In MoE, routing is token-driven. Each token passes through a router (gate network) that scores candidate experts and selects the top-k. Experts that aren't selected are bypassed entirely for that token.

The first recurring tension this creates is load imbalance. Similar tokens tend to prefer the same experts, so within a batch, certain experts get flooded while others sit idle. In extreme cases this degenerates into expert collapse — some experts receive no tokens at all. Switch Transformer (Fedus et al., 2022) reported that applying top-1 routing without an auxiliary loss causes exactly this kind of collapse.

The standard fix is to add an auxiliary load balancing loss during training:

$$\mathcal{L}{aux} = \alpha \cdot N \sum{i=1}^{N} f_i \cdot P_i$$

where $f_i$ is the fraction of tokens routed to expert $i$ and $P_i$ is the fraction of router probability mass assigned to expert $i$. The loss nudges the two distributions to align. The catch: this only operates at training time. If the input distribution at inference differs from the training distribution, routing bias re-emerges.

The second problem is broken gradient flow. Top-k hard routing — pick k, discard the rest — is non-differentiable. Experts that aren't selected receive no gradient. Over repeated training steps, frequently selected experts grow stronger while the rest atrophy, creating polarization. Workarounds like STE (Straight-Through Estimator) and noisy top-k exist, but they don't change the underlying structure.

Flipping the Routing Direction — Mixture of Tokens

Mixture of Tokens (MoT) (Antoniak et al., NeurIPS 2024) flips who drives routing. Instead of tokens choosing experts, each expert decides which combination of tokens from the entire batch it processes.

The mechanism works as follows. A controller computes a weight matrix $\Phi \in \mathbb{R}^{E \times T}$ over all tokens in the batch for each expert ($E$ = number of experts, $T$ = total tokens in batch). The input to expert $e$ is not a single token but a linear combination of all tokens:

$$\tilde{x}e = \sum{t} \Phi_{e,t} \cdot x_t$$

And symmetrically, the update for token $t$ is a linear combination of all expert outputs:

$$y_t = \sum_{e} \Phi_{e,t} \cdot \text{Expert}_e(\tilde{x}_e)$$

This is why the paper calls it "cross-example aggregation": tokens from different sequences (examples) in the batch are mixed together before being fed to a single expert.

This soft aggregation structurally addresses both tensions described above. All experts are always active, so gradient flows uniformly across all of them. Because $\Phi$ is softly normalized, the situation where tokens pile up on one expert simply cannot arise — load imbalance is structurally prevented without any auxiliary loss. The entire computation is differentiable, which improves training stability. The paper reports a 3× training speedup over a dense Transformer while matching conventional MoE quality.

What Each Approach Actually Addresses

The auxiliary loss in MoE doesn't fully close the gap between training and inference. This is especially pronounced in prefill/decode disaggregation setups: during the decode phase, where token count per batch is tiny, imbalance worsens. Some expert GPUs are overloaded; others are nearly idle.

MoT simply doesn't have this tension. The caveat is that cross-example mixing degrades when batch sizes are small. With only one sequence in the batch, the benefit of mixing tokens across different examples disappears — you're only mixing tokens within that one sequence.

Serving Perspective — The Real Cost of All-to-All Communication

For serving engineers, the hard part of MoE isn't the parameter count — it's the communication overhead from expert parallelism. When experts are distributed across multiple GPUs, you need a dispatch step to send tokens to the right GPU and a combine step to gather results back. Both are implemented as all-to-all collective operations. That's two all-to-all calls per forward pass, and four per MoE layer during training (including the backward pass).

Measured numbers make this concrete. When running DeepSeek-V2-Lite inference on SGLang with 8-GPU expert parallelism, up to 59.2% of the MoE layer forward pass latency came from expert parallelism communication. More time spent shuffling tokens between GPUs than doing the actual expert computation.

MegaBlocks (Gale et al., 2022) achieves 1.8×–2.4× training throughput improvement through GPU kernel-level optimizations for sparse MoE, but it doesn't eliminate all-to-all communication — it reduces the cost by overlapping communication with compute. vLLM handles this by supporting multiple all-to-all backends (NCCL EP, DeepEP HT/LL, etc.) as pluggable options.

How this changes with MoT at large serving scale isn't yet well-characterized empirically. Cross-example aggregation can be implemented as dense matrix operations similar to attention, which opens the possibility of reusing FlashAttention-style kernels instead of dispatching tokens to other GPUs. If that pans out, it would be a structural bypass of MoE's chronic all-to-all overhead.

The difference also shows up in continuous batching. In MoE, if one expert GPU becomes a straggler, the entire batch latency is gated by that expert. MoT distributes mixed tokens evenly across experts by design, so this problem doesn't arise structurally.

Current Limitations and Practical Judgment

MoT's barriers to production deployment aren't just a lack of real-world case studies. Because cross-example aggregation mixes tokens from different examples within the same batch, implementing prefill-decode separation with KV cache reuse gets complicated. The paper states compatibility with autoregressive generation, but how it integrates with vLLM's PagedAttention and continuous batching optimizations requires separate validation.

Memory is another consideration. Soft aggregation requires maintaining the weight matrix $\Phi$ across the entire batch, so activation memory scales with batch size. Hard-routing MoE only needs to maintain the k active expert paths, so its activation memory footprint is much smaller.

If you're running MoE today and want to measure expert utilization, per-layer token distribution histograms are the most direct signal. Enable expert parallelism in vLLM with --enable-expert-parallel, then either collect router logits via a custom hook or pull per-expert token counts from SGLang's internal statistics to see actual imbalance in your serving environment.

How quickly MoT absorbs the serving optimization ecosystem — PagedAttention, FlashAttention, continuous batching — will determine whether the idea becomes practical.