💬 NLP

EMO: Pretraining Mixture of Experts for Emergent Modularity

EMO introduces a novel Mixture-of-Experts architecture that leverages document boundaries during pretraining to enable emergent, semantically specialized expert modularity, allowing for selective expert deployment with minimal performance degradation compared to standard monolithic or conventional MoE models.

Original authors: Ryan Wang, Akshita Bhagia, Sewon Min

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ryan Wang, Akshita Bhagia, Sewon Min

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

The Problem: The "All-Or-Nothing" Giant

Imagine you have a massive, all-knowing library (a Large Language Model) that contains every book ever written, from advanced quantum physics to cooking recipes and coding manuals.

Currently, if you want to ask the library a simple question about how to bake a cake, you have to open the entire building, turn on every light, and hire every single librarian to stand by, even though you only need the one librarian who knows about baking. This is incredibly expensive and slow.

Some people tried to fix this with Mixture of Experts (MoE) models. Think of this as a library with hundreds of specialized librarians. When you ask a question, the system picks just a few librarians to help you.

The Catch: In standard MoE models, the system is chaotic. If you ask a question about baking, the system might accidentally wake up the librarian who knows about grammar, the one who knows about history, and the one who knows about coding. Because the "wrong" librarians are still active, you can't just fire the coding librarian to save money; the system breaks if you try to use only the baking experts. You still have to keep the whole team on payroll.

The Solution: EMO (The Organized Library)

The authors introduce EMO, a new way to train these models so they naturally organize themselves into neat, independent teams.

The Secret Sauce: The "Document" Rule
The key idea is simple: Everything in one document usually belongs to the same topic.

If you are reading a document about coding, every sentence in that document is about coding.
If you are reading a document about medicine, every sentence is about medicine.

EMO uses this fact as a rule during training. It tells the system: "For this specific document, pick a small group of experts (a 'pool') and force every single word in that document to use only experts from that group."

It's like telling a group of students: "For this specific essay, you must only use the resources from the Science Section. Do not go to the History Section."

Because the model follows this rule for millions of documents, it learns to naturally group its experts. The "coding" experts learn to stick together, the "math" experts stick together, and the "medical" experts stick together. They stop mixing with unrelated topics.

The Results: Cutting the Team Without Breaking the System

The authors built a massive model (14 billion parameters total, but only 1 billion active at a time) and tested it.

Full Team Performance: When they use the whole team of experts, EMO works just as well as standard models. It's smart and accurate.
The "Cut" Test: This is where EMO shines. They tried to run the model using only 25% of the experts (just the math experts, for example).
- Standard Model: If you try to use only 25% of a standard model, it crashes. Performance drops by 10–15% because the remaining experts are a random, confused mix.
- EMO: If you use only 25% of EMO, performance drops by only 1%. The model stays smart because that 25% is a perfectly organized, specialized team.

The Analogy:

Standard MoE: Like a random assortment of tools in a toolbox. If you take out half the tools, you might lose the hammer and the screwdriver, leaving you with only a wrench and a saw. You can't build a house.
EMO: Like a toolbox where the tools are sorted into labeled drawers. If you only need to fix a leak, you open the "Plumbing" drawer. You have every tool you need for that job, and you don't need to carry the "Gardening" or "Electrical" drawers with you.

What the Experts Actually Learned

The researchers looked inside the model to see what the experts were doing.

Old Models: The experts were learning low-level tricks, like "I handle the word 'the'" or "I handle commas." This is like a librarian who only knows how to shelve books by the color of their spine.
EMO: The experts learned high-level topics. One group became the "Math Team," another became the "Coding Team," and another became the "Medical Team." This is like having a librarian who actually knows the subject matter.

Why This Matters

This paper shows that we can build huge, powerful AI models that are modular. Instead of carrying a giant, heavy backpack (the full model) everywhere you go, you can carry a small, specialized kit for the specific task you are doing right now.

Memory Efficiency: You don't need to load the whole model into memory if you only need to do math.
Flexibility: You can mix and match these expert groups.
No Human Labels: The model figured this out all by itself. The researchers didn't have to tell it, "You are the math expert." The model discovered the math experts on its own just by following the "Document Rule."

In short, EMO turns a monolithic, expensive giant into a set of Lego blocks that can be snapped together or taken apart depending on what you need to build, without losing the ability to build a castle.

Technical Summary: EMO - Pretraining Mixture of Experts for Emergent Modularity

Problem Statement

Large Language Models (LLMs) are typically deployed as monolithic systems, requiring the full model to be loaded even when an application requires only a narrow subset of capabilities (e.g., code generation or mathematical reasoning). This paradigm incurs unnecessary computational costs and memory usage. While Mixture-of-Experts (MoE) architectures offer a potential solution by activating only a subset of experts per input, standard MoEs fail to support modular deployment. In existing MoEs, tokens within a single input activate diverse experts, often driven by low-level lexical patterns (e.g., prepositions, punctuation) rather than high-level semantic domains. Consequently, restricting inference to a specific subset of experts for a given domain leads to severe performance degradation, rendering MoEs impractical for memory-constrained settings where only a fraction of the model can be loaded.

Methodology: EMO

The authors introduce EMO, an MoE architecture designed to induce modularity as a first-class objective during pretraining. The core goal is to train a model where expert subsets can be used independently for specific domains and composed into a strong general-purpose model, without relying on human-defined domain labels.

Key Mechanism: Document-Level Shared Pools

The central intuition is that tokens within the same document usually share a semantic domain. EMO leverages document boundaries as a weak supervisory signal to enforce this structure:

Shared Expert Pool: For each document, the router selects a shared pool of experts ( $D$ ) from the total set of routed experts.
Restricted Routing: All tokens within that document are constrained to select their active experts only from this shared pool.
Dynamic Granularity: To ensure flexibility across different memory budgets, the size of the expert pool ( $d$ ) is treated as a random variable sampled from a uniform distribution $U\{k, n_r\}$ during pretraining, where $k$ is the number of active experts per token and $n_r$ is the total routed experts.

Training Objective and Technical Considerations

Objective: The model is trained using standard autoregressive language modeling loss ( $L_{CE}$ ) combined with auxiliary losses.
Load Balancing: A critical technical contribution is the adoption of global load balancing. Standard micro-batch load balancing conflicts with the document-level routing constraint (forcing tokens in a micro-batch to spread across experts). EMO aggregates routing statistics across data-parallel groups to compute load balancing loss, ensuring uniform expert utilization across the entire dataset while maintaining consistency within documents.
Architecture: The authors trained a 1B-active, 14B-total parameter model with 128 total experts (127 routed, 1 shared) on 1 trillion tokens. The architecture incorporates improvements over prior baselines, including pre-norm, the removal of QK-norm, and the inclusion of a shared expert.

Key Contributions

Emergent Modularity: EMO demonstrates that modular structure can emerge directly from data using document boundaries as the only signal, eliminating the need for predefined domain labels or human priors.
Selective Expert Use: The model enables the deployment of small expert subsets (e.g., 12.5% to 25% of experts) that retain near-full-model performance on domain-specific tasks, a capability standard MoEs lack.
Semantic Specialization: Unlike standard MoEs that specialize in low-level syntactic patterns, EMO experts organize into coherent groups specializing in high-level semantic domains (e.g., math, code, biomedical, politics).
Pareto Frontier Improvement: EMO expert subsets outperform memory-matched models trained from scratch (both dense and standard MoE), establishing a new Pareto frontier for the memory-accuracy trade-off.

Results

Full-Model Performance: As a complete model, EMO matches the performance of a standard MoE trained on the same data (1T tokens), achieving competitive results on benchmarks like MMLU, MMLU-Pro, and GSM8K.
Subset Performance:
- When retaining 25% of experts, EMO incurs only a 1% absolute drop in MMLU accuracy, whereas standard MoEs degrade by over 10%.
- When retaining 12.5% of experts, EMO incurs a 3% drop, while standard MoEs drop by 15%.
- EMO subsets consistently outperform dense models and standard MoEs trained from scratch with the same parameter count.
Robustness: The performance of EMO is robust across different expert selection methods (router-based vs. Easy-EP pruning) and remains effective even with extremely small validation sets for expert selection.
Analysis: Clustering analysis reveals that EMO tokens from the same document are grouped into the same semantic clusters (e.g., "health, medical & wellness"), whereas standard MoE tokens are dispersed across clusters based on surface-level features.

Significance and Claims

The paper claims that EMO provides a path toward modular, memory-efficient deployment of large, sparse models. By enabling the independent use and composition of expert subsets, EMO addresses the memory bottlenecks associated with scaling MoEs. The authors highlight several implications:

Accessible Deployment: Users can serve only the experts relevant to a specific task or domain, significantly reducing memory requirements.
Fine-Grained Control: Modularity allows for the selective enabling or disabling of capabilities (e.g., excluding experts associated with spam or adult content) at inference time.
Modular Development: The structure suggests a potential paradigm for modular pretraining, where task-specific expert subsets could be trained independently and later reintegrated.
Monitorability: The distinct routing patterns provide structured signals for auditing and debugging model behavior.

The authors conclude that modularity can be built into pretraining itself, moving away from monolithic systems toward composable architectures that are easier to deploy, adapt, and inspect.