Routing in Mixture of Experts (MoE) is a core mechanism in modern Machine learning architectures that dynamically assigns each input token to a subset of specialized expert networks. Instead of activating all parameters for every input, a router, or gating network, learns to direct tokens to the most relevant experts, balancing model capacity with computational efficiency. This approach allows Large language models to scale to trillions of parameters while keeping inference costs manageable, as only a fraction of the network is active per token.
The concept originates from the early 1990s work on adaptive mixtures of local experts, but it gained prominence in the 2010s with the introduction of sparsely-gated MoE layers in deep learning. Modern implementations, such as those used in models from Google DeepMind and OpenAI, rely on learned routing functions that output probability distributions over experts, often incorporating auxiliary losses to ensure balanced usage. Routing has become a critical design choice in state-of-the-art generative AI systems, influencing both training efficiency and final model quality.
Historical Development
The idea of routing inputs to specialized components predates deep learning, with roots in ensemble methods and modular neural networks. In 1991, Michael I. Jordan and Robert Jacobs published foundational work on hierarchical mixtures of experts, where a gating network weighted the outputs of several feedforward networks. This early framework established the core principle: learn to partition the input space so that different experts handle different regions.
In 2017, researchers at Google DeepMind (then Google Brain) introduced the sparsely-gated MoE layer in a paper led by Noam Shazeer. This design used a softmax gating function to select the top-k experts for each token, with k typically set to 1 or 2. The authors demonstrated that MoE layers could scale to hundreds of experts, achieving state-of-the-art results on language modeling and machine translation benchmarks while reducing computational cost per example. This work directly influenced later architectures like the Switch Transformer (2021), which simplified routing to select only one expert per token, and the GShard framework for massively parallel training.
Routing Mechanisms
Routing functions can be categorized into several types. The most common is token-choice routing, where each token independently selects its top-k experts based on a learned score. This is implemented as a linear projection of the token's hidden state followed by a softmax over expert indices. The selected experts then process the token, and their outputs are weighted by the routing probabilities and summed.
An alternative is expert-choice routing, where each expert selects the top-k tokens from a batch, ensuring that all experts receive a minimum load. This approach, popularized in the Switch Transformer and later refined in models like Mixtral, addresses load imbalance issues but requires careful handling of token-to-expert assignments.
Other variants include hashing-based routing, which uses deterministic hash functions to assign tokens without learned parameters, and hierarchical routing, where a two-level gating system first selects a group of experts and then a specific expert within that group. Each method trades off flexibility, computational overhead, and training stability.
Load Balancing and Auxiliary Losses
A major challenge in MoE routing is load imbalance: the router may learn to always select a few popular experts, leaving others undertrained. To mitigate this, most implementations add an auxiliary load-balancing loss that penalizes uneven distribution of tokens across experts. The standard formulation, introduced in the 2017 Shazeer paper, computes the fraction of tokens routed to each expert and encourages uniformity via a cross-entropy term.
More recent approaches, such as the Switch Transformer's load balancing loss, use a simpler coefficient that multiplies the dot product between the router's average probability and the actual token count per expert. Some systems, like GShard, employ capacity factors that limit how many tokens each expert can process, forcing the router to distribute work. These techniques are essential for stable training and preventing expert collapse, where some experts become dead weights.
Architectural Integration
Routing is typically applied within transformer blocks, replacing the feedforward network (FFN) with an MoE layer. In a standard Transformer (architecture), each token passes through a multi-head attention sublayer followed by a position-wise FFN. In an MoE transformer, the FFN is replaced by a set of expert FFNs, each with its own parameters, and a router that selects which experts to activate.
This integration allows models to have a large total parameter count while keeping the active parameter count per token constant. For example, a model with 64 experts, each with 1 billion parameters, has 64 billion total parameters but only activates 2 billion per token if k=2. This property is crucial for scaling to trillion-parameter models, as seen in systems like the Mixture of Experts models from Anthropic and other labs.
The router itself is a small neural network, often a single linear layer followed by a softmax, with its own parameters trained jointly with the experts. Some architectures use separate routers for different layers or attention heads, and recent work explores learned routing policies that adapt based on token type or position.
Training Dynamics
Training MoE models with routing introduces unique challenges. The router's discrete decisions (selecting top-k experts) are non-differentiable, so gradients flow only through the selected experts' outputs weighted by the routing probabilities. This creates a moving target problem: as experts improve, the router's preferences shift, potentially causing instability.
To address this, researchers use techniques like noisy top-k gating, where Gaussian noise is added to routing logits during training to encourage exploration. Another approach is to use a straight-through estimator for the routing decision, treating the selection as a hard assignment but passing gradients through the probability weights. Additionally, some methods anneal the temperature of the softmax during training to gradually sharpen routing decisions.
Load balancing losses are typically weighted with a small coefficient (e.g., 0.01) to avoid dominating the primary task loss. In practice, training MoE models requires careful hyperparameter tuning, and many systems employ auxiliary losses for both load balance and router confidence.
Applications in Large Language Models
Routing in MoE has become a standard technique in large language models. Notable examples include:
- The Switch Transformer (2021) from Google, which scaled to 1.6 trillion parameters with a sparse MoE architecture.
- GShard (2020), which demonstrated efficient training of MoE models across thousands of TPU cores.
- Mixtral 8x7B (2023) from Mistral AI, which uses 8 experts with top-2 routing, achieving performance comparable to larger dense models.
- DeepSeek-V3 (2024), which employs a fine-grained MoE with 256 experts and top-8 routing.
- Several models from OpenAI and Anthropic are reported to use MoE layers, though exact details are often proprietary.
These models demonstrate that routing enables cost-effective scaling, as the compute per token remains constant even as total parameters grow. This has made MoE a preferred choice for serving large models in production environments, including cloud platforms like Amazon Web Services, Microsoft Azure, and Google Cloud.
Efficiency and Hardware Considerations
Routing introduces communication overhead in distributed training and inference, as tokens must be sent to the appropriate expert, which may reside on a different device. This all-to-all communication pattern can become a bottleneck, especially with many experts. Hardware vendors like NVIDIA (though not listed, implied) and AMD have developed optimized kernels for MoE operations, and specialized AI accelerators from Groq and SambaNova are designed to handle sparse activation patterns efficiently.
To reduce communication, some systems use expert parallelism, where experts are replicated across devices, or hybrid approaches that combine data and expert parallelism. The choice of routing granularity (token-level vs. block-level) also affects efficiency. Token-level routing offers finer control but more communication, while block-level routing groups tokens to reduce overhead.
Challenges and Future Directions
Despite its success, routing in MoE faces several open problems. One is the trade-off between expert specialization and generalization: overly specialized experts may not transfer well to new tasks. Another is the difficulty of fine-tuning MoE models, as the router may need to adapt to new data distributions. Research on adaptive routing, where the number of active experts varies per token, and on learned routing policies that consider context beyond the current token, is ongoing.
Future directions include combining routing with other efficiency techniques like Model Pruning and Quantization, and developing routers that can handle multi-modal inputs. As models continue to scale, routing will remain a key mechanism for balancing capacity and compute, with potential applications beyond language to vision and reinforcement learning.
See Also
- Transformer (architecture)
- Large language model
- Multi-Head Attention
- Model Pruning
- Google DeepMind
- OpenAI
- Anthropic
References
- Shazeer, N., et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.
- Fedus, W., et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.
- Lepikhin, D., et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.
- Jiang, A. Q., et al. (2023). Mixtral of Experts.
- DeepSeek-AI (2024). DeepSeek-V3 Technical Report.