Switch Transformer › Distillation and Fine-tuning

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.