PaperlayerView exact revision on arXiv

Source-linked AI summary

Taming Sparsely Activated Transformer with Stochastic Experts

Simiao Zuo, Xiaodong Liu, Jian Jiao, Young Jin Kim, Hany Hassan, Ruofei Zhang, Tuo Zhao, Jianfeng Gao

arXiv:2110.04260v3cs.CLcs.LG

TL;DR

SAMs offer computationally efficient scaling, but their larger parameter counts do not reliably improve performance and routing remains problematic. THOR uses randomly activated experts with consistency-regularized training, achieving stronger machine-translation results, including a 2 BLEU multilingual gain over an equal-sized Switch model and parity with an MoE 18 times larger.

  • Problem

    Sparsely activated models can scale parameters without comparable computational growth, but they are reported to be parameter inefficient and commonly rely on problematic gating-based routing.

  • Method

    THOR randomly activates experts during training and inference and uses consistency regularization so experts learn from one another as teachers.

  • Results

    THOR consistently outperforms Transformer and MoE models across machine-translation settings, including a 2 BLEU multilingual improvement over an equal-sized Switch model and parity with an MoE 18 times larger.

  • Takeaways & Limitations

    THOR is more parameter-efficient, less prone to overfitting, and produces more consistent predictions across the evaluated translation settings.

  • Takeaways & Limitations

    Inference remains stochastic: different random seeds can produce a 0.5 BLEU difference on a Czech-to-English dataset, and the paper focuses on NLP, especially machine translation.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Sparsely activated models reduce the computational cost of scaling neural networks but can be parameter inefficient. THOR replaces learned gating with stochastic experts and consistency regularization, achieving strong machine-translation results across resource settings.

  • 1 INTRODUCTION: Sparsely activated models select only a subset of parameters for each input, enabling much larger models without a significant increase in computational cost.Switch-large has 35 times more parameters than T5-large while requiring the same forward FLOPs.
  • 1 INTRODUCTION: Despite their scale and efficiency, SAMs can be parameter inefficient, with Switch-large only slightly exceeding T5-large on GLUE and underperforming it on ARC.The reported GLUE scores are 88.5 versus 87.8, while ARC scores are 66.0 versus 68.8.
  • 1 INTRODUCTION: Existing SAM research mainly improves input routing through gating networks, but these mechanisms can suffer from load imbalance when weights collapse toward one expert.Proposed mitigations include noisy gating, expert capacity limits, and load-balancing losses.
  • 1 INTRODUCTION: THOR randomly activates experts without gating and trains them with cross-entropy plus consistency regularization from other experts as teachers.The objective encourages all experts to make consistent predictions while learning from training data.
  • 1 INTRODUCTION: THOR outperforms state-of-the-art MoE models by an average of 2 BLEU on twelve low-resource translation tasks and matches a 5.5-billion-parameter MoE in multilingual translation with 18 times fewer parameters.In multilingual translation, the 300-million-parameter THOR model improves over an equal-sized MoE by 2 BLEU.

2 BACKGROUND

Sparsely activated models use parallel expert networks and route each input to only a small subset of them. This keeps computation roughly independent of the total number of experts, although learned gates can collapse and create load imbalance.

  • 2 BACKGROUND: An expert layer contains an attention mechanism and multiple parallel feed-forward networks, with each feed-forward network functioning as an expert.A gating mechanism determines which expert or experts receive each input.
  • 2 BACKGROUND: The gating mechanism computes routing scores from the input using a trainable weight matrix, then selects the top-K experts for activation.The activated set T contains K experts selected from the total N experts.
  • 2 BACKGROUND: Activating K experts instead of N keeps forward-pass FLOPs from increasing with the total number of experts when K is much smaller than N.The paper gives K = 2 and N = 2048 in GShard as an example.
  • 2 BACKGROUND: Learned gate weights can collapse so that nearly all inputs are routed to one expert, creating load imbalance.Existing approaches use heuristics such as Gaussian noise, capacity limits, and balancing losses to mitigate this issue.

3 ANALYSIS OF SPARSELY ACTIVATED MODELS

Experiments compare MoE and Switch Transformer routing behaviors against Transformer baselines, finding similar performance across gating and random-routing variants. The analysis identifies load imbalance and effectively random gating behavior, while noting the paper’s NLP and translation scope.

  • MoE(dec) and MoE(tok) perform similarly, and neither significantly outperforms the Transformer on validation loss or BLEU.
  • Gating can collapse toward one expert, whereas another observed regime distributes inputs nearly equally with almost identical routing probabilities.
  • Switch Transformer variants perform similarly whether routing by token, sentence, or randomly, and do not outperform the vanilla Transformer after load imbalance is alleviated.
  • The paper focuses on natural language processing, particularly neural machine translation, and notes that other fields may yield different conclusions.
  • The analysis identifies two gating behaviors: load imbalance and random routing.

4 THOR: TRANSFORMER WITH STOCHASTIC EXPERTS

THOR replaces gating-based routing with stochastic expert selection and consistency-regularized training. It addresses inference variability through consistency while simplifying routing and avoiding load imbalance and extra gating parameters.

  • Routing design: Random expert selection removes the load-imbalance problem and eliminates the need for load-balancing heuristics or additional gating parameters.Each expert receives a fair chance to be trained across iterations, and THOR adds no parameters for routing.
  • Limitation: Random expert selection can make inference results vary across random seeds, producing a 0.5 BLEU difference on a Czech-to-English translation dataset.This variability is the main limitation of stochastic inference and motivates the consistency regularizer.
  • THOR training: THOR randomly selects expert pairs for each layer during training, then trains the paired predictions with cross-entropy and consistency regularization.Experts learn from training labels and from each other as teachers, encouraging consistent predictions while retaining stochastic specialization.
  • Inference: During inference, THOR can randomly select an expert per layer and input, or use alternatives such as expert ensembling.The inference procedure follows the stochastic-selection design used during training, with ensemble selection discussed as another option.
  • Relation to dropout: THOR is conceptually similar to dropout but can be used during inference and yields a more structured sparse model suitable for distributing experts across GPUs.Unlike dropout, THOR selects expert components and supports serving-time sparsity.

5 EXPERIMENTS

THOR is evaluated against Transformer, Switch Transformer, and regularized baselines across low-resource, rich-resource, and multilingual translation settings. It consistently improves translation quality, with the strongest parameter-efficiency gains in multilingual translation.

  • 5.2 LOW-RESOURCE TRANSLATION: More than 1.0 BLEU: THOR improves over Switch Transformer on all 12 low-resource translation tasks, while Switch matches Transformer at 26.3 average BLEU.THOR also significantly outperforms SMART and R3F under the same comparison.
  • 5.3 RICH-RESOURCE TRANSLATION: 0.3 BLEU on En-De and 0.1 BLEU on En-Fr: THOR sets new state-of-the-art results on the two rich-resource WMT benchmarks without data augmentation or pretrained language models.The comparison uses BLEU and sacreBLEU scores.
  • 5.4 MULTILINGUAL TRANSLATION: 17 of 20 tasks: THOR outperforms the multilingual baseline, with gains of 4.7 BLEU on Gu-En and 6.7 BLEU on Hi-En.Improvements are generally larger on smaller datasets, while larger datasets still show gains such as +0.9 BLEU on Cs-En and +1.1 on Fi-En.
  • 5.5 ABLATION EXPERIMENTS: Consistency regularization is crucial: removing it causes the main performance loss, whereas removing either cross-entropy term produces only a small BLEU decrease.The two cross-entropy terms play similar roles in training.
  • 5.5 ABLATION EXPERIMENTS: THOR remains robust as regularization increases above α > 2.0, makes more consistent predictions than Switch Transformer, and avoids the overfitting seen in larger Transformer FFNs.THOR’s prediction variance is below 0.002, and BLEU consistently improves as experts increase from 2 to 8.

6 CONCLUSION

The paper concludes that THOR provides an efficient alternative to classic sparsely activated models. Across translation settings, it combines stochastic experts with consistent predictions and stronger generalization than Transformer and Switch Transformer baselines.

  • 6 CONCLUSION: THOR randomly activates experts and trains them with consistency regularization so experts learn from both data and one another.This design removes the need for a gating mechanism while encouraging consistent predictions.
  • 6 CONCLUSION: THOR preserves the efficient training and serving profile of classic MoE models while improving parameter efficiency, prediction consistency, and resistance to overfitting.The conclusion attributes these properties to results across the evaluated settings.
  • 6 CONCLUSION: Across low-resource, rich-resource, and multilingual translation, THOR significantly outperforms both vanilla Transformer and Switch Transformer.The conclusion summarizes a comprehensive machine-translation evaluation.

A.1 TRAINING DETAILS

The paper compares THOR with alternative expert formulations and specifies the multilingual data preprocessing and optimization setup. The multilingual data combines German-English and Vietnamese-English corpora with separate language dictionaries.

  • A.1 TRAINING DETAILS: MoE(dec) uses a separate Transformer decoder per expert, whereas MoE(tok) represents experts with distinct learned input tokens.For MoE(tok), routing replaces the beginning-of-sentence token with an expert-specific token.
  • A.1 TRAINING DETAILS: The multilingual data combines preprocessed De-En and Vi-En datasets using BPE with 10,000 merge operations.Separate dictionaries are learned for English and the combined German-plus-Vietnamese languages.
  • A.1 TRAINING DETAILS: Training uses Adam with learning rate 0.001 and an effective batch size of 64k tokens, while inference uses beam size 5 and length penalty 1.0.The setup uses 8k tokens per GPU across 8 GPUs.

A.2 ADDITIONAL RESULTS

Additional Switch Transformer analyses find that gating behaves like random routing, while load balancing changes expert utilization without changing this routing behavior.

  • The examined figures compare average routing confidence and expert load for sentence-level and token-level Switch variants.Figures 10–13 cover configurations with and without load-balancing loss.
  • Without load balancing, Switch(s) routes nearly randomly despite a 10% expert-load difference between 55% and 45%.After about 10k iterations, both experts have routing confidence around 0.60.
  • Without load balancing, Switch(t) also exhibits random routing, with both experts’ confidence converging to around 0.55 and imbalanced loads.Token-level routing allows different tokens in one sentence to reach different experts.
  • Load-balancing loss reduces expert-load imbalance to below 0.4% for Switch(s) and below 0.2% for Switch(t), versus about 10% without it.Routing confidence still converges to about 0.60 for Switch(s) and 0.55 for Switch(t).

B DATASETS

The paper evaluates translation using low-resource and multilingual datasets drawn from IWSLT, Europarl, and WMT sources.

  • Low-resource experiments use English-Vietnamese, English-German, and English-French data from IWSLT, plus English-Romanian, English-Latvian, and English-Czech data from Europarl and WMT.The latter languages use Europarl training data and WMT’17 validation and testing data.
  • Table 6 summarizes the statistics of the low-resource translation datasets used in the experiments.
  • The multilingual translation task uses datasets whose statistics and sources are summarized in Table 7.English is the other language in all listed multilingual translation tasks.

C TRAINING DETAILS

Training uses dataset-specific preprocessing, optimization, regularization, batching, and decoding settings for the reported translation experiments.

  • IWSLT and WMT data use joint source-target dictionaries with BPE vocabularies built from 10,000 and 40,000 split operations, respectively.Other preprocessing follows Fairseq.
  • The low-resource setup uses α = 5.0, 32k-token batches, Adam, learning rate 0.0015, and 40k training steps.The model with the highest validation BLEU is tested, using beam size 5 for decoding.
  • Another translation setup uses α = 2.0, Adam with learning rate 0.001, and beam settings of 4 for En-De and 10 for En-Fr.The reported metrics include BLEU and sacreBLEU.
  • A larger training configuration uses 1.6 million-token batches, 24 GPUs, RAdam, learning rate 0.05, dropout 0.1, label smoothing 0.1, and α = 4.0.

D ADDITIONAL EXPERIMENTS

Additional experiments show that increasing experts improves Switch and THOR performance, with THOR outperforming both baselines even with two experts, while routing overhead slows inference.

  • Switch Transformer trains faster than the vanilla Transformer, and this scaling advantage becomes more pronounced as the number of experts increases.The comparison uses training perplexity against wall-time measured in GPU hours.
  • With more experts, Switch Transformer validation BLEU improves beyond the vanilla Transformer at the same FLOPs, whereas 2 experts perform slightly worse.The result indicates that Switch requires substantially more parameters to surpass dense performance.
  • THOR significantly outperforms the vanilla Transformer and Switch Transformer even with 2 experts, and its performance improves as experts increase.The experiment uses Transformer-IWSLT on WMT’16 En-De, with model sizes ranging from 48M to 456M parameters.
  • Switch Transformer and THOR infer more slowly than the vanilla Transformer because input routing introduces computation and communication overhead.The overhead becomes more noticeable as the number of experts grows.
Loading 2110.04260v3…