EMO: Allen Institute and UC Berkeley's AI Redefines Efficiency

A Major Advancement for Language Models

The Allen Institute for AI, in collaboration with UC Berkeley, has unveiled an innovative language model named EMO. This model stands out due to its modular structure, allowing its internal components to specialize in specific domains such as medicine or politics, while maintaining impressive overall performance.

EMO employs a unique approach by setting document boundaries during training. This method enables the modules to develop expertise in distinct content areas rather than focusing solely on linguistic structures. Remarkably, even when the model is reduced to just 25% of its modules, its performance decreases by only about one percentage point. This offers significant advantages in terms of storage space efficiency and targeted control over the content domains covered.

The Concept of Mixture of Experts

Mixture of Experts (MoE) architectures have become common in language models, as demonstrated by DeepSeek-V4 and Qwen3.5. These models activate only a few experts per token, allowing for expansion to hundreds of billions of parameters without excessively increasing computational costs. However, the complete model must still be loaded into memory, as different tokens within a task call upon various experts.

The article highlights that in standard MoEs, experts often focus on superficial linguistic models, responding to elements such as prepositions or punctuation, rather than more complex domains like mathematics or code. This complicates the creation of a useful subset.

Using Document Boundaries as Training Signals

EMO addresses this issue with a simple yet effective approach. Instead of sorting training data into fixed domains like mathematics or biology, the authors use document boundaries. Tokens within a document generally belong to the same domain, allowing the model to choose its active experts from a shared pool.

The model determines which experts belong to this pool by averaging its router preferences for all tokens in a document and retaining the most frequently selected ones. This method trains modularity as a primary goal, enabling the selection of an arbitrary subset of experts for a given domain without harming the model's overall performance.

Adjustments for Stable Training

To ensure training stability, two adjustments were necessary. First, the authors stopped calculating load balancing locally by training batch, preferring a global calculation across many documents. This avoids conflicts between training objectives, one grouping tokens within a document and the other distributing them across a maximum number of experts.

Second, the researchers randomly varied the size of the document pool during training. This teaches the model to work with subgroups of experts of different sizes during inference.

Performance with a Fraction of Experts

The team trained an MoE with 1 billion active parameters and 14 billion parameters in total, spread across 128 experts, with eight active per token, on a pre-training corpus of 1 trillion tokens. As a complete model, EMO matches a standard MoE trained in the same manner and outperforms OLMoE despite using five times more data.

By reducing the number of experts, the researchers found that with only 25% of the remaining experts (32 out of 128), EMO loses about one percentage point of absolute performance on average across several benchmarks. At 12.5% (16 experts), the drop is about three points. A standard MoE, in the same configuration, loses between 10 and 15 percentage points.

Analysis of Expert Learning

To understand how EMO operates, the researchers examined the distribution of tokens to experts internally. For each token, they recorded the probability that the router would send it to each expert, creating a sort of fingerprint for each token. These fingerprints were then clustered.

Unlike a standard MoE where each token independently chooses its experts, EMO enforces consistent use of experts by defining a shared pool per document. This promotes domain specialization.

Practical Applications and Beyond

The most obvious application of this technology is running models in memory-constrained environments, where only the experts relevant to the domain are loaded. EMO's subgroups of experts match or surpass a standard MoE with 32 experts and a dense model with eight active parameters.

The researchers also envision fine-tuning models in real-time. For example, a children's application could disable clusters responding to inappropriate content. In an initial test, a subgroup of 32 experts from EMO was retrained and reintegrated into the complete model of 128 experts, enhancing the overall model without reaching the level of the standalone subgroup.

Finally, EMO could facilitate monitoring, as the experts make visible which parts of the model a given input utilizes. Ai2 has released the EMO model, a comparably trained MoE benchmark, along with the training code on Hugging Face and GitHub. The researchers have also made available an interactive demo of token activations. However, questions remain regarding the optimal selection and combination of expert subgroups, the retraining of individual modules for specific tasks, and the use of the modular structure to make models more interpretable.