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

TL;DR

This work tackles the scalability limits of discrete Mixture of Experts (MoE) by introducing ∞-MoE, which treats experts as a continuous space and activates only a sparse subset per token. The method defines a router over a continuous index $z$, yielding $y = \int p(z|x) f(x,z)\,dz$ and approximating it with Monte Carlo sampling; a Gaussian router with mean $\mu(x)$ and diagonal covariance $\Sigma(x)$ guides sampling, while each sampled $z$ modulates an FFN via a top-$N\%$ mask. The FFN is constructed as $f(x,z) = W_2(\mathrm{Act}(W_1 x) \odot \mathrm{mask}(z))$, enabling effectively infinite experts with per-token sparsity. Empirically, ∞-MoE outperforms Dense, Switch Transformer, and standard MoE on GPT-2 Small/Medium zero-shot benchmarks, with the inference-time sampling rate $K$ providing a controllable speed–accuracy tradeoff; ablations show data scale, index dimension, and custom kernels further shape performance and efficiency. Overall, the approach offers a parameter-efficient route to vastly larger expressive capacity and adaptable runtime behavior for large language models.

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.

Figures (4)

• Figure 1: Overview of the proposed Infinite Mixture of Experts ($\infty$-MoE). The router outputs a continuous distribution over the expert space, and each sample selects a unique expert.
• Figure 2: Comparison of MoE and $\infty$-MoE models on several tasks while varying the number of experts $K \in \{1,2,3,4,8\}$. For GPT-2 Small, $K=2$ yields 124 million active parameters.$\infty$-MoE consistently achieves strong accuracy across a wide range of $K$, even with fewer experts.
• Figure 3: Accuracy on HellaSwag as a function of training data size (in billions of tokens). $\infty$-MoE is compared against a MoE baseline (GPT-2 Small backbone).
• Figure 4: FFN forward latency vs. active rate. The custom kernel achieved 1.25-fold speedup.