$\infty$-MoE: Generalizing Mixture of Experts to Infinite Experts

TL;DR

This paper introduces $ abla$-MoE, a novel approach that extends the Mixture of Experts model to an infinite number of experts by selecting parameter subsets based on continuous values, improving training stability and efficiency.

Contribution

It proposes a continuous expert selection mechanism enabling infinite experts, enhancing training stability and flexibility in model size and inference trade-offs.

Findings

Achieves comparable performance to larger dense models with fewer parameters.

Allows flexible adjustment of experts at inference for accuracy-speed trade-off.

Improves accuracy by up to 2.5% over traditional MoE.

Abstract

The Mixture of Experts (MoE) selects a few feed-forward networks (FFNs) per token, achieving an effective trade-off between computational cost and performance. In conventional MoE, each expert is treated as entirely independent, and experts are combined in a discrete space. As a result, when the number of experts increases, it becomes difficult to train each expert effectively. To stabilize training while increasing the number of experts, we propose $\infty$-MoE that selects a portion of the parameters of large FFNs based on continuous values sampled for each token. By considering experts in a continuous space, this approach allows for an infinite number of experts while maintaining computational efficiency. Experiments show that a GPT-2 Small-based $\infty$-MoE model, with 129M active and 186M total parameters, achieves comparable performance to a dense GPT-2 Medium with 350M parameters. Adjusting the number of sampled experts at inference time allows for a flexible trade-off between accuracy and speed, with an improvement of up to 2.5\% in accuracy over conventional MoE.