Modularity revealed in Mixture-of-Experts models

Ai2 released EMO. This Mixture-of-Experts (MoE) model is 1B-active, 14B-total in scale, with 8 active experts and 128 total experts. It was trained on 1 trillion tokens, and was designed to maintain strong general performance when using all experts together, while also ensuring performance doesn't collapse significantly even when only a subset of experts is selected.

The key is using document boundaries as a weak supervision signal. Tokens from the same document are routed only within a single shared pool of experts, and the router averages preferences across the whole document to select the most frequently used group of experts. As a result, experts tend to cluster around semantic domains like medical, news, politics, and film, rather than surface patterns like prepositions or punctuation.

For training stability, global load balancing was applied at a scale larger than the document level, and document pool size was not fixed but randomly sampled. Applying load balancing at the local microbatch level would spread tokens from the same document across multiple experts, conflicting with EMO's objective, but balancing globally allows both intra-document consistency and overall expert utilization to be achieved together.

On benchmarks, EMO performed similarly to standard MoE when using all experts, and remained robust even under conditions where only a subset of experts was kept.

• Keeping only 25% of experts (32) resulted in an overall average performance drop of about 1pp.
• Keeping only 12.5% of experts (16) resulted in a drop of about 3pp.
• These results held both after pretraining and after fine-tuning.

Expert selection is also cheap. Using just a single example with few-shot demonstrations, it was possible to find a group of experts that matched the performance of a module chosen using the entire validation set, and this combined well with existing expert-pruning methods such as Easy-EP.

When router activations were clustered based on the first 100 tokens of 12,000 pretraining documents, EMO formed semantically grounded clusters such as Health, Medical & Wellness, News Reporting, US Politics & Elections, and Film & Music. In contrast, standard MoE was drawn more toward surface-level token patterns such as Prepositions, Proper Names, Copula Verbs, and Definite Articles.

Ai2 released the EMO base model, a standard MoE baseline trained on the same data, and the training code. They emphasized that this is an early step toward a modular language model that works even when a small expert subset is selected and combined.