Understanding the Mixture of Experts (MoE) Architecture:
Introduction to MoE
Mixture of Experts (MoE) is a neural network architecture that consists of multiple sub-networks (called “experts”) and a gating mechanism that routes each input to the most appropriate expert. In essence, an MoE model is like an ensemble of specialists: each expert network is trained to handle a different kind of input or feature, and a gate (or router) learns to choose which expert(s) should be “consulted” for each input. Only a subset of the experts is activated for any given input, which makes MoEs very efficient for scaling to large model sizes
Think of it this way: instead of having one gigantic model handle everything, MoE breaks the model into many smaller expert models, each specializing in certain inputs. For example, one expert might specialize in processing questions about math, another expert might focus on language translation, another on code, etc. When a new input comes in, the gating mechanism will dynamically select the expert(s) best suited for that input. This means each input is handled by a mixture of expert outputs, rather than a one-size-fits-all model. The core idea is that by dividing the problem among experts, the network can learn more specialized and efficient behaviors for different types of inputs, all while being part of one unified model.
Technical Details of MoE
In an MoE model, the architecture is structured to include a pool of expert networks and a gating network that assigns work to those experts. Concretely, an MoE layer might replace a standard layer (for instance, the feed-forward layer in a Transformer) with multiple parallel expert layers plus a learned router. Each expert is itself a neural network (often a feed-forward network) with its own parameters, and the gating network determines which expert’s output to use for a given input. The experts all receive the same input, but the gate decides, based on the input, how to weight each expert’s contribution or even to select a single expert.
Illustration of a Switch Transformer MoE layer (light blue section) replacing a standard feed-forward network. In this example, the token “More” is routed to Expert 2 (with gate score 0.65) and “Parameters” to Expert 1 (with gate score 0.8). Only the selected expert’s feed-forward computation is used for each token (solid lines), making the layer sparse and efficient. Other tokens in the batch may be routed to different experts in parallel (dotted lines).
Reference: https://arxiv.org/pdf/2101.03961.pdf
Gating mechanism:
The gating network usually produces a set of scores (or probabilities) for the experts based on the input. A common implementation is to use a small neural network (often just a single linear layer) that outputs one logit per expert, followed by a Softmax function to turn these into probabilities.
In practice, MoE models often use a sparse gating strategy – meaning the gate will assign most of the weight to one or a few top experts and zero to the rest. For example, the gate might pick the single best expert (this is called Top-1 gating) or the top two experts (Top-2 gating) for each input token. If an expert’s weight is zero, the model can skip computing that expert altogether, which saves a lot of computation. The Switch Transformer is an MoE model that uses Top-1 gating – each token is routed to only one expert – making the computation very efficient.
Expert networks:
Each expert in an MoE is a normal neural sub-network (for instance, a feed-forward layer or an MLP). These experts can be as simple or complex as needed. In many MoE implementations for Transformers, each expert is just a feed-forward network similar to the ones used in standard Transformer layers. However, one could also design more complex experts (for example, an expert itself could be a mini-neural network or even another MoE, leading to hierarchical mixtures of experts!). During training, all experts are trained simultaneously, but for each training sample, only the experts selected by the gate get updated (since they are the ones used to produce the output for that sample).
Types of MoE Models
1. Hard MoE:
- The gating network selects a few experts (often one) for each input.
- Efficient in large-scale models but non-differentiable.
2. Soft MoE:
- The gating network assigns continuous probabilities to all experts, combining their outputs.
- Fully differentiable and suitable for gradient-based optimization.
3. Hierarchical MoE:
Uses multiple layers of expert models, where each level refines the expert selection.
How MoE differs from traditional neural networks:
In a traditional dense neural network (or layer), the same parameters are used for every single input. For instance, if you have a regular feed-forward layer with 100 million parameters, all 100 million of those parameters will be activated and used to compute the output for each input example. MoE is different because the parameters used can vary depending on the input – only a small fraction of the model’s parameters (the ones belonging to the selected experts) are active for any given input. This means an MoE model can have a very large number of total parameters, but the effective computation per input is much smaller. For example, Google’s Switch Transformer (a language model using MoE) has about 26 billion total parameters, but only ~700 million of those are active during inference for any given token. In contrast, a dense model of 26 billion parameters would use all 26 billion for every input. Thus, MoE achieves the best of both worlds: it can increase model capacity (parameters) significantly without a proportional increase in computation for each sample. This conditional computation is the key difference – it makes MoE models sparse (only part of the network fires) whereas traditional models are dense (the entire network fires on every input).
Applications of MoE
MoE architectures have become especially popular in large-scale AI models because they offer a way to scale model size and complexity in a compute-efficient manner. Here are a few notable applications and use cases of MoE in different domains:
• Natural Language Processing (Large Language Models): One of the flagship examples is Google’s Switch Transformer, which was among the first models to exceed a trillion parameters by using MoE. The Switch Transformer (introduced in 2021) used 2048 experts in certain layers, reaching a total of 1.6 trillion parameters, yet it trains faster than smaller dense models because at any given time only one expert per input is active. In fact, the authors reported a 4× speed-up in pre-training compared to a dense Transformer of similar quality. This model demonstrated that we could match the performance of extremely large models while significantly cutting down computational cost. Another example from Google is GShard (2020), which applied MoE to a massive multilingual translation model. GShard enabled scaling a Transformer to over 600 billion parameters using MoE and distributed training, achieving state-of-the-art translation quality across 100 languages. This was trained on Google’s TPU pods with experts spread across devices, showing MoE’s power in multi-lingual NLP. More recently, Google’s GLaM model used MoE to match GPT-3 level performance on language tasks with only about one-third of the energy cost of GPT-3 by activating far fewer parameters per input and in Mistral’s Mixtral 8x7B each layer is composed of 8 feedforward blocks—that is, experts—each of which has 7 billion parameters. For every token, at each layer, a router network selects two of those eight experts to process the data. It then combines the outputs of those two experts and passes the result to the following layer. The specific experts selected by the router at a given layer may be different experts from those selected at the previous or next layer.
• Computer Vision: MoEs have also been explored in computer vision. Researchers have created vision transformers that incorporate MoE layers (sometimes called V-MoE models). For example, Google’s research on “Scaling Vision with Sparse MoEs” introduced models where certain transformer blocks use expert layers, allowing the model to specialize in different image patches or features. In practice, this means a vision MoE can allocate different experts to focus on different visual patterns or object types. Such models have been used to train extremely large vision models with less computational cost than a dense model of equivalent size. While MoE is not yet as common in vision as in NLP, these experiments show that the concept can generalize beyond text, potentially improving efficiency in image recognition or multi-task vision systems.
• Recommendation Systems and Multi-Task Learning: MoE is very useful in systems that must handle multiple objectives or tasks. A prominent example is Google’s MMoE (Multi-gate Mixture-of-Experts) model used for recommendation and advertising systems. In recommendation systems, one often needs to predict several things at once (for instance, predicting whether a user will click on an item and whether they will purchase it). Traditional single-task models can falter when forced to handle conflicting objectives. The MMoE approach uses a set of shared experts and multiple gating networks – essentially one gate per task – so that each task can dynamically prioritize different experts. This allows the model to explicitly model task relationships and share learnings when beneficial, while still keeping certain expertise separate for each objective. The result is a more efficient multi-task model that often outperforms training separate models or naive multi-task networks. In summary, MoE helps recommendation systems scale to handle many goals at once, improving both accuracy and efficiency by allocating specialized “experts” for different tasks.
Key benefits of MoE: Across these applications, MoE architectures have shown several clear advantages:
• Efficiency in Training: MoE models can reach a target performance with significantly less computation than dense models. For example, the Switch Transformer (an MoE model) was able to achieve the same accuracy as a dense Transformer in 7× fewer training steps (with the same computing resources). This faster training is because each training sample only utilizes a small fraction of the model’s weights, so effectively the model learns faster given a fixed budget of FLOPs.
• Scalability in Model Size: Because adding more experts increases the total parameters but not the per-input computation, MoEs make it feasible to scale models to hundreds of billions or trillions of parameters. This scalability allows capturing more knowledge or diversity in the model. For instance, by scaling up the number of experts, researchers saw consistent quality improvements on tasks without increasing the inference cost proportionally. In one case, increasing experts from 128 to 2048 (as done in GShard) dramatically boosted translation quality. MoEs thus offer a path to extremely large models that would be impractical as dense models.
• Reduced Inference Cost per Parameter: MoEs decouple model capacity from runtime cost. A model like GLaM or Switch has a gigantic number of parameters in total, but since only a few are used for any input, the inference computation (and by extension, energy cost) is much lower than a dense model of equal size. Google’s GLaM model is a great example: it achieved performance on par with the 175B-parameter GPT-3 model while using only one-third of the energy during training, thanks to MoE’s sparse activation. At inference time, MoE models also benefit – e.g., Switch Transformer with 1.6T parameters only needs to activate a network equivalent to a few hundred million parameters for each token, which can make it faster than a dense trillion-parameter model. (Note: A trade-off is that all those parameters still need to be stored in memory – MoEs require loading the full model weights, which can be memory-intensive, even if not all are used every time.)
In summary, MoE architectures have enabled AI models that are far larger and more specialized than previously possible, across NLP, vision, and recommendation systems. Major tech companies like Google have leveraged MoEs in cutting-edge projects (Switch Transformer, GShard, etc.), and other organizations (e.g. Meta’s large translation MoE, Microsoft’s DeepSpeed MoE for large models) are also exploring this approach to build efficient, scalable AI systems.
Recent Advancements in MoE
MoE has evolved significantly since its original conception in the 1990s. Early work introduced the idea of a gating network choosing between expert models, but it wasn’t until the deep learning era that MoEs were scaled up to extreme sizes. In 2017, researchers at Google reimagined MoEs for large deep networks with the Sparsely-Gated MoE layer, enabling what they called “outrageously large” neural networks (hundreds of billions of parameters) by activating only a small fraction of the network for each input. Since then, a series of advancements have addressed key challenges like training stability and efficient routing, making MoEs more practical and powerful. Recent SoTA models like Google’s multi-modal Gemini 1.5 and IBM’s enterprise-focused Granite 3.0 are MoE models. DeepSeek R1, which has comparable performance to GPT-4o and o1, is an MoE architecture with 671B total and 37B activated number of parameters and 128 experts
• Load balancing and expert utilization: One challenge in MoE training is ensuring that all experts are utilized reasonably, rather than some experts doing most of the work while others are rarely used. If the gating router sends too many inputs to a few experts and almost none to others, those few become overloaded (and get disproportionately trained) while the idle experts don’t learn well at all. This imbalance can cause a form of model collapse where effectively the model isn’t really using its full capacity. It also creates routing bottlenecks – if one expert gets, say, ten times more tokens than another, it becomes a slow hotspot that can hold up the whole batch during processing. To tackle this, researchers introduced several solutions. One is adding random noise to the gating probabilities (a technique introduced in the 2017 MoE paper) so that the selection isn’t greedy and inputs get spread out more evenly. Another common solution is an auxiliary load-balancing loss: during training, the model is given a small extra penalty whenever the distribution of tokens across experts is too uneven. This encourages the gate to route some traffic to less-used experts. The Switch Transformer, for example, uses a load-balancing loss to keep any single expert from taking too large a fraction of the tokens. Additionally, MoE implementations often enforce a capacity limit per expert – meaning each expert can only process up to a certain number of tokens from a batch. If more tokens than that are routed to the expert, the excess will either be dropped or routed to a second-choice expert. This prevents one expert from becoming a severe bottleneck and helps maintain throughput in distributed settings. Together, these strategies (noisy gating, balance losses, capacity caps) greatly alleviate the expert imbalance problem, allowing training to remain efficient even as the number of experts grows.
• Stability and training improvements: Large MoE models initially suffered from training instabilities – for instance, the model could become too dependent on a few experts or have difficulty converging, and fine-tuning an MoE on new tasks sometimes led to overfitting. Recent research has made MoEs more stable. One notable advancement is the router z-loss (introduced in the ST-MoE paper by Google in 2022), which adds a small penalty on the gating network’s logits to discourage them from becoming extremely large. In simpler terms, this keeps the gate’s output distribution a bit smoother (not too close to a one-hot assignment), which was found to reduce training instabilities without hurting model quality. Another trick was using higher precision for the gating computations: the Switch Transformer team discovered that doing the router’s Softmax in full precision (while the rest of the model used lower precision like bfloat16) helped avoid numerical issues that were causing training to blow up. Moreover, techniques like increased regularization on expert layers (e.g. applying higher dropout inside experts to prevent overfitting) have been used during fine-tuning to improve generalization. All these improvements address the unique challenges of training such a sparse and conditional model, making MoEs more robust in practice.
• Simplified and improved routing strategies: As mentioned, the original MoE designs often allowed routing to multiple experts per input (which gives the model flexibility but increases complexity). The Switch Transformer simplified this by using only one expert per token (Top-1 gating). This change had several benefits: it reduced the router computation (no need to calculate and combine outputs from two or more experts), cut down communication overhead between devices (each token’s data goes to just one expert machine instead of potentially two or more), and allowed larger batch sizes per expert since the tokens weren’t split among as many experts. Remarkably, this simplification did not hurt model quality much – Switch Transformer maintained accuracy while greatly streamlining the MoE mechanism. This showed that sometimes less is more in MoE routing: a simpler routing can be easier to train and scale. Other research has explored alternative routing mechanisms too, such as hashing-based routing (to remove the need for a Softmax entirely) and more complex mixture setups like multiple layers of experts or experts assigned to specific data domains. The field is actively experimenting with how to best route inputs to experts in a way that is both effective and computationally cheap.
• Scalable infrastructure and sparsity optimization: Because MoE models often have to distribute different experts across different hardware (GPUs/TPUs) for parallelism, a lot of innovation has gone into making this distribution efficient. Frameworks like GShard (for TensorFlow) and DeepSpeed MoE (for PyTorch) were developed to automate the process of splitting experts across devices and handling the communication under the hood. One concept is expert parallelism, where each expert runs on a separate device (or group of devices), and during a forward pass, tokens are sent to whatever device hosts their assigned expert. Naively, this could be slow due to network communication, but these frameworks use smart scheduling and batching of communication (sometimes an all-to-all operation) to minimize overhead. The result is that very large MoE models can be trained and inferred on clusters almost as efficiently as smaller models. For example, GShard’s automatic sharding allowed Google to train the 600B model on 2048 TPU cores with high utilization. Researchers have also introduced grouped routing strategies – e.g., grouping experts by machine and even adding a small penalty if a machine’s experts as a whole get too much traffic – to ensure the load is balanced not just per expert but also per machine. All these advancements address the so-called “routing bottleneck,” ensuring that routing computations and data exchanges do not become the limiting factor as we scale up the number of experts. In short, the MoE community has made great progress in the software and algorithms needed to support massive MoE models, from improved optimizers and initialization methods to libraries that handle the complex parallelism behind the scenes.
MoE research is very active, and new ideas continue to improve the paradigm. We’ve seen work on hierarchical MoEs (experts made of smaller experts), dynamic capacity (adjusting how many experts to use on the fly), and even distilling MoE models into smaller dense models for deployment. Each of these is aimed at leveraging the strengths of MoE (huge capacity, efficient use of compute) while mitigating downsides (complex training dynamics, memory use). The evolution from the early MoE models to today’s sophisticated systems like Switch Transformers demonstrates how these challenges can be overcome one by one.
Key takeaways: MoEs are a form of conditional computation that lets us build bigger and smarter models without necessarily making them slower. The gating mechanism is crucial – it learns to send each input to the best expert, using something like a softmax-based scoring of experts. The output is a weighted combination (often mostly from one expert) of the expert networks’ outputs. This dynamic routing is what differentiates MoEs from ordinary neural networks that treat every input the same way. The result is a flexible, scalable model architecture that can allocate its “brain power” where needed – much like how in a large organization, a query gets routed to a specific department of experts rather than bothering everyone.
References:
Switch Transformers: Scaling to Trillion Parameter Models
with Simple and Efficient Sparsity: https://arxiv.org/pdf/2101.03961.pdf
