EMO: How Constraining Expert Selection to Document-Level Pools Causes Semantic Modularity to Emerge During MoE Pretraining
Mixture-of-Experts (MoE) models have become the dominant architecture for scaling large language models: by activating only a small subset of experts (feed-forward networks) per token, they increase total parameters without proportionally increasing compute. Yet a fundamental problem persists — even though MoEs activate few experts per token, most or all experts are used over the course of a task, because routing decisions are made independently per token and tend to specialize on low-level lexical features (prepositions, punctuation, token identity) rather than semantic domains. This means you cannot deploy a small subset of experts for a specific domain — the full model is always required, creating severe memory bottlenecks as models grow to hundreds of billions of parameters.
EMO, a collaboration between UC Berkeley and the Allen Institute for AI (Ai2), addresses this with a strikingly simple idea: constrain all tokens within a document to select their active experts from a shared pool, while allowing different documents to use different pools. This document-level pooling constraint, applied throughout pretraining, causes coherent expert groupings to emerge from the training data — without any human-defined domain labels.
The Core Result
A 1B-active / 14B-total EMO pretrained on 1T tokens matches standard MoE performance as a full model. Critically, it enables selective expert use: retaining only 25% of experts incurs just a ~1% absolute performance drop, and 12.5% incurs ~3%. Standard MoEs degrade by 10–15% under the same constraint.
| Expert Retention | Standard MoE Drop | EMO Drop |
|---|---|---|
| 25% (32 of 128) | ~10% absolute | ~1% absolute |
| 12.5% (16 of 128) | ~15% absolute | ~3% absolute |
| 6.2% (8 of 128) | Below dense baseline | Still outperforms dense baseline |
After fine-tuning, 12.5% expert retention perfectly recovers full-model performance on GSM8K. EMO’s expert subsets even outperform models trained from scratch at matched memory budgets — a single checkpoint serves as a “model family” at multiple deployment tiers.
Figure 1 — EMO’s two-level routing mechanism. Standard MoE (left) routes each token independently. EMO (right) first selects a document-level expert pool D, then constrains all tokens to route within D. A shared expert (always active, outside the pool) handles universal computation.
The Mechanism: How Document-Level Pooling Works
In a standard MoE, each token independently computes p_t = softmax(router(x_t)) and selects its top-k experts. There is no cross-token constraint. EMO inserts a single modification: a document-level pooling step before per-token routing.
For a document with T tokens, the document expert pool D is the top-d experts from the average routing distribution:
D = TopK( (1/T) Σ p_t , d )All tokens are then masked and renormalized to route only within D. The hyperparameter d controls granularity: d = k forces all tokens to use the same experts (maximum modularity); d = n_r recovers standard MoE. During training, d is sampled uniformly per document from {k, ..., n_r}, exposing the model to all pool sizes and enabling flexible deployment at any subset size.
A shared expert (always active, not subject to pool masking) provides a baseline of always-available computation, ensuring general-purpose capability even when the pool is small.
Why Semantic Specialization Emerges
The central question: why does averaging routing distributions across a document shift expert specialization from syntactic to semantic?
Figure 2 — Standard MoE experts specialize in low-level lexical patterns (prepositions, copula verbs). EMO experts specialize in high-level semantic domains (film reviews, health, news, politics). This shift directly enables modularity.
The answer has three parts:
Averaging as a low-pass filter. The document average (1/T)Σp_t ≈ softmax(P · (1/T)Σh_t) suppresses high-frequency syntactic noise (token-identity features that are uncorrelated across tokens) while preserving low-frequency semantic signal (domain-level features that are correlated across all tokens in a document). The signal-to-noise ratio improves approximately linearly with document length.
The geometric explanation. Wang et al. (2026) proved that MoE routers are linear projections, so routing is determined by hidden state geometry. In standard MoEs, early-layer hidden states are dominated by token identity, making lexical features the most discriminative. EMO’s document-level averaging projects out token-identity components, leaving the document’s semantic domain as the dominant signal.
Global load balancing. EMO uses global load balancing (aggregating routing statistics across data-parallel groups) instead of micro-batch load balancing. This prevents all documents from selecting the same pool: the pool constraint ensures expert consistency within documents, while global load balancing encourages diverse expert usage across documents. The two become complementary rather than conflicting.
Why doesn’t the model route by the most common token type? Four mechanisms prevent this: (1) shared experts absorb common syntactic processing, (2) the pool must serve ALL tokens including rare domain-specific ones, (3) global load balancing prevents pool homogenization, and (4) the averaged routing distribution is a learned abstracted signal, not a token-frequency count.
The Training Pipeline: Making It Stable
The most critical training challenge is a tension between load balancing and the document-pool constraint. Standard micro-batch load balancing pushes tokens within the same micro-batch to spread across many experts — directly opposing EMO’s constraint. The result: severe training instability.
The solution is global load balancing (following Qiu et al. 2025 / Qwen3): routing frequency is computed via all-reduce across all data-parallel groups, while per-group routing probabilities are retained. This makes the two objectives complementary.
EMO also improves on its OLMoE baseline architecturally: pre-norm instead of post-norm, removed QK-norm, and a shared expert. These changes make EMO’s baseline trained on 1T tokens outperform OLMoE trained on 5T tokens — a 5× training efficiency improvement.
Training from scratch is important: annealing a standard MoE with the EMO objective for 50B tokens produces partial modularity but underperforms. The routing constraint must shape representations from initialization.
Deployment Implications
EMO enables static expert subset deployment — a fundamentally different paradigm from dynamic expert offloading. Rather than swapping experts in and out of GPU memory at inference time, EMO lets you permanently load only a domain-specific subset and serve it as a standalone model.
Figure 3 — EMO’s expert subsets maintain near-full-model performance even at aggressive retention ratios, while standard MoEs degrade sharply. EMO pushes the Pareto frontier for the memory-accuracy trade-off.
The deployment workflow is simple: collect 5–100 validation examples for the target domain, forward-pass through the router to aggregate routing probabilities, rank experts, and retain the top-d per layer. Optionally fine-tune (1 epoch, lr=5e-5). EMO is robust to the selection method — both router-based ranking and Easy-EP (a state-of-the-art pruning method) produce strong results.
This is complementary to, not competing with, dynamic serving systems like PreMoe, SwapMoE, and MoE-Infinity. For known, narrow tasks, EMO’s static approach is strictly better (no runtime I/O overhead). For diverse, unpredictable queries, dynamic approaches or the full EMO model remain necessary.
Important caveat: When the deployment task is general (e.g., MMLU’s “Other” catch-all category), EMO’s expert subsets underperform models trained from scratch at matched memory budgets. EMO is best suited for domain-specific serving rather than general-purpose chatbots at reduced memory.
The Broader Landscape: Where EMO Fits
MoE research has historically split into two branches: scalable architectures (Shazeer → Switch → Mixtral → OLMoE → DeepSeekMoE → DeepSeek-V3 → Qwen3) focusing on activating fewer experts per token to scale cheaply, and modular architectures (Branch-Train-Merge → ModuleFormer → Branch-Train-MiX → FlexOlmo) focusing on making expert subsets independently usable. EMO is the first to unify both goals.
Prior modular approaches fell short in different ways. FlexOlmo and BTX require predefined domain labels (ambiguous, inject human bias). ModuleFormer uses a mutual information objective but EMO’s authors found it doesn’t outperform standard MoEs and degrades below 40% expert retention. DeepSeekMoE aims for ultimate specialization through fine-grained segmentation but still produces syntactic-level specialization.
The key insight: EMO demonstrates that the same architecture produces radically different specialization depending solely on the training objective. Modularity is not primarily an architectural property but a training-objective property.
Implications for Safety and Control
EMO’s semantic modularity opens a path to inference-time safety modulation that connects to GRAM (Gradient-Routed Auxiliary Modules), a framework from Anthropic. The paper envisions disabling expert clusters for “spam, gambling, or adult content” in child-facing applications — the unsupervised analogue of GRAM’s supervised capability isolation. This convergence suggests modular pretraining is becoming a central paradigm for AI safety, not just efficiency.
The vision extends to modular development: training expert subsets independently and reintegrating them, rather than retraining the entire system. A preliminary experiment shows improvement but not full recovery — integration challenges remain.
Limitations and Open Questions
EMO is not without significant limitations:
- Scale is untested. 1B-active / 14B-total is 50–100× below frontier models like DeepSeek-V3 (671B). No evidence the method transfers to that scale.
- General tasks break modularity. The “Other” category failure means mixed-domain queries don’t benefit from selective deployment.
- Document boundaries aren’t universal. The assumption that tokens within a document share a domain breaks for multi-file code repositories, multi-turn conversations, and multi-topic documents. No code generation benchmark is evaluated.
- The “no human priors” claim is partially misleading. Deployment requires a validation set, subset size choice, and optionally fine-tuning — all human-in-the-loop.
- Modular insertion doesn’t work yet. The reintegration experiment shows improvement but not full recovery.
- Community traction is modest. 43 GitHub stars, ~212 downloads, no independent reproduction as of mid-2026.
Conclusion
EMO represents a paradigm shift — from MoE-as-scaling-device to MoE-as-composable-architecture. The document-level pool constraint is architecturally minimal yet fundamentally changes the routing optimization landscape, causing semantic expert groupings to emerge self-supervised. The result: a single pretrained model that can be deployed at multiple memory tiers with graceful degradation, pushing the Pareto frontier for the memory-accuracy trade-off.
The implications extend beyond efficiency. EMO’s semantic modularity enables capability control (disabling dangerous experts), modular maintenance (training and reintegrating expert subsets), and structured monitorability (expert activations as interpretable signals). Whether this approach scales to frontier models, handles code data effectively, and achieves seamless modular insertion remain the critical open questions for the future of composable AI architectures.