💬 NLP

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

The paper proposes $\infty$ -MoE, a novel architecture that generalizes Mixture of Experts to an infinite space by selecting continuous portions of large feed-forward network parameters per token, thereby enabling effective training with a massive number of experts while achieving performance comparable to much larger dense models and offering a flexible accuracy-speed trade-off.

Original authors: Shota Takashiro, Takeshi Kojima, Shohei Taniguchi, Yusuke Iwasawa, Yutaka Matsuo

Published 2026-01-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shota Takashiro, Takeshi Kojima, Shohei Taniguchi, Yusuke Iwasawa, Yutaka Matsuo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, high-tech library where every book represents a piece of knowledge. In a traditional library (a standard AI model), you have a few very large, general-purpose librarians. When you ask a question, every librarian tries to answer it, and their answers are blended together. This is accurate but slow and expensive because you're paying all of them to work on every single question.

To fix this, researchers invented the Mixture of Experts (MoE). Instead of having everyone work, you have a "Head Librarian" (called a router) who looks at your question and picks just two specific librarians to answer it. This is much faster and cheaper. However, there's a catch: in this traditional setup, you have to decide exactly how many librarians to hire (say, 16 or 100). If you hire too many, the Head Librarian gets confused about who to pick, and training the whole system becomes a nightmare. It's like trying to manage a team of 1,000 people when you can only talk to two at a time; it's hard to keep everyone in sync.

The New Idea: ∞-MoE (Infinite Mixture of Experts)

The authors of this paper, from the University of Tokyo, asked a bold question: What if we didn't have a fixed number of librarians at all? What if we had an infinite number of them?

They call their new system ∞-MoE. Here is how it works, using a simple analogy:

1. The Continuous "Expert Space"

Instead of having Librarian #1, Librarian #2, and Librarian #3 sitting in a row, imagine the librarians are spread out across an infinite, continuous map.

In the old system, you had to pick specific points on a map (like "Librarian at coordinate 5").
In the new system, the Head Librarian draws a cloud on the map for every question you ask. This cloud represents a range of experts that might be helpful.

2. Sampling Instead of Picking

When you ask a question, the Head Librarian doesn't just pick two specific people. Instead, they take a "snapshot" (a sample) from that cloud.

They might pick a point at coordinate 5.2.
They might pick a point at coordinate 5.3.
Because the map is continuous, every single time they take a snapshot, they get a unique expert that has never existed before.

3. The "Mask" (Turning Neurons On and Off)

How do you have infinite experts without building infinite computers? The secret is in the mask.
Think of the AI's brain as a giant grid of lightbulbs (neurons).

In a standard AI, all the lightbulbs are on.
In the old MoE, the router turns on a specific block of lightbulbs for Expert #1 and a different block for Expert #2.
In ∞-MoE, the "sample" (the coordinate on the map) acts like a stencil. The system takes the giant grid of lightbulbs and uses the stencil to turn only the top 25% of the brightest bulbs on for that specific question.
Because the stencil is based on a continuous number, every question gets a slightly different pattern of lightbulbs turned on. It's like having a unique, custom-made tool for every single sentence you type, but you only build the tool out of parts you already have.

Why is this a big deal?

The paper tested this on two models (GPT-2 Small and Medium) and found some impressive results:

Better Performance with Less "Active" Brain:
The ∞-MoE model with 129 million active parameters (the parts actually working) performed just as well as a standard, dense model with 350 million parameters. It's like getting the intelligence of a giant brain while only using the energy of a small one.
Flexibility:
In the old MoE, if you wanted to go faster, you had to retrain the whole model. In ∞-MoE, you can just change the number of "samples" (snapshots) you take at the moment you ask a question.
- Need speed? Take fewer samples (fewer experts).
- Need maximum accuracy? Take more samples.
- The paper shows you can get a 2.5% accuracy boost over standard MoE just by tweaking this setting, without changing the model itself.
Stability:
One of the biggest problems with adding more experts to old systems is that they become unstable and hard to train. Because ∞-MoE treats experts as a smooth, continuous space rather than a jumpy list of separate people, the training stays stable even as the "number" of experts grows toward infinity.

The Catch (Limitations)

The authors are honest about what they haven't solved yet:

Scaling Up: They tested this on medium-sized models. They aren't 100% sure how it will behave on the massive models used by companies today (like GPT-4 scale), though they expect it to work.
Hardware Speed: While the math says it should be fast, the actual computer code is tricky. Because the "lightbulbs" turned on change for every single word, it's hard for standard computer chips to process them all at once efficiently. They had to write special code to make it run fast, and there is still room for improvement.

Summary

∞-MoE is like upgrading from a fixed team of specialists to a shape-shifting, infinite team. Instead of hiring 100 specific people, you have a magical machine that can instantly create a unique, perfect specialist for every single question you ask, using only a fraction of the computer power. It allows AI to be smarter and more flexible without getting bogged down by the cost of having a massive, static brain.

Technical Summary: ∞-MoE: Generalizing Mixture of Experts to Infinite Experts

1. Problem Statement

Large Language Models (LLMs) have achieved significant performance gains through scaling parameters, data, and compute. However, simply increasing model size incurs substantial computational and memory overheads. Mixture of Experts (MoE) architectures address this by routing each token to a sparse subset of specialized experts, expanding parameter capacity without a linear increase in per-token compute costs.

Despite these benefits, conventional MoE faces limitations when scaling the number of experts:

Discrete Constraints: Experts are treated as entirely independent entities in a discrete space.
Training Instability: As the number of experts increases, training each expert effectively becomes difficult.
Router Overhead: Aggressively increasing the expert pool (e.g., to millions) can lead to router overhead growth or over-compressed experts that degrade accuracy.

The central question posed is whether performance can be further improved by moving beyond a finite set of experts to an effectively infinite number, while maintaining computational efficiency.

2. Methodology: ∞-MoE

The authors propose ∞-MoE, a framework that generalizes MoE from a discrete set of experts to a continuous, potentially uncountably infinite expert space.

2.1 Continuous Expert Space

Unlike standard MoE, where a router selects a discrete index $i$ from a set $Z = \{1, \dots, n\}$ , ∞-MoE defines the expert space as a continuous domain $Z \subseteq \mathbb{R}^{d_z}$ .

Router: The router outputs a probability density $p(z|x)$ over the continuous space, modeled as a Gaussian distribution $\mathcal{N}(z | \mu(x), \Sigma(x))$ .
Sampling: During training and inference, the model samples $K$ continuous expert indices $z^{(k)} \sim p(z|x)$ .
Integration: The output is approximated via Monte Carlo sampling: $y \approx \frac{1}{K} \sum_{k=1}^K f(x, z^{(k)})$ .

2.2 Expert Design via Dynamic Masking

To realize an infinite number of experts within a fixed parameter budget, the authors modulate a single Feed-Forward Network (FFN) using the sampled continuous index $z$ .

Mechanism: The sampled $z$ is used to generate a mask that dynamically selects a subset of neurons in the intermediate layer of the FFN.
Implementation:
1. Compute intermediate values $\hat{m} = W_z z$ .
2. Apply a top- $N\%$ operator to retain only the largest $N\%$ of nodes, setting the rest to zero.
3. The final expert output is $f(x, z) = W_2 (\text{Act}(W_1 x) \odot \text{mask}(z))$ .
Effect: Each distinct sampled $z$ activates a unique subset of the FFN's parameters, effectively creating a distinct "expert" from a continuous continuum while maintaining the sparsity characteristic of traditional MoE.

3. Key Contributions

Theoretical Generalization: The paper introduces a formulation that extends MoE from a discrete expert pool to a continuous domain, theoretically allowing for an unbounded number of experts.
Efficient Implementation: By combining continuous sampling with dynamic neuron masking, the method achieves the capacity of infinite experts while preserving the computational efficiency of sparse routing (only a small fraction of parameters are active per token).
Flexible Inference: The approach allows for a flexible trade-off between accuracy and speed at inference time by adjusting the number of samples $K$ .
Custom Kernel Optimization: The authors developed a fused CUDA kernel to handle the masking and reduction operations efficiently, achieving a 1.25-fold speedup over standard PyTorch implementations.

4. Experimental Results

The authors evaluated ∞-MoE on GPT-2 Small and GPT-2 Medium architectures across various zero-shot NLP benchmarks (BoolQ, HellaSwag, WinoGrande, ARC, OpenBookQA, RACE).

Performance vs. Dense Models: A GPT-2 Small-based ∞-MoE (129M active, 186M total parameters) achieved an average score of 39.8, outperforming the dense GPT-2 Medium (350M parameters) which scored 39.8 (tied in average, but ∞-MoE used fewer active parameters). Specifically, ∞-MoE outperformed the dense GPT-2 Small (38.5) and standard MoE variants.
Performance vs. Standard MoE: ∞-MoE consistently outperformed Switch Transformer (Top-1) and standard MoE (Top-2) across all benchmarks. For GPT-2 Small, it improved the average score from 38.8 (MoE) to 39.8.
Scaling with $K$ : Increasing the number of samples $K$ during inference yielded additional accuracy gains. Even with minimal sampling ( $K=1$ ), the model maintained a 2.5% accuracy improvement over standard MoE.
Stability: Routing entropy measurements indicated that the model quickly stabilized during training, with high normalized routing entropy ( $\approx 0.9$ ), suggesting uniform utilization of the expert space.
Ablation Studies:
- Dataset Size: Performance on HellaSwag scaled positively with training data size up to 100 billion tokens.
- Index Dimension: Increasing the dimension of the continuous index $z$ improved performance up to $d_z=128$ , with no further gains observed at higher dimensions.

5. Significance and Claims

The paper claims that ∞-MoE successfully bridges the gap between the high capacity of dense models and the efficiency of sparse MoE. By treating experts as a continuous space, the method allows for "infinite" specialization without the training instability associated with massive discrete expert pools.

The authors emphasize that the framework offers a flexible speed-accuracy trade-off: users can adjust the number of sampled experts ( $K$ ) at inference to meet specific latency or hardware constraints without retraining.

6. Limitations and Future Work

The authors explicitly note several limitations and open challenges:

Scaling to Larger Models: The behavior of ∞-MoE on models larger than GPT-2 Medium (e.g., GPT-3 scale) is not yet fully understood, particularly regarding the trade-offs between total parameters and active parameters.
Router Distributions: The current implementation uses a unimodal Gaussian router. More complex distributions (e.g., mixtures of Gaussians) could offer better coverage but introduce training variance and sampling complexity.
Domain Generalization: Applicability to non-NLP domains (e.g., Vision Transformers or Multimodal models) remains unexplored.
Hardware Efficiency: While the theoretical FLOPs are low, practical throughput can be hindered by the dynamic nature of the active neuron sets, which complicates batching and the use of fixed-shape GEMMs. The authors note that their current implementation prioritizes throughput via masked dense matrix multiplications, but closing the gap between theoretical compute savings and practical speed requires further system-level optimizations.