Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

CompeteSMoE -- Statistically Guaranteed Mixture of Experts Training via Competition

Nam V. Nguyen, Huy Nguyen, Quang Pham, Van Nguyen, Savitha Ramasamy, Nhat Ho

TL;DR

The paper tackles the challenge that traditional SMoE routing can be suboptimal because the router does not fully align with the experts' computations. It introduces a competition mechanism where all experts compute outputs and a winner-take-all rule selects the top responders, leading to better sample efficiency than softmax routing. Building on this, CompeteSMoE trains a router via a distillation objective to approximate the competition policy, augmented with a diversity loss and a carefully designed scheduling scheme to keep overhead low. Theoretical guarantees for Gaussian MoE show improved convergence and density-estimation rates, while empirical evaluations on vision-language fine-tuning and language pre-training demonstrate faster convergence and stronger zero-shot performance with modest training overhead. Overall, the approach offers a scalable, statistically grounded path to more effective SMoE training for large models with practical applicability.

Abstract

Sparse mixture of experts (SMoE) offers an appealing solution to scale up the model complexity beyond the mean of increasing the network's depth or width. However, we argue that effective SMoE training remains challenging because of the suboptimal routing process where experts that perform computation do not directly contribute to the routing process. In this work, we propose competition, a novel mechanism to route tokens to experts with the highest neural response. Theoretically, we show that the competition mechanism enjoys a better sample efficiency than the traditional softmax routing. Furthermore, we develop CompeteSMoE, a simple yet effective algorithm to train large language models by deploying a router to learn the competition policy, thus enjoying strong performances at a low training overhead. Our extensive empirical evaluations on both the visual instruction tuning and language pre-training tasks demonstrate the efficacy, robustness, and scalability of CompeteSMoE compared to state-of-the-art SMoE strategies. We have made the implementation available at: https://github.com/Fsoft-AIC/CompeteSMoE. This work is an improved version of the previous study at arXiv:2402.02526

CompeteSMoE -- Statistically Guaranteed Mixture of Experts Training via Competition

TL;DR

The paper tackles the challenge that traditional SMoE routing can be suboptimal because the router does not fully align with the experts' computations. It introduces a competition mechanism where all experts compute outputs and a winner-take-all rule selects the top responders, leading to better sample efficiency than softmax routing. Building on this, CompeteSMoE trains a router via a distillation objective to approximate the competition policy, augmented with a diversity loss and a carefully designed scheduling scheme to keep overhead low. Theoretical guarantees for Gaussian MoE show improved convergence and density-estimation rates, while empirical evaluations on vision-language fine-tuning and language pre-training demonstrate faster convergence and stronger zero-shot performance with modest training overhead. Overall, the approach offers a scalable, statistically grounded path to more effective SMoE training for large models with practical applicability.

Abstract

Sparse mixture of experts (SMoE) offers an appealing solution to scale up the model complexity beyond the mean of increasing the network's depth or width. However, we argue that effective SMoE training remains challenging because of the suboptimal routing process where experts that perform computation do not directly contribute to the routing process. In this work, we propose competition, a novel mechanism to route tokens to experts with the highest neural response. Theoretically, we show that the competition mechanism enjoys a better sample efficiency than the traditional softmax routing. Furthermore, we develop CompeteSMoE, a simple yet effective algorithm to train large language models by deploying a router to learn the competition policy, thus enjoying strong performances at a low training overhead. Our extensive empirical evaluations on both the visual instruction tuning and language pre-training tasks demonstrate the efficacy, robustness, and scalability of CompeteSMoE compared to state-of-the-art SMoE strategies. We have made the implementation available at: https://github.com/Fsoft-AIC/CompeteSMoE. This work is an improved version of the previous study at arXiv:2402.02526
Paper Structure (42 sections, 5 theorems, 87 equations, 8 figures, 14 tables)

This paper contains 42 sections, 5 theorems, 87 equations, 8 figures, 14 tables.

Table of Contents

  1. Introduction
  2. CompeteSMoE
  3. Background
  4. Routing via Competition
  5. Scheduled Training of the Router
  6. Router Loss
  7. Diversity Loss
  8. Router Training Schedule
  9. The CompeteSMoE Algorithm
  10. Statistical Guarantee of the Competition Mechanism
  11. Related Work
  12. Sparse Mixture of Experts
  13. Competitive Learning
  14. Experiment
  15. Experimental Settings
  16. ...and 27 more sections

Key Result

Proposition 3.1

With the MLE defined in equation (eq:MLE), the convergence rate of the density estimation $p_{\widehat{G}_n}(Y|X)$ to the ground-truth density $p_{G_*}(Y|X)$ is given by: Above, we denote $V(p_1,p_2):=\frac{1}{2}\int|p_1-p_2|\mathrm{d} m$ as the Total Variation distance between two probability density functions $p_1,p_2$ dominated by the Lebesgue measure $m$.

Figures (8)

  • Figure 1: The evolution of zero-shot performance averaged over nine visual instruction tuning tasks throughout training of various SMoE algorithms using a 5.1B parameters backbone.
  • Figure 2: An illustrative of the interleaved learning phases in CompeteSMoE: (a) activating all experts for the router to learn the competition policy; and (b) normal routing using the router.
  • Figure 3: Comparison of expert change rates at different training stages. Lower values are better.
  • Figure 4: Performance comparison of different activation functions used within the Competition Mechanism over 9 benchmarks.
  • Figure 5: Learning performance of $\mathcal{L}_{\mathcal{D}}$ and $\mathcal{L}_{{\mathcal{D}}_{\text{wo-reg}}}$ measured by the Level Learning metric at every 20% of training steps on the MMBench-EN benchmark.
  • ...and 3 more figures

Theorems & Definitions (6)

  • Proposition 3.1
  • Theorem 3.2
  • Theorem K.1
  • Lemma L.1: Theorem 7.4,Vandegeer-2000
  • Lemma L.2
  • proof : Proof of Lemma \ref{['lemma:upper_bounds']}