EMO: Training Modular Language Models Through Document-Level Expert Pooling

Paper: 2605.06663 Authors: Ryan Wang, Akshita Bhagia, Sewon Min Categories: cs.CL

The Gap

Standard MoEs activate sparse experts per token, promising efficiency. But when you try to deploy only a subset of experts for a specific domain (say, just the “math experts” for a calculator app), performance collapses. The model was never trained to work with partial expert sets—it learned entangled dependencies across all experts. This blocks memory-efficient deployment: you still need to load the full 14B model even if your app only needs 3B worth of capabilities.

Prior work either accepts this limitation (deploy everything) or uses human-defined expert assignments (code experts, math experts), which requires manual annotation and doesn’t scale. The gap: how do you train an MoE where expert subsets are independently functional, without human priors?

Problem: Standard MoE experts are entangled
   |
   v
Observation: Tokens in same document share domain
   |
   v
Method: Force document tokens to use shared expert pool
   |
   v
Evidence: 25% experts retained -> 1% performance drop
   |
   v
Conclusion: Emergent modularity enables subset deployment

The Increment

One sentence: Before EMO, deploying a subset of MoE experts meant broken models; after EMO, you can ship 25% of the experts and retain 99% of performance.

Core Mechanism

EMO adds one constraint during training: all tokens within a document must select experts from the same small pool (e.g., 4 experts out of 64 total). Different documents can use different pools, but within a document, the expert choice is shared. This forces the model to organize experts by domain—since documents are typically domain-coherent (a math paper uses math throughout, a code file uses code throughout), tokens from similar domains naturally converge on similar expert pools.

The architecture is otherwise standard: 8 MoE layers, each with 64 experts, top-2 routing per token. The key difference is the routing mechanism. Instead of each token independently picking 2 experts from all 64, the model first samples a pool of 4 experts for the entire document, then each token picks 2 from that pool. The pool is sampled using an aggregated routing score across all document tokens, ensuring coherence.

During inference, you can identify which expert subsets handle which domains (by analyzing routing patterns on domain-specific data), then deploy only those subsets. A code-focused app loads the 16 experts that activate on code; a math app loads a different 16. The model was trained to function this way, so performance degrades gracefully rather than collapsing.

Standard MoE:                    EMO:
                                 
Token1 ---> [64 experts]         Document boundary
Token2 ---> [64 experts]              |
Token3 ---> [64 experts]              v
  (independent routing)          [Sample pool: 4/64]
                                      |
                                      v
                                 Token1 ---> [4 experts]
                                 Token2 ---> [4 experts]
                                 Token3 ---> [4 experts]
                                   (shared pool)

Think of it like a restaurant kitchen. In a standard MoE, every dish (token) can call any chef (expert) from the entire kitchen—flexible but chaotic, and you can’t send half the chefs home without ruining service. EMO is like assigning a small team of chefs to each table (document). A table ordering Italian gets the Italian team; a table ordering sushi gets the sushi team. Each dish at the table still picks from its assigned team, but the team itself is coherent. Now you can run an Italian-only night with just the Italian chefs, and service works fine. The constraint (team per table) forces specialization to emerge naturally based on what customers order together.

Key Concepts

• Document-level expert pooling: Instead of letting each token independently choose from all experts, EMO first selects a small subset of experts for the entire document, then restricts all tokens in that document to choose from that subset. This is implemented by computing an aggregated routing score across all document tokens, sampling a pool (e.g., 4 experts), then performing standard top-k routing within that pool. The key insight: documents are domain-coherent, so this constraint aligns expert selection with semantic domains without explicit labels.

• Emergent modularity: The property that expert subsets become independently functional without being explicitly trained for independence. In standard MoEs, experts co-adapt—expert A learns to rely on expert B’s outputs downstream, creating entanglement. EMO’s document pooling breaks this: since different documents use different pools, an expert can’t rely on another expert always being present. This forces each expert to be more self-contained, and groups of experts that co-occur frequently (because they handle the same domain) naturally form coherent modules.

• Semantic vs syntactic specialization: Standard MoEs show syntactic specialization—one expert handles punctuation, another handles capitalization, another handles specific token positions. These are low-level patterns that don’t correspond to human-meaningful categories. EMO’s experts show semantic specialization—one group handles math, another handles code, another handles legal text. This emerges because document pooling groups experts by the domains that co-occur within documents, which are semantic units. Semantic specialization is what enables practical subset deployment: you can ship “the math experts” because they form a coherent functional unit.

Framework Shift

Before (standard MoE):           After (EMO):

[All 64 experts]                 [All 64 experts]
       |                                |
   per-token                      per-document
    routing                         pooling
       |                                |
       v                                v
  Token uses                    Document uses
  any 2/64                      shared 4/64
                                       |
  (entangled,                    per-token
   not modular)                   routing
                                       |
                                       v
                                  Token uses
                                  any 2/4
                                       
                                  (modular,
                                   subset-deployable)

From token-level independence to document-level coherence, the core shift is trading routing flexibility for emergent structure.

Expert Assessment

Problem choice: Real gap. As models scale to 100B+ parameters, memory constraints become the bottleneck for deployment. The inability to use MoE subsets is a known pain point, not a manufactured problem. This sits at the intersection of efficiency and modularity, both hot topics.

Method maturity: Elegant insight, not brute force. The document pooling constraint is simple—almost trivial—but the emergent behavior is non-obvious. The risk: it relies on documents being domain-coherent, which is true for pretraining corpora but may not hold for all applications (e.g., multi-topic conversations). The paper doesn’t explore failure modes deeply.

Experimental integrity: Baselines are fair. The comparison to standard MoE is apples-to-apples (same architecture, same data, same compute). The subset deployment experiments are convincing: 25% retention with 1% drop is a clean result. However, the paper only tests on one model size (1B active, 14B total) and one dataset (1T tokens). Scaling behavior is unknown. The domain specialization analysis (Section 4.3) is qualitative—they show examples of expert routing on math/code, but don’t quantify purity or overlap rigorously.

Writing quality: The core idea is clear, but the paper front-loads motivation and delays the method until Section 3. Section 2 (related work) is dense and could be trimmed. The biggest weakness: the paper doesn’t discuss when document pooling might fail. What if documents are multi-domain? What if the corpus has poor document boundaries? These are practical concerns left unaddressed. Rewriting Section 5 (limitations and future work) to be more critical would elevate the paper.

Verdict: weak accept — Solves a real problem with a simple, effective method, but limited experimental scope and insufficient failure mode analysis prevent a strong accept.

Takeaways

Constraint-driven emergence: When you want a model to learn structure (modularity, specialization, composability), don’t annotate it—constrain the training dynamics to make that structure advantageous. EMO’s document pooling is a forcing function: it makes modular experts more efficient than entangled ones. This principle transfers: if you want disentangled representations, constrain what can interact; if you want hierarchical structure, constrain information flow across levels.

Semantic units as inductive bias: Documents are semantic units (mostly). By aligning model constraints with semantic boundaries, EMO gets semantic specialization for free. This is cheaper than human annotation and more robust than clustering embeddings post-hoc. Generalizes to other domains: use sentence boundaries for fine-grained tasks, use user sessions for personalization, use code functions for program synthesis.

Graceful degradation through modularity: The 25% → 1% drop result is the key practical win. If you’re building systems that need to scale down (edge deployment, cost optimization, A/B testing), design for subset functionality from the start. Retrofitting modularity onto a monolithic model is hard; training for it is easy if you pick the right constraints.