General • • ardatugsat@gmail.com

Mixture-of-Experts in Large Language Models: Routing Mechanisms, Capacity Constraints, and the Load Balancing Problem

Abstract

Mixture-of-Experts (MoE) architectures have emerged as one of the most computationally compelling approaches to scaling large language models without proportional increases in inference cost. By conditionally activating only a subset of model parameters per forward pass, MoE models achieve a favorable tradeoff between total parameter count and floating-point operations (FLOPs) per token. This post provides a rigorous technical analysis of sparse MoE systems as applied to transformer-based language models. We examine the routing mechanism in depth — from the original softmax-gating formulation to the token-choice and expert-choice variants — and analyze the central challenge of load imbalance, wherein a small number of experts absorb disproportionate routing mass. We survey auxiliary loss formulations designed to enforce load balance, discuss capacity collapse as a failure mode, and evaluate the empirical performance of recent MoE LLMs including Switch Transformer, GLaM, Mixtral 8x7B, and DeepSeekMoE. The post concludes with open research questions around expert specialization, routing interpretability, and the tension between sparsity and training stability.

1. Introduction

The prevailing wisdom in large language model scaling has been articulated most clearly by Kaplan et al. (2020): model performance improves as a power law with respect to parameter count, dataset size, and compute budget. However, this scaling paradigm assumes dense models — ones where all parameters participate in every forward pass. This assumption is computationally expensive: doubling the parameter count doubles inference cost.

Mixture-of-Experts architectures offer an escape from this constraint. The core idea is elegant: replace each dense feed-forward sublayer in a transformer with a collection of $N$ specialized sub-networks (experts), and route each input token through only $k \ll N$ of them. The activated fraction $k/N$ determines the sparsity level and, consequently, the ratio of total parameters to active parameters per forward pass.

This architectural choice has profound implications. A model with 8 experts per layer and $k=2$ active experts can achieve the total parameter count of an 8� larger dense model while using only 2� the FLOPs per token during inference. The promise is significant: sub-linear scaling of compute with respect to parameters.

Yet MoE models are not simply “free lunches.” They introduce a cluster of interrelated challenges: load imbalance (some experts receive far more tokens than others), capacity overflow (tokens may be dropped if an expert’s buffer is full), routing collapse (a degenerate state where routing concentrates on very few experts), and distributed training complexity (experts must be sharded across devices, requiring all-to-all communication patterns).

In this post, I analyze MoE architectures at the level of their mathematical formulation, examine the failure modes that arise in practice, and survey the mitigation strategies proposed in the literature. The goal is to provide a technically grounded understanding of why MoE models behave as they do — not to advocate for or against a particular implementation.

2. Related Work

The conceptual foundations of mixture-of-experts trace back to Jacobs et al. (1991), who proposed gating networks for combining specialized submodels in a supervised learning context. Jordan and Jacobs (1994) extended this to a hierarchical MoE framework and provided an EM-based training algorithm. These early formulations were applied to relatively shallow models and small datasets, and the approach lay relatively dormant until the modern deep learning era.

The key paper that revived MoE for deep learning at scale is Shazeer et al. (2017), “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.” This work introduced the sparsely-gated MoE layer into LSTMs, using a noisy top-$k$ gating mechanism and an auxiliary load-balancing loss. The paper demonstrated that MoE layers could achieve better perplexity than dense baselines at equal inference cost, establishing the fundamental efficiency tradeoff that motivates all subsequent work.

Fedus et al. (2022) — Switch Transformer — simplified the gating to top-1 routing (i.e., each token routes to exactly one expert) and scaled the approach to transformer architectures, achieving competitive performance with significantly fewer FLOPs than T5 dense models. The simplification from top-$k$ to top-1 brought both benefits (reduced all-to-all communication) and costs (less redundancy, higher routing sensitivity).

Du et al. (2022) introduced GLaM (Generalist Language Model), a 1.2 trillion parameter MoE model with top-2 routing that outperformed GPT-3 on multiple few-shot benchmarks while using approximately one-third of the energy consumption during inference. GLaM represents perhaps the most compelling demonstration of the efficiency proposition of MoE at scale.

Lepikhin et al. (2021) — GShard — addressed the distributed systems challenge of MoE, introducing an expert-parallel sharding strategy that enabled scaling to 600 billion parameters across 2048 TPU cores. This work is notable for the engineering formalism it brings to MoE: explicit capacity factors, auxiliary loss terms, and dispatch mechanisms are all specified with the precision needed for large-scale implementation.

Most recently, Jiang et al. (2024) released Mixtral 8x7B, a sparse MoE model that achieves performance competitive with Llama 2 70B while activating only 12.9 billion parameters per forward pass out of a total of 46.7 billion. Mixtral demonstrated that open-weight MoE models are viable at the level of quality previously associated only with large dense proprietary systems.

3. Technical Analysis

3.1 The Routing Mechanism

Consider a transformer layer with feed-forward sublayer replaced by an MoE module. Given a token representation $\mathbf{x} \in \mathbb{R}^d$, the router computes a probability distribution over $N$ experts:

$$G(\mathbf{x}) = \text{Softmax}(\mathbf{x} W_g)$$

where $W_g \in \mathbb{R}^{d \times N}$ is the gating weight matrix. Under top-$k$ sparse routing, only the $k$ experts with highest gate values are activated:

$$\hat{G}(\mathbf{x}) = \text{TopK}(G(\mathbf{x}), k)$$

The output of the MoE layer is then:

$$\text{MoE}(\mathbf{x}) = \sum_{i \in \mathcal{T}(\mathbf{x})} \hat{G}_i(\mathbf{x}) \cdot E_i(\mathbf{x})$$

where $\mathcal{T}(\mathbf{x})$ is the set of top-$k$ expert indices and $E_i$ is the $i$-th expert network (typically a two-layer MLP). The gate values serve both as a weighting mechanism for the expert outputs and as a differentiable routing signal that propagates gradients back to $W_g$.

An important subtlety: the $\text{TopK}$ operation is not differentiable with respect to which experts are selected. Gradients only flow through the continuous gate values $\hat{G}_i$, not through the discrete selection. This means the router cannot directly learn from the downstream quality of its routing decisions — it can only adjust the magnitude of selected experts’ contributions, not the selection itself. This is a fundamental limitation of the straight-through-estimator-free formulation used in most MoE implementations.

3.2 Load Imbalance and the Collapse Problem

The most persistent challenge in training MoE models is load imbalance. Formally, define the load $L_i$ on expert $i$ across a batch of $T$ tokens as:

$$L_i = \sum_{t=1}^{T} \mathbf{1}[i \in \mathcal{T}(\mathbf{x}_t)]$$

In a balanced system, $L_i \approx kT/N$ for all $i$. In practice, routing tends to concentrate: once an expert begins receiving more tokens, it processes more data, updates its weights more frequently, and often becomes better at a certain type of input — which makes the router send even more tokens to it. This positive feedback loop leads to a regime called routing collapse, where one or a small number of experts handles the majority of tokens while the rest receive nearly zero traffic and consequently fail to train.

Routing collapse is particularly insidious because it is stable: once collapsed, the gradient signal from the auxiliary loss may be insufficient to overcome the inertia of a highly trained specialist expert versus undertrained generalist ones.

3.3 Auxiliary Load-Balancing Loss

Shazeer et al. (2017) proposed an auxiliary importance loss to combat imbalance. Define the soft importance of expert $i$ as:

$$\text{Imp}(i) = \sum_{t=1}^{T} G_i(\mathbf{x}_t)$$

The coefficient of variation squared, $\text{CV}(\text{Imp})^2$, provides a scalar measure of imbalance. Minimizing this term encourages all experts to receive similar total routing probability. However, this soft importance measure can diverge from actual load (hard token counts), since high-probability tokens that aren’t selected still contribute to soft importance.

Fedus et al. (2022) proposed a more direct formulation. Define the fraction of tokens dispatched to expert $i$:

$$f_i = \frac{1}{T} \sum_{t=1}^{T} \mathbf{1}[i \in \mathcal{T}(\mathbf{x}_t)]$$

and the mean gate probability assigned to expert $i$:

$$P_i = \frac{1}{T} \sum_{t=1}^{T} G_i(\mathbf{x}_t)$$

The Switch Transformer auxiliary loss is:

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

where $\alpha$ is a hyperparameter controlling the strength of the regularization. This loss has a minimum when $f_i = 1/N$ for all $i$ (perfect balance). Crucially, $f_i$ is non-differentiable but $P_i$ is, so the loss provides a differentiable signal that encourages $P_i$ to be uniform, which in turn nudges token assignments toward balance.

The choice of $\alpha$ is delicate. Too small and collapse occurs; too large and the auxiliary loss dominates the language modeling objective, harming model quality. Empirically, Switch Transformer found $\alpha = 10^{-2}$ to work well for top-1 routing, while top-2 routing with GShard required $\alpha = 10^{-3}$.

3.4 Capacity Factors and Token Dropping

In distributed MoE training, each expert has a fixed capacity — the maximum number of tokens it can process in a single forward pass. Define the capacity $C$ as:

$$C = \left\lfloor \frac{k \cdot T}{N} \cdot \text{cf} \right\rfloor$$

where $\text{cf} \geq 1$ is the capacity factor. A capacity factor of 1.0 corresponds to the theoretical uniform allocation; factors above 1.0 provide headroom for imbalance. When a token is routed to an already-full expert buffer, it is dropped: its contribution to the layer output is simply its residual stream value, unmodified by any expert computation.

Token dropping is a significant concern. Dropped tokens represent computation that was budgeted but not performed, effectively reducing the model’s capacity to process certain inputs. The frequency of token dropping is directly tied to load imbalance: in a well-balanced system with $\text{cf} = 1.25$, dropping is rare. In a collapsed system, dropped tokens can become a significant fraction of the batch.

3.5 Expert-Choice Routing

Zhou et al. (2022) proposed an inversion of the standard routing paradigm: rather than having tokens choose their experts (token-choice), have experts choose their tokens (expert-choice). Each expert selects the top-$C$ tokens most relevant to it, based on the same gating scores:

$$\mathcal{S}_i = \text{TopC}\left(\{G_i(\mathbf{x}_t)\}_{t=1}^T, C\right)$$

Expert-choice routing guarantees perfect load balance by construction — every expert processes exactly $C$ tokens. However, it introduces a different pathology: heterogeneous token coverage. Some tokens may be selected by multiple experts; others by none. For language modeling, where every token must produce an output, unselected tokens again fall back to passing through the residual stream unchanged.

The expert-choice formulation also changes the semantics of the model: under token-choice, a token has agency in selecting its processing pathway; under expert-choice, an expert has agency in selecting which tokens it processes. Neither formulation is strictly superior, and the choice has downstream consequences for what the experts learn to specialize in.

3.6 DeepSeekMoE and Fine-Grained Expert Decomposition

Dai et al. (2024) introduced DeepSeekMoE, which proposes a fine-grained expert segmentation strategy. Rather than having a small number of large experts ($N$ small, $d_{\text{ffn}}$ large), DeepSeekMoE uses a large number of small experts ($N$ large, $d_{\text{ffn}}$ proportionally small), and activates more of them per token while preserving the total active parameter count.

The key insight is that finer-grained experts can achieve more precise combinations of knowledge, reducing redundancy between experts and improving parameter utilization. If each expert is smaller, the routing mechanism can compose more nuanced combinations without increasing active FLOPs. DeepSeekMoE-16B was shown to match the performance of LLaMA2-7B dense models with approximately 40% of the active parameters.

4. Discussion

4.1 Expert Specialization: Myth or Reality?

A natural hypothesis is that MoE experts develop semantic specializations — one expert handling syntax, another semantics, a third factual recall, and so on. The empirical evidence is mixed. Fedus et al. (2022) found that Switch Transformer experts do exhibit positional preferences (some experts handle tokens at specific positions more frequently) and token-type preferences (punctuation versus content words). However, the extent to which these preferences constitute meaningful semantic specialization versus superficial statistical regularities is unclear.

Interpretability work on MoE models is in its early stages. The routing decision is made by a linear projection of the residual stream, which means that routing is determined by linear features of the token representation at that layer. Given what we know about how information is structured in transformer residual streams (Elhage et al., 2022), this suggests that routing specialization is likely to reflect the kinds of linear features that are prominent at each layer — predominantly syntactic at earlier layers, more semantic at later layers.

4.2 Training Stability and the Role of Noise

Shazeer et al. (2017) proposed adding tunable Gaussian noise to the gating logits before applying TopK:

$$\tilde{G}(\mathbf{x}) = \text{Softmax}\left(\mathbf{x} W_g + \epsilon \cdot \text{Softplus}(\mathbf{x} W_{\text{noise}})\right)$$

where $\epsilon \sim \mathcal{N}(0, 1)$. The noise serves as an exploration mechanism: it prevents the router from early on committing to a fixed set of experts and gives lower-ranked experts occasional opportunities to process tokens and update their parameters. At inference time, the noise is typically removed. The relationship between this training-time noise injection and the final quality of the trained model remains an active area of research.

4.3 MoE vs. Dense: When Does Sparsity Help?

The efficiency advantage of MoE over dense models is clearest when comparing models at equal inference FLOPs. At equal FLOPs, MoE models consistently outperform dense models because they have access to more total parameters — more knowledge storage capacity — for the same computation. The empirical scaling law for MoE models appears to shift the dense scaling curve favorably: for a given FLOPs budget, an MoE model tends to achieve lower loss than a dense model trained to the same compute budget.

However, the picture is more nuanced when considering training efficiency. MoE models with $N$ experts per layer require roughly $N$ times more memory for expert parameters, and the all-to-all communication pattern in expert-parallel training introduces significant overhead. The compute-to-memory bandwidth ratio of modern accelerators means that MoE models are often memory-bound rather than compute-bound in practice, partially negating the FLOPs savings.

4.4 Open Problems

Several fundamental questions remain unresolved:

5. Conclusion

Mixture-of-Experts architectures represent one of the most practically important ideas in modern large language model design. The central promise — decoupling total parameter count from per-token compute — has been validated empirically across a range of scales, from Switch Transformer to Mixtral 8x7B. However, the realization of this promise requires careful engineering: load-balancing losses, capacity factors, and routing formulations that avoid collapse while enabling genuine specialization.

The routing mechanism is the heart of MoE, and it remains underspecified by current theory. We know empirically that top-$k$ routing with auxiliary load-balancing losses works, but we do not have a satisfying theoretical account of what the trained routing function represents or how to design it from first principles. As MoE models become more prevalent — both in research and in production — this gap between empirical practice and theoretical understanding will become increasingly important to close.

For practitioners working with MoE models, the key takeaways are: (1) monitor expert load statistics during training and treat imbalance as an early warning sign; (2) tune $\alpha$ (the auxiliary loss coefficient) carefully and consider decaying it over training; (3) capacity factors above 1.0 reduce dropped tokens at the cost of increased memory, and the right tradeoff depends on the degree of imbalance; and (4) expert-choice routing eliminates load imbalance by construction but introduces uneven token coverage, which may be preferable or harmful depending on the task.

References

Speculative Decoding for Large Language Model Inference: Mechanisms, Theory, and Empirical Tradeoffs
In-Context Learning as Implicit Bayesian Inference: Theory, Evidence, and Open Problems

Leave a Comment

Your email address will not be published. Required fields are marked *