PaperlayerView exact revision on arXiv

Source-linked AI summary

StableMoE: Stable Routing Strategy for Mixture of Experts

Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, Furu Wei

arXiv:2204.08396v1cs.LGcs.CL

TL;DR

Existing learning-to-route MoE methods can change the expert assigned to an input during training, even though inference uses one expert, which tends to harm sample efficiency. StableMoE learns and distills a routing strategy, then freezes the resulting router; across language modeling and multilingual machine translation, it outperforms existing MoE methods in convergence speed and performance.

  • Problem

    Existing learning-to-route MoE methods may change an input’s target expert during training, although inference activates only one, tending to harm sample efficiency.

    Source: §1 · p. 1

  • Method

    StableMoE learns a balanced, cohesive routing strategy, distills it into a lightweight router decoupled from the backbone, then freezes it for token assignment in stage two.

    Source: §1 · p. 1, §3.1 · p. 3

  • Results

    Across language modeling and multilingual machine translation, StableMoE outperforms existing MoE methods in convergence speed and performance.

    Source: §6 · pp. 8–9

  • Takeaways & Limitations

    StableMoE has a stable, cohesive, and balanced routing strategy, while the other three MoE methods cannot meet all three properties simultaneously.

    Source: §3.4 · p. 4

Abstract

from arXiv · show

The Mixture-of-Experts (MoE) technique can scale up the model size of Transformers with an affordable computational overhead. We point out that existing learning-to-route MoE methods suffer from the routing fluctuation issue, i.e., the target expert of the same input may change along with training, but only one expert will be activated for the input during inference. The routing fluctuation tends to harm sample efficiency because the same input updates different experts but only one is finally used. In this paper, we propose StableMoE with two training stages to address the routing fluctuation problem. In the first training stage, we learn a balanced and cohesive routing strategy and distill it into a lightweight router decoupled from the backbone model. In the second training stage, we utilize the distilled router to determine the token-to-expert assignment and freeze it for a stable routing strategy. We validate our method on language modeling and multilingual machine translation. The results show that StableMoE outperforms existing MoE methods in terms of both convergence speed and performance.

1 Introduction

The introduction identifies changing token-to-expert assignments as a sample-efficiency problem and proposes StableMoE’s two-stage routing strategy. It reports that routing changes persist late in training and that StableMoE outperforms existing MoE methods in convergence speed and performance.

1 Introduction

  • MoE lets Transformers expand their parameter scale by activating only a few experts per token, keeping computational and memory overhead tolerable.Source: §1 · p. 1
  • Existing learning-to-route methods can assign the same input to different experts during training, although inference activates only one, tending to harm sample efficiency.The same input can update different experts even though only one is finally used.Source: §1 · p. 1
  • 40.9% of tokens had a last fluctuation step after 20% of training, 29.1% after halfway, and 15.4% after 80%.The measurements examine validation tokens in BASE Layer.Source: §1 · p. 1
  • StableMoE learns a balanced, cohesive routing strategy, distills it into a lightweight router, then freezes that router for token assignment in stage two.The sigmoid gating mechanism propagates task-objective supervision to the routing strategy, while a balance loss encourages balanced assignment.Source: §1 · pp. 1–2
  • Experiments on language modeling and multilingual machine translation show StableMoE outperforms existing MoE methods in convergence speed and performance.Source: §1 · p. 2

2 Background: Mixture-of-Experts for Transformers

The section describes Transformer layer computation and how sparse expert gating is incorporated into Transformer blocks. It also includes a figure illustrating StableMoE’s two-stage routing procedure.

2 Background: Mixture-of-Experts for Transformers

  • A standard L-layer Transformer maps an input sequence of T tokens through its layers; each token has a hidden state, and the layer computation uses self-attention and a feed-forward network.Source: §2 · p. 2
  • MoE layers insert a set of FFN experts between neighboring Transformer blocks, with sparse gating activating only a few or one expert per token.The gating function is sparse for computational efficiency.Source: §2 · p. 2
  • StableMoE’s illustrated procedure learns a routing strategy, distills it into a lightweight router, then freezes that router for stable routing in stage two.Source: Figure 3 · p. 3

3 Method

STABLEMOE addresses routing fluctuation by learning balanced, cohesive routing, distilling it into a lightweight router, and freezing that router for stage 2 and inference. The comparison presents the two-stage design as combining stable, cohesive, and balanced routing.

3.1 Training Stage 1: Learn Routing Strategy

  • Stage 1 learns balanced assignments with a balance loss and cohesive routing through sigmoid-gate feedback, while distilling the learned strategy into a lightweight router.Tokens are assigned to their highest-affinity expert during this stage.Source: §3.1 · p. 3

3.2 Training Stage 2: Learn with Stable Routing Strategy

  • In stage 2, the frozen distilled router fixes token-to-expert assignments, and training uses the task loss without balance or distillation losses.Source: §3.2 · p. 4

3.3 Inference

  • Inference also uses the frozen distilled router, keeping routing consistent with stage 2.Source: §3.3 · p. 4

3.4 Comparison with Existing MoE Methods

  • STABLEMOE combines stable, cohesive, and balanced routing, while the other three compared MoE methods do not meet all three properties simultaneously.Source: §3.4 · p. 4

4 Experiments

Experiments evaluate StableMoE on language modeling and multilingual machine translation, where it outperforms other MoE methods and converges faster in language modeling. The section also describes the datasets and training setups used for these evaluations.

4.3.1 Language Modeling

  • StableMoE lowers language-model perplexity by 0.3-0.8 versus existing MoE methods under the base setting and consistently outperforms them in both base and large settings.It also achieves about 3.7 lower perplexity than the standard Transformer and about 1.3 higher perplexity than the deeper larger model in the base setting.Source: Table 2 · p. 5, §4.3.1 · p. 6

4.1 Tasks and Datasets

  • Language modeling uses the English subset of CC100; translation uses 32.5 million WMT parallel data for English and nine other languages, augmented with 180 million backtranslation data; runs train 60K and 352K steps, respectively.Source: §4.1 · p. 5, §4.2 · p. 5

4.2 Experimental Setup

  • Both tasks use two training stages; the language-modeling setup inserts one MoE layer after the middle Transformer block, while translation inserts layers in the encoder and decoder.The respective stage-1 allocations are 6K and 30K steps.Source: §4.2 · p. 5

4.3.2 Multilingual Machine Translation

  • At the same FLOPs, StableMoE beats the dense model by 1.22 test BLEU; MoE expands parameters by 523% while increasing FLOPs by 9.3%.StableMoE also achieves the best average test BLEU among the compared MoE methods on WMT X→En translation.Source: Table 3 · p. 6, §4.3.2 · p. 6

4.3.1 Language Modeling

  • With validation perplexity plotted against training wall time, MoE models usually converge faster than dense models, and StableMoE is fastest among the MoE methods.Larger dense models achieve better final validation perplexity.Source: Figure 4 · p. 6, §4.3.1 · p. 6

4.4 Analysis

The analysis tests routing stability and several design choices, including expert count, expert parameters, MoE-layer placement, stage allocation, and distilled-router variants. It finds that fixed routing improves performance and prevents routing fluctuation during StableMoE’s second training stage.

4.4.1 Effects of Hyperparameters

  • As expert counts increase, validation perplexity tends to fall; StableMoE with 16 experts outperforms BASE Layer with 32, and 32 experts achieves similar perplexity to BASE Layer with 64.StableMoE performs best across the tested expert counts.Source: §4.4.1 · pp. 6–7
  • Increasing expert parameters from 454M to 1.51B improves performance, and StableMoE performs best under both settings.The comparison varies the number of expert sublayers.Source: Figure 6 · p. 7, §4.4.1 · p. 7
  • Stacking three MoE sublayers in the middle Transformer position performs better than placing them at the bottom, top, or uniformly scattered.The comparison covers four MoE-layer placements.Source: Table 4 · p. 7, §4.4.1 · p. 7

4.4.2 Effects of the Fixed Routing Strategy

  • Fixed routing lowers BASE Layer validation perplexity by 0.63, while disabling StableMoE’s fixed-routing stage raises it by 0.20 versus the full model.The authors report that both comparisons support a performance benefit from fixed routing.Source: Table 5 · p. 7, §4.4.2 · p. 7

4.4.1 Effects of Hyperparameters

  • With a word-embedding distilled router, assigning 6K steps, or 10% of total training, to stage 1 is reported as a good balance point.The authors speculate that the embedding is learned quickly, making longer stable training more important.Source: §4.4.1 · p. 7

4.4.3 Variants of Distilled Routers

  • Among tested distilled routers, word embeddings perform best, while the 3-layer Transformer performs poorly; the authors recommend the simpler word-embedding router.They suggest the 32-category classification signal may be insufficiently informative for learning a complex router.Source: Table 6 · p. 8, §4.4.3 · p. 8

4.4.4 Analysis of Routing Fluctuations

  • BASE Layer shows notable routing fluctuation, whereas StableMoE has none during stage 2 because its routing strategy is fixed.The analysis tracks validation-token assignments every 500 steps and defines each token’s last fluctuation step relative to its final-step expert.Source: Figure 7 · p. 8, §4.4.4 · p. 8

5 Related Work

The related-work discussion contrasts dynamically learned token routing with static hash routing and positions StableMoE as combining learned assignments with a frozen router.

Learning-to-route methods

Static routing

  • Hash Layer (Roller et al., 2021) uses a static, pre-designed token-level hash table, but its randomly determined assignments limit its performance ceiling.Source: §5 · p. 8

StableMoE

  • StableMoE aims to retain learned balanced and cohesive routing while preventing fluctuation with a frozen lightweight router that mimics the original strategy.Source: §5 · p. 8

6 Conclusion

The conclusion presents StableMoE as a two-stage response to routing fluctuation: it learns and distills a routing strategy, then freezes the lightweight router. The paper reports evaluation on language modeling and multilingual machine translation, with better performance and convergence speed than existing MoE methods.

Conclusion

  • StableMoE first learns a balanced, cohesive routing strategy and distills it into a lightweight router decoupled from the backbone, then freezes the router for remaining training.Source: §6 · p. 8
  • The paper evaluates StableMoE on language modeling and multilingual machine translation.Source: §6 · p. 8
  • The reported results show that StableMoE outperforms existing MoE methods in both convergence speed and performance.Source: §6 · p. 9

A Hyper-parameters for Language Modeling

This section directs readers to the language-modeling hyper-parameters for STABLEMOE's base and large settings in Table 7.

A Hyper-parameters for Language Modeling

B Hyper-parameters for Multilingual Machine Translation

This section points to the hyper-parameters for STABLEMOE in multilingual machine translation and includes a representative visualization of expert-assigned tokens.

B Hyper-parameters for Multilingual Machine Translation

C Visualization of the Fixed Routing Strategy of STABLEMOE

The section visualizes STABLEMOE's fixed routing strategy through frequent validation-set tokens assigned to experts. The examples show that tokens assigned to the same expert usually share common features.

C Visualization of the Fixed Routing Strategy of STABLEMOE

Loading 2204.08396v1…