EP023: Scaling Switch Transformers to Trillion Parameters

The paper "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity" introduces the Switch Transformer, a sparsely-activated model designed to address the computational barriers of scaling deep learning models.Here is a summary of the key points:• Core Innovation: The authors simplify the Mixture-of-Experts (MoE) routing algorithm. Unlike previous MoE models that route tokens to multiple experts, the Switch Transformer uses a Switch layer that routes each token to a single expert (k=1). This simplification reduces routing computation and communication costs while preserving model quality.• Efficiency and Speed: The architecture allows for massive parameter counts without a corresponding increase in computational cost per example. It achieves up to a 7x increase in pre-training speed compared to the T5-Base model while using the same computational resources.• Stability and Training Techniques: To overcome the training instabilities common in large sparse models, the paper introduces several techniques:    ◦ Selective Precision: Using float32 precision specifically for the router mechanism while keeping the rest of the model in bfloat16.    ◦ Improved Initialization: Reducing the initialization scale of weight matrices.    ◦ Expert Dropout: A regularization technique that increases dropout specifically within the expert layers during fine-tuning to prevent overfitting.• Scaling to Trillions: By combining data, model, and expert parallelism, the authors successfully trained models with up to 1.6 trillion parameters. The 1.6T parameter model (Switch-C) showed no training instability.• Downstream Performance: The model demonstrated superior scaling on diverse natural language tasks. It improved over the mT5 baseline across all 101 languages tested in a multilingual setting.• Distillation: The authors showed that large sparse models can be distilled into smaller dense models, reducing model size by up to 99% while retaining approximately 30% of the quality gains from the large teacher model.