CompeteSMoE -- Statistically Guaranteed Mixture of Experts Training via Competition

TL;DR

This paper introduces CompeteSMoE, a novel training method for sparse mixture of experts that uses a competition-based routing mechanism, improving efficiency, robustness, and scalability in large language and vision models.

Contribution

It proposes a new competition-based routing mechanism for SMoE, with theoretical guarantees and an effective training algorithm for large models.

Findings

Better sample efficiency than softmax routing

Strong performance on language and vision tasks

Robustness and scalability demonstrated

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