Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

StableMoE: Stable Routing Strategy for Mixture of Experts

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

TL;DR

This work identifies routing fluctuation as a key inefficiency in learning-to-route MoE Transformers, where token-to-expert assignments vary during training but inference uses a single activated expert. It introduces StableMoE, a two-stage approach that first learns a balanced, cohesive routing and distills it into a lightweight router, then freezes this router to provide stable routing throughout training and inference. Empirical results on language modeling and multilingual machine translation show StableMoE achieves faster convergence and improved perplexity/BLEU over prior MoE methods, with robust performance across hyperparameters and router variants. The method effectively combines the benefits of learning-to-route strategies with a fixed, stable routing in deployment, enabling scalable, efficient training of large MoE-augmented transformers.

Abstract

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.

StableMoE: Stable Routing Strategy for Mixture of Experts

TL;DR

This work identifies routing fluctuation as a key inefficiency in learning-to-route MoE Transformers, where token-to-expert assignments vary during training but inference uses a single activated expert. It introduces StableMoE, a two-stage approach that first learns a balanced, cohesive routing and distills it into a lightweight router, then freezes this router to provide stable routing throughout training and inference. Empirical results on language modeling and multilingual machine translation show StableMoE achieves faster convergence and improved perplexity/BLEU over prior MoE methods, with robust performance across hyperparameters and router variants. The method effectively combines the benefits of learning-to-route strategies with a fixed, stable routing in deployment, enabling scalable, efficient training of large MoE-augmented transformers.

Abstract

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.
Paper Structure (37 sections, 9 equations, 7 figures, 9 tables)

This paper contains 37 sections, 9 equations, 7 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Background: Mixture-of-Experts for Transformers
  3. Method
  4. Training Stage 1: Learn Routing Strategy
  5. Balance Loss
  6. Distilled Router
  7. Training Objective
  8. Training Stage 2: Learn with Stable Routing Strategy
  9. Inference
  10. Comparison with Existing MoE Methods
  11. Assignment Algorithm
  12. Gating Function
  13. Balance Loss
  14. Experiments
  15. Tasks and Datasets
  16. ...and 22 more sections

Figures (7)

  • Figure 1: Illustration of the routing fluctuation problem. The same input is assigned to different experts along with training. However, during inference, only one expert is sparsely activated for the input. The routing fluctuation tends to harm sample efficiency because the same input updates different experts while only one is used.
  • Figure 2: Cumulative token percentage with regard to the last fluctuation step of tokens for BASE Layer base. A substantial portion of tokens still change their target experts even if the training is nearing the end.
  • Figure 3: Illustration of two training stages in StableMoE. In training stage 1, we learn a routing strategy and distill it into a lightweight router. Then, we freeze the distilled router for stable routing in training stage 2.
  • Figure 4: Convergence speed of different models. TRM is a shorthand for Transformer.
  • Figure 5: Comparison of MoE-based Transformers with different numbers of experts. Lower perplexity indicates better performance.
  • ...and 2 more figures