Research map / Architectures and efficiency / Mixture of experts
FlexOLMo: Open Language Models for Flexible Data Use
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces FlexOlmo, a MoE LM whose experts are trained independently on private data (with a frozen public expert as anchor) and merged at inference via domain-informed router embeddings, enabling data opt-in/opt-out without joint training.
Problem
Training language models on distributed, closed datasets without data sharing, while allowing data to be flexibly included or excluded at inference time without retraining.
Approach
Each data owner trains a FFN expert independently alongside a frozen public expert and shared attention, learning a router embedding initialized from domain embeddings (plus a negative bias); experts are then concatenated into one MoE at inference, so experts can be added or removed without further training.
Key results
- FlexOlmo (8 experts, 37B total / 20B active) averages 52.4 vs. 42.4 for the public-only model and 47.6 for BTM (top-2) across 31 tasks, about 10% relative over the best baseline.
- On the 4-expert setup, FlexOlmo averages 47.8 vs. 36.9 for the public model (about 41% relative gain) and 43.4 for BTM.
- Training-data extraction rate on math data is 0.7% with the math expert included, vs. 1.6% for a dense math-trained model and 0.1% for the public model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| combine independently trained domain experts using FlexOlmo's MoE routing | OLMo 2 7B public model · Four-set evaluation (24 tasks) | average: 36.9 -> 46.7 | Helped |
| use FlexOlmo instead of the strongest merging baseline | OLMo 2 7B public model · Four-set evaluation (24 tasks) | average: 43.4 -> 46.7 | Helped |
| combine eight independently trained experts using FlexOlmo | OLMo 2 7B public model · Eight-set evaluation (31 tasks) | average: 47.6 -> 52.4 | Helped |
| remove router bias (ablation) | FlexOlmo · Four-set evaluation (24 tasks) | average: 46.7 -> 45.8 | Hurt |
| remove domain embedding initialization and router bias (ablation) | FlexOlmo · Four-set evaluation (24 tasks) | average: 46.7 -> 44.4 | Hurt |
| remove training to coordinate (ablation) | FlexOlmo · Four-set evaluation (24 tasks) | average: 46.7 -> 38.8 | Hurt |
| initialize router embeddings with GRIT embeddings rather than public model embeddings (ablation) | FlexOlmo · Four-set evaluation (24 tasks) | average: 43.5 -> 46.7 | Helped |
| train router embeddings on proxy data after merging (ablation) | FlexOlmo · Eight-set evaluation (31 tasks) | average: 51.3 -> 52.4 | Helped |
| extend OLMo-2 7B with independently trained math and code experts | OLMo-2 7B 4T-token checkpoint · Math2 | score: 42.6 -> 51.0 | Helped |
| extend OLMo-2 7B with independently trained math and code experts | OLMo-2 7B 4T-token checkpoint · Code4 | score: 13.3 -> 18.9 | Helped |
| continue pretraining with rewritten Reddit data rather than filtered Reddit data | OLMo 7B · MC9 | score: 0.74 to 0.76 | Helped |
| continue pretraining with rewritten Reddit data rather than filtered Reddit data | OLMo 7B · MMLU | score: 0.62 to 0.66 | Helped |
Surprising findings
- Continued pretraining on closed domain sets improves some tasks but suffers catastrophic forgetting and inconsistent performance, whereas the merged model does not.
- Naively adding a public expert plus a math and code expert to OLMo-2 7B-style training yields a 2.5x inference-FLOP model that outperforms OLMo-2 7B at equal training FLOPs.
Limitations
- Extraction attacks show a nonzero, though small, leakage of training data from shared expert weights, which may need DP for sensitive data.
- Closed datasets are simulated (except for Reddit-style constraints) rather than real private data owners' sets.
- Performance is slightly below the data-controlled unrestricted MoE (2x FLOPs, 1x data) on some tasks; coding performance drops in some configurations.
Experimental setup
Datasets: DCLM-Baseline StarCoder Dolma Reddit FineMath4+ Dolmino Math Mix FlexMix PushShift Reddit
Benchmarks: MC9 Gen5 MMLU MMLU-Pro AGIEval BBH Math2 Code4 SciRIFF5 NewsG PoemG
Metrics: accuracy exact match pass@1 F1 LM-judge score extraction rate
Baselines: Model soup Branch-Train-Merge (BTM) BTX Prompt-based routing Unrestricted MoE Individual experts
Builds on: Branch-Train-Mix (BTX) Branch-Train-Merge (BTM) DEMix OLMo 2 GritLM
Closest papers by meaning
- DataComp-LM: In search of the next generation of training sets for language models (NeurIPS 2024 · similarity 0.88)
- RedPajama: an Open Dataset for Training Large Language Models (NeurIPS 2024 · similarity 0.87)
- Knowledge Fusion of Large Language Models (ICLR 2024 · similarity 0.87)
- Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.87)
- OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models (ICML 2024 · similarity 0.87)
- RegMix: Data Mixture as Regression for Language Model Pre-training (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.