The development of large language models (LLMs) has traditionally required gathering massive datasets into a centralized repository. Many of those collections include sensitive information, copyrighted text or data with strict usage limits. Such requirements block organizations with extensive proprietary or regulated data from contributing to model creation. FlexOlmo, presented by researchers at the Allen Institute for AI and partner institutions, offers a method for training and running models under these governance constraints.
Standard workflows merge all datasets into one corpus, locking in inclusion choices. There is no way to remove data once the model is built. That setup breaks rules such as HIPAA and GDPR, and it cannot respect license terms for commercial or attribution-restricted content or handle private code and clinical records properly.
FlexOlmo’s design splits training into modular units on local data and supports toggling specific dataset contributions during inference without needing to retrain the public model.
The system uses a Mixture-of-Experts (MoE) design, where each expert is a separate feedforward network (FFN) module. A public model, referred to as M_pub, acts as a shared core. Data holders train their own expert M_i on local data D_i, while all attention layers and other shared parameters remain fixed.
Sparse activation triggers only some experts for each token. Expert routing uses a matrix based on domain embeddings so tokens go to the right module without joint training. A bias term keeps no expert from dominating. This layout lets clients include or exclude modules at run time.
Each expert M_i is trained to match M_pub in a hybrid MoE setup. During that phase, M_pub and attention layers remain fixed. Only M_i’s FFN and its router embedding r_i are updated. To initialize r_i, samples from D_i are encoded by a pretrained encoder, and their average forms the initial embedding. Optional tuning on proxy public data refines r_i with minimal compute overhead.
Researchers built a training collection called FLEXMIX. It contains a public mix of general web data and seven closed subsets: News, Reddit, Code, Academic Text, Educational Text, Creative Writing and Math. Each expert sees just one of those partitions, simulating the real-world case where organizations cannot share data because of legal, ethical or operational restrictions.
The FlexOlmo architecture was tested on 31 benchmarks spanning ten task categories. Those include general language understanding tests such as MMLU and AGIEval, generative question answering with the GEN5 dataset, code generation challenges like Code4, and several mathematical reasoning tasks drawn from Math2. It was also assessed on summarization, translation, commonsense reasoning, reading comprehension, entity linking and dialogue acts.
Baseline strategies include Model soup, which averages fine-tuned model weights across multiple specialist runs to capture different domain nuances; Branch-Train-Merge (BTM), which ensembles output probabilities using optimized weights; BTX, which turns separate dense models into a MoE by transferring their parameters into expert slots; and a prompt-based router that trains an instruction-tuned classifier to send new queries to the correct expert.
FlexOlmo achieved a 41 percent relative gain over the base public model and improved by 10.1 percent compared to the best merging baseline, BTM. The benefits were strongest on tasks linked to the closed domains, confirming that dedicated expert modules add real value.
Tests highlight each design choice’s role. Skipping expert-public coordination harms scores. Random router embeddings weaken module separation. Removing the bias term skews expert calls, mainly when over two modules are active. Token-level patterns show math queries pick the math expert at deeper stages, and common tokens use M_pub, confirming more nuanced behavior than simple single-expert routing. Early layers still route many tokens to M_pub, preserving a stable base across all experiments.
One standout feature is deterministic opt-out at inference. If an expert is removed from the router matrix, its influence vanishes entirely. In experiments, omitting the News expert reduced scores on a news task but left other benchmarks unchanged, illustrating the precise scope of each module’s effect. That feature can support audit trails in regulated settings by providing a clear record of which data influences each inference.
Data extraction risks were assessed with known attacks. The public-only model had a 0.1 percent extraction rate; a dense math-trained model hit 1.6 percent; FlexOlmo with the math expert showed 0.7 percent. Teams may add differential privacy at the expert level for stronger protection. The architecture supports privacy methods or encrypted training without altering its core. Those results show that splitting data by domain can reduce leakage compared to monolithic dense models.
The team applied the FlexOlmo process to a strong baseline, OLMo-2 7B, which was pretrained on four trillion tokens. Adding two experts—Math and Code—boosted the average benchmark score from 49.8 to 52.8 without touching the central model weights. That result shows the approach scales and fits into existing training pipelines.
FlexOlmo delivers a structured method for building modular LLMs under tight governance rules. It supports distributed training on local holdings and gives organizations a clear way to switch data contributions on or off at run time. Empirical findings place its performance on par with both unified models and ensemble-based systems.
This framework is well suited for environments that require data to remain on premises, that enforce dynamic usage policies or that must comply with strict regulations. FlexOlmo demonstrates an operational path to high-performance language models that respect real-world data constraints. All credit for this research goes to the project’s original authors and contributors.