EMO, Document-Level MoE Pretraining

AllenAI released EMO. It pretrained an MoE with 1B active / 14B total parameter scale on 1T tokens, using an architecture of 8 active experts / 128 total experts.

It used document boundaries as a weak supervision signal so that tokens from the same document are routed only within the same expert pool. Unlike standard MoE where tokens scatter based on surface-level patterns, this design allows document-level semantic modules to form.

• Global load balancing reduced expert imbalance while maintaining routing consistency within documents.
• Rather than fixing the per-document pool size, it used random sampling to handle various subset sizes.
• EMO's routing clusters grouped into domains like Health, News Reporting, US Politics, and Film & Music, while standard MoE concentrated on surface features like Prepositions, Proper Names, Copula Verbs, and Definite Articles.

On overall benchmarks, it maintained performance similar to standard MoE, and using all experts together preserved general-purpose model performance. However, performance degradation was small even when reducing experts, making it advantageous for selective use.

• With a 25% expert subset, the average drop was about 1%p.
• With a 12.5% subset, the average drop was only about 3%p.
• Standard MoE under the same conditions collapsed sharply as subsets shrank, falling to near-random levels at small subsets.
• Task-specific module selection was possible with a small amount of validation data or even just 1 example, and it was also combined with Easy-EP.

Released alongside this are the full EMO model, a standard-MoE baseline trained on the same data, and the training code.