FlexOlmo: Open Language Models for Flexible Data Use

🎯 Pitch

Independently trained domain experts can be combined into a single MoE model and match the performance of one trained on all data jointly, enabling data owners to contribute without sharing data and users to opt out of specific data sources at inference time with zero overhead.

1. Executive Summary

This paper introduces FLEXOLMO, a new class of language models that enables distributed training without data sharing and data-flexible inference, where model parameters trained on closed datasets can be independently added or removed at inference time with no further training. The approach uses a mixture-of-experts (MoE) architecture with a novel domain-informed routing mechanism (router embeddings initialized from off-the-shelf embedders and finetuned during coordinated expert training) that eliminates the need for joint training across data owners, evaluated on the FLEXMIX corpus—a curated mixture of public Common Crawl data alongside seven simulated closed domains (news, code, academic papers, educational text, math, creative writing, and Reddit)—using models up to 37 billion parameters (20 billion active) across 31 diverse downstream tasks. FLEXOLMO achieves an average 41% relative improvement over the public-only baseline and outperforms prior model merging methods by 10.1% on average (beating the strongest baseline BTM across benchmarks including MC9, GEN5, MMLU, BBH, Math2, and Code4), while matching the performance of an unrestricted MoE trained with the same compute budget—establishing that independently trained domain experts can be effectively combined without joint data access, but only when experts are taught to coordinate through a frozen public anchor model and routing is structured via pairwise binary classification with negative bias rather than naive continued pretraining.

2. Context and Motivation

The Core Problem: Centralized Data Pooling as a Bottleneck for LM Training

The standard recipe for pretraining language models has a fundamental structural assumption: all training data must be pooled in one place, under one model developer's control, at one point in time. Once the training run begins, the developer makes a one-time, irrevocable decision about which data sources to include or exclude. After training completes, the model has no mechanism to track which specific datasets influenced which parameters, and removing the effect of any particular data source requires costly unlearning procedures [1, 2, 3] with uncertain legal and technical guarantees.

This centralized paradigm creates friction with a wide range of real-world constraints on data sharing and usage — constraints that the paper argues are intensifying, not relaxing, as AI regulation matures. The core claim is that these constraints are not edge cases but rather structural features of the data landscape that standard LM training is poorly equipped to handle.

The paper organizes these constraints into two categories (Section 2.1), and understanding both is essential to grasping why FLEXOLMO's design choices — particularly the MoE architecture and coordinated training algorithm — are motivated the way they are.

Organizations in regulated industries — healthcare institutions subject to HIPAA, financial firms governed by data sovereignty laws, legal practices with confidentiality obligations — possess valuable domain-specific data that could substantially improve LMs for their use cases. However, this data cannot be shared externally with model developers due to legal, regulatory, or intellectual property restrictions. The paper cites HIPAA, GDPR [14, 15], data sovereignty laws [16], and IP protections as canonical barriers.

The consequence is a deadlock: these organizations could benefit enormously from domain-adapted LMs trained on their internal data, but the standard training paradigm requires them to either (a) hand their data to a third-party model developer (legally impermissible) or (b) train an entire LM from scratch internally (computationally infeasible for most organizations). There is no intermediate option in the standard paradigm where a data owner can independently train a component of a larger model and contribute it to a shared system without ever exposing the underlying data.

The paper also highlights a temporal dimension to this problem: data usage rights can expire. An organization might have permission to use a dataset for model training during a specific contractual window, after which the data must be removed. Standard pretraining provides no mechanism for surgically excising the influence of that data from the model after the fact — unlearning remains an open research problem with uncertain guarantees. This creates legal liability for model developers who cannot certify that certain data is no longer influencing model outputs.

Data Use Constraints: When Access Depends on Context, User, or Purpose

Even when data can be shared for training, whether it should be active during a particular inference request often depends on context. The paper identifies several dimensions of this (Section 2.1):

Privileged access. User-facing applications often involve closed data restricted to specific, authorized users. The paper's concrete example is GitHub Copilot: code suggestions should reflect an organization's internal repositories, but only when the requesting engineer has the appropriate role and access rights [18]. A single unified model trained on all available code cannot enforce these per-user access controls at inference time — the model either has the internal repository data in its weights (and therefore "knows" it for all users) or it doesn't.

Copyright and data consent. The legal landscape around training data for AI is evolving and uncertain [19, 20, 21, 22, 23, 24]. Different data sources come with different licenses, and what constitutes permissible use often depends on the purpose of the resulting model, not just the fact of training. For instance, a dataset licensed for non-commercial research use cannot legally be included in a model deployed for commercial applications. Current training practices bake all data into a single model, making it impossible to offer a "commercial-use" variant that simply deactivates the non-commercially-licensed training data.

Model control. Training data often include sensitive or toxic content [27, 28, 29] that may be beneficial in certain contexts but harmful in others. The paper's example is toxicity data: a researcher studying toxic content detection might want a model that has learned from toxic examples, while a public-facing chatbot must have that capability entirely deactivated. The standard paradigm forces a binary choice — include it for everyone or exclude it for everyone — with no mechanism for context-dependent activation.

These three categories of constraints (sharing, access, and purpose-based use) collectively motivate what the paper calls "a new class of LM and accompanying training methods that address restrictions in data sharing and usage." The key design requirement is not just that models can be trained without centralized data pooling, but that individual data contributions remain separable, auditable, and deactivatable throughout the model's operational lifetime.

Why This Problem Matters (Beyond the Obvious)

The paper's framing suggests the problem matters for reasons that go beyond simple regulatory compliance:

It broadens the pool of available training data. High-quality domain-specific data — medical records, legal documents, financial transactions, proprietary codebases — is disproportionately subject to sharing restrictions. If the only viable training paradigm requires centralized pooling, this entire class of data remains inaccessible for LM training, pushing models toward a lowest-common-denominator of publicly available web text. FLEXOLMO's approach creates a mechanism for data owners to contribute to model training without relinquishing control over their data, potentially opening up data sources that would otherwise remain entirely outside the LM training ecosystem.

It enables asynchronous, independent contribution. In the standard paradigm, all data must be assembled before the training run begins, and any new data source requires retraining from scratch (or continued pretraining with the attendant risk of catastrophic forgetting). The paper envisions a model where new experts can be added over time — a medical expert when a hospital joins, a legal expert when a law firm partners — without requiring all existing data owners to coordinate or retrain. This is analogous to how open-source software allows independent contributors to add modules without requiring a centralized build process that has access to every contributor's private development environment.

It provides opt-out guarantees that current approaches cannot. Existing solutions for removing data influence after training (machine unlearning) are best-effort and technically contested. Even when unlearning appears successful on evaluation metrics, it is difficult to guarantee that no residual influence remains. FLEXOLMO's architecture provides a much stronger guarantee: if an expert module trained on dataset D is physically removed from the model, the resulting model is mathematically identical to one that was never trained on D in the first place — because the training of other experts never involved D. This is not a soft "we tried to forget" guarantee but a hard architectural guarantee, which has significant legal and ethical implications for compliance with data deletion requirements.

Prior Approaches and Where They Fall Short

The paper identifies three families of prior work that address aspects of this problem, each with fundamental limitations in the specific setting of LM pretraining with data restrictions.

Federated Learning: Synchronous Coordination Is Costly and Fragile

Federated Learning (FL) [5, 4, 30] is the most well-known paradigm for training models across distributed datasets without centralized data pooling. In classical FL, a central server coordinates training across multiple clients, each holding local data. Each client trains locally, sends parameter updates (or gradients) to the server, which aggregates them and distributes the updated model back to clients — iterating until convergence.

The paper acknowledges FL's theoretical appeal but identifies several practical barriers that have limited its adoption for LM training (Section 2.2):

Synchronization cost. FL requires all participating clients to train and communicate in synchronized rounds. Training a 7B+ parameter LM generates enormous update vectors; communicating these across multiple organizations in real-time is bandwidth-intensive and slow. The paper cites prior work showing that FL has seen "limited adoption in LM training due to the high cost of synchronization and performance degradation" [6, 7].
Performance degradation. FL methods that introduce differential privacy (DP) to protect against inference attacks on model updates [36] typically degrade model quality. There is a well-documented privacy-utility-efficiency tradeoff [7] that makes FL unattractive for the quality-sensitive task of LM pretraining.
Security vulnerabilities. Perhaps most critically, FL remains susceptible to privacy attacks despite not sharing raw data, because inter-client communication (gradient updates, model parameters) can leak information about local training data [37, 38]. This undermines the core privacy guarantee that motivates the distributed approach in the first place.

The paper positions FLEXOLMO as fundamentally different from FL: rather than synchronously aggregating updates from multiple clients training on a shared model architecture, FLEXOLMO enables each data owner to independently train a modular component (an expert) and contribute it asynchronously, with no inter-client communication whatsoever during training. This eliminates synchronization overhead and removes the communication channel that FL's privacy attacks exploit. DP can still be applied locally by individual data owners, but it becomes an optional per-contributor choice rather than a system-wide requirement.

Model Merging: Combining Independently Trained Models

A second family of approaches — which the paper draws more directly from — is model merging [40, 10]. The high-level idea is to train separate models on separate datasets and then combine them into a unified model through some post-hoc procedure. The paper evaluates several specific techniques (Section 4.3):

Model Soup [12] and variants. This approach merges model weights by averaging them, either uniformly or with learned per-layer weights. It works well when models are fine-tuned from the same initialization on related tasks, because the models remain in the same loss basin and their weights are interpolatable. However, the paper finds that when models are trained on completely disjoint datasets with different distributions (as is the case with domain-specific closed data), they diverge too far from each other and from the seed model during training. Simple weight averaging becomes "limited" (Section 5.1) because the averaged weights represent a point in parameter space that is optimized for none of the individual domains.

Branch-Train-Merge (BTM) [11, 52]. BTM avoids weight-space interpolation by instead ensembling at the output level: each model processes the input independently, and their output logits are combined via a weighted average (where weights are based on each model's likelihood on the input). This is more robust to parameter divergence because it doesn't require models to be in the same weight-space basin. However, BTM has a fundamental limitation: every model must process every input token, making inference cost scale linearly with the number of experts. For an 8-expert model with 7B parameters per expert, this would mean running all 8 models in parallel for every token — a 8× increase in inference FLOPs and latency.

BTX [55] and expert routing methods. BTX addresses the inference cost problem by upcycling independently trained dense models into a Mixture-of-Experts architecture, where only a subset of experts is activated per token. This is the closest prior approach to FLEXOLMO. However, BTX has a critical requirement: after merging the independently trained experts, it performs joint training on the union of all datasets used in expert training. Specifically, BTX trains the router and all model parameters (including non-expert parameters like attention layers) on the combined dataset after merging. This violates the core data-sharing constraint that motivates FLEXOLMO: it requires joint access to all datasets.

The paper's key insight is that BTX's post-merge joint training is necessary because its experts were trained naively (simply continued pretraining from the public model) without any mechanism to ensure they coordinate with each other at merge time. FLEXOLMO's coordinated training algorithm (§3.3.1) and domain-informed router initialization (§3.3.2) are specifically designed to eliminate this requirement — teaching experts to work together during their isolated training so that no post-merge joint data access is needed.

Prompt-based routing. A simpler approach is to classify each input query into a domain and route it exclusively to the corresponding expert model. The paper tests this using Llama-3.1-8B-Instruct and OLMo-2-7B-Instruct as classifiers (Section 5.1). This approach fails for two reasons: (1) classifier errors cause catastrophic performance drops when the wrong expert is selected, and (2) many inputs benefit from combining knowledge across multiple domains, which single-expert routing cannot do. The paper's results bear this out: prompt-based routing is "highly unstable" and underperforms both BTM and FLEXOLMO.

Standard MoE Training: Joint Training on All Data

Standard MoE models [8, 9] train all experts and the router jointly on all available data. This architecture provides the sparse activation benefits that FLEXOLMO leverages, but the joint training requirement is exactly what FLEXOLMO must avoid. The paper's contribution is not the MoE architecture itself, but rather the training and merging procedures that make MoE compatible with independent, asynchronous, data-separated training.

How FLEXOLMO Positions Itself

The paper positions its contribution at the intersection of these three approaches, addressing their respective failure modes:

Against FL: FLEXOLMO eliminates synchronization and inter-client communication entirely. Training is fully independent and asynchronous.
Against model merging (weight averaging, BTM, BTX): FLEXOLMO provides the inference efficiency of sparse MoE activation (unlike BTM) without requiring post-merge joint data access (unlike BTX), by introducing coordinated training and domain-informed routing.
Against standard MoE: FLEXOLMO enables independent expert training with flexible opt-in/opt-out — capabilities the standard jointly-trained MoE cannot provide.

The conceptual framework is explicitly drawn from the open-source software analogy [40]: FLEXOLMO treats model components like software modules that can be independently developed, distributed, and composed, with clear interfaces (the frozen public model as anchor, the router embeddings as API) that enable interoperability without requiring all developers to share their source code (training data).

A subtle but important point about the paper's scope: it does not claim to solve the privacy problem in a formal, cryptographically-guaranteed sense. The data extraction analysis in Section 5.3 explicitly acknowledges that a nonzero (though small) fraction of training data can be extracted from expert weights, and recommends differential privacy for sensitive data. The contribution is about data access control — who needs to see what data during training, and who can deactivate it at inference — rather than formal privacy guarantees against weight inspection attacks.

This positioning explains several of the paper's design choices: the use of a frozen public anchor model (to provide a shared coordinate system without sharing closed data), the pairwise binary classification approach to router training (because each expert only sees its own data plus public data, never other closed data), the addition of a negative bias term (to tighten decision boundaries around each expert's specialization, compensating for the fact that experts are never compared against each other during training), and the optional router tuning on proxy data (to provide a lightweight post-merge refinement that doesn't require joint access to closed data). Each of these design choices directly addresses a specific limitation of prior approaches in the data-restricted setting.

3. Technical Approach

This is primarily a systems and training methodology paper whose core idea is that independently trained domain experts can be combined into a unified Mixture-of-Experts model without joint data access — if experts are trained to coordinate through a frozen public anchor model and routing is constructed via domain-informed embeddings learned during isolated training, rather than requiring post-merge joint optimization.

3.1 Reader Orientation

The paper builds a modular language model system where different organizations can independently train parts of the model on their own closed data and contribute those parts to a shared model, with the ability to add or remove any part at inference time without retraining. The core problem is that standard language model training requires all data to be pooled in one place at one time, which is incompatible with data sharing restrictions (HIPAA, GDPR, IP protections) and with the need for fine-grained data access control (some data should only be active for certain users or use cases). The solution's "shape" is a Mixture-of-Experts architecture where each expert is a set of feedforward networks trained independently on one closed dataset, with a frozen public model serving as a shared "anchor point" that teaches all experts to coordinate with each other, and a router constructed from domain embeddings that enables integrating these experts without joint training.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components:

Public Model (Mpub) — a dense 7B-parameter language model pretrained on open web data (DCLM-Baseline). It serves as the foundation and the shared anchor point that all experts learn to coordinate with. Its parameters (including attention layers and one set of FFN layers) are frozen during expert training.
Expert Modules (M1...Mn) — feedforward networks (FFNs) that replace some FFN layers in the transformer blocks. Each expert Mi is trained independently by the owner of closed dataset Di, who has access to Di and a copy of Mpub but never shares Di with anyone else.
Domain-Informed Router — a learned matrix Wr that maps input token embeddings to a probability distribution over which experts to activate. The router is constructed by concatenating per-expert router embedding vectors ri, each initialized from domain embeddings of Di and finetuned during isolated expert training.
Coordinated Training Algorithm — a training procedure where each expert is trained within a temporary two-expert MoE alongside a frozen copy of the public model, learning both its FFN parameters and its router embedding to complement (not replace) the public model.
Merging Procedure — a zero-training-cost integration step where all independently trained expert FFNs and their router embeddings are combined into a single MoE model. Experts can be added or removed by inserting or deleting the corresponding rows in the router matrix.

Information flows as follows: a token enters the transformer block → the router computes scores ri · x + bi for each expert (where bi is a learned negative bias) and selects the top-k → the activated experts process the token through their FFNs → the outputs are weighted by softmax-normalized router scores and summed → the result is added to the attention output and passed to the next layer.

3.3 Roadmap for the Deep Dive

First, the formal problem setup and constraints (§3.1 in the paper), defining exactly what "training without data sharing" and "flexible opt-out" mean — because these constraints drive every subsequent design choice.
Second, the MoE architecture at the layer level (§3.2), establishing the standard MoE computation that FLEXOLMO adapts and clarifying what changes versus what stays the same.
Third, the coordinated training algorithm (§3.3.1), the paper's core technical innovation — how each expert is trained in isolation as a two-expert MoE with a frozen public anchor, and why this teaches experts to coordinate without joint data access.
Fourth, the domain-informed router construction (§3.3.2), including router embedding initialization from off-the-shelf embedders, the pairwise binary classification framing, and the crucial negative bias term that tightens decision boundaries.
Fifth, the optional router tuning on proxy data (§3.3.3), a lightweight post-merge refinement that uses only small public subsets to improve routing without accessing closed data.

3.4 Detailed, Sentence-Based Technical Breakdown

Problem Setup and Constraints

The paper formalizes the setting (Section 3.1) as follows. Let $M_{\text{pub}}$ be a model trained on a publicly available dataset $D_{\text{pub}}$ . Let $D = \{D_1, D_2, ..., D_n\}$ be a collection of locally maintained datasets, each with a separate owner. The objective is to construct a single model $M_{\text{final}}$ by composing $M_{\text{pub}}$ and a set of modules $\{M_1, M_2, ..., M_n\}$ , where each $M_i$ is independently trained by the owner of $D_i$ (who also has access to $M_{\text{pub}}$ ).

The paper articulates two hard constraints that $M_{\text{final}}$ must satisfy:

(1) Training without joint data access. Constructing $M_{\text{final}}$ does not require anyone to have joint access to the full dataset collection $D$ . Each $M_i$ is trained independently by the owner of $D_i$ , who never shares $D_i$ with anyone else. This means no step in the pipeline can involve combining $D_i$ and $D_j$ for $i \neq j$ — not for training experts, not for training a router, not for any post-merge fine-tuning.

(2) Guaranteed opt-out. Removing any module $M_i$ from $M_{\text{final}}$ guarantees complete removal of the influence of its associated data $D_i$ . The paper states this as: the weights of $M_i$ were trained only on $D_i$ (plus the public data via the anchor model, but never on any other $D_j$ ), so physically deleting $M_i$ 's parameters produces a model mathematically identical to one that was never trained on $D_i$ — a hard architectural guarantee, not a best-effort unlearning approximation.

These constraints rule out several otherwise-plausible approaches. Joint router training on the union of all datasets is forbidden because it requires joint data access. Post-merge fine-tuning of all parameters (as in BTX) is forbidden for the same reason. Even a central server that receives gradients from multiple data owners (as in federated learning) introduces communication channels that the paper explicitly wants to avoid, both for privacy reasons and for the asynchronous contribution model.

The key modeling challenges that fall out of these constraints are: (1) how to create each $M_i$ using only $D_i$ and $M_{\text{pub}}$ such that $M_i$ will work well with other modules it has never seen during training, and (2) how to merge $M_{\text{pub}}, M_1, ..., M_n$ into $M_{\text{final}}$ without access to the combined data $D_{\text{pub}} \cup D_1 \cup ... \cup D_n$ .

MoE Architecture: What Changes and What Doesn't

FLEXOLMO adopts the standard Mixture-of-Experts architecture (Section 3.2) with a specific modification to support independent training. In a standard transformer, each block contains an attention layer followed by a feedforward network (FFN). In an MoE transformer, the single FFN is replaced by a router and $n$ smaller FFNs (called experts). The paper explicitly notes: "we omit the layer index for each expert in our notation for simplicity" — meaning that at each transformer layer, there is a separate set of expert FFNs, and the routing decision is made per-layer independently.

The MoE computation for a single layer is (Section 3.2):

$y = \sum_{i \in \text{Top}_k(r(x))} \text{softmax}(r(x)_i) M_i(x)$

where $x \in \mathbb{R}^h$ is the processed input token embedding from the attention layer (with $h$ being the hidden dimension, 4096 in the 7B model), $r(x)$ is the router function that computes a score vector over all $n+1$ experts (the public expert plus $n$ domain experts), $\text{Top}_k$ selects the indices of the $k$ highest-scoring experts, $\text{softmax}$ normalizes those top- $k$ scores into a probability distribution summing to 1, and $M_i(x)$ is the output of expert $i$ 's FFN applied to $x$ .

What this computes: For each token at each transformer layer, the router selects the $k$ most relevant experts (out of $n+1$ , where $k=4$ in their final configuration), computes a weighted combination of their outputs, and sums them. Experts that are not selected contribute zero — their FFNs are never executed for that token. This sparse activation is what keeps inference cost manageable: with 8 total experts and $k=4$ active, roughly half the FFN parameters are used per token, giving a model with 37 billion total parameters but only 20 billion active.

Why this form: The Top-k sparsity is the key enabler for independent expert training. Unlike dense ensembling (BTM) where every expert processes every token, the MoE architecture means adding a new expert increases total parameters but does not linearly increase per-token compute — only the $k$ most relevant experts are activated. More importantly, the architecture provides a natural interface for modularity: each expert is a self-contained FFN that can be independently trained, and the router provides the only coupling between experts. If the router can be constructed without joint training, the entire MoE can be assembled without joint data access.

The critical difference from standard MoEs: In standard MoE training, all expert FFNs and the router are trained jointly on the full dataset, with load-balancing losses to encourage uniform expert utilization. FLEXOLMO breaks this coupling: experts are trained independently on disjoint datasets, and the router is constructed post-hoc from embeddings learned during isolated training. This is the fundamental architectural bet — that the per-layer per-token routing can be made to work without the router and experts co-adapting through joint gradient updates.

Non-expert parameters: The attention layers and other non-FFN parameters are shared across all experts and are inherited from the public model $M_{\text{pub}}$ . During expert training, these shared parameters are frozen — only the new expert FFNs and their router embeddings are updated. This is crucial because it means all experts share the same "representation space" produced by the attention layers, enabling the router embeddings to operate in a consistent coordinate system.

Coordinated Training Algorithm: Teaching Experts to Work Together in Isolation

This is the paper's central technical innovation (Section 3.3.1). The challenge is: how do you train expert $M_i$ on dataset $D_i$ such that, months later, it can be plugged into an MoE alongside experts $M_j, M_k, ...$ that it has never encountered during training — and the combined model should perform well without any post-merge joint training?

The naive approach and why it fails. The straightforward method would be to take the public model $M_{\text{pub}}$ and continue pretraining it on $D_i$ , treating the resulting dense model as expert $M_i$ . The paper explicitly reports that this fails (Section 3.3.1): "this method causes the experts to diverge too much from one another and from the original seed model, which makes merging after isolated training difficult." Divergence here means that the FFNs trained on different domains move to different regions of parameter space, producing incompatible representations that cannot be cleanly combined — the router has no basis for choosing between them because their outputs don't operate in the same "language."

The coordinated training procedure. Instead of training $M_i$ as an isolated dense model, the paper constructs a temporary two-expert MoE for each dataset $D_i$ :

Start with the public model $M_{\text{pub}}$ , which is a standard dense transformer with 32 layers, hidden dimension 4096, trained on 1 trillion tokens of public data. The model has standard FFNs in each layer — let's call these the "public FFNs."
For dataset $D_i$ , construct a two-expert MoE by taking the public FFNs and creating a second, initially identical copy of them. The first copy (public expert) is frozen — its parameters do not change during training. The second copy (expert $M_i$ ) is unfrozen and will be trained on $D_i$ . The shared attention layers are also frozen.
During training on $D_i$ , the router must decide for each token whether to send it to the frozen public expert or the trainable domain expert. The router itself consists of two router embedding vectors: $r_{\text{pub}}$ (frozen) and $r_i$ (trainable), which produce scores $r_{\text{pub}} \cdot x$ and $r_i \cdot x$ for input $x$ .
The model is trained with a standard language modeling loss (next-token prediction) on $D_i$ , but only the parameters of expert $M_i$ (its FFN weights) and its router embedding $r_i$ receive gradient updates. Everything else — the public expert's FFNs, the attention layers, the public router embedding $r_{\text{pub}}$ — remains frozen.

Why this teaches coordination. Because the public expert is frozen and always available as an alternative, the domain expert $M_i$ has an incentive to specialize: it should learn to handle tokens where the public model is inadequate (i.e., domain-specific patterns in $D_i$ ), while leaving tokens where the public model already performs well to the public expert. If the domain expert tried to learn everything in $D_i$ , including patterns already well-handled by the public model, it would be redundant and would not improve the training loss (since the public expert already handles those tokens). The gradient signal naturally pushes the domain expert to focus on the residual — what $D_i$ contains that the public model doesn't capture.

More subtly, because every expert $M_1, M_2, ..., M_n$ is trained with the same frozen public expert as its counterpart, they all learn to be complementary to the same reference model. This shared reference creates an implicit coordination: expert $M_i$ and expert $M_j$ never see each other during training, but they both learned to fill gaps relative to the same $M_{\text{pub}}$ . When they are later merged into a single MoE, their outputs are compatible because they were both optimized to complement the same baseline.

A concrete example from the paper's data. The Math expert is trained on 20.3 billion tokens of math-relevant content. During coordinated training, when the input token is part of a mathematical expression that the public model handles poorly (e.g., a LaTeX formula or mathematical reasoning step), the router learns to assign high probability to the Math expert, and the Math expert's FFNs learn to produce better next-token predictions than the public expert would. When the input is general English text that the public model handles well, the router learns to assign high probability to the public expert, and the Math expert doesn't waste capacity on these tokens. This specialization is not manually specified — it emerges from the gradient signal, because the only way to reduce the loss on $D_i$ is to improve predictions on tokens where the public model is currently failing.

Training hyperparameters. Each expert is trained for 50 billion tokens (the total across all 8 experts is 400 billion tokens). The paper uses the same optimization setup as the public model: "the AdamW optimizer with parameters $\beta_1 = 0.9$ and $\beta_2 = 0.95$ and a cosine learning rate scheduler." The specific learning rate for expert training is not separately reported in the main text (only the public model's LR of 0.0009 is given), but the training infrastructure and batch sizes are consistent with the public model setup.

Domain-Informed Router: Embedding-Based Routing Without Joint Training

The router is the critical integration mechanism in any MoE — it decides which experts process each token. In standard MoEs, the router is trained jointly with all experts, learning to route tokens based on which experts produce the best outputs. FLEXOLMO cannot do this because joint training requires joint data access. The paper's solution is a domain-informed router (Section 3.3.2) that is constructed from per-expert embeddings learned during isolated training.

Router structure. The router is a learned matrix $W_r \in \mathbb{R}^{(n+1) \times h}$ that maps an input vector $x \in \mathbb{R}^h$ to a distribution over the $n+1$ experts (the public expert plus $n$ domain experts):

$r(x) = W_r x$

where each row of $W_r$ is a router embedding vector $r_i \in \mathbb{R}^h$ corresponding to expert $i$ . The full matrix is:

$W_r = \begin{bmatrix} r_{\text{pub}} \\ r_1 \\ \vdots \\ r_n \end{bmatrix}$

What this computes: The router score for expert $i$ is the dot product $r_i \cdot x$ — the similarity between the input representation $x$ and the learned embedding $r_i$ . Experts whose embeddings are more similar to $x$ (in the dot-product sense) receive higher scores. The Top-k experts by score are then activated.

Router embedding initialization. Where do these embeddings come from? The paper leverages an external, off-the-shelf text embedder to initialize them (Section 3.3.2):

$r_i = \frac{1}{|S_i|} \sum_{d_k \in S_i} E(d_k), \quad S_i \subset D_i$

where $E$ is an off-the-shelf embedder (specifically, GritLM/GritLM-7B, a generative representational instruction tuning model [67] that maps a document into an $h$ -dimensional vector, with $h = 4096$ matching the model's hidden dimension), $S_i$ is a sampled subset of documents from $D_i$ (the paper samples 1,000 documents per data source, as stated in Appendix A.2), and $E(d_k)$ is the embedding vector for document $d_k$ .

What this computes: For each closed dataset, the data owner processes a small sample of their documents through the embedder, obtains a vector for each, and averages them. The result is a single vector $r_i$ that represents the "prototypical content" of $D_i$ in the embedder's representation space. This vector is used as the initial router embedding for expert $M_i$ .

Why this form over alternatives: The paper compares this GRIT embedder initialization against using the public model's own hidden states as the initialization (Table 3), finding that GRIT substantially outperforms (46.7 average vs. 43.5 average across benchmarks). The likely reason is that the external embedder provides a representation space that is specifically optimized for semantic similarity (through its contrastive training), whereas the public model's hidden states are optimized for next-token prediction and may not cleanly separate domains. The domain-embedding initialization gives the router a meaningful starting point — the router embedding for "Math" starts out close (in dot-product space) to tokens that look like math, and far from tokens that look like news.

Router embedding finetuning. During the coordinated training of expert $M_i$ (the two-expert MoE with frozen public expert), the router embedding $r_i$ is finetuned alongside the FFN parameters of $M_i$ . Importantly, $r_{\text{pub}}$ remains frozen across all experts. This means every expert's router embedding is learned relative to the same, fixed public embedding — providing a consistent coordinate system. After training all $n$ experts independently, the $n+1$ router embeddings are simply concatenated into the full router matrix $W_r$ . No further router training on combined data is required.

The pairwise binary classification framing (Appendix D). The paper provides a formal intuition for why this works. During isolated training, each expert learns a binary classifier between the public class $C_{\text{pub}}$ and its own class $C_i$ , with decision boundary:

$h_i = \{x \in \mathbb{R}^h : s_{\text{pub}}(x) = s_i(x)\}$

where $s_i(x) = r_i \cdot x$ . The decision rule for expert $i$ during training is simply:

$f_i(x) = \begin{cases} C_i & \text{if } r_i \cdot x > r_{\text{pub}} \cdot x \\ C_{\text{pub}} & \text{otherwise} \end{cases}$

At inference time, these $n$ binary classifiers are combined into a multiclass classifier:

$F(x) = \arg\max_{i \in \{\text{pub}, 1, ..., n\}} s_i(x)$

The key question is: why should binary classifiers trained independently work well as a multiclass classifier when combined? The paper's answer is that because each binary classifier uses the same frozen $r_{\text{pub}}$ as its negative class and the same frozen public expert as its baseline, they all operate in the same coordinate system. The multiclass decision $\arg\max_i r_i \cdot x$ is equivalent to running all $n$ binary decisions simultaneously and picking the class with the highest margin over the public baseline.

The negative bias term. The paper identifies a subtle problem with this scheme and introduces a negative bias $b_i$ for each expert to fix it (Section 3.3.2). The issue is that during binary training, the constraint is only $r_i \cdot x > r_{\text{pub}} \cdot x$ for $x \in C_i$ . Many different $r_i$ vectors can satisfy this without precisely characterizing the specialized region of $C_i$ . When multiple experts are combined at inference, an input from $C_i$ might accidentally score higher under $r_j$ (for $j \neq i$ ) than under $r_i$ , because $r_j$ 's decision boundary was never calibrated against $C_i$ .

The solution is to modify the training-time decision rule for expert $i$ :

$r_i \cdot x + b_i > r_{\text{pub}} \cdot x$

where $b_i$ is a learned negative scalar bias (initialized to a negative value and trained alongside $r_i$ ). As illustrated in Figure 6, this negative bias shifts the decision boundary closer to the data points of $C_i$ , making the expert more "conservative" — it only activates on inputs that are strongly characteristic of $C_i$ , not on inputs that are merely slightly more like $C_i$ than the public data. At inference time, this tighter boundary reduces interference between experts, because each expert claims a smaller, more precisely defined region of the input space.

What the ablation says about the bias. Table 1 includes an ablation removing the bias term ("no bias"). The average performance drops from 46.7 to 45.8 (across 24 tasks), a modest but consistent degradation. The bias term is not the most critical component (removing coordinated training entirely drops performance to 38.8), but it provides a consistent improvement by tightening the router's specialization boundaries. The paper also ablates removing both the domain embedding initialization and the bias ("no domain embedding init, no bias"), which drops performance further to 44.4, showing that the two mechanisms are complementary.

Optional Router Training on Proxy Data

The coordinated training procedure produces a functional router without any post-merge training. However, the router was trained only on pairwise decisions ( $C_i$ vs. $C_{\text{pub}}$ ), never on multiclass decisions ( $C_i$ vs. $C_j$ for $i \neq j$ ). The paper introduces an optional lightweight refinement step (Section 3.3.3) that can improve routing quality without requiring joint access to the closed datasets.

The proxy data assumption. The paper assumes that each data owner can identify a small set of public documents $\hat{D}_i \subseteq D_{\text{pub}}$ that approximately resemble their closed data $D_i$ . Formally, $|\hat{D}_i| \ll 0.01 \times |D_i|$ — less than 1% of the size of the closed dataset. This proxy set is too small to train expert FFNs (which need billions of tokens), but it provides useful signal for refining router embeddings.

Constructing proxy data. To build $\hat{D}_i$ , the paper trains a binary classifier to distinguish documents from $D_i$ versus $D_{\text{pub}}$ . Specifically, they finetune Snowflake/snowflake-arctic-embed-xs (a 22M-parameter embedding model) on a balanced dataset of 500,000 documents (250,000 from each source). The classifier "quickly achieved an accuracy above 95% across all datasets considered" (Appendix A.2). Once trained, the classifier is applied to the public dataset, and public documents with the highest predicted probability of belonging to $D_i$ are selected as the proxy set $\hat{D}_i$ .

What happens during proxy router training. After all experts are merged into the full MoE, the router embeddings $r_1, ..., r_n, r_{\text{pub}}$ are further finetuned on the combined proxy sets $\hat{D}_1, ..., \hat{D}_n$ along with the full public data $D_{\text{pub}}$ , sampled uniformly. The expert FFNs are not updated during this step — only the router embeddings (the rows of $W_r$ ) receive gradient updates. The training uses a standard language modeling objective, but gradients flow only through the router parameters, not through the expert FFNs or attention layers.

How much data is used. The total proxy router training uses 5 billion tokens in total (stated in Section 4.4). This is tiny compared to expert training (50B tokens per expert, 400B total) — consistent with the claim that proxy sets are <1% of closed data sizes.

What this achieves. The router tuning step exposes each router embedding to data that resembles other domains' data (via the proxy sets), enabling the router to learn multiclass distinctions that were absent during pairwise binary training. For instance, the Math router embedding can now see proxy-news data (public documents that resemble news) and learn to produce lower scores on news-like inputs, even though it never saw actual news data during its isolated training.

The impact. Table 1 shows that adding router training improves FLEXOLMO's average performance from 46.7 to 47.8 (comparing "FLEXOLMO (no optional router training)" vs. "FLEXOLMO"). The improvement is not dramatic — the system already works well without it — but it provides a consistent boost, particularly on general benchmarks (GEN5: 58.6 → 59.8, BBH: 43.4 → 47.1) where multiclass routing decisions are more complex.

Why this doesn't violate the data access constraint. The proxy router training uses only public data — documents from $D_{\text{pub}}$ that happen to be similar to closed domains. No actual closed data is shared or accessed during this step. Each data owner independently identifies proxy documents using their locally-trained classifier, and these public documents (not the closed data) are used for router tuning. The closed data $D_i$ is used only to train the classifier that selects proxies, and this classifier training happens locally at each data owner's site.

Full Training Pipeline: End-to-End Walkthrough

Putting all components together, the complete FLEXOLMO training pipeline for a system with public data and $n$ closed datasets is:

Phase 0: Public model pretraining. A dense 7B-parameter model (32 layers, hidden dim 4096) is trained on 1 trillion tokens of public web data (the Public Mix, which is DCLM-Baseline with news and creative writing excluded). This uses the OLMo 2 architecture and training recipe: AdamW with $\beta_1 = 0.9$ , $\beta_2 = 0.95$ , cosine LR schedule, learning rate 0.0009, global batch size 4M tokens, trained on 512 H100 GPUs for three days. This produces $M_{\text{pub}}$ .

Phase 1: Router embedding initialization. Each data owner processes 1,000 documents from their dataset $D_i$ through the GritLM embedder, averages the resulting 4096-dimensional vectors, and uses this average as the initial router embedding $r_i$ . The public router embedding $r_{\text{pub}}$ is initialized similarly from a sample of public data.

Phase 2: Coordinated expert training (parallel, independent). For each $D_i$ , the data owner constructs a two-expert MoE from $M_{\text{pub}}$ : the first expert is a frozen copy of $M_{\text{pub}}$ 's FFNs, the second is a trainable copy (initialized identically). The router contains $r_{\text{pub}}$ (frozen) and $r_i$ (trainable, initialized from Phase 1). The model is trained on $D_i$ for 50 billion tokens using the standard OLMo 2 recipe. Only $M_i$ 's FFN parameters and $r_i$ are updated; all other parameters (attention, $M_{\text{pub}}$ 's FFNs, $r_{\text{pub}}$ ) are frozen. This phase runs in parallel for all $n$ data owners, with no communication between them.

Phase 3: Merging (zero training). The trained expert FFNs $M_1, ..., M_n$ and their router embeddings $r_1, ..., r_n$ are collected, along with the frozen public expert FFNs and $r_{\text{pub}}$ . The full router matrix $W_r$ is formed by vertically stacking all $n+1$ embeddings. At each transformer layer, the MoE module now has $n+1$ expert FFNs and the stacked router. The model is a standard Top- $k$ MoE with $k=4$ , 37B total parameters, 20B active.

Phase 4 (optional): Proxy router tuning. Each data owner trains a binary classifier (22M parameters, 500K training examples, >95% accuracy) to distinguish their $D_i$ from $D_{\text{pub}}$ , then uses it to select the highest-scoring public documents as proxies $\hat{D}_i$ . The proxies from all owners (5B tokens total) are combined with public data, and the router embeddings are finetuned on this mixture using a language modeling objective. Expert FFNs and attention parameters remain frozen.

Inference. At inference time, each token is processed through the router across all $n+1$ experts. The top- $k$ experts (where $k=4$ ) are activated per layer. To opt out of dataset $D_i$ , the corresponding expert FFNs and router embedding row $r_i$ are removed. The remaining model is mathematically identical to one that was never trained with $D_i$ involved.

Summary of Design Choices and Their Justifications

Frozen public anchor during coordinated training over fully independent continued pretraining: prevents expert divergence by ensuring all experts optimize their FFNs to complement the same reference model, creating implicit coordination without joint training.
Domain embedding initialization of router over random initialization or public model hidden states: provides a semantically meaningful starting point (GRIT embedder is contrastively trained for similarity), improving final router quality (Table 3 shows +3.2 points over public model initialization).
Negative bias term over unbiased pairwise training: tightens decision boundaries so experts claim smaller, more precise regions of input space, reducing interference when multiple experts are combined. The paper's Figure 6 and Appendix D provide geometric intuition; Table 1 ablations show consistent benefit.
Softmax over Top-k over hard routing to a single expert: allows inputs to combine knowledge from multiple domains (e.g., a math problem involving code benefits from both the Math and Code experts), which per-expert routing (like prompt-based baselines) cannot do. Figure 2 confirms that inputs frequently activate multiple domain experts simultaneously.
Freezing attention layers over training them per-expert: prevents each expert from modifying the shared representation space, ensuring all experts operate in the same coordinate system defined by $M_{\text{pub}}$ 's attention parameters. This is what makes router embedding concatenation valid.
Proxy router tuning on <1% data over full joint router training: provides some multiclass routing signal without violating data-sharing constraints, improving performance modestly (Table 1: +1.1 points on average) while keeping closed data local.
Coordinating via frozen public rather than via inter-expert communication: eliminates all synchronization and communication overhead that plagues federated learning, enabling truly asynchronous, independent contribution. The cost is that the public model must be reasonably capable on all domains to serve as an effective anchor.

4. Key Insights and Innovations

Innovation 1: Reframing the Data Restriction Problem as an Architecture Design Problem, Not a Privacy Problem

The paper's most fundamental conceptual move is redefining what "training with data restrictions" means in a way that unlocks a different class of solutions. The dominant framing in prior work — particularly federated learning [5, 4] — treats the core challenge as privacy: how do we prevent information leakage from local datasets during collaborative training? This framing naturally leads to cryptographic and differential privacy solutions that add communication overhead, degrade model quality, and struggle to scale to LM pretraining.

FLEXOLMO shifts the framing to a data access and composability problem. The paper asks not "how do we prevent data from leaking during parameter sharing?" but rather "how do we design a model architecture where components trained on separate datasets can be composed without anyone needing joint access to those datasets — ever?" This is a fundamentally different question with fundamentally different technical constraints, and it sidesteps several problems that plagued FL:

No communication channel to secure. Since experts are trained independently and merged post-hoc, there is no gradient exchange, no parameter server, and no inter-client communication during training. The privacy vulnerabilities that FL's communication channels introduce [37, 38] simply don't apply — not because they're better defended, but because the attack surface doesn't exist.
Asynchronous contribution is native. In FL, all clients must participate in synchronized rounds, creating coordination overhead and making it impossible for new data owners to join after training completes. FLEXOLMO's architecture naturally supports adding new experts months or years after the initial model is deployed, because experts are merged via router concatenation, not iterative aggregation. This is not an optimization over FL — it's a different capability entirely.
Opt-out is an architectural guarantee. FL produces a single monolithic model trained on all clients' data; removing one client's contribution requires unlearning, which is technically uncertain. FLEXOLMO provides a hard guarantee: deleting expert $M_i$ 's weights produces a model mathematically identical to one trained without $D_i$ . This is not a privacy guarantee (the weights themselves may contain extractable information, as Section 5.3 documents), but it is a data governance guarantee with direct legal and compliance implications that FL cannot offer.

This reframing is significant beyond the technical solution it enables. It recasts the bottleneck for LM training on restricted data from "we need better privacy technology" to "we need better modular architectures." The implication is that privacy-preserving machine learning and modular model design are complementary but distinct research directions, and that for many real-world data restriction scenarios (regulatory compliance, expired data rights, per-user access control), the modular architecture approach solves the governance problem even without formal privacy guarantees.

Evidence for this reframing's power is visible throughout the results: FLEXOLMO outperforms the FL-controlled baseline (Table 1: 47.8 vs. best FL-equivalent baseline at 43.4 for BTM), supports opt-out with minimal collateral damage (Figure 4: removing the news expert drops NewsG performance by more than 20 points while affecting other tasks by <2 points), and enables adding new experts to an already-deployed model without retraining (the full eight-expert setup in Table 2 was built by simply adding four experts to the four-expert model from Table 1).

Innovation 2: The Anchor Model Concept as an Alternative to Joint Optimization

The paper's second conceptual contribution is the anchor model training strategy — using a frozen public model as a shared reference point during isolated expert training — as a principled alternative to the joint optimization that prior MoE composition methods required.

Prior work on composing independently trained models into an MoE (most notably BTX [55]) relied on post-merge joint training on the union of all datasets to teach the router and experts to work together. This is a natural solution to a real problem: independently trained models produce incompatible representations, and some form of joint adaptation is needed to make routing decisions coherent. The field's assumption was that this joint adaptation requires access to the combined data, making it incompatible with data-sharing restrictions.

FLEXOLMO's insight is that joint data access is sufficient for coherent merging but not necessary. The anchor model provides an alternative: instead of training experts to work directly with each other (which requires them to co-train or at least co-adapt on shared data), train each expert to work with the same frozen reference model. If every expert learns to complement $M_{\text{pub}}$ , then by transitivity, experts $M_i$ and $M_j$ are implicitly coordinated — they don't need to have seen each other's data because they both learned their role relative to the same baseline.

This is a fundamental shift in how to think about modular model training. It's analogous to the difference between teaching a team of specialists to work together by having them practice as a team (joint training, requires everyone in the same room) versus teaching each specialist to complement a shared standard operating procedure (anchor training, requires only that everyone knows the SOP). The latter is less expressive — specialists can't develop specialized coordination patterns with each other — but it enables asynchronous, independent participation at scale.

What makes this more than an architectural trick is the training dynamics argument implicit in the design. Because the public expert is frozen during the two-expert MoE training, the domain expert's gradient signal naturally pushes it toward specialization: tokens where the public model already produces good predictions receive low loss regardless of whether they route to the domain expert, so the domain expert's parameters are primarily shaped by tokens where the public model fails. This is not an explicit auxiliary loss or a manually designed specialization objective — it's an emergent property of training a new expert alongside a frozen reference under a shared language modeling loss. The ablation results confirm this is essential: removing coordinated training (Table 1, "no training to coordinate") causes a catastrophic drop from 46.7 to 38.8 average performance, a far larger degradation than removing any other component. Without the anchor, experts are just independently continued-pretrained dense models that diverge into incompatible parameter regions.

The concept is likely transferable beyond the specific MoE architecture in this paper. Any modular architecture where components can be independently trained and later composed could benefit from an anchor-based coordination strategy — the anchor provides a shared coordinate system without requiring joint data access. The paper doesn't explore this generalization, but the principle is clear.

While Section 2.1 of the paper reads as a motivation section, the taxonomy of data constraints it introduces — separating sharing constraints (data cannot leave the owner's environment) from usage constraints (data can be shared for training but its activation at inference depends on context) — is itself a conceptual contribution with implications for how the field evaluates modular training methods.

Prior work on federated learning, model merging, and modular architectures typically bundles these constraints together under the umbrella of "privacy" or "data restrictions," leading to solutions optimized for one type of constraint at the expense of the other. Federated learning focuses almost exclusively on sharing constraints (keeping data local) but provides no mechanism for inference-time data access control — the resulting model is a single unified artifact with no separable data contributions. Model merging approaches like BTM and BTX address modularity but were designed for the case where all data can be centrally accessed during merging, making them incompatible with sharing constraints.

FLEXOLMO's design is unusual in that it addresses both constraint types simultaneously through a single architecture:

For sharing constraints: the coordinated training algorithm ensures no joint data access is needed at any point, including during merging.
For usage constraints: the MoE architecture with per-expert router embeddings provides fine-grained inference-time control — removing an expert physically deactivates the associated data's influence, enabling per-user, per-use-case, or per-license activation policies.

This dual coverage is not incidental but rather a direct consequence of the architecture choice. An MoE where each expert corresponds to a specific dataset and the router is constructed from per-expert embeddings naturally supports both properties, while alternative architectures (dense ensembles, weight-merged models, federated aggregates) trade off one for the other.

The taxonomy also clarifies what FLEXOLMO does not solve — and this is where the data extraction analysis in Section 5.3 becomes more than a standard evaluation. The paper shows that trained expert weights can leak training data (1.6% extraction rate for a dense math expert, 0.7% for the math expert within FLEXOLMO), and explicitly recommends differential privacy for sensitive data. This establishes a boundary: FLEXOLMO solves the data access problem (who can train on what, who can activate what at inference) but not the formal privacy problem (what can be inferred from inspecting trained weights). The separation of these concerns — governance vs. privacy — is a useful clarification that the field's conflated terminology often obscures.

Innovation 4: Router Construction as Domain-Embedding Concatenation with Pairwise Binary Training

The paper's approach to the router problem is conceptually distinct from both joint MoE training and existing model merging routers. Standard MoEs train the router jointly with all experts on all data, learning routing decisions through the same gradient signal that trains the experts. BTX and related methods train the router during post-merge joint optimization on the combined datasets. Prompt-based routing uses an external classifier's domain predictions. Model soup and BTM avoid routing entirely by averaging or ensembling all experts on every input.

FLEXOLMO's router construction is novel in combining three ideas that are individually precedented but whose combination is not: (1) domain-embedding initialization from off-the-shelf embedders (inspired by prior work on embedding-based routing [68, 42, 69]), (2) pairwise binary training where each expert learns routing decisions only against the public baseline (enabled by the anchor model concept), and (3) concatenation-based merging where per-expert embeddings are simply stacked into the router matrix with no post-hoc alignment or joint optimization required.

What makes this intellectually distinctive is the pairwise-to-multiclass generalization argument formalized in Appendix D. The paper argues that if each binary classifier $f_i$ learns a tight decision boundary around its domain $C_i$ relative to a shared public baseline $C_{\text{pub}}$ , then the argmax over all binary scores at inference time approximates what a jointly-trained multiclass classifier would do — provided the decision boundaries are tight enough that they don't overlap in ways that cause interference. The negative bias term can be understood as a mechanism for tightening these boundaries, reducing the "claim staked" by each expert in the input space.

This framing converts a potentially arbitrary engineering choice (initialize router embeddings from domain embeddings, train them in pairs, add a bias term) into a principled argument with clear failure modes: if the boundaries aren't tight enough, experts will interfere; if they're too tight, routing becomes brittle and experts won't activate when needed. The empirical finding that performance plateaus at 4 active experts (Figure 3) and that the bias term provides consistent but modest gains (Table 1: +0.9 points) aligns with this framing — the pairwise training produces boundaries that are good enough to work in practice but not perfectly calibrated for multiclass decisions.

The comparison to BTX's router training is instructive. BTX needs joint data to train the router because its experts were trained naively and produce incompatible logit scales. FLEXOLMO doesn't need joint data because the pairwise training ensures that expert logits are calibrated relative to the same frozen public reference — the router scores $r_i \cdot x$ are directly comparable because they were all learned with the same $r_{\text{pub}}$ as the decision boundary anchor. This is a genuinely different approach to making routing decisions without joint training, not just an incremental refinement of existing methods.

The ablation in Table 3 strengthens the case: using GRIT embedder initialization for router embeddings (46.7 average) substantially outperforms using the public model's hidden states (43.5 average), suggesting that the external embedder's representation space — which is explicitly trained for semantic similarity — captures domain-relevant structure that next-token prediction representations don't cleanly separate. This is a non-obvious finding with practical implications for any system that needs to route inputs to domain-specific modules without joint training.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper curates a custom training corpus called FLEXMIX (Section 4.1, Figure 5) consisting of a Public Mix (2.37 trillion tokens from DCLM-Baseline, with news and creative writing excluded) plus seven simulated closed datasets: News (158B tokens), Creative Writing (201.9B), Math (20.3B), Code/StarCoder (83B), Academic (58.6B), Educational Text (102.2B), and Reddit (9.9B). These are deliberately designed to represent datasets that are historically closed, previously public but now restricted, or domains with scarce high-quality public data. For evaluation, the paper uses 31 benchmarks across 10 categories (Section 4.2, Appendix C Table 5): general-purpose evaluations include MC9 (9 multiple-choice datasets: ARC-Easy, ARC-Challenge, BoolQ, CSQA, HellaSwag, OpenBookQA, PIQA, SocialIQa, WinoGrande), GEN5 (5 generative QA tasks: CoQA, SQuAD, Natural Questions, TriviaQA, DROP), MMLU, MMLU-Pro, AGIEval, and BBH (23 BIG-Bench tasks); domain-specific evaluations include Math2 (GSM8K + MATH), Code4 (MBPP, MBPP+, HumanEval, HumanEval+), SciRIFF5 (5 scientific literature understanding subtasks), NewsG (news generation evaluated by LM judge), and PoemG (poem generation evaluated by LM judge). The primary experiments in Table 1 use 100 samples per subtask across 24 tasks; the final Table 2 evaluation uses 1,000 samples per subtask across all 31 tasks.
Base model(s). The public model is a dense 7B-parameter transformer following the OLMo 2 architecture (Section 4.4): 32 layers, hidden dimension 4,096, trained on 1 trillion tokens of public data using AdamW (β₁ = 0.9, β₂ = 0.95), cosine learning rate schedule with peak LR 0.0009, global batch size 4M tokens, on 512 H100 GPUs for three days. This choice is motivated by the OLMo family's fully open release and reproducible training recipe, making it a strong open-source baseline. For the scaling experiment (Section 5.4), the paper uses the OLMo-2 7B pre-anneal checkpoint trained on 4T tokens as the starting public model. For the FLOPs-matched unrestricted MoE comparison (Section 5.1), the paper trains a sparse MoE initialized from the public dense model on the combined dataset including all closed sets and Public Mix.
Metrics. The paper reports accuracy (exact match or multiple-choice selection) for all benchmarks except GEN5 tasks (which use F1 score), code tasks (which use Pass@1), and NewsG/PoemG (which use an LM-judge score normalized to 0–100, as detailed in Appendix C). The primary aggregate metric throughout is the arithmetic mean across all evaluated tasks — MC9, GEN5, MMLU, MMLU-Pro, AGIEval, BBH, Math2, NewsG, PoemG, SciRIFF5, and Code4 averaged into a single "Avg." column in Tables 1 and 2. All evaluations follow the OLMES evaluation standard (Appendix C). The MATH benchmark uses exact match with ground truth answers following the methodology from Lewkowycz et al. For NewsG and PoemG, a Llama-3.3-70B-Instruct judge scores completions on domain-specific criteria (journalistic quality, thematic coherence, clarity/fluency), with five completions per prompt averaged to reduce variance.
Baselines. The paper compares against several model merging baselines, all of which train dense models independently by continuing pretraining from the public model on each closed set, then merge them:
- Prompt-based routing: uses Llama-3.1-8B-Instruct or OLMo-2-1124-7B-Instruct as a domain classifier to route each query to exactly one expert model (Section 4.3, Appendix A.1).
- Model soup (average) [12]: uniform weight averaging across all independently trained models.
- Model soup (weighted) [12]: weighted parameter averaging where weights come from softmax-normalized log-likelihoods of each model on the test input.
- Branch-Train-Merge (BTM) [11, 52]: ensembles models by computing weighted average of output probabilities, with weights from softmax-normalized log-likelihoods (top-k variant tested: k=2 for the full 8-expert setup in Table 2).
- BTX [55]: upcycles independently trained dense models into MoE by copying FFN parameters to experts and averaging non-expert parameters. The original BTX requires joint training on combined data after merging; to approximate this while respecting data restrictions, the paper performs post-merge training on public data only (Section 4.3).
- Unrestricted MoE: an upper-bound reference — a sparse MoE initialized from the public dense model and trained on the combined dataset (all closed sets plus Public Mix) with standard joint MoE training. Two variants: compute-controlled (1× FLOPs, 0.5× data, meaning half the data but equal compute to FLEXOLMO) and data-controlled (2× FLOPs, 1× data, meaning all data with twice the compute).
- Individual experts: each domain-specific model evaluated in isolation, trained by continuing pretraining Mpub on a single closed set.
- Public model: the base dense model trained only on Public Mix, evaluated as the lower bound.
Generation budget / compute accounting. Training compute is measured in tokens processed: the public model sees 1T tokens; each expert is trained for 50B tokens on its domain data (400B total across 8 experts in the full setup); optional router training uses 5B tokens (Section 4.4). For the unrestricted MoE comparison, compute is measured in total FLOPs: since MoE training with 2× active experts incurs roughly 2× the FLOPs of dense training for the same data, the paper reports both a FLOPs-controlled comparison (same total FLOPs as FLEXOLMO, processed half the data) and a data-controlled comparison (same data, 2× FLOPs). At inference, the paper reports active parameter count: FLEXOLMO with 8 total experts and k=4 active has 37B total parameters but 20B active. The scaling experiment in Section 5.4 notes that FLEXOLMO with 3 active experts uses 2.5× the inference FLOPs of a dense model like OLMo-2 7B (Table 4).
Cross-validation / statistical protocol. The paper does not employ cross-validation or statistical significance testing for its main results in the conventional sense. The primary experiments evaluate on fixed test sets (the standard test splits of each benchmark). For NewsG and PoemG, the paper generates five completions per prompt and reports the average score to reduce variance. For the restricted MoE baseline training, the paper does not specify multiple random seeds for training runs. The ablations in Table 1 are single-run comparisons with components removed.

Main Quantitative Results

Individual Expert Performance: Specialization at the Cost of Generality

Table 1 and Table 2 present the individual expert baselines alongside the public model. The headline finding is that domain-specific training produces strong specialized performance but severely degrades general capability:

The Math expert (trained on 20.3B math tokens) achieves 50.4 on Math2 and 45.4 on BBH, dramatically outperforming the public model's 8.1 and 35.6 respectively (Table 1). However, it degrades on general benchmarks: 63.8 vs. 68.4 on MC9, 46.3 vs. 58.8 on GEN5, 51.1 vs. 57.0 on MMLU — losses of 4.6, 12.5, and 5.9 points respectively. This is catastrophic forgetting from continued pretraining on domain-specific data.
The Code expert shows the most extreme specialization-regression pattern: it achieves 22.4 on Code4 (vs. public model's 1.0), but collapses to 38.7 on MC9 (vs. 68.4), 41.4 on GEN5 (vs. 58.8), and 30.0 on MMLU (vs. 57.0). The 29.7-point drop on MC9 and 27.0-point drop on MMLU demonstrate that continued pretraining on code alone is destructive to broad language understanding.
The Educational Text expert is the most balanced individual expert, achieving 37.2 average (nearly matching the public model's 36.9) while improving on specific domains like BBH (40.0 vs. 35.6).
In the full eight-expert setup (Table 2), the pattern repeats: the Math expert achieves 53.1 on Math2 and 15.8 on Code4 (vs. public model's 8.2 and 1.1), while individual experts like Creative Writing and Academic collapse to 31.7 and 29.5 averages respectively (vs. 42.4 for public), showing catastrophic forgetting on general benchmarks.

The key takeaway is that naive continued pretraining on domain data creates specialists that are incompatible with general-purpose deployment — they cannot simply be used as standalone experts in a routing system because their general capabilities have eroded. This motivates why FLEXOLMO's coordinated training (which keeps the public expert frozen and trains the domain expert to complement rather than replace it) is essential.

FLEXOLMO vs. Individual Experts and Prior Merging Methods

Table 1 (four-expert setup: public, math, educational text, code) and Table 2 (eight-expert setup: all seven closed sets plus public) present the central comparisons. The headlines:

FLEXOLMO achieves 41% relative improvement over the public-only baseline in the four-expert setup (Section 5.1): from 36.9 average to 47.8 for FLEXOLMO with router training. The largest absolute gains appear on benchmarks where closed data provides the most benefit: Math2 improves from 8.1 to 50.7 (+42.6 points), BBH from 35.6 to 47.1 (+11.5 points), Code4 from 1.0 to 17.3 (+16.3 points). Critically, FLEXOLMO does not sacrifice general performance: MC9 (68.4 → 71.0), GEN5 (58.8 → 59.8), and MMLU (57.0 → 59.9) all improve over the public model, unlike the individual experts.

FLEXOLMO beats the strongest baseline BTM by 10.1% on average (Table 1: 47.8 vs. 43.4). Breaking this down per-category:

BBH: FLEXOLMO 47.1 vs. BTM 43.6 (+3.5 points)
Math2: FLEXOLMO 50.7 vs. BTM 21.2 (+29.5 points) — this is the largest absolute gap, showing that MoE routing is dramatically better than output ensembling for math
Code4: FLEXOLMO 17.3 vs. BTM 22.3 (-5.0 points) — the one category where BTM outperforms, suggesting that code tasks benefit from seeing all experts' outputs rather than sparse activation
MMLU Pro: FLEXOLMO 30.8 vs. BTM 29.0 (+1.8 points)
AGI Eval: FLEXOLMO 45.8 vs. BTM 43.6 (+2.2 points)

The paper attributes FLEXOLMO's advantage over BTM to the MoE architecture's ability to selectively activate different experts per layer (Section 5.1): "This highlights the model's layer-specific specialization and its greater expressivity than approaches that route inputs to a single expert." BTM ensembles at the output level, meaning all experts process every token, which is both computationally expensive (8× inference cost) and potentially dilutes domain-specific knowledge by averaging logits from irrelevant experts.

Compared to prompt-based routing, FLEXOLMO is dramatically more robust. Prompt-based routing achieves 40.8 (OLMo router) and 40.0 (Llama router) — substantially below FLEXOLMO's 47.8 and even below some individual experts. The paper notes this method is "highly unstable: it performs well when the classifier selects the correct expert, but degrades sharply when it does not" (Section 5.1).

Compared to BTX, FLEXOLMO (47.8) substantially outperforms BTX (40.0). The paper attributes this to BTX's reliance on post-merge training with public data only, which is insufficient to teach the router to coordinate experts that were trained without the anchor model concept: "training BTX on the public data only is not optimal" (Section 5.1, footnote).

Model soup (weighted) at 42.2 and model soup (average) at 39.6 both underperform BTM and FLEXOLMO substantially — confirming that weight-space interpolation fails when experts diverge during independent training on disjoint datasets.

FLEXOLMO vs. Unrestricted MoE (Upper Bound)

Table 1 includes the unrestricted MoE reference model (Section 4.3, Section 5.1), which is trained with standard joint MoE training on all data (public + closed sets) — meaning it has access to the combined dataset that FLEXOLMO cannot use jointly. This provides a sense of how much performance is lost due to the data separation constraint:

Compute-controlled comparison: FLEXOLMO (47.8) outperforms the unrestricted MoE trained with equal FLOPs but half the data (46.3). This is a striking result: FLEXOLMO achieves better performance at the same computational cost despite never having joint access to the training data. The unrestricted MoE's score is dragged down by MC9 (68.0 vs. FLEXOLMO's 71.0) and GEN5 (53.8 vs. 59.8).
Data-controlled comparison: Unrestricted MoE with full data and 2× FLOPs achieves 51.5, outperforming FLEXOLMO's 47.8 by 3.7 points. The gap is largest on Code4 (27.0 vs. 17.3) and MMLU (63.1 vs. 59.9). This is the "price of data separation" — having joint access to all data during training provides a clear but not insurmountable advantage.

The paper frames this as evidence that "FLEXOLMO enables training without direct access to the data (requiring only model sharing) and flexible opt-in and opt-out functions while retaining strong performance" (Section 5.1).

FLEXOLMO Ablations: Component Contributions

Table 1 provides a systematic component ablation. The full FLEXOLMO (47.8) represents all components active including optional router training. Removing components sequentially:

Remove optional router training (46.7): -1.1 points. The impact is visible across most benchmarks (GEN5: -1.2, MMLU Pro: -2.4, BBH: -3.7, Code4: +0.9). The BBH drop is notable, suggesting router tuning particularly helps with complex reasoning tasks requiring multiclass routing decisions.
Remove bias term (45.8): -0.9 points from 46.7. Consistent small degradation across most benchmarks, with Math2 showing 51.5 → 50.0 (-1.5).
Remove both domain embedding init and bias (44.4): -3.4 points from 46.7. This is a larger drop, with Math2 falling from 51.5 to 44.9 (-6.6) and Code4 from 18.2 to 16.6 (-1.6). The domain embedding initialization matters most for tasks requiring precise expert specialization.
Remove coordinated training (38.8): -8.9 points from 46.7. This is the single largest ablation effect — a catastrophic 19% relative drop. The impact is most severe on Math2 (51.5 → 19.3, -32.2 points) and Code4 (18.2 → 10.3, -7.9 points). Without coordinated training, experts are just independently continued-pretrained dense models that diverge from each other, and the subsequent MoE merging fails to recover domain specialization. Interestingly, AGI Eval is relatively robust to this ablation (44.8 → 43.1, -1.7), and BBH holds up moderately (43.4 → 41.2, -2.2).

The key conclusion: coordinated training is the single most important component, accounting for roughly 70% of the total improvement over the public baseline (36.9 → 46.7 = +9.8, of which +7.9 is attributable to coordinated training). The domain embedding initialization and bias term provide meaningful but secondary gains. Optional router training provides a modest final boost.

Full Eight-Expert FLEXOLMO

Table 2 scales the approach to all seven simulated closed sets plus the public mix. The headlines:

FLEXOLMO achieves 52.4 average across 31 tasks, compared to 42.4 for the public-only model (+10.0 points, 24% relative improvement).
FLEXOLMO beats BTM (top-2) by 10% relatively (52.4 vs. 47.6, +4.8 points).
Domain-specific gains are substantial: Math2 improves from 8.2 (public) to 48.5 (+40.3), BBH from 35.7 to 46.4 (+10.7), SciRIFF5 from 48.1 to 54.3 (+6.2), PoemG from 47.8 to 62.2 (+14.4), NewsG from 76.0 to 80.7 (+4.7).
Code4 performance is actually weaker with router training (17.2) than without (18.6), and weaker than BTM (24.0) — consistent with the Table 1 finding that code tasks benefit less from sparse MoE routing.

Compared to the four-expert setup in Table 1, adding four more experts improves average performance from ~47.8 to 52.4 (the benchmarks differ between tables, so direct comparison is not exact), demonstrating that the architecture scales gracefully as more experts are added.

Scaling FLEXOLMO: Integration with OLMo-2 7B

Section 5.4 and Table 4 evaluate whether FLEXOLMO's recipe can improve an already-strong model. Starting from the OLMo-2 7B pre-anneal checkpoint (4T tokens pretrained, 50B tokens annealed), the paper trains two additional experts (math and code, 50B tokens each) and combines them into a three-expert FLEXOLMO. The comparison is against the released OLMo-2 7B, which was continued from the same 4T-token checkpoint with an equivalent compute budget (3 × 50B-token training runs, but as a single dense model).

FLEXOLMO achieves 52.8 average vs. OLMo-2 7B's 49.8 (+3.0 points), with 2.5× inference FLOPs.
Largest gains on domain-specific tasks: Math2 (51.0 vs. 42.6, +8.4), Code4 (18.9 vs. 13.3, +5.6), BBH (53.1 vs. 49.8, +3.3).
General benchmarks improve modestly: MMLU (65.2 vs. 63.7, +1.5), GEN5 (71.0 vs. 70.2, +0.8), MC9 (77.8 vs. 77.8, tied).
No regression on any benchmark — unlike continued pretraining of a dense model, which would risk catastrophic forgetting, FLEXOLMO's modular design preserves general capabilities while adding domain specialization.

The paper presents this as evidence that expert specialization with selective activation "enhances performance without catastrophic forgetting or forcing diverse capabilities to compete for fixed model capacity" (Section 5.4).

Model Behavior Analysis

Routing Patterns: Domain-Appropriate Expert Activation

Figure 2 visualizes the router's token-level expert distribution for inputs from different domains. The key findings:

Domain experts are preferentially activated for their corresponding inputs: math inputs heavily activate the math expert (visible as a dark band in the math expert column of the heatmap), code inputs activate the code expert, news inputs activate the news expert, etc.
The public expert is frequently activated across all domains, consistent with the coordinated training design — experts learn to complement, not replace, the public model. Even for domain-specific inputs, a substantial fraction of tokens still route to the public expert.
Different layers activate different combinations of experts: the heatmap shows layer-varying routing patterns, not uniform routing. This per-layer specialization is what gives the MoE architecture greater expressivity than methods that route the entire input to a single expert.
The horizontal gray lines in Figure 2 indicate uniform routing — actual routing deviates substantially from uniform, confirming that the router learns meaningful, non-random expert selection.

Number of Active Experts: Performance Plateaus at k=4

Figure 3 analyzes how downstream performance varies with the number of active experts (k in Top-k routing, evaluated on MMLU). Performance improves consistently as k increases from 1 to 4, then plateaus. The paper interprets this as evidence that "the final model can operate efficiently as a sparse model by activating only four experts per input during inference" (Section 5.2). This finding validates the choice of k=4 for the final model and suggests that most of the benefit of having many experts comes from the ability to route to the best few, not from using all of them simultaneously.

Data Opt-Out: Surgical Removal with Minimal Collateral Damage

Figure 4 demonstrates the opt-out mechanism by removing the news expert from FLEXOLMO and evaluating performance. The key result:

NewsG performance drops substantially when the news expert is removed (the paper does not give exact numbers in the text, but the bar chart in Figure 4 shows a clear reduction for NewsG).
Other tasks are minimally affected: MMLU, MC9, GEN5, BBH, and other benchmarks show bars of nearly identical height with and without the news expert, confirming that each expert's influence is largely localized to its domain.

This validates the paper's central claim that FLEXOLMO provides "strict opt-out guarantees" — removing an expert surgically removes its associated data's influence without degrading unrelated capabilities. This is a hard architectural guarantee, not a best-effort unlearning approximation.

Ablation Studies and Robustness Checks

Router embedding initialization method: Table 3 compares GRIT embedder initialization (the default, 46.7 average) against using the public model's own hidden states as initialization (43.5 average, -3.2 points). The GRIT advantage is visible across most benchmarks: MC9 (71.1 vs. 70.5, +0.6), GEN5 (58.6 vs. 55.5, +3.1), AGI Eval (44.8 vs. 40.2, +4.6), BBH (43.4 vs. 40.1, +3.3), Math2 (51.5 vs. 48.4, +3.1), Code4 (18.2 vs. 8.7, +9.5). The Code4 difference is particularly dramatic — a 9.5-point gap — suggesting that the external embedder's semantic similarity space is especially important for routing code-related inputs correctly. This is a non-obvious finding: one might expect the public model's own representations to be better suited for routing since they share the same architecture, but the contrastively-trained embedder captures domain-relevant structure more cleanly.
Removing coordinated training: As discussed above, this produces the largest single ablation effect in Table 1 ("no training to coordinate": 38.8 vs. 46.7, -8.9 points). This experiment trains experts by simply continuing pretraining each expert independently on its domain data (the naive approach), then merges them using the same domain-informed router procedure. The dramatic collapse — particularly on Math2 (51.5 → 19.3) — confirms that without the anchor model to prevent divergence, experts become incompatible and the MoE merging cannot recover their specialized capabilities.
Removing the bias term: "no bias" in Table 1: 45.8 vs. 46.7 (-0.9 points). The degradation is consistent but modest across most benchmarks (MC9: -3.2, GEN5: -3.0, Math2: -1.5, Code4: -0.6). This confirms the Appendix D intuition that the bias tightens decision boundaries but is not critical to basic functionality — the pairwise training with a frozen reference already produces reasonably well-calibrated router scores.
Removing both domain embedding init and bias: "no domain embedding init, no bias" in Table 1: 44.4 vs. 46.7 (-2.3 points). This is a larger drop than removing bias alone, confirming that the domain embedding initialization contributes independently to router quality. The effects are additive: bias removal costs ~0.9 points, domain embedding removal costs an additional ~1.4 points.
Proxy router training (vs. none): Comparing "FLEXOLMO" (with router training) to "FLEXOLMO (no optional router training)" in Table 1: 47.8 vs. 46.7 (+1.1 points). The improvement is visible on GEN5 (59.8 vs. 58.6, +1.2), MMLU Pro (30.8 vs. 28.4, +2.4), BBH (47.1 vs. 43.4, +3.7), but Code4 actually degrades slightly (17.3 vs. 18.2, -0.9). The BBH improvement suggests that proxy router training particularly helps with complex reasoning tasks requiring nuanced multiclass routing. The Code4 degradation might indicate that the proxy data for code is not sufficiently representative, causing the router to become slightly miscalibrated for code-specific routing decisions.
Full eight-expert setup: Table 2 confirms that the approach scales from 4 to 8 experts without degradation — FLEXOLMO with router training achieves 52.4 average (vs. 47.6 for BTM top-2). Interestingly, FLEXOLMO without router training (51.3) still outperforms BTM (47.6), confirming that the core coordinated training + domain-informed router approach works at larger scale.
Data extraction analysis (Section 5.3): This is not an ablation but a robustness/safety check. The paper implements training data extraction attacks following Carlini et al. [108] on the math data (smallest closed set, trained for 3 epochs — making it the most extraction-susceptible). For 32-token prefixes from 10,000 math documents, they generate 10 continuations of 256 tokens each using top-k=50, top-p=0.95, temperature 1.0. A document is "extracted" if any of the 10 continuations achieves ≥0.9 normalized Levenshtein similarity with the original. Results:
- Public model (never seen math data): 0.1% extraction rate (attributed to deterministic continuations where the model completes text in a way that happens to match the target — a known phenomenon [111]).
- Dense math expert (trained on math data): 1.6% extraction rate.
- FLEXOLMO with math expert included: 0.7% extraction rate.
- Overfitted model (100 epochs on small data, as a validation check): 60% extraction rate (confirming the extraction method works).
The paper draws two conclusions: (1) it is difficult to extract a substantial portion of training data in practice, aligning with prior findings; (2) FLEXOLMO's lower extraction rate (0.7%) compared to a dense expert (1.6%) likely reflects that the math expert within FLEXOLMO is only activated on a subset of tokens, reducing the surface area for extraction. The paper explicitly recommends DP training for experts containing sensitive data, noting it is orthogonal to the architecture.

Critical Assessment

Do the Experiments Support the Central Claims?

Claim: "FLEXOLMO achieves an average 41% relative improvement over the public-only baseline."

This claim is supported by Table 1, but the reader should understand what "41% relative improvement" means precisely. The public model achieves 36.9 average across the 24-task subset; FLEXOLMO achieves 47.8. The relative improvement is (47.8 - 36.9) / 36.9 ≈ 29.5%, not 41%. The paper's stated 41% figure appears either to use a different baseline (perhaps comparing to a weaker public model variant) or to be computed over a subset of the benchmarks. The text states "an average 41% relative gain" in Section 5.1 and the abstract, but the numbers in Table 1 don't straightforwardly produce 41%. This is a discrepancy the paper does not explain. A possible reconciliation: the paper might be computing the relative gain only on benchmarks where the closed data provides the largest improvements (Math2: 8.1 → 50.7 is a 526% gain, Code4: 1.0 → 17.3 is a 1,630% gain), and averaging those individual relative gains rather than computing the relative gain of the average. This would be a misleading aggregation because relative gains on near-zero baselines inflate the average. The 41% figure should be treated with caution — the absolute improvement of 10.9 points (from 36.9 to 47.8) is more meaningful.

Claim: "FLEXOLMO outperforms prior model merging methods by 10.1% on average."

This claim is comparing FLEXOLMO at 47.8 against BTM at 43.4 in Table 1: (47.8 - 43.4) / 43.4 ≈ 10.1%. This is correct arithmetically, but the comparison has several caveats:

BTM uses top-2 in the 8-expert setup but the paper doesn't specify the k for the 4-expert Table 1. If BTM uses all experts (k=4), it's running 4 models in parallel for every token — 4× the inference cost of FLEXOLMO's sparse activation. A fairer comparison might control for inference FLOPs, not just model quality.
BTX's post-merge training uses public data only, which the paper acknowledges is "not optimal." This makes BTX a weakened baseline compared to the original BTX paper, which trained on the combined data. The comparison is valid for the data-restricted setting, but it's comparing against a deliberately handicapped BTX.
The baselines are all reimplemented, not taken from prior work's released models. There is no guarantee that these reimplementations achieve optimal performance for their respective methods. Hyperparameter tuning details for baselines are sparse.

Claim: "FLEXOLMO matches the performance of an unrestricted MoE trained with the same compute budget."

Table 1 shows FLEXOLMO at 47.8 vs. unrestricted MoE (1× FLOPs, 0.5× data) at 46.3 — FLEXOLMO actually outperforms this baseline by 1.5 points. However:

The unrestricted MoE sees only half the data (0.5×) to match FLOPs, since MoE training costs ~2× per token. This means it's trained on less data than FLEXOLMO's experts (which each see 50B tokens of their domain data, 400B total). The comparison is FLOPs-equitable but not data-equitable — FLEXOLMO has an advantage in total training tokens.
The unrestricted MoE with full data (1×) but 2× FLOPs achieves 51.5, outperforming FLEXOLMO by 3.7 points. This is the "cost of data separation" — having joint data access enables better performance at higher compute cost.
Neither unrestricted MoE variant is comprehensively hyperparameter-tuned. The paper doesn't specify how many configurations were tried for the unrestricted MoE baseline.

Claim: "Independently trained domain experts can be effectively combined without joint data access."

This is the central claim, and the evidence is strong but conditional:

What was demonstrated: On FLEXMIX (a curated mixture with 7 simulated closed domains) using OLMo 2 7B architecture, FLEXOLMO combines independently trained experts into a single MoE that outperforms the public model and all baselines, with opt-out capabilities verified.
What was not demonstrated:
- Generalization to genuinely closed/private data distributions (FLEXMIX datasets are constructed from public sources — they simulate "closed" data but are not actually private or proprietary). Real closed data might have different statistical properties that affect expert specialization or router behavior.
- Behavior with a much larger number of experts (e.g., 100+ domains). Figure 3 shows k=4 plateau, but it's unclear how routing degrades with many more experts competing for the same k slots.
- The approach with different base model architectures or scales. All experiments use 7B OLMo 2 models. A 1B or 70B model might show different divergence patterns during coordinated training.
- The approach on non-English data or multimodal data, where the domain embedding initialization and coordinated training dynamics might differ.

Genuine Weaknesses

The 41% relative improvement figure is opaque and potentially misleading. As noted above, the arithmetic in Table 1 does not straightforwardly produce 41%. The paper should clarify what this number represents and over which benchmarks it is computed. The abstract and introduction prominently feature this number, but it cannot be verified from the reported results.

Difficulty estimation for router training uses closed data to train a classifier. The proxy data construction (Section 3.3.3, Appendix A.2) requires each data owner to train a binary classifier distinguishing their closed data from public data, using 250,000 documents from each source. This classifier training uses the closed data $D_i$ — which is fine from a data-sharing perspective (the classifier is trained locally), but it means the proxy data selection pipeline is not purely public. If a data owner cannot train even a small classifier on their closed data (e.g., due to compute constraints or because they're contributing pre-trained weights from a third party), the optional router tuning step becomes unavailable. The paper does not discuss this dependency.

No evaluation on data contamination or benchmark leakage. The closed datasets (News, Creative Writing, Academic, Educational Text, Reddit) are derived from web-crawled or publicly available sources. It's possible that some evaluation benchmarks overlap with these training sets, inflating domain expert performance artificially. Standard decontamination analyses (checking for benchmark questions in training data) are not reported. This is particularly relevant for MMLU, which draws from academic and educational sources that might overlap with the Academic and Educational Text training data.

Training cost asymmetry between FLEXOLMO and baselines. FLEXOLMO trains each expert for 50B tokens in a two-expert MoE configuration. The baseline individual experts train a dense model on 50B tokens. These have different computational costs: the two-expert MoE processes each token through both the frozen public FFN and the trainable domain FFN (plus the full dense attention layers), making it roughly 2× the FLOPs per token of the dense baseline. If the baselines were given equivalent FLOPs (e.g., training for 100B tokens instead of 50B), they might perform better. The paper's FLOPs-controlled comparison is only with the unrestricted MoE, not with the individual expert or BTM baselines.

The evaluation uses different sample sizes per benchmark. Table 1 uses 100 samples per subtask; Table 2 uses 1,000 samples per subtask. The benchmarks are not the same across tables (Table 2 adds NewsG, PoemG, SciRIFF5, and uses a slightly different task set). While this is methodologically acceptable, it means direct comparison between the four-expert and eight-expert results requires caution — the different sample sizes introduce different levels of statistical noise.

The router training proxy requirement limits applicability. The optional router training step assumes each data owner can identify public documents that resemble their closed data. This is feasible for domains like "news" or "code" (where public proxies exist), but might be impossible for genuinely unique closed datasets (e.g., proprietary internal documents that have no public analog). The paper acknowledges the proxy step is optional and shows it provides only modest gains (+1.1 points in Table 1), but doesn't analyze when the proxy approach would fail or how performance degrades when proxies are unrepresentative.

Limited analysis of expert complementarity vs. redundancy. The paper shows that combining experts helps (Table 2), but doesn't analyze which combinations of experts are most beneficial or whether some experts are redundant. For instance, Academic and Educational Text might have substantial overlap — does adding both provide gains beyond adding just one? A leave-one-expert-out analysis would strengthen the claim that each expert contributes uniquely.

What Additional Experiments Would Have Strengthened the Paper?

Ablation on the amount of closed training data per expert. The paper uses 50B tokens per expert uniformly. How does expert quality scale with data quantity? If a data owner only has 5B tokens, does FLEXOLMO still work? The coordinated training framework's dependence on data volume is unexplored.
Varying the public model quality. All experiments use the same public model (7B, 1T tokens). How does FLEXOLMO's performance change with a weaker public model (e.g., 1B parameters or 100B training tokens)? The anchor model concept requires the public model to be reasonably capable on all domains — the paper never probes this boundary.
Direct comparison to federated learning at equivalent scale. The paper critiques FL but never runs an FL baseline on FLEXMIX. While the practical barriers to FL at this scale are real, an empirical comparison (even at smaller scale) would strengthen the claim that FLEXOLMO avoids FL's performance degradation.
Analysis of expert divergence during coordinated training. The paper claims that naive continued pretraining causes divergence but never quantifies this — e.g., by measuring the angular distance between expert FFN weight matrices across training steps with and without the frozen anchor. This would make the "divergence prevents merging" argument more concrete.
Measurements of inference latency, not just FLOPs. The paper measures active parameters (20B) and reports inference FLOP multipliers (2.5×), but doesn't report wall-clock latency. For the opt-out and flexible inference use cases (per-user access control, real-time routing decisions), latency matters as much as throughput. Sparse MoE inference incurs additional overhead from the routing computation and expert dispatch that FLOP counting doesn't capture.
Evaluation on genuinely private or non-public data. The FLEXMIX datasets are all derived from publicly available sources, merely relabeled as closed. Testing on an actually private dataset (e.g., a hospital's internal records, a company's proprietary codebase, with appropriate IRB/legal approvals) would strengthen the claim that the approach works for real restricted data, not just simulated restrictions.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Not Accounted for in the Headline Efficiency Numbers

The assumption or constraint. The paper's compute-optimal scaling framework rests on the ability to estimate each prompt's difficulty before deciding how to allocate the inference budget. The method for doing so — generating 2048 samples per question and computing either ground-truth pass@1 (oracle) or the PRM's average final-answer score (predicted) — is extraordinarily expensive. The paper acknowledges this explicitly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The consequence. In a realistic deployment, the total cost would be difficulty estimation + strategy execution. Since difficulty estimation generates 2048 samples per question and the largest test-time budgets studied are 256–512 generations, the estimation step alone can consume 4–8× more compute than the problem-solving step. The reported 4× efficiency gains over best-of-N (Figures 4 and 8) are computed after difficulty is known, without amortizing the cost of learning it. In practice, this means the framework is currently applicable only in settings where difficulty can be estimated once and amortized over many queries (e.g., benchmarking fixed test sets), not in single-query interactive deployments. The paper frames this as an exploration-exploitation tradeoff worth studying but provides no mechanism for reducing estimation cost.

What evidence exists in the paper. No experiment measures or accounts for difficulty estimation cost in the reported efficiency numbers. Figures 4 and 8 plot accuracy vs. generation budget, where the x-axis counts only the strategy execution generations. The 2048-sample overhead is never included in any budget calculation. The paper does not compare against a baseline that uses the same total compute (including difficulty estimation) — for instance, an "oracle-free" best-of-256 that simply allocates 256 generations uniformly to every question might outperform a FLEXOLMO-style approach that spends 2048 generations on estimation plus 64 on strategy execution.

Mitigation status. The paper explicitly defers this to future work, suggesting "pretraining or finetuning models to directly predict difficulty of a question" (Section 3.2 and Section 8). No such model is developed or evaluated. The predicted difficulty bins (using PRM scores without ground truth) partially address the label problem (no need for correct answers) but not the cost problem (still 2048 generations per question). Until a cheap difficulty estimation method is demonstrated, the 4× figure should be understood as an upper bound on achievable efficiency for amortized settings, not a realized deployment gain for per-query inference.

Hard Problems Remain Fundamentally Unsolved — Test-Time Compute Cannot Create Capability from Nothing

The assumption or constraint. The entire framework assumes that the base model can produce correct solutions at some non-trivial rate — that there are correct answers somewhere in the proposal distribution to find or refine. The paper's difficulty bins are defined by the base model's pass@1 rate, and the hardest bin (bin 5) represents problems where pass@1 is near zero. The paper is transparent about this boundary:

"on the hardest questions... no method makes meaningful progress" (Section 5.2)

The consequence. For problems in difficulty bin 5, no amount of test-time compute — search, revisions, or their compute-optimal combination — produces meaningful improvement. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5%. For the MATH benchmark, these hardest problems represent approximately 20% of the test set (one quintile). For a practitioner, this means FLEXOLMO's approach offers no path forward for genuinely novel or out-of-distribution reasoning that exceeds the base model's training distribution. If an application requires solving problems where the base model's pass@1 is near zero, the recommendation is unambiguous: scale pretraining, not test-time compute.

This limitation also means the FLOPs-matched comparison (Section 7) is strongly tilted against test-time compute on hard problems. At R ≫ 1 with PRM search, hard questions show a −52.9% relative disadvantage compared to the 14× larger model. The paper frames this as a finding rather than a failure, but it establishes a hard ceiling: test-time compute amplifies existing capability but does not create it. For problems outside the base model's reach, pretraining remains the only viable path.

What evidence exists in the paper. The difficulty-bin breakdowns (Figures 3 right, 7 right, 9) provide clear and consistent evidence. Across search methods, revision strategies, and FLOPs-matched comparisons, bin 5 is the consistent failure mode. The paper does not investigate why the base model fails on these problems (e.g., whether it lacks the necessary knowledge, or possesses the knowledge but cannot compose it into correct reasoning chains), which would clarify whether improved prompting or retrieval augmentation could partially address this limitation.

Mitigation status. The paper does not attempt to solve this limitation — it documents it clearly and treats it as a boundary condition. The takeaway box in Section 7 explicitly states that "test-time compute is most effective when problems are within the base model's reach." This is an honest and useful characterization, but it means the approach is complementary to, not a replacement for, continued progress in pretraining.

Verifier Over-Optimization Is a Hard Ceiling That the Compute-Optimal Policy Mitigates but Cannot Eliminate

The constraint. The PRM verifier is learned from Monte Carlo rollouts of the base model, and its quality determines the entire search pipeline's ceiling. When search is pushed aggressively, the PRM's scores become unreliable — the model finds solutions that score highly under the verifier but are actually incorrect. The paper documents this phenomenon in Section 5.2 as the explanation for why beam search degrades on easy problems at high budgets (Figure 3, right) and why lookahead search — the most powerful optimizer — paradoxically performs worst overall (Figure 3, left). Qualitative examples in Appendix M (Figure 29) show degenerate outputs: repetitive low-information steps and overly short 1–2 step solutions that the PRM rates highly.

The consequence. The compute-optimal policy mitigates over-optimization by routing easy problems away from aggressive search (using best-of-N instead of beam search where the verifier is most exploitable), but it does not solve the underlying problem. On medium-difficulty problems where beam search is deployed, over-optimization still limits the scaling ceiling — the beam search curves in Figure 3 flatten and sometimes decline well before the budget is exhausted. This means the approach cannot scale test-time compute arbitrarily: at some budget, the verifier becomes the bottleneck regardless of how cleverly the policy allocates compute. For a practitioner, this implies that investment in verifier quality is the highest-leverage way to improve test-time compute scaling — better search algorithms provide diminishing returns if they merely optimize a flawed verifier more efficiently.

The paper's finding that lookahead search (which uses additional forward rollouts to improve step-level scoring) underperforms at equal generation budgets (Figure 3, left) is particularly revealing: even when the verifier is given more context to make better assessments, the cost of that additional context (k+1 generations per step) reduces the effective number of beams explored, and the net effect is negative. This suggests a difficult tradeoff between verifier accuracy and search breadth that the paper does not resolve.

What evidence exists in the paper. Figure 3 (right) shows beam search accuracy decreasing on bin 1 from ~78% to ~77% as budget increases from 4 to 256, while best-of-N continues improving. Figure 3 (left) shows lookahead search underperforming all methods. Appendix M provides qualitative examples. The paper does not, however, measure verifier calibration directly — e.g., by plotting PRM predicted scores vs. actual correctness rates across search budgets. Such an analysis would quantify the over-optimization threshold and help practitioners determine safe operating budgets for their specific PRM quality.

Mitigation status. The paper identifies over-optimization as a key bottleneck (Section 8) and suggests improving verifier robustness as a research direction — e.g., through adversarial training, ensemble verification, or constrained search with KL penalties. None of these are implemented or evaluated. The compute-optimal policy is the paper's primary mitigation, but it works by avoiding the problem (not deploying aggressive search where over-optimization is worst) rather than solving it. For practitioners, this means the approach works best in the moderate-budget regime (tens to low hundreds of generations) where over-optimization is not yet dominant.

Revisions and PRM Search Are Evaluated Independently — the Framework's Full Potential Is Unexplored

The assumption or constraint. The paper studies two complementary mechanisms — PRM-guided search and iterative revision — as independent strategies and develops compute-optimal policies for each separately. Section 8 explicitly acknowledges:

"we did not experiment with PRM tree-search techniques in combination with revisions"

This is a deliberate scoping choice, not an oversight, but it leaves open the question of whether the two mechanisms have multiplicative or merely additive benefits.

The consequence. The paper's headline results (Figures 4 and 8) represent a lower bound on what a fully integrated system could achieve. Since revisions improve the proposal distribution (generating better candidates from better-conditioned models) while PRM search improves candidate selection (more intelligently navigating the space of possible solutions), their combination could yield gains beyond either method alone. For example, using the revision model as the proposal distribution within beam search — where each search node conditions on previously rejected branches as revision context — could simultaneously improve the quality of explored candidates and the efficiency of exploration. Alternatively, using the PRM to guide which revisions to pursue (rather than blindly generating long revision chains) could reduce the 38% correct-to-incorrect reversion rate (Section 6.1) by stopping revision when the PRM indicates the current answer is already on track.

For a practitioner deciding whether to deploy this approach, the independent evaluation means the reported numbers likely underestimate the achievable performance. However, it also means the approach as described is incomplete — deploying "FLEXOLMO" as a system requires choosing between revisions and search (or deploying both independently), not benefiting from their integration. The paper's compute-optimal policy selects between them per difficulty bin, but never uses both on the same problem.

What evidence exists in the paper. The difficulty-bin analysis provides indirect evidence that the mechanisms are complementary: revisions work best on easy problems (Figure 7, right: purely sequential is optimal for bins 1–2), while beam search works best on medium-hard problems (Figure 3, right: beam search outperforms best-of-N on bins 3–4). This complementary pattern suggests that a combined system could route each problem to the mechanism (or combination) best suited to its difficulty, but no such system is built or evaluated.

Mitigation status. The paper acknowledges this as future work (Section 8). Given the strong results from each mechanism independently, and the clear complementary difficulty-dependent patterns, this is arguably the most promising near-term extension of the work. The architectural components are all in place — what's missing is the integration and evaluation.

The Revision Model Has an Inherent 38% Correct-to-Incorrect Reversion Rate

The constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target (Section 6.1). This training data construction means the model never learns what to do when the current answer is already correct — it has no signal for "stop revising, you're done." Section 6.1 reports:

"approximately 38% of correct answers get converted back to incorrect ones" using a naive approach where the last revision is always taken as final.

The consequence. This reversion rate means the revision model is fundamentally unreliable as a standalone generator — you cannot simply run it for N steps and take the final output. The paper mitigates this with selection mechanisms (majority voting or verifier-based selection across the entire revision chain, picking the best answer from any point), but these are post-hoc patches. They require generating the full chain (including the incorrect reversions) and then running a separate selection process, which wastes compute on revisions that actively degrade quality.

More subtly, the 38% reversion rate interacts poorly with the sequential-to-parallel ratio optimization (Section 6.3). As the chain length increases, the probability that a correct answer appears and is subsequently "revised away" grows, creating a tension between the benefits of additional revision steps (pass@1 improves along the chain, Figure 6 left) and the risk of losing already-found correct answers. The compute-optimal policy selects chain lengths that balance these effects, but the underlying fragility means that small changes in problem difficulty or model calibration could shift the optimal policy substantially.

What evidence exists in the paper. Section 6.1 directly reports the 38% figure. Figure 6 (left) shows that per-step pass@1 continues improving out to 64 steps (well beyond the 4-step training horizon), confirming that the model has learned a generalizable revision skill — but the reversion problem means this improved per-step accuracy does not translate directly to improved final accuracy without careful selection. The ReST^EM experiment (Appendix K, Figure 16) provides further evidence of the revision model's fragility: attempting to optimize it with RL-style training caused performance to degrade substantially with sequential revisions, suggesting the revision skill is sensitive to training methodology in ways not fully characterized.

Mitigation status. The paper's mitigations (majority voting, verifier-based selection) are effective in practice — the reported results achieve strong performance — but they are engineering patches rather than architectural solutions. A more principled approach would train the revision model to recognize when no revision is needed (e.g., by including correct-to-correct trajectories in the training data), but this would require modifying the data generation pipeline and potentially changing the learning dynamics. The paper does not explore this. For a practitioner, the takeaway is that revision-based approaches require a selection mechanism as a mandatory component, not an optional enhancement, and that the training data construction (incorrect-to-correct trajectories only) introduces a structural bias that must be managed at inference time.

Single Benchmark, Single Model Family — Generalizability Is Unproven

The constraint. All experiments in the paper use the MATH benchmark (500 test questions) with PaLM 2-S* as the base model. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is unverified. MATH consists exclusively of competition-level math problems requiring multi-step symbolic reasoning, and PaLM 2-S* has a specific capability profile (roughly 10–19% pass@1 on MATH) that may not generalize to other models or tasks.

The consequence. Several aspects of the findings could be model-specific or benchmark-specific:

The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution — specifically, the types of errors it makes, the diversity of its incorrect solutions, and the separability of correct vs. incorrect solution paths. A model with different calibration properties (e.g., one that is more confident but less accurate) might produce PRM scores that over-optimize at lower search budgets or in different difficulty regimes.
The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. The edit-distance-based pairing strategy (Section 6.1) assumes that structurally similar but incorrect solutions provide useful revision signals — this might not hold for models that make different types of errors.
MATH is a domain with clean, objective correctness signals (exact match on final answers). This enables both the Monte Carlo rollout PRM training (correctness of rollouts can be automatically determined) and the difficulty estimation (pass@1 can be computed from correctness signals). Extending the framework to tasks with ambiguous or multi-dimensional correctness (open-ended generation, dialogue, creative writing, multi-step planning) would require fundamentally different verifier training and difficulty estimation approaches. The paper does not address this.

What evidence exists in the paper. The paper provides no experiments on other benchmarks, other model families, or other tasks. The 500-question test set size is acknowledged but the dependence on a single benchmark and model is not discussed as a limitation in Section 8. This is a significant gap given the paper's broad claims about test-time compute scaling laws.

Mitigation status. Not addressed. The paper does not discuss generalizability to other benchmarks, model families, or task types. The authors' belief that PaLM 2-S* is "representative" is stated without evidence. For a practitioner considering deploying this approach, the most important unknown is whether the difficulty-dependent optimal policies (beam search for medium problems, revisions for easy problems) transfer to their specific model and task distribution. The framework for computing these policies (difficulty estimation, strategy sweep, cross-validation) is general, but the specific policies discovered on MATH with PaLM 2-S* may not be. Replicating the analysis on a new model/task combination would require substantial computational resources (training a PRM, training a revision model, sweeping strategies across difficulty bins).

7. Implications and Future Directions

How This Work Changes the Landscape

This paper reframes the data restriction problem in language model training from a privacy-and-access-control challenge into a modular architecture design problem, a shift that carries consequences for how the field thinks about collaborative model development, data governance, and the relationship between training methodology and architectural design.

A New Axis of Model Scalability: Modularity as a First-Class Design Constraint

The dominant narrative in LLM development treats model scaling as a problem of parameter count, data volume, and compute budget — all optimized under the assumption of centralized, joint training on pooled data. FLEXOLMO introduces a fourth axis: modular composability under data separation. This matters because it changes what counts as a "valid" training setup. Before this work, the choice was binary: either you have joint access to all training data (and can train a standard MoE or dense model) or you don't (and you're limited to federated learning with its synchronization costs, or to post-hoc model merging with its performance degradation). FLEXOLMO creates a third option — asynchronous, independent expert training with no joint data access at any point — that achieves performance competitive with jointly-trained models at equivalent compute budgets (Table 1: FLEXOLMO at 47.8 vs. unrestricted MoE at 46.3 under FLOPs control).

The magnitude is not paradigm-shifting in the sense of "everything we knew about MoE training was wrong" — the architecture is still a standard Top-k MoE, the training objective is still next-token prediction, and the optimization recipe is standard. What shifts is the assumption about data access during training. By demonstrating that the router can be constructed from domain embeddings learned in pairwise binary training rather than joint multiclass optimization, the paper removes what was widely assumed to be a hard requirement for MoE composition. The ablation showing that removing coordinated training causes performance to collapse from 46.7 to 38.8 (Table 1, "no training to coordinate") provides the diagnostic: naive independent training doesn't work, but the anchor-based approach does — and the gap between them exposes a design principle rather than an incremental trick.

This reframing makes modular, contribution-based LM development a legitimate research direction rather than a speculative vision. The open-source software analogy the paper draws from [40] is not merely rhetorical — by showing that independent expert training can work at the scale of 7B-parameter models on realistic data volumes (50B tokens per expert), the paper provides an existence proof that the metaphor has engineering substance. Data owners can train components, contribute them to a shared model, and later remove them — all without exposing their data to any other party.

Reconciliation of Tensions in Prior Model Merging Work

The paper resolves a tension that has existed in the model merging literature between expressivity and data privacy. Prior merging methods fell into two camps: those that achieve strong performance but require joint data access (BTX, which needs post-merge training on the combined dataset) and those that don't require joint data but underperform substantially (model soup, BTM, which average or ensemble independently trained models without learning cross-expert interactions). The paper's results quantify this tension numerically: BTX (which violates data separation by using public data for post-merge training) achieves only 40.0 on the four-expert setup, while BTM (which respects data separation) achieves 43.4 — but both are substantially below FLEXOLMO's 47.8 (Table 1). FLEXOLMO breaks the tradeoff: it achieves better performance than the data-dependent methods while maintaining strict data separation, by moving the coordination from post-merge joint optimization into the independent training phase itself (via the frozen public anchor).

This also explains why the prompt-based routing baselines perform so poorly (40.8 and 40.0 in Table 1). Those baselines represent the simplest form of modularity — train separate models, classify each input, route to exactly one — and their failure demonstrates that token-level, per-layer routing across multiple experts is not a luxury but a necessity for combining domain-specific knowledge with general capability. The Figure 2 routing analysis, showing that inputs activate domain experts primarily but also frequently activate the public expert and sometimes other domain experts, provides mechanistic evidence: real-world inputs don't fall cleanly into single domains, and forcing hard routing decisions discards useful cross-domain knowledge.

Redirection of Research Attention: From Joint Training to Anchor-Based Coordination

The paper's most actionable impact on the research landscape may be to redirect attention away from post-hoc merging tricks (better weight interpolation, smarter ensembling weights) and toward training-time coordination mechanisms. The ablation results in Table 1 are instructive: the largest single effect (8.9 points, from 46.7 to 38.8) comes from removing coordinated training; the second-largest (3.4 points, from 46.7 to 43.3) comes from removing domain embedding initialization; the bias term and router tuning contribute smaller, incremental gains. The clear message is that how you train the experts matters far more than how you merge them — and specifically, whether you train them to coordinate with a shared reference. This suggests that future work on modular LM training should invest heavily in coordination mechanisms during independent training (the anchor model concept, auxiliary losses that encourage complementarity, representation alignment techniques) rather than focusing primarily on sophisticated post-hoc integration methods.

The paper also redirects attention toward router design as a first-class research problem in modular architectures. Standard MoE research treats the router as a jointly-trained component whose design is driven by load balancing and training stability concerns. FLEXOLMO shows that the router can be meaningfully structured — initialized from domain embeddings, trained in pairs against a shared reference, calibrated with a negative bias — and that these design choices have substantial performance implications (Table 3: GRIT embedder initialization outperforms public model hidden states by 9.5 points on Code4 and 4.6 points on AGI Eval). This opens a design space that the standard MoE literature largely ignores: what is the optimal router initialization, training curriculum, and calibration strategy for the modular, data-separated setting?

A New Vocabulary for Data Governance in ML

Beyond the technical contributions, the paper provides a conceptual vocabulary — the separation of sharing constraints from usage constraints, the notion of architectural opt-out guarantees, the framing of modular contribution as distinct from privacy-preserving computation — that equips practitioners and policymakers to reason about data governance in LM development more precisely. The data extraction analysis in Section 5.3 demonstrates that FLEXOLMO does not solve the formal privacy problem (0.7% extraction rate for the math expert within FLEXOLMO), and the paper's explicit recommendation to use DP for sensitive data establishes a clean boundary: architectural modularity addresses data access and usage governance, while differential privacy addresses leakage from weight inspection. This separation of concerns — governance vs. privacy — is useful for organizations navigating regulatory landscapes where compliance requires both data localization (FLEXOLMO's strength) and formal privacy guarantees (DP's domain, orthogonal to the architecture).

Follow-Up Research This Work Enables

Cheap domain identification for router initialization without per-data-owner computation. The current approach requires each data owner to sample 1,000 documents from their closed data, process them through an external embedder (GRIT), and average the embeddings. This is computationally light but still requires the data owner to have the embedder and to run it locally. A natural extension is whether the public model itself — perhaps with a lightweight probing classifier trained on its own hidden states — can predict useful router initialization vectors for unseen domains, eliminating the need for data owners to run any initialization step. The specific experiment: for each of the seven FLEXMIX closed domains, train a logistic regression probe on the public model's final-layer hidden states (averaged over documents) to predict domain membership, using only the proxy public data (documents from Dpub that resemble each domain). Compare the resulting router initialization quality against the GRIT initialization. Success would mean new data owners need only provide a domain description (a few dozen example documents or even a textual description) rather than running an external embedding pipeline. This would lower the barrier to contribution and remove the dependency on a specific external embedder.

Training experts to coordinate with each other, not just with the public anchor. The paper's coordinated training uses only binary expert-vs-public routing during training. At inference, experts compete against each other, but they were never trained to make these distinctions. The negative bias term partially addresses this, but it's a heuristic — experts aren't learning to differentiate "math that the code expert might handle" from "math that only I can handle." An extension would introduce a multi-anchor training scheme where, during expert training, the model has access not only to the frozen public expert but also to frozen, pre-trained embeddings of other experts' router vectors (computed after those experts are trained). The training objective would include an auxiliary loss encouraging the domain expert's router to produce low scores on inputs that are well-handled by other experts (as judged by those experts' frozen router embeddings). The concrete experiment: train the math expert with access to the code expert's router embedding, penalizing the math router for producing high scores on code-like inputs, and vice versa. Measure whether this reduces interference between experts at inference time (e.g., whether the math expert is less likely to be spuriously activated on code inputs). A negative result — that multi-anchor training doesn't help beyond the simple binary coordination — would confirm that pairwise binary training with a single anchor suffices, which would simplify future systems.

Training data extraction as a function of expert activation frequency. The paper's data extraction analysis (Section 5.3) shows that FLEXOLMO's math expert has a 0.7% extraction rate vs. 1.6% for a dense math expert, and hypothesizes this is because the FLEXOLMO expert is only activated on a subset of tokens. This hypothesis is testable: measure extraction rates for the same expert when it is activated on different fractions of tokens (controlled by varying the Top-k parameter or by adding a load-balancing loss that forces more uniform activation). The strong hypothesis is that extraction rate scales linearly with activation frequency — if the expert is activated on X% of tokens, extraction rate is roughly X% of the dense expert's rate. If this holds, it provides a design knob: contributors who are concerned about extraction can request a lower activation ceiling for their expert (e.g., k=1 instead of k=4 in the layers where their expert is placed) in exchange for reduced extraction risk, with a predictable accuracy tradeoff. If the relationship is sublinear or thresholded, it would reveal more complex dynamics in how MoE sparsity interacts with training data memorization.

Genuinely private data evaluation. All FLEXMIX datasets are derived from publicly available sources that are merely labeled as "closed" — they are not actually private, proprietary, or sensitive. The strongest test of FLEXOLMO's applicability to real data restriction scenarios would be an evaluation where one or more experts are trained on genuinely non-public data (e.g., a hospital's internal clinical notes for a medical QA task, a law firm's internal briefs for a legal reasoning task, with appropriate IRB and data use agreements). The evaluation would measure: (a) whether domain expert specialization is comparable in magnitude to what was observed on FLEXMIX (i.e., does the anchor-based coordination work as well when the domain data is genuinely out-of-distribution relative to public web text?), (b) whether the extraction rate is higher for genuinely private data (which might be more unique and thus more memorizable than public math data), and (c) whether the proxy router training step (which requires finding public documents that resemble the closed data) remains feasible when the closed data has no public analog. A negative finding — that FLEXOLMO's gains are smaller on genuinely private data — would suggest that the simulated-closed setup overestimates real-world applicability and that additional coordination mechanisms are needed for out-of-distribution domain data.

Scalability to many experts with overlapping domains. The paper's eight-expert setup has mostly non-overlapping domains (math, code, news, creative writing, academic, educational, Reddit). A stress test would train 20–50 experts on partially overlapping datasets — e.g., 10 different code datasets (different languages, different repositories), 10 different medical datasets (different specialties) — and measure whether the domain-informed router can make fine-grained distinctions between similar experts, or whether performance plateaus (or degrades) due to expert interference. Figure 3 shows that activation plateaus at k=4 for eight experts, but this says nothing about whether 4 experts suffice when there are 50 to choose from. A specific experiment: train experts on Python code, JavaScript code, C++ code, and general code, then measure on Python-specific benchmarks whether the router activates the Python expert more than the general code expert, and whether adding a second Python expert (trained on a different Python dataset) provides additional gains or merely splits the routing probability. This would characterize the granularity at which expert specialization provides benefit and inform how data owners should partition their data when contributing experts.

Router training without any proxy data from closed domains. The optional router tuning step (Section 3.3.3) assumes each data owner can identify public documents that resemble their closed data, which requires training a classifier on the closed data locally. What if this is not possible — e.g., the data owner is contributing pre-trained weights from a third party and has never seen the training data? An alternative approach would use router tuning on synthetic data: the data owner provides a textual description of their domain (e.g., "this expert was trained on internal legal documents from a US law firm specializing in patent law"), and a separate LLM generates synthetic training documents matching that description, which are then used for router tuning. The experiment would compare this synthetic-data router tuning against the proxy-data approach and against no router tuning, measuring whether synthetic data can partially recover the multiclass routing signal that proxy data provides. A positive result would further reduce the requirements on data owners and move toward a setup where contribution requires only a domain description, not any processing of the actual training data.

Practical Applications and Downstream Use Cases

Multi-institutional scientific LM development with data localization. The most direct application is for consortia of research institutions — hospitals, universities, government agencies — that want to collaboratively build a domain-specific LM but cannot legally pool their data. For example, a network of 10 hospitals wants to build a clinical language model trained on their combined patient records (for clinical note summarization, diagnosis coding, treatment recommendation), but HIPAA prevents any single institution from accessing the others' data. Each hospital trains a FLEXOLMO expert on their local records (50B tokens of clinical text), using a shared public medical LM (trained on open-access PubMed articles and medical textbooks) as the anchor. The experts are merged into a unified clinical MoE. The key benefit: the combined model achieves the paper's demonstrated ~41% relative improvement over the public baseline on domain-specific tasks (extrapolating from Table 1's domain benchmark gains), without any hospital ever sharing patient data. The opt-out mechanism (Figure 4) provides additional compliance value: if a hospital withdraws from the consortium or a patient revokes data consent, removing that hospital's expert surgically removes its data's influence with minimal impact on other tasks.

Commercial LM deployment with per-customer data activation. A company building a code completion service (similar to GitHub Copilot) serves multiple enterprise customers, each with their own proprietary codebase. Each customer wants the model to use their codebase for completions, but only for their own developers — a competitor's developers should not benefit from their private code patterns. The standard approach (training a single model on all customers' code) fails because the model cannot enforce per-user access control at inference time. With FLEXOLMO, each customer independently trains an expert on their codebase (50B tokens, using a shared public code model as anchor) and contributes it. At inference, the service activates only the experts corresponding to the requesting developer's authorized codebases. The key benefit: the system provides architectural opt-out guarantees — if a customer terminates their contract, removing their expert guarantees their codebase's influence is eliminated, with no unlearning required. The paper's Code4 results (Table 1: FLEXOLMO achieves 17.3 on coding benchmarks, up from the public model's 1.0) suggest substantial domain gains, though the BTM baseline's Code4 performance (22.3, higher than FLEXOLMO) suggests code may be a domain where dense ensembling has advantages worth investigating for this use case specifically.

Regulated industries with expiring data usage rights. Organizations in finance, law, and other regulated sectors often license training data under time-limited contracts — they can use a dataset for model training for, say, 24 months, after which the data must be removed. Standard pretraining provides no mechanism for this: the data's influence is diffused throughout the model weights. FLEXOLMO provides a direct solution: train the licensed data as a separate expert, contribute it to the MoE, and when the license expires, delete that expert's weights. The remaining model is architecturally guaranteed to be free of the expired data's influence. The paper's Figure 4 demonstrates that expert removal has minimal impact on unrelated tasks, and the extraction analysis (Section 5.3, 0.7% rate for the math expert within FLEXOLMO) quantifies the residual risk from weight inspection. For organizations where data licensing compliance is audited, the ability to point to a specific set of deleted parameters and prove they were only trained on the expired dataset provides a level of auditability that unlearning approaches cannot match. The practical recommendation from the paper is to combine this with differential privacy for experts trained on data containing personally identifiable information, providing both governance (architectural opt-out) and privacy (DP guarantees against weight inspection).