[Paper Review] Taming Sparsely Activated Transformer with Stochastic Experts

Simiao Zuo, Xiaodong Liu+6 more|arXiv (Cornell University)|Oct 8, 2021

Topic Modeling33 references27 citations

TL;DR

THOR random-expert activation with a consistency regularizer outperforms standard Transformer and Switch MoE models across low-resource, rich-resource, and multilingual MT tasks, showing improved parameter efficiency.

ABSTRACT

Sparsely activated models (SAMs), such as Mixture-of-Experts (MoE), can easily scale to have outrageously large amounts of parameters without significant increase in computational cost. However, SAMs are reported to be parameter inefficient such that larger models do not always lead to better performance. While most on-going research focuses on improving SAMs models by exploring methods of routing inputs to experts, our analysis reveals that such research might not lead to the solution we expect, i.e., the commonly-used routing methods based on gating mechanisms do not work better than randomly routing inputs to experts. In this paper, we propose a new expert-based model, THOR (Transformer witH StOchastic ExpeRts). Unlike classic expert-based models, such as the Switch Transformer, experts in THOR are randomly activated for each input during training and inference. THOR models are trained using a consistency regularized loss, where experts learn not only from training data but also from other experts as teachers, such that all the experts make consistent predictions. We validate the effectiveness of THOR on machine translation tasks. Results show that THOR models are more parameter efficient in that they significantly outperform the Transformer and MoE models across various settings. For example, in multilingual translation, THOR outperforms the Switch Transformer by 2 BLEU scores, and obtains the same BLEU score as that of a state-of-the-art MoE model that is 18 times larger. Our code is publicly available at: https://github.com/microsoft/Stochastic-Mixture-of-Experts.

Motivation & Objective

• Motivate why sparsely activated models (SAMs) may be parameter inefficient despite large parameter counts.
• Investigate whether gating-based routing (top-k experts) outperforms random routing in MoE-style architectures.
• Propose THOR, a SAM where experts are randomly activated with a consistency regularizer to align predictions across experts.
• Evaluate THOR on low-resource, rich-resource, and multilingual machine translation to assess parameter efficiency and generalization.

Proposed method

• Analyze gating-based MoE models and identify load imbalance and random routing behaviors.
• Introduce THOR where a pair of experts is randomly activated per layer per training iteration and during inference.
• Opt a consistency regularizer that minimizes cross-entropy losses from two random expert selections plus a KL-based consistency term.
• Train THOR with a dual-teacher-like setup where experts learn from each other to produce consistent predictions.
• Evaluate THOR under low-resource, rich-resource, and multilingual MT settings using standard benchmarks.

Experimental results

Research questions

• RQ1 Do sparsely activated models inherently underperform compared to densely activated models of similar size?
• RQ2 Is gating-based routing necessary for MoE-like models to achieve gains, or can random expert activation be effective?
• RQ3 Can a consistency regularizer enable robust training and inference when experts are randomly activated?
• RQ4 What is THOR’s performance advantage over Transformer and Switch Transformer across MT tasks and settings?

Key findings

• THOR consistently surpasses the vanilla Transformer and the Switch Transformer across all three settings.
• In low-resource MT, THOR improves average BLEU by over 1.0 point vs Switch and outperforms SMART and R3F baselines.
• In rich-resource MT, THOR sets new state-of-the-art on En-De and En-Fr without data augmentation or pretraining.
• In multilingual MT, THOR with 300M parameters matches the BLEU of a 5.5B-parameter Switch MoE, demonstrating 18x parameter efficiency.
• THOR exhibits higher prediction consistency and reduced variance compared to Switch Transformer, with less overfitting as model size grows.