Switch Transformer › Scaling Results
Language AI HandbookAll chapters Switch Transformer › Scaling ResultsOne section · Sections in this chapter8 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
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.
Speed vs. Quality Tradeoffs
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.
Scaling to Trillion Parameters
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.
def compare_model_costs(
d_model: int,
d_ff: int,
num_layers: int,
num_experts: int,
seq_len: int = 512,
) -> dict:
"""Compare parameter counts and FLOPs for dense vs Switch."""
# Dense model: each layer has one FFN
dense_ffn_params = d_model * d_ff + d_ff * d_model # Up + down projection
dense_total_ffn = dense_ffn_params * num_layers
# Switch model: each layer has num_experts FFNs
switch_total_ffn = dense_ffn_params * num_layers * num_experts
# Router parameters (small)
router_params = d_model * num_experts * num_layers
# FLOPs per token (both activate same amount)
dense_flops = 2 * dense_ffn_params # Forward pass FLOPs
switch_flops = 2 * dense_ffn_params # Same! Only one expert activates
return {
"dense_params": dense_total_ffn,
"switch_params": switch_total_ffn + router_params,
"param_ratio": (switch_total_ffn + router_params) / dense_total_ffn,
"flops_ratio": switch_flops / dense_flops,
}# Compare at different scales: (Name, d_model, d_ff, layers, experts)
configs = [
("Base", 768, 3072, 12, 128),
("Large", 1024, 4096, 24, 128),
("XL", 2048, 8192, 24, 256),
]
comparison_results = []
for name, d_model, d_ff, layers, experts in configs:
res = compare_model_costs(d_model, d_ff, layers, experts)
comparison_results.append((name, d_model, experts, res))Comparison: Dense vs Switch Transformer ================================================================= Base (d_model=768, 128 experts): Dense FFN params: 56,623,104 Switch total params: 7,248,936,960 Parameter ratio: 128.0x more parameters FLOPs ratio: 1.0x (same compute!) Large (d_model=1024, 128 experts): Dense FFN params: 201,326,592 Switch total params: 25,772,949,504 Parameter ratio: 128.0x more parameters FLOPs ratio: 1.0x (same compute!) XL (d_model=2048, 256 experts): Dense FFN params: 805,306,368 Switch total params: 206,171,013,120 Parameter ratio: 256.0x more parameters FLOPs ratio: 1.0x (same compute!)
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.
Sample Efficiency
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.
