PaperlayerView exact revision on arXiv

Source-linked AI summary

EMO: Pretraining Mixture of Experts for Emergent Modularity

Ryan Wang, Akshita Bhagia, Sewon Min

arXiv:2605.06663v2cs.CL

TL;DR

EMO addresses the difficulty of using only capability-relevant parts of MoE models by training experts into coherent, reusable groups from document boundaries. It matches standard MoE performance as a full model while retaining performance with only 25% or 12.5% of experts, dropping 1% or 3% respectively.

  • Problem

    Existing MoEs require the full model because tokens in one task activate diverse experts, limiting independent expert-subset use in memory-constrained deployment.

  • Method

    EMO restricts all tokens in each document to a shared expert pool, allowing coherent modular groups to emerge without human-defined domains.

  • Results

    Retaining 25% (12.5%) of experts causes only a 1% (3%) absolute performance drop, while EMO’s subsets specialize by semantic domains rather than low-level syntax.

  • Takeaways & Limitations

    EMO demonstrates a path toward modular, memory-efficient deployment of large sparse models with expert subsets that remain effective independently and in composition.

  • Takeaways & Limitations

    Using a fixed expert-pool size during training limits inference-time flexibility because performance degrades for subsets of other sizes.

Abstract

from arXiv · show

Large language models are typically deployed as monolithic systems, requiring the full model even when applications need only a narrow subset of capabilities, e.g., code, math, or domain-specific knowledge. Mixture-of-Experts (MoEs) seemingly offer a potential alternative by activating only a subset of experts per input, but in practice, restricting inference to a subset of experts for a given domain leads to severe performance degradation. This limits their practicality in memory-constrained settings, especially as models grow larger and sparser. We introduce EMO, an MoE designed for modularity-the independent use and composition of expert subsets-without requiring human-defined priors. Our key idea is to encourage tokens from similar domains to rely on similar experts. Since tokens within a document often share a domain, EMO restricts them to select experts from a shared pool, while allowing different documents to use different pools. This simple constraint enables coherent expert groupings to emerge during pretraining using document boundaries alone. We pretrain a 1B-active, 14B-total EMO on 1T tokens. As a full model, it matches standard MoE performance. Crucially, it enables selective expert use: retaining only 25% (12.5%) of experts incurs just a 1% (3%) absolute drop, whereas standard MoEs break under the same setting. We further find that expert subsets in EMO specialize at semantic levels (e.g., domains such as math or code), in contrast to the low-level syntactic specialization observed in standard MoEs. Altogether, our results demonstrate a path toward modular, memory-efficient deployment of large, sparse models and open new opportunities for composable architectures.

1 Introduction

EMO addresses the deployment limits of standard MoEs by training document-level shared expert pools that yield modular, selectively usable subsets while preserving full-model performance. Its semantic domain specialization supports memory-efficient and composable deployment without human-defined domain priors.

  • Motivation: Standard MoEs remain effectively monolithic because tokens in one task activate diverse experts specializing in low-level lexical patterns.This makes expert subsets unusable independently and limits deployment in memory-constrained settings.
  • Method: EMO restricts all tokens within each document to select experts from a shared pool, allowing modular structure to emerge from document boundaries alone.The approach avoids human-defined domain priors while encouraging similar-domain tokens to use similar experts.
  • Results: 1B-active, 14B-total parameters and 1 trillion tokens define the EMO pretraining setup.As a full model, EMO matches the overall performance of a standard MoE.
  • Results: 25% retained experts incur a 1% absolute drop, while 12.5% retained experts incur a 3% absolute drop.Standard MoEs break under the same selective-expert setting, whereas EMO supports effective subset composition.
  • Analysis: EMO expert subsets specialize in higher-level semantic domains and topics, unlike standard MoE experts’ lower-level syntactic specialization.The paper identifies examples including math and code for semantic specialization, versus prepositions and punctuation for syntactic specialization.
  • Implications: The results establish modularity as a built-in property of large language models and enable targeted extension training plus more interpretable, debuggable components.The authors release EMO and a matched baseline trained on the same data for reproducibility and further study.

2 Related Work

Prior work uses MoE sparsity to scale Transformers, but inactive experts still create inference-memory bottlenecks. This paper positions modular expert subsets as a complementary solution, contrasting it with existing specialization, structured-expert, domain-partitioning, and ModuleFormer approaches.

  • Mixture-of-Experts as Scalable Architectures: MoE architectures scale Transformers by activating only a subset of experts per input, with newer systems increasing expert counts and sparsity.DeepSeek-V3 is cited as using hundreds of experts per layer while activating only a small subset per token.
  • Mixture-of-Experts as Scalable Architectures: Memory bottlenecks arise because inactive experts must remain in VRAM during inference, motivating scaling laws, efficient serving, and expert pruning.These approaches address memory constraints without the selective expert-subset use introduced here.
  • Mixture-of-Experts as Scalable Architectures: This work proposes selective expert-subset use for downstream tasks as a complementary way to alleviate memory bottlenecks in large, sparse MoEs.The contribution is framed as a new deployment strategy alongside prior memory-focused approaches.
  • Specialization and Modularity of Existing MoEs: Prior MoE studies often find specialization based on surface-level or low-level lexical patterns, while some report specialization concentrated in only a tiny subset of experts.The supplied passage lists token IDs, prepositions, and punctuation as examples of such cues.
  • Specialization and Modularity of Existing MoEs: Standard MoEs trained conventionally do not support meaningful small expert subsets for downstream domains, motivating modularity as a first-class training objective.The paper states that modular training enables selective subset use and that this behavior is robust across subset-selection methods.
  • Training MoEs with Structured or Specialized Experts: ModuleFormer shares the goal of standalone expert subsets, but its evaluation omits standard MoEs and its performance degrades significantly when less than 40% of experts are retained.ModuleFormer uses an objective maximizing mutual information between tokens and experts; the authors report no improvement over standard MoEs in their reproduction.

3 Modular Mixture of Experts (EMO)

EMO makes modularity a first-class MoE objective by enabling expert subsets to operate independently while preserving a strong full-model composition. It induces domain-coherent expert pools from document boundaries without explicit semantic labels, while addressing load balancing and variable subset-size deployment.

  • EMO targets expert subsets that work in isolation for downstream domains while their composition remains a strong general-purpose model.
  • Document-level routing selects a shared expert pool for each document and constrains all its tokens to route within that pool, allowing modular groups to emerge without domain labels.Different documents may use different pools, using document boundaries as weak supervision.
  • The subset-granularity hyperparameter d trades specialization against expressivity: d = k forces shared experts, whereas d = nr recovers standard MoE routing.Smaller d yields more specialized subsets; larger d provides greater flexibility but weaker modular structure.
  • Global load balancing resolves the conflict between uniform expert utilization and document-level consistency that arises under standard micro-batch balancing.Aggregating routing statistics across data-parallel groups makes the objectives largely complementary and supports stable training.
  • Sampling d independently for each document from U{k, . . . , nr} prevents overfitting to one pool size and supports expert subsets with varying capacities.This exposes the model to different expert pool sizes during pretraining for flexible selective use at inference.

4 Experimental Setup

The experiments compare EMO with a standard MoE using matched architecture and training, evaluating both full-model performance and selective activation of domain-specific expert subsets. Models use 1B active and 14B total parameters, with 128 experts and 8 experts activated per token.

  • Experimental Configuration: 1B active and 14B total parameters are distributed across 128 experts, with 127 routed experts, 1 shared expert, and 8 experts activated per token.The architecture includes n = 128, n_r = 127, n_s = 1, and k = 8.
  • Experimental Configuration: The baseline MoE and EMO share the same architecture and differ only in their training objectives.Both models are trained from scratch on the OLMoE pretraining corpus.
  • Training: 1 trillion tokens plus a 50B-token linear annealing phase are used for main training, while ablations use 130B tokens and include dense and smaller-MoE comparisons.The main models are trained from scratch on the OLMoE pretraining corpus.
  • Evaluation Settings: The evaluation covers full-model zero-shot performance and selective expert use for task-specific domains.Selective evaluation uses coarse-grained MMLU and MMLU-Pro domains, including math, physics, health, philosophy, and history, plus GSM8K.
  • Expert Selection: A small validation set identifies relevant experts per domain using either aggregated routing probabilities or Easy-EP.Even a single few-shot example can be sufficient to select an effective expert subset.

5 Results and Analysis

EMO matches standard MoE performance as a full model while preserving performance under aggressive expert-subset deployment. Its learned modularity reflects semantic domain specialization, unlike the lexical or syntactic specialization of standard MoEs.

  • Full-model performance: EMO matches standard MoE performance in both 1T- and 130B-token settings while outperforming matched-active-parameter dense models in the latter.The baseline MoE is competitive with OLMOE despite using 1T rather than 5T tokens.
  • Selective expert use: Over 10% drops at 25% expert retention show that standard MoEs fail under subset deployment, even with fine-tuning.At 25% retention, experts decrease from 128 to 32, and performance falls below a matched-active-parameter dense model.
  • Selective expert use: About 1% and 3% drops at 25% and 12.5% retention show that EMO preserves performance and can outperform dense baselines with only 6.2% of experts.After fine-tuning, GSM8K subsets retaining up to 12.5% of experts perfectly recover full-model performance.
  • Selective expert use: EMO’s 32-expert and 8-expert subsets match or outperform memory-matched standard-MoE and dense models trained from scratch.These comparisons use subsets of the larger pretrained EMO model.
  • Expert selection robustness: Standard MoEs degrade sharply as subsets shrink under Easy-EP or router-based selection, whereas EMO remains robust and largely selection-scheme insensitive.Random selection quickly converges to random performance, while EMO performs strongly under both evaluated nonrandom strategies.

6 Future Directions

EMO’s emergent modularity suggests future directions for accessible sparse-model deployment, inference-time capability control, modular updates, and improved monitoring. These opportunities arise from selectively managing semantically organized expert subsets.

  • Accessible Deployment of Large Sparse MoEs: Selective expert use offers an orthogonal path to more accessible deployment and adaptation of large, highly sparse MoEs.This approach uses small expert subsets for a given domain rather than the full model.
  • Fine-grained Control: Semantic expert organization could enable inference-time control by selectively enabling or disabling domain-specific subsets.Examples include excluding spam, gambling, or adult-content clusters in child-facing applications and conditionally exposing biomedical knowledge.
  • Fine-grained Control: Capability isolation at inference time could provide an alternative to dataset-level filtering by managing which domains are exposed in different scenarios.This direction follows from selectively controlling semantically specialized expert clusters.
  • Modular Development and Maintenance: Modular pretraining could train task- or domain-specific expert subsets independently and later reintegrate them into the full model.In a preliminary test, a finetuned 32-expert subset was inserted into the original 128-expert model; the result improved over the original full model but did not match the standalone subset.
  • Higher Degrees of Monitorability: Expert activations could make models easier to monitor and audit by revealing which model components are used for each input.Activation mismatches, such as a math answer strongly using creative-writing or low-quality-web experts, could prompt closer inspection.

7 Conclusion

EMO makes modularity emerge during pretraining by routing tokens from each document through a shared expert pool. This produces expert subsets specialized to high-level capabilities without human-defined domains or task labels, while matching standard MoEs as a full model.

  • 7 Conclusion: EMO is designed to make modularity emerge during pretraining.The model is a mixture-of-experts architecture.
  • 7 Conclusion: Shared expert pools for tokens within each document induce subsets specializing to high-level tasks and capabilities.This mechanism does not rely on human-defined domains or task labels.
  • 7 Conclusion: As a full model, EMO matches standard MoEs in general performance.The conclusion states that the induced modular structure does not come at the cost of general performance.

A Architectural & Training Details and Ablations … B Selective expert use Details

The ablations show that EMO’s document-level routing requires global load balancing for stable training, while dynamic expert-subset training and shared experts improve selective-use flexibility and performance. Additional tuning identifies effective learning-rate, load-balancing, and normalization choices for EMO and standard MoEs.

  • A.1 Load Balancing: Global balancing aggregates routing statistics across data-parallel groups instead of enforcing uniform expert usage within each micro-batch.The modified objective uses aggregated routing frequencies while retaining per-group routing probabilities, approximating global routing over a larger sequence set.
  • A.1 Load Balancing: Global load balancing produces more stable EMO pretraining than standard micro-batch-level balancing.Micro-batch balancing is unstable because it conflicts with document-level routing, whereas global statistics preserve consistent routing within documents while diversifying utilization across documents.
  • A.2 How to Choose d: Training with a distribution over d yields robust performance across a wide range of inference-time expert subset sizes.Fixed-d models perform well at their training subset size but degrade at other sizes, limiting flexible selective use.
  • A.3 Shared Experts: Shared experts improve EMO performance on MMLU, consistent with prior observations in DeepSeek-MoE.Table 2 compares configurations without shared experts, with fixed d = 32, and with dynamic d sampled from U(8, 128).
  • A.4 Tuning LR and LB: lr = 4e −3 and lb = 1e −1 are the selected hyperparameters for the standard MoE baseline.The learning rate 4e −3 gave the best results, while increasing load balancing from 1e −2 to 1e −1 improved training stability.
  • A.4 Tuning LR and LB: lr = 4e −3 gave EMO the strongest performance by 3000 training steps with lb = 1e −1.Small loss spikes at this learning rate were reduced by applying global load balancing.
  • A.5 Prenorm vs ReorderedNorm: Prenorm with removed QK-norm achieves lower training loss than ReorderedNorm on both EMO and standard MoEs.These comparisons excluded global load balancing, shared experts, and dynamic d.

B.1 130B Token Experiments

On 130B-token training runs, EMO matches or outperforms memory-matched baselines for selective expert use across MMLU, MMLU Pro, and GSM8K. These results improve the memory–accuracy Pareto frontier.

  • Experimental setup: The experiments evaluate models trained on 130B tokens and include memory-matched comparisons.The reported comparisons follow the selective-expert-use evaluation described in §5.2.
  • Selective expert use: EMO matches or outperforms all baselines in selective expert use across MMLU, MMLU Pro, and GSM8K.The comparisons include a standard 32-expert MoE and a dense model trained from scratch.
  • Selective expert use: The selective-use results push the Pareto frontier in memory–accuracy trade-offs.The comparison covers models trained from scratch, including a standard MoE with 32 experts and a dense model.

B.2 Ablations on expert subset initialization · B.3 Selective Expert Use on "Other" Category in MMLU and MMLU Pro · B.4 Annealing Standard MoEs to be Modular.

EMO’s expert selection remains effective with very little validation data, though prompt format and task characteristics influence subset quality. Selective use is weaker on broad catch-all tasks, and inducing modularity during pretraining outperforms applying the modularity objective only during annealing.

  • B.2 Ablations on expert subset initialization: EMO shows little degradation when few-shot validation data decreases to a single example, indicating sample-efficient expert selection.Few-shot demonstrations may provide sufficient token-level signals for selecting experts.
  • B.2 Ablations on expert subset initialization: With zero-shot prompts, performance declines modestly as validation size decreases, even with only 5 validation examples.This supports EMO’s effectiveness under highly data-constrained conditions.
  • B.2 Ablations on expert subset initialization: Validation-set effects are task-dependent: on GSM8K, performance improves as validation-set size decreases.Smaller sets may produce more focused estimates of expert relevance, while aggregation can smooth signals and reduce specialization.
  • B.2 Ablations on expert subset initialization: Few-shot validation prompts can outperform zero-shot prompts even when test examples use zero-shot formats, showing that validation content and structure shape expert selection.The interaction between validation and evaluation formats remains unresolved.
  • B.3 Selective Expert Use on "Other" Category in MMLU and MMLU Pro: On general MMLU Other and MMLU Pro categories, EMO subsets of 32 and 8 experts struggle to match Reg MoE @32 and Dense @8 trained-from-scratch baselines.These catch-all tasks are broad rather than localized to a narrow capability.
  • B.4 Annealing Standard MoEs to be Modular: EMO outperforms EMO-anneal across most tasks and expert subset sizes when both models receive 1T-token pretraining and 50B-token annealing.EMO-anneal is a standard MoE converted using the document-level expert pool constraint only during annealing.
  • B.4 Annealing Standard MoEs to be Modular: Applying the document-level expert pool objective only after pretraining causes EMO-anneal to underperform EMO on most benchmarks across most subset sizes.The result indicates that pretraining from scratch is important for realizing selective-use gains.

B.5 Generations from Selective Expert Use on GSM8K

Zero-shot GSM8K generations show that EMO remains coherent under selective expert use, whereas regular MoE subsets often produce incorrect or degraded reasoning. Across examples, regular MoE performance improves inconsistently with larger subsets and can still fail at 128 experts.

  • Evaluation setup: The evaluation uses EMO and Regular MoE models trained on 1T tokens, tested zero-shot across different expert subset sizes without finetuning.The section specifically compares selective expert use on GSM8K generations.
  • Overall comparison: 8-expert subsets of EMO consistently produce coherent outputs, while regular MoE subsets deteriorate.This is the section’s qualitative comparison across selective expert sizes.
  • GSM8K examples: For the 9-hour sleep problem, Regular MoE fails with 8 and 32 experts but gives the correct 3-hour answer with 128 experts.The 8-expert output answers 7, the 32-expert output answers 3 after an incorrect intermediate claim, and the 128-expert output correctly computes 9 - 6 = 3.

C Evaluation Details · D Token Clustering Details

The evaluation uses category-based expert selection and benchmark-specific scoring protocols, while token-clustering analysis finds a dedicated first-token expert cluster that quickly transitions to stable document-level activation patterns.

  • C Evaluation Details: 40% of MMLU and MMLU-Pro examples are sampled for validation/training, with the remaining 60% reserved for evaluation.MMLU’s 57 subjects are grouped into 17 broader categories to ensure sufficient examples per category.
  • C Evaluation Details: Expert selection processes all subjects within each MMLU category using a single expert group, while aggregate MMLU and MMLU-Pro results exclude “other.”The evaluation categories and subjects are listed in Table 4.
  • C Evaluation Details: Multiple-choice benchmarks report raw accuracy from the highest-log-likelihood answer choice, while generation benchmarks report Gen5 recall and GSM8K exact match.These protocols define the reported metrics across the listed evaluation tasks.
  • C Evaluation Details: During selective expert use, the same examples serve as validation data for expert selection and training data for finetuning.For non-MMLU tasks, original train and validation splits are merged for both purposes; finetuning masks prompt inputs and optimizes output tokens.
  • D Token Clustering Details: EMO consistently activates a dedicated expert cluster for the first token of each document.The first token is processed without context, motivating this distinct activation behavior.
  • D Token Clustering Details: After observing one token, EMO usually transitions to a specific expert activation pattern that often remains stable throughout the document.The first-token cluster does not extend to the second token.
Loading 2605.06663v2…