Switch Transformer: Top-1 Routing at Scale

The previous chapters in this part established the fundamentals of Mixture of Experts: how expert networks specialize, how gating mechanisms route tokens, and how auxiliary losses maintain load balance. Yet despite these elegant solutions, MoE models remained notoriously difficult to train at scale. Training instability, communication overhead, and complex gradient interactions meant that even the most carefully designed MoE architectures struggled to reliably improve over dense baselines when pushed to the sizes that truly matter for language modeling. The Switch Transformer, introduced by Fedus, Zoph, and Shazeer in 2022, changed this by making a counterintuitive choice: route each token to just one expert instead of two or more. This radical simplification, combined with careful engineering around capacity limits and training stability, enabled scaling to over a trillion parameters while maintaining computational efficiency.

The Switch Transformer demonstrated that complexity is not always the path to better performance. By removing the interpolation between multiple experts that previous MoE designs required, Switch reduced communication overhead, simplified gradient computation, and paradoxically improved model quality. This might seem backwards at first: consulting more experts should give richer representations, right? The Switch paper showed empirically that this intuition breaks down at scale, because the benefits of simpler, cleaner optimization outweigh the representational gains of multi-expert interpolation. Understanding why this works reveals something deep about how sparse models behave during training.

Think of top-1 routing like assigning each student in a class to exactly one specialist tutor for a given question, rather than having them consult two tutors and average the advice. At first glance, two tutors sounds better. But in practice, assigning clear responsibility to one tutor means the tutor learns the student's needs and can be held accountable for the outcome. Two tutors can disagree, dilute each other's signal, or simply duplicate effort. The clarity of single-expert assignment can matter more than the breadth of multi-expert consultation.

This chapter examines the design decisions behind Switch Transformer: the Switch layer that replaces standard feed-forward networks, the top-1 routing mechanism and why it outperforms top-2, the capacity factor system that handles routing imbalances, the training stability techniques that make trillion-parameter training feasible, and the scaling results that established MoE as a path toward larger language models. We also examine fine-tuning challenges, distillation strategies, and limitations that subsequent architectures have addressed.

The core architectural innovation of Switch Transformer is the Switch layer, which replaces the standard feed-forward network (FFN) in each transformer block with a sparse, routed alternative. To understand why this substitution matters, we need to first recall what happens in a standard transformer and then see how the Switch layer turns this computation into something far more powerful.

Recall from our discussion of feed-forward networks in Part XV that a standard transformer block processes each token through a two-layer MLP. This feed-forward network applies the same transformation to every token in the sequence, using shared parameters regardless of what the token is or what context it appears in. The network has learned everything it knows about language compressed into a single set of weight matrices, and every token, whether it is a mathematical symbol, a proper noun, or a common preposition, passes through the exact same computation:

where:

• $x$: the input vector to the layer, typically of dimension $d_{\text{model}}$
• $W_1$: learnable weight matrix projecting from $d_{\text{model}}$ to the larger hidden dimension $d_{ff}$
• $W_2$: learnable weight matrix projecting back from $d_{ff}$ to $d_{\text{model}}$
• $b_1, b_2$: learnable bias vectors
• $\text{activation}$: the non-linear activation function (typically ReLU or GeLU)

This standard FFN architecture has served transformers well, but it embodies a basic limitation: every token receives identical processing. A token representing a mathematical concept receives the same transformation as a token representing a cooking ingredient. The network must somehow compress all of its knowledge about syntax, semantics, world facts, and linguistic patterns into a single set of parameters. At small scales this is manageable; at large scales it creates a capacity bottleneck where the dense parameters struggle to simultaneously stand for all the heterogeneous knowledge the model needs.

In a Switch layer, we break free from this constraint. Instead of one FFN, we have $E$ expert networks, each with the same architecture as the original FFN but with independent parameters. Each expert can specialize in different types of tokens or different aspects of language understanding. A router network decides which single expert should process each token, letting the model to deploy specialized computation where it matters most. The router's decision is content-dependent: a token representing a mathematical formula might consistently route to expert 7, while tokens that are common function words might distribute across experts 2, 4, and 9.

This architectural choice creates a natural division of labor. One expert might become skilled at processing mathematical notation, another at handling named entities, and yet another at processing syntactic function words. The key insight is that this specialization emerges from training rather than being imposed by design. You do not have to tell the model what each expert should specialize in; the combination of routing decisions and gradient updates naturally drives experts toward complementary specializations. The router and the experts co-evolve during training, with the router learning to recognize which tokens should go where, and the experts learning to excel at the tokens they receive.

Prior MoE models, including the influential GShard architecture, used top-2 routing: each token was sent to two experts, with the outputs combined as a weighted average. The reasoning was intuitive: combining perspectives from multiple specialists should produce richer representations. After all, if one expert knows about syntax and another knows about semantics, shouldn't we benefit from consulting both? The original Shazeer MoE paper from 2017 used top-2 routing as a default, and subsequent work largely inherited this choice without questioning it deeply. Switch Transformer was specific partly for its willingness to challenge this received wisdom empirically rather than theoretically.

Switch Transformer challenged the top-2 assumption with empirical evidence that contradicted conventional wisdom. The authors found that routing to a single expert reaches comparable or better performance while giving several practical advantages that compound at scale. The finding held across multiple model sizes, pre-training objectives, and evaluation benchmarks. Its consistency indicated that top-1 routing reflects MoE optimization rather than a dataset-specific artifact.

The effectiveness of top-1 routing can be understood through several lenses, each revealing a different facet of why simpler is better in this context. These arguments reinforce each other, and together they explain why the empirical results so consistently favor single-expert assignment over the seemingly more powerful two-expert approach.

Reduced communication cost. In distributed training, each expert typically lives on a different device. This means that routing a token to an expert requires sending that token's representation across the network to wherever the expert resides, processing it there, and returning the result. This all-to-all communication is a basic bottleneck in distributed MoE training: it requires every device to communicate with every other device, and the total data transferred scales with the number of tokens per batch times the dimensionality of the token representations. Top-2 routing means every token must participate in two such communication rounds, effectively doubling the network traffic. Top-1 cuts this communication cost in half, which becomes increasingly important as models scale across hundreds or thousands of accelerators. At the scale of Switch-C, which used 2048 experts across a correspondingly large device mesh, this factor-of-two communication reduction translates to a major practical advantage.

Simpler gradient flow. With top-2 routing, the gradient for a token flows through two expert networks weighted by their respective probabilities. This creates interdependencies that complicate optimization: the gradient with respect to one expert's parameters depends on what the other expert produced, because the router probability determines how much each expert's output contributes to the final result. When both experts influence the output, neither expert receives a clean signal telling it how to improve. Top-1 gives cleaner gradients, where each expert receives full responsibility for the tokens it processes. When expert 7 handles a token and the model produces a poor output, the gradient flows cleanly back through expert 7 and through the router's probability for expert 7. This clarity in the optimization objective helps the model converge more reliably and prevents situations where two experts compete for the same tokens without either fully committing to specialization.

Higher effective capacity per expert. If each token visits two experts, and each expert has a fixed capacity (more on this shortly), then top-2 routing effectively halves the number of tokens any expert can accommodate. Top-1 routing means twice as many tokens can be processed per expert for the same capacity setting. This translates to larger effective batch sizes per expert, which improves the quality of gradient estimates and allows each expert to see more diverse examples during training. Better gradient quality per expert means faster and more reliable specialization.

Empirical validation at multiple scales. The Switch Transformer paper demonstrated that top-1 routing with the same total compute reaches better perplexity than top-2 routing across multiple model scales. This wasn't a marginal improvement; the gains were consistent and significant, validating the theoretical arguments with concrete evidence. The perplexity gap was largest at the largest scales, consistent with the hypothesis that the communication and gradient benefits compound as scale increases.

The mathematical formulation for top-1 routing is straightforward, which is part of its elegance. Given a token representation $x$, the router computes a probability distribution over all available experts by first projecting the input into a logit vector and then applying softmax:

where:

• $p(x)$: the probability distribution over experts for the input token, with $p_i(x)$ representing the probability of routing to expert $i$
• $W_r$: the learnable router weight matrix of shape $[d_{\text{model}} \times E]$, projecting the input to a vector of expert logits
• $x$: the input token representation, a vector of dimension $d_{\text{model}}$

This probability distribution is the router's assessment of how well-suited each expert is for processing this particular token. The selected expert is simply the one with the highest probability: $e^* = \arg\max_i p_i(x)$.

Why does this formula make sense? Notice that the linear projection $W_r \cdot x$ produces a scalar score for each expert, and the softmax converts these scores into a proper probability distribution that sums to 1. The expert with the highest score gets the highest probability, and the argmax selects that expert as the routing destination. The simplicity of this computation is what allows the router to be extremely fast, adding minimal overhead to the forward pass.

Once the expert is selected, the output of the Switch layer combines the expert's processing with a scaling factor based on the router's probability:

where:

• $y$: the output of the Switch layer for the token, which replaces the dense FFN output in the transformer block
• $e^*$: the index of the selected expert (the one with the highest probability)
• $p_{e^*}(x)$: the probability assigned to the selected expert, acting as a multiplicative gating factor
• $E_{e^*}(x)$: the output of the selected expert network after processing the input token through its feed-forward layers

The multiplicative weighting by $p_{e^*}(x)$ serves an important purpose beyond just combining the expert output. It gives a differentiable signal that allows gradients to flow back through the router, letting the router to learn which experts work best for which tokens. Without this multiplicative term, the argmax operation would be non-differentiable: the gradient would have no information about how to adjust the router's parameters. By preserving the router probability in the output, we create a path through which gradients can flow during backpropagation, letting end-to-end training of both the router and the experts.

A potential concern with top-1 routing is that hard assignment creates discontinuities. Small changes in a token's representation could flip it to a different expert, potentially causing unstable training dynamics. If a token is hovering on the decision boundary between two experts, tiny perturbations might cause it to oscillate between them during training. This could create a noisy gradient signal where the expert receiving a token changes unpredictably from step to step, preventing stable specialization.

However, the multiplicative weighting smooths this effect in an elegant way. If a token is marginally assigned to expert 3 with probability 0.51 (vs. 0.49 for expert 2), the output is scaled by 0.51. This scaling gives a soft transition that prevents gradient spikes. When the token is near a decision boundary, the output magnitude is reduced, naturally dampening the impact of near-ties. As the router becomes more confident about a token's assignment, the probability approaches 1.0 and the full expert output is used.

This design ensures that tokens near decision boundaries contribute less strongly to the gradient signal, letting the model to focus its learning on tokens where the routing decision is clear and real. Think of it as the model saying: when I am not sure which expert should handle a token, I will use a softer, more conservative output rather than committing fully to an uncertain choice. This prevents spurious gradient updates driven by borderline routing decisions and encourages the router to become increasingly confident over the course of training, pushing probabilities toward 0 or 1 for most tokens as the experts develop clear specializations.

Even with perfect load balancing losses pushing the router toward uniform distribution, routing can never be perfectly uniform in practice. The discrete nature of expert selection means that random fluctuations will always cause some experts to receive more tokens than others in any given batch. This creates a practical problem that cannot be ignored: how many tokens should each expert be prepared to handle? If we allocate too few slots, important tokens get dropped. If we allocate too many, we waste memory that could be used for more experts or larger expert networks.

The capacity factor ($C$) is Switch Transformer's solution to this challenge. It defines how much buffer capacity each expert has beyond the perfectly balanced allocation. This gives a mechanism to handle the inevitable routing imbalances without catastrophic failure. The capacity factor is one of the most practically important hyperparameters in a Switch Transformer, because it directly governs the tradeoff between memory usage and the frequency of token dropping.

Think of the capacity factor like reserving seats at a restaurant. If you know on average 10 people will want to sit at each of your 8 tables, but arrivals are random, you do not set every table for exactly 10 guests. You set each table with a few extra chairs to handle groups that show up together. The capacity factor tells each expert how many extra chairs to set out. Setting $C = 1.25$ is like setting 12.5 chairs at each table: you can accommodate bursts of extra demand without anyone being turned away, at the cost of a modest amount of wasted seating capacity when demand is below peak.

If a batch contains $n$ total tokens and we have $E$ experts, perfect balance would assign $n/E$ tokens to each expert. In an ideal world, the router would reach this exact distribution, and we could set the capacity to exactly $n/E$ without any waste. But in reality, some experts will be more popular than others for any given batch, driven by the content distribution of that particular batch. The capacity factor scales the allocation to accommodate this variance:

where:

• $n$: the total number of tokens in the batch, equal to batch size times sequence length
• $C$: the capacity factor (a scalar multiplier greater than 1.0 for typical use)
• $E$: the number of experts available
• $\lfloor \cdot \rfloor$: the floor function, rounding down to ensure an integer number of token slots

Understanding this formula requires thinking about what happens at different values of $C$. With $C = 1.0$, each expert has exactly enough slots for perfect balance. Any deviation from uniform routing causes some tokens to be dropped. With $C = 1.25$, each expert can handle 25% more tokens than the perfectly balanced allocation, so routing variance must be quite substantial before any tokens are dropped. With $C = 2.0$, experts can absorb twice the expected load, in effect eliminating dropping at the cost of using twice the memory per expert.

Why does this formula make sense? Notice that $n/E$ is the perfectly balanced allocation per expert. Multiplying by $C$ scales this up proportionally. If $C = 1.5$ and each expert would receive 100 tokens under perfect balance, each expert instead gets capacity for 150 tokens, leaving a 50-token buffer for routing imbalance. The floor function ensures integer capacity, since you cannot process a fractional token.

A capacity factor of $C = 1.0$ means experts have exactly enough slots for perfect balance. In practice, routing is never perfect, so some experts overflow. The Switch solution is simple and pragmatic: tokens that exceed an expert's capacity skip the expert entirely and pass through via the residual connection.

This design creates a spectrum of tradeoffs that you can select based on your constraints:

• $C = 1.0$: Minimal memory usage, but many tokens may be dropped (skipped)
• $C = 1.25$: 25% buffer handles typical load variance well
• $C = 2.0$: Generous buffer, fewer drops, but double the memory per expert

When an expert reaches capacity, additional tokens routed to it are "dropped": they bypass the expert entirely and only the residual connection preserves their information. This might seem problematic at first glance. After all, if the router determined that a token should go to a particular expert, doesn't skipping that expert harm the model's performance? The answer is fine-grained: token dropping does have a cost, but several design choices mitigate this cost enough that moderate dropping is acceptable and the overall system works well.

Understanding why token dropping is tolerable requires considering what information is lost when a token gets dropped. The token's representation entering the Switch layer is not destroyed; it continues unchanged through the residual connection and remains available for all subsequent transformer layers. What the token misses is the specific expert processing in this one layer. Given that a large Switch Transformer might have 12 or 24 such layers, missing one expert processing in one layer is a relatively small fraction of the total computation applied to that token.

Several factors mitigate the impact of token dropping, making this design more reliable than it initially appears:

• Load balancing reduces drops. The auxiliary balancing loss (covered in Chapter 6 of this part) incentivizes the router to distribute tokens evenly across experts. A well-trained router rarely overloads any single expert severely, minimizing the frequency of capacity overflow. During training, the model learns routing strategies that naturally keep expert loads within bounds, so capacity overflow becomes increasingly rare as training progresses.

• Residual connections preserve information. Dropped tokens still retain their original representation through the residual connection; they simply do not receive the expert's specialized processing for that layer. The token's information is not lost, just unchanged by that particular expert. This is a graceful degradation rather than a catastrophic failure.

• Drops are distributed across tokens. With proper load balancing, drops rarely concentrate on specific tokens across multiple layers. A token might skip one expert in layer 5 but be processed normally in layers 1 through 4 and 6 through 12. The impact of a single dropped expert processing is diluted across the many layers of the network, preventing any single token from being systematically disadvantaged.

• Training adapts to token dropping. The model learns to function with some level of token dropping, developing representations that do not critically depend on every expert processing at every layer. This adaptation means that occasional capacity overflow does not catastrophically harm performance. The model has seen token dropping during training and has learned to work around it.

Let's assemble the full Switch layer, incorporating routing, capacity management, and the load balancing loss we discussed in Chapter 6. This implementation ties together all the concepts we have explored: the router that assigns tokens to experts, the capacity mechanism that handles overflow, and the auxiliary loss that encourages balanced utilization. Reading through this code carefully is one of the best ways to develop a concrete mental model of how all the pieces fit together in a working Switch layer.

The output shows that tokens are distributed across experts, with the auxiliary loss encouraging this balance. As training progresses, the balance loss pushes the router toward a more uniform distribution while allowing specialization. The residual connection at the end of the forward pass saves dropped tokens from receiving zero information. Without output = output + x, any token that did not get processed by an expert would contribute a zero vector to subsequent layers, destroying that token's representation entirely. The residual connection is a necessary safety net for the capacity-management system.

The Switch Layer has four hyperparameters that you will need to configure for any practical deployment, and the tradeoffs they govern are nontrivial. Understanding what each parameter controls helps you make principled choices when adapting Switch Transformer to a new scale or training environment.

The key parameters for the Switch Layer are:

• num_experts: Number of expert networks. More experts increase total parameter capacity without increasing compute per token, since only one expert activates. However, more experts also means more communication overhead in distributed settings and more routing complexity. Typical values in the Switch paper range from 128 to 2048, depending on the number of devices available.
• capacity_factor: Multiplier defining the expert buffer size ($C$). Controls the tradeoff between memory and token dropping. The recommended default of 1.25 works well for most training configurations, but you should increase it to 1.5 or 2.0 if you observe high drop rates during early training.
• balance_loss_weight: Coefficient for the auxiliary load balancing loss to prevent routing collapse. If set too high, it forces uniform routing at the cost of specialization. If set too low, routing collapse can occur. The paper found 0.01 to be a reliable default across a range of configurations.
• d_ff: Hidden dimension of the expert feed-forward networks. This should match the FFN dimension of the dense layers in your baseline transformer to ensure a fair parameter count comparison for each individual expert.

Switch Transformers, like other MoE models, can be unstable during training. The sparse routing mechanism introduces discontinuities that dense models do not face, and the discrete expert selection can create optimization challenges that compound over the course of a long training run. The authors identified several strategies necessary for stable training at scale, each addressing a specific source of instability. These strategies are not just engineering tricks; they stand for real insights about the pathological training dynamics that arise in sparse models and how to counteract them.

One counterintuitive finding was that using lower precision (bfloat16 or float16) in most of the model while keeping the router in float32 materially improved stability. This selective precision strategy shows a deep understanding of where numerical precision matters most in the network.

The router's softmax operation is particularly sensitive to numerical precision because small differences in logits get amplified exponentially. When computing $e^{z_i}$ for logits $z_i$, even small rounding errors can cause some experts to receive disproportionately high or low probabilities. In lower precision, these errors accumulate and can cause the router to make inconsistent or degenerate routing decisions. A rounding error of 0.001 in a logit becomes a rounding error of $e^{0.001} \approx 1.001$ in the exponential, but if the logit magnitude itself is large (say 20), then small absolute errors correspond to large relative errors in the resulting probabilities.

By keeping the router computation in float32 while letting the rest of the model to use memory-efficient bfloat16, Switch Transformer reaches both numerical stability and computational efficiency. The router is a tiny fraction of total computation (a single linear layer followed by a softmax), so the precision overhead is minimal in terms of memory and compute. The expert networks, which dominate the parameter count and computation, can safely use lower precision because the relative smoothness of their FFN computations makes them less sensitive to rounding errors.

As we discussed in Chapter 7, the router z-loss penalizes large router logits, preventing the softmax from becoming too peaked. This is particularly important for Switch Transformer because top-1 routing already creates hard decisions. Extremely peaked softmax distributions worsen the discontinuity, making tiny input changes flip routing decisions sharply. The z-loss is a direct intervention against this failure mode.

The z-loss gives a gentle pressure against this behavior. To understand what it penalizes, we first need to understand the log-sum-exp function, which is closely related to the softmax denominator. For a batch of $n$ tokens with logit vectors $z^{(i)}$ each containing $E$ elements, the z-loss is:

where:

• $L_z$: the auxiliary router z-loss, added to the total training objective with a small coefficient
• $n$: the total number of tokens in the batch
• $E$: the total number of experts
• $z_j^{(i)}$: the logit (pre-softmax score) for expert $j$ given token $i$
• $\sum_{j=1}^{E} e^{z_j^{(i)}}$: the softmax partition function (denominator), which captures the total magnitude of the logit vector

This formula penalizes large logits to maintain numerical stability. Each component serves a specific purpose:

1. Sum of Exponentials: The inner sum $\sum e^{z_j}$ captures the aggregate magnitude of the logits. When logits are large, this sum grows exponentially. This gives a sensitive detector for logit magnitude.
2. Logarithm: The $\log$ scales this value back to the linear domain, in effect computing something close to $\max_j z_j$ plus a small correction term. This makes the penalty proportional to the maximum logit rather than exponential in it.
3. Square: Squaring the result penalizes large positive values strongly while being smooth near zero, creating a strong incentive to keep logits bounded without overly constraining small logit values.

The log-sum-exp is tightly related to the maximum logit value. For a vector with one dominant entry $z_{\max}$ and all others much smaller, $\log \sum_j e^{z_j} \approx z_{\max}$. Minimizing $(\log \sum e^{z_j})^2$ is therefore approximately equivalent to minimizing $z_{\max}^2$, directly penalizing large maximum logits. Squaring the log-sum-exp creates a smooth, differentiable penalty that applies strong pressure when logits are large.

Without this loss, the model could increase logits indefinitely. Due to softmax's translation invariance, logits of [5, 3, 2] and [500, 498, 497] produce identical probability distributions. But the second case risks floating-point overflow and creates extremely peaked distributions that harm optimization. The z-loss prevents this degenerate behavior.

The Switch paper recommended several initialization and regularization tweaks that collectively improve training stability. These recommendations emerged from extensive empirical experimentation and address specific failure modes that appeared during early Switch Transformer training runs.

Smaller router initialization. The router weights should be initialized with a smaller standard deviation than the default, typically by using a scale factor of 0.1 relative to the default initialization. This prevents the router from making extremely confident routing decisions early in training, before it has seen enough data to form reliable preferences. With large initial weights, the router commits strongly to certain routing patterns from the very first batches, which can create path dependency: early routing decisions determine which tokens experts see, which determines how experts specialize, which in turn reinforces the original routing decisions. Starting with small weights keeps the router more agnostic early on, letting more exploratory routing before specialization solidifies.

Increased dropout in experts. Higher dropout rates of 0.1 to 0.2 inside expert networks help prevent overfitting to specific routing patterns. Because each expert only sees the subset of tokens routed to it, the effective training sample for each expert is smaller than the full batch. Higher dropout compensates for this reduced sample diversity and prevents experts from memorizing the specific tokens they receive rather than learning general patterns.

Expert dropout as structured regularization. Occasionally dropping entire experts during training encourages the model to develop redundancy and prevents over-reliance on any single expert. When an expert is dropped during a training step, its tokens must be handled by the residual connection alone, which trains the rest of the network to function reliably without that expert's contribution. This technique is analogous to dropout at the expert level rather than the neuron level. This gives structured regularization that improves generalization.

These techniques work together to prevent premature specialization, where the model commits too strongly to certain routing patterns before it has seen enough data to make informed decisions. The goal is to keep the training dynamics exploratory during the early phase, letting the router and experts to co-evolve gradually toward a stable and beneficial division of labor rather than locking into a suboptimal pattern early and never recovering.

Scaling experiments provide the strongest evidence for Switch Transformer's effectiveness. The authors compared Switch Transformer against dense T5 models with equivalent compute budgets. The sparse architecture showed consistent benefits across a wide range of scales. These comparisons were carefully controlled: the dense and sparse models were trained for the same number of FLOPs, on the same data, with the same optimizer. This allows architectural differences to be isolated as the cause of performance differences.

The scaling results painted a clear picture: Switch Transformer consistently outperformed dense baselines at every model size tested, and the gap tended to grow with scale. This is exactly what you would hope for in a scaling architecture: it works and gets relatively better as the scale increases, suggesting that MoE architectures align with the scaling properties of language modeling in a basic way.

The key insight is that Switch models reach better quality at the same training compute. A Switch-Base model with 7B total parameters (but only activating approximately 100M per token) outperformed T5-Base while using the same FLOPs per training step. This was not achieved by throwing more compute at the problem; it emerged from the more efficient use of parameters that sparsity lets. The same compute budget, applied to a model with more total capacity, produces a better model.

The improvements were substantial across multiple dimensions:

• Pre-training speed: Switch-Base achieved T5-Base quality in approximately one-seventh the training time. This is a large speedup: if T5-Base takes a week to train to a given quality level, Switch-Base can reach that same quality in roughly one day using the same hardware.
• Same-compute quality: Given equal compute, Switch models consistently achieved lower perplexity, with larger gaps at larger model scales.
• Scaling efficiency: The gap widened at larger scales, with Switch-XXL showing large improvements over T5-XXL. This suggests that the benefits of expert sparsity compound as the total parameter count grows.

The Switch Transformer demonstrated that MoE lets scaling to very large parameter counts. The largest model, Switch-C, contained 1.6 trillion parameters distributed across 2048 experts. This scale was simply not achievable with dense models given the computational resources available at the time; training a dense 1.6 trillion parameter model would require roughly 2048 times the compute and memory resources of Switch-C's training run. This makes it impractical with existing hardware.

Despite this large parameter count, the computational cost per token remained manageable because only one expert activates per token. A 1.6 trillion parameter model might have the knowledge capacity of its full parameter count, but the inference cost of a model roughly 2000 times smaller. This decoupling of stored capacity from active computation is the central promise of MoE architectures, and Switch-C was the first demonstration of it at truly extreme scale.

This comparison illustrates the key advantage: Switch models can have 128x or more parameters while using identical compute per forward pass. The additional parameters give more capacity for learning without proportionally increasing training cost. This decoupling of parameters from compute was the central insight that enabled scaling to trillion-parameter models. You pay for the experts in memory and communication, but not in FLOPs: each token is processed by exactly one expert regardless of how many total experts exist.

Beyond raw speed, Switch Transformers showed improved sample efficiency. The models achieved better quality with fewer training tokens, suggesting that the sparse expert structure lets more effective use of training data. This aligns with the intuition that specialized experts can learn more from each example in their domain.

When a token about mathematics routes to a mathematics-specialized expert, that expert receives a concentrated signal for updating its parameters. In contrast, a dense model must update all parameters regardless of the token's content, diluting the learning signal across parameters that may not be relevant. The sparsity of MoE routing creates a natural form of gradient focusing: the gradient for each parameter flows only from the tokens most relevant to that parameter's specialization, rather than averaging across the full diversity of the training corpus.

This improved sample efficiency has a practical implication: Switch models can reach a given quality level faster in wall clock time and with fewer training examples. For organizations with limited access to training data, this means MoE architectures may extract more value from a fixed data budget than dense alternatives.

One challenge with large MoE models is deployment: serving a 1.6T parameter model requires distributing experts across many devices, increasing latency and infrastructure complexity. A model that needs 2048 devices to run inference is impractical for most organizations, even if training it was feasible. The Switch paper explored distillation as a solution, showing that much of the knowledge learned by large MoE models can be transferred to smaller, more deployable architectures.

The core idea behind MoE distillation is the same as standard knowledge distillation: train a smaller student model to match the outputs of a larger teacher model. The student benefits from the teacher's soft probability distributions over the vocabulary, which encode more information than hard one-hot labels. For language modeling, this means the student is trained to produce the same token probability distributions as the Switch teacher, capturing the teacher's uncertainty and confidence patterns rather than just matching the highest-probability prediction.

The authors found that Switch Transformers could be distilled into dense models that retain much of the quality gain. A Switch-Base model distilled to a T5-Base architecture achieved 30% of the original quality improvement while eliminating the MoE complexity entirely. This gives a practical path for organizations that want MoE's training benefits without its deployment challenges: train with MoE to get fast, high-quality pre-training, then distill to a dense model for efficient inference.

The distillation process is straightforward: train the dense student model to match the Switch teacher's outputs, using a combination of hard labels (for tasks with discrete targets) or soft targets (for language modeling). The student learns to approximate the ensemble behavior of the sparse experts using its single, dense feed-forward network. While 30% retention might sound modest, it is a free improvement over training the dense model from scratch: you get some of the benefits of a much larger model at the cost of running one MoE training run.

The distillation framing also reveals something interesting about what MoE models learn. The fact that some of the quality improvement survives distillation into a dense model suggests that MoE training produces better representations and token predictions, not just better routing. The sparse expert training procedure appears to teach the model things about language that dense training at the same compute budget misses, and those things are expressible in a dense architecture, albeit at reduced strength.

Fine-tuning MoE models presents unique challenges that you must address carefully. These challenges differ from the fine-tuning challenges of dense models and arise specifically from the sparse routing mechanism and the way experts specialize during pre-training.

Expert collapse during fine-tuning. During fine-tuning on narrow domains, some experts may receive very few tokens, causing their parameters to drift or become useless. The fine-tuning dataset may not cover the full range of content that the pre-trained router learned to handle. If you fine-tune a Switch model on a medical question-answering dataset, for example, the routing patterns might concentrate heavily on the experts that handle medical and scientific terminology, leaving other experts starved of gradient signal. Over time, these underutilized experts can deteriorate, reducing the model's versatility. Increasing the load balancing loss weight during fine-tuning helps maintain expert diversity by keeping routing more uniform even when the domain is narrow.

Capacity tuning across dataset sizes. Fine-tuning datasets are often orders of magnitude smaller than pre-training corpora, changing the optimal capacity factor. With smaller batches, the variance in routing decisions increases, because each batch contains a less representative sample of the full token distribution. Higher routing variance means more overflow, potentially requiring higher capacity factors to avoid excessive token dropping. The authors recommended re-tuning the capacity factor for each downstream task, particularly when moving from web-scale pre-training to task-specific fine-tuning with small curated datasets.

Transfer gaps on specific tasks. While Switch models excelled at pre-training perplexity, the gains sometimes diminished after fine-tuning. Dense models occasionally caught up on specific tasks, suggesting that MoE's advantages are most pronounced for general language modeling rather than task-specific optimization. This finding motivates the distillation approach: use MoE for pre-training to get better representations, then either distill to dense for fine-tuning and deployment, or accept the infrastructure complexity of MoE for tasks where the pre-training quality gap persists through fine-tuning.

Let's visualize how tokens get routed to experts in a trained Switch layer. This visualization helps build intuition about what the router learns and how it distributes tokens across the available experts. We will trace through a full forward pass with concrete numbers to make the routing mechanism tangible.

The heatmap shows each token's probability distribution over experts. Notice how most tokens have one dominant expert (high probability in one column, near-zero in the others). This shows the router's confident routing decisions even with random initialization. This confidence emerges naturally from the softmax: when one logit is substantially higher than the others, the softmax output is a near-one-hot vector. The bar chart shows the actual token distribution, with the red line showing perfect balance. With random initialization, we see large variance in expert popularity. During actual training, the auxiliary load balancing loss would apply pressure to even out this distribution, driving the bar heights toward the red line while letting real departures driven by token content.

To fully understand how a single token moves through a Switch layer, let's trace through the complete computation manually with concrete numbers. This numerical walkthrough makes the abstract formulas concrete and reveals exactly what happens at each step.

Consider a simple setting: we have 4 experts ($E = 4$), a model dimension of $d_{\text{model}} = 4$ (tiny, for illustration), and a batch with 8 tokens total. The capacity factor is $C = 1.5$.

Step 1: Compute capacity. We apply the capacity formula:

Each expert can accept at most 3 tokens before it is considered full and subsequent tokens are dropped.

Step 2: Router produces logits. For one specific token with representation $x = [1.2, -0.3, 0.8, 0.5]$, the router's linear layer produces logits $z = W_r \cdot x$. Suppose this gives $z = [2.1, 0.5, 1.8, -0.2]$.

Step 3: Softmax converts logits to probabilities. We compute:

The sum of exponentials is $8.17 + 1.65 + 6.05 + 0.82 = 16.69$. Dividing each by the sum:

Step 4: Select the top expert. The highest probability is $p_0 = 0.49$ for expert 0, so $e^* = 0$. This token will be processed by expert 0.

Step 5: Check capacity and route. We check whether expert 0 has room. If fewer than 3 tokens have already been assigned to expert 0, we process this token. The expert runs its FFN on the token representation and returns output $E_0(x)$.

Step 6: Scale by router probability. The Switch layer output for this token is:

The factor 0.49 scales the output. This shows the router's moderate confidence. If the router were completely certain (probability 1.0), the expert output would pass through unchanged. If the router were uncertain (probability near 0.25 for each expert), the output would be substantially attenuated.

Step 7: Add residual. The final output adds the original input: $y_{\text{final}} = 0.49 \times E_0(x) + x$. This residual connection ensures the token retains its full original representation even if the expert's contribution is small.

This trace reveals why the router probability matters as a multiplicative weight: it gives a differentiable path for gradients from the output loss back to the router's linear projection weights, letting end-to-end training. Without this multiplicative weighting, the routing decision would be a non-differentiable argmax with no gradient signal for the router.

Despite its successes, Switch Transformer has important limitations that influenced subsequent work and that you should understand before deploying these models. MoE is not a free lunch: the benefits in parameter efficiency come with added complexity and memory use. Communication costs and operational fragility also matter. Being clear-eyed about these limitations is needed for making good architectural decisions.

MoE models, including Switch Transformer, exhibit more training instability than dense counterparts. Routing decisions create discontinuities in the loss surface, and router collapse (where all tokens route to one or a few experts) can occur if load balancing fails or if the router's initial weights are poorly set. Selective precision and z-loss help substantially but do not eliminate these issues. Training large Switch models requires careful monitoring, thoughtful hyperparameter choices, and sometimes manual intervention when instability appears. Experience training dense language models does not transfer completely because MoE failure modes differ qualitatively.

The instability problem also interacts with training duration. Switch models sometimes show good early progress followed by a sudden degradation, where the router enters a degenerate state from which it cannot recover within normal training dynamics. This catastrophic forgetting of good routing patterns is more common in MoE than in dense models, and it makes checkpoint selection and training monitoring especially important.

While Switch models excel at pre-training, the advantages often shrink after fine-tuning. On some tasks, dense models with equivalent compute match or exceed Switch performance after task-specific training. This suggests MoE's benefits may be most pronounced for general-purpose language modeling rather than task-specific optimization. The expert specializations that emerge during pre-training may not always align with the demands of a specific downstream task, and the routing may need to reorganize substantially during fine-tuning. If fine-tuning disrupts the carefully developed routing patterns, some of the pre-training quality advantage can be lost. You should evaluate whether MoE is appropriate for your specific use case and evaluate dense baselines before committing to the infrastructure complexity of MoE deployment.

Deploying MoE models requires expert parallelism, where different experts live on different devices. This introduces all-to-all communication patterns that standard data parallelism or tensor parallelism do not require. Every forward pass involves gathering tokens from all devices, routing them to the appropriate expert devices, processing them there, and scattering the results back. This communication pattern is harder to optimize, more sensitive to network latency, and more prone to load imbalances than standard parallelism strategies. The infrastructure burden limited early adoption of Switch-style models, though frameworks like DeepSpeed, Megatron-LM, and FairScale have since added MoE support, materially reducing the barrier to experimentation.

The deployment complexity also affects serving infrastructure. A dense model can often be served from a single high-memory GPU or a small number of GPUs using standard tensor parallelism. A MoE model with hundreds or thousands of experts may require a dedicated serving cluster with carefully orchestrated routing logic. This operational burden is a real cost that must be weighed against the quality benefits when making deployment decisions.

While FLOPs per token remain constant, the total parameter memory scales linearly with expert count. A 128-expert model requires 128 times the FFN memory compared to a dense model with the same per-token computation, even though only 1 out of 128 experts activates per token. This memory overhead dominates at inference time, where batch sizes may be small and where the large inactive expert weights still consume device memory. The ratio of active parameters to total parameters, sometimes called the activation ratio, is 1/E for a Switch model with E experts, meaning that the large majority of parameters sit idle during any given forward pass. This is by design and is what creates the parameter efficiency advantage, but it also means that raw memory requirements scale with total parameters rather than active parameters.

The memory-compute tradeoff differs substantially between training and inference. During training, large batches amortize the memory cost over many tokens, and communication overhead is a one-time cost per step. During inference, individual requests may involve just a handful of tokens, making memory inefficiency more pronounced relative to the compute performed.

Though residual connections preserve dropped tokens' information, consistent dropping on the same tokens across multiple layers does have a measurable performance cost. For tasks requiring precise token-level processing, such as sequence labeling, named entity recognition, or tasks where rare but important tokens carry necessary information, the capacity factor must be tuned carefully. In practice, this means MoE models often need larger batch sizes to ensure adequate expert utilization and low drop rates. Applications with strict latency requirements and small batch sizes may find dense models more suitable, because the capacity factor cannot be increased arbitrarily without increasing memory requirements proportionally.

Despite these limitations, Switch Transformer established that simplified MoE architectures could scale efficiently. The design choices, top-1 routing, capacity factor, and training stabilization strategies, became foundational for subsequent work. Mixtral, which we will explore in the next chapter, builds directly on these foundations while introducing innovations for open-source deployment, including using top-2 routing with a smaller number of experts and optimizing the implementation for practical use on commodity hardware.

Switch Transformer demonstrated that radical simplification can enable MoE scalability. Routing each token to one expert rather than two halved communication overhead, cleaned gradient flow, and increased effective batch sizes per expert. The empirical results supported these theoretical arguments, and consistent gains across scales indicated that top-1 routing was not a dataset-specific trick.

The key innovations include:

• Top-1 routing: Each token goes to exactly one expert, halving communication cost and simplifying gradients compared to top-2 routing, while counterintuitively achieving better performance
• Capacity factor: A tunable buffer ($C$, typically 1.25) determines how many tokens each expert can handle, with overflow tokens skipping to the residual path in a graceful degradation pattern
• Selective precision: Keeping the router in float32 while using lower precision elsewhere improves training stability by preventing rounding errors from corrupting routing decisions
• Load balancing with z-loss: The auxiliary balance loss from prior MoE work combines with router z-loss to maintain balanced expert utilization and prevent logit magnitude from growing without bound
• Initialization and dropout strategies: Smaller router initialization, higher expert dropout, and occasional expert-level dropout prevent premature specialization and improve generalization

The scaling results were strong: Switch-Base achieved T5-Base quality in one-seventh the training time, and the architecture scaled to 1.6 trillion parameters. These results established MoE as a practical path toward larger language models, showing that parameter count and computational cost can be decoupled. The capacity-factor mechanism sets a hard limit on tokens per expert and drops overflow to residual connections, handling imperfect routing without catastrophic failures. Combined with aggressive load balancing, this made MoE training reliable enough for production scale.

The limitations, training instability, fine-tuning gaps, infrastructure complexity, and memory overhead, are real and have motivated substantial follow-on research. Mixtral and other successors addressed several of these challenges, but all retain the core Switch Transformer insight: simpler routing, properly stabilized, lets the parameter scaling that makes modern large language models possible.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about Switch Transformer architecture and its innovations.