Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ReXMoE: Reusing Experts with Minimal Overhead in Mixture-of-Experts

Zheyue Tan, Zhiyuan Li, Tao Yuan, Dong Zhou, Weilin Liu, Yueqing Zhuang, Yadong Li, Guowei Niu, Cheng Qin, Zhuyu Yao, Congyi Liu, Haiyang Xu, Boxun Li, Guohao Dai, Bo Zhao, Yu Wang

TL;DR

ReXMoE decouples expert dimensionality from per-layer budgets, enabling richer expert combinations without sacrificing individual expert capacity or inflating overall parameters, and is confirmed as new design paradigm for parameter-efficient and scalable MoE-based LLMs.

Abstract

Mixture-of-Experts (MoE) architectures have emerged as a promising approach to scale Large Language Models (LLMs). MoE boosts the efficiency by activating a subset of experts per token. Recent works show that fine-grained experts substantially enriches the combinatorial flexibility of active experts and enhances model expressiveness. However, such a design is fundamentally limited by the layer-local routing mechanism: each layer is restricted to its own expert pool. This requires a careful trade-off between expert dimensionality and routing diversity given fixed parameter budgets. We describe ReXMoE, a novel MoE architecture that improves routing beyond the existing layer-local approaches by allowing routers to reuse experts across adjacent layers. ReXMoE decouples expert dimensionality from per-layer budgets, enabling richer expert combinations without sacrificing individual expert capacity or inflating overall parameters. To this end, we propose a new progressive scaling routing (PSR) strategy to gradually increase the candidate expert pool during training. As a result, ReXMoE improves both language modeling and downstream task performance. Extensive experiments on models ranging from 0.5B to 7B parameters across different architectures demonstrate that ReXMoE consistently improves performance under fixed architectural dimensions, confirming ReXMoE as new design paradigm for parameter-efficient and scalable MoE-based LLMs.

ReXMoE: Reusing Experts with Minimal Overhead in Mixture-of-Experts

TL;DR

ReXMoE decouples expert dimensionality from per-layer budgets, enabling richer expert combinations without sacrificing individual expert capacity or inflating overall parameters, and is confirmed as new design paradigm for parameter-efficient and scalable MoE-based LLMs.

Abstract

Mixture-of-Experts (MoE) architectures have emerged as a promising approach to scale Large Language Models (LLMs). MoE boosts the efficiency by activating a subset of experts per token. Recent works show that fine-grained experts substantially enriches the combinatorial flexibility of active experts and enhances model expressiveness. However, such a design is fundamentally limited by the layer-local routing mechanism: each layer is restricted to its own expert pool. This requires a careful trade-off between expert dimensionality and routing diversity given fixed parameter budgets. We describe ReXMoE, a novel MoE architecture that improves routing beyond the existing layer-local approaches by allowing routers to reuse experts across adjacent layers. ReXMoE decouples expert dimensionality from per-layer budgets, enabling richer expert combinations without sacrificing individual expert capacity or inflating overall parameters. To this end, we propose a new progressive scaling routing (PSR) strategy to gradually increase the candidate expert pool during training. As a result, ReXMoE improves both language modeling and downstream task performance. Extensive experiments on models ranging from 0.5B to 7B parameters across different architectures demonstrate that ReXMoE consistently improves performance under fixed architectural dimensions, confirming ReXMoE as new design paradigm for parameter-efficient and scalable MoE-based LLMs.
Paper Structure (35 sections, 5 equations, 9 figures, 9 tables)

This paper contains 35 sections, 5 equations, 9 figures, 9 tables.

Table of Contents

  1. introduction
  2. Related Works
  3. Mixture-of-Experts.
  4. Parameter Reusing in Transformers.
  5. Method
  6. Review of TopK Routing Mixture-of-Experts
  7. Expanding Candidate Expert Pool
  8. Discussion.
  9. Progressive Scaling Routing Strategy
  10. Experiments
  11. Experimental Setup
  12. Training environment.
  13. Training data.
  14. Evaluation metrics.
  15. Main Results
  16. ...and 20 more sections

Figures (9)

  • Figure 1: Overview of ReXMoE. Compared to vanilla MoE, ReXMoE enables more flexible expert combinations by reusing experts from adjacent layers. The only additional overhead comes from the router, which learns to route tokens to the expanded candidate pool. During training, ReXMoE progressively scale the candidate pool by gradually reducing the number of masked experts until all experts are available.
  • Figure 2: Comparison of prefill and decode goodput between base MoE and ReX models. Numbers above the bars indicate the relative speedup over the base MoE.
  • Figure 3: Average accuracy during training, distribution of load balance violations (degree of expert load imbalance), and distribution of under-utilized experts ratios (larger values indicate more inactive experts) under different cross-layer expert reuse sizes.
  • Figure 4: Loss curves for different strategies.
  • Figure 5: Activate ratio of MoE-SE and ReX-SE-R4 across layers in different tasks. The gray dashed lines indicate uniform distribution. ReXMoE shows stronger ability in task specialization.
  • ...and 4 more figures