FlexOLMo: Open Language Models for Flexible Data Use

We thank the reviewer for their time and constructive feedback. We appreciate that you like the core idea of FlexOlmo. We address each of your specific comments below.

The training procedure of the expert model and route embedding are not clearly explained, particularly regarding how their parameters are jointly updated across the dataset, in Section 3.3.1 and 3.3.2.

The training process of Flexolmo is as follows (as shown in Figure 1):

1. First, a public model is trained on data without any restrictions.
2. Then, each data owner trains their local expert alongside a frozen copy of the public model. This public model acts as a shared "anchor," ensuring that all independently trained experts learn to operate within a common functional space. During this local training stage, the public model and its public router embedding remain frozen. Only the local expert's parameters (FFN layers) and its corresponding router embedding are updated.
3. Finally, in the merging step, we combine the frozen public expert and all the independently trained local experts by simply concatenating their learned router embeddings into a joint table.

We thank the reviewer for this feedback and will make the explanation more clear in the final version of the paper.

For Section 3.3.3, the distribution of private data may also be privacy-sensitive.

We agree with the reviewer that the selection of public data based on the private distribution may carry the risk of indirectly revealing information about it. We clarify in the paper that this is an optional step a data owner would only undertake if they are comfortable with this risk.

I'm not fully convinced how FlexOlmo can handle restrictions including copyright restrictions. In certain scenarios, take toxicity classification as an example, it would be helpful if the authors could provide a more detailed explanation of how FlexOlmo addresses this challenge and why it is preferable to conventional methods, such as directly fine-tuning on datasets that contain toxic content.

Training a single FlexOlmo model allows for deploying different combinations of experts (and their associated data) to different users based on specific access criteria such as permissions, policies, or preferences. Here are two concrete examples:

• Copyrighted Data: Suppose two newspapers contribute experts trained on their proprietary data. They could restrict access so that only users with an active subscription to their respective papers have the option to activate the associated expert, or activate both experts if the user subscribes to both, while others cannot.
• Content Moderation: The FlexOlmo architecture allows for the on-demand deployment of specialized experts, such as one for toxicity detection. This expert can be activated specifically for sensitive applications like toxicity detection, while remaining inactive in other use cases where child-safety is important In summary, we train a single FlexOlmo model, but it can function as if it were trained on many different data combinations, customizing its usage to match specific deployment requirements and constraints.

Lack of domain-specific evaluation. Specifically, there are no corresponding evaluations for Reddit and Textbook domains.

Thank you for this suggestion. We include many domain-specific evaluations as detailed in Section 4.2. For datasets like Textbook, which consists of broad texts from STEM domains, its benefits are reflected in improved performance on STEM-related benchmarks like BBH, Math, and Code. Similarly, some datasets like Reddit are deemed to be generally helpful for improving the model's world knowledge and conversational abilities and are not limited to specific domain evaluations.

I notice that the training of FlexOlmo requires substantial computational resources in Section 4.4. In practical decentralized training scenarios, however, data owners may not have access to such extensive resources for continual training. This limitation should be discussed.

The substantial computational resources mentioned in Section 4.4 (512 H100s for 3 days) are for training the base public model on 1T tokens. Training the individual local experts requires significantly less compute (e.g., 50B tokens for one expert). Our framework is applicable to cross-silo federated learning, where clients are typically organizations with access to reasonable compute resources. Furthermore, as we will discuss in the limitations section, future work can explore extending our model to accommodate clients with varying compute budgets by allowing them to train different-sized local experts.

From experiments, the majority of FlexOlmo's performance gains appear to come from "Math2". If the math domain is excluded, does FlexOlmo still outperform BTM and BTX?

Yes, it does. After excluding the Math2 results, FlexOlmo's average score is still better than the baselines. The average scores are: FlexOlmo (47.4), BTM (46.5), and BTX (43.3).

What is the motivation behind processing and filtering the original Reddit dataset? Why not directly use the same version adopted in previous studies?

In the case of the Reddit data, we drew on data from a parallel line of work, which was motivated by the observation that prior versions of Reddit data showed limited performance gains on downstream tasks, particularly on knowledge-based QA tasks. The hypothesis of that work was that Reddit's valuable, specialized knowledge could be better leveraged by more targeted filtering and data refinement. The process to arrive at the final data involved (1) constructing thread contexts inspired by QA structure, (2) filtering for high-quality, academically relevant subreddits, and (3) rewriting the content to reduce noise and match standard QA formats. The effectiveness of this approach was validated experimentally: when continuing pretraining on OLMo 7B, the rewritten Reddit data yielded substantial improvements over simply using the filtered data – MMLU performance improved from 0.62 to 0.66, and MC9 improved from 0.74 to 0.76. This Reddit data is released to facilitate future research.

What is the rationale for selecting these eight domains? Are there any references or prior work supporting the inclusion of domains such as "Creative Writing" or "Reddit"? Furthermore, is there any semantic relationship or overlap among these domains? The core contribution and motivation of the constructed dataset FlexMix is unclear. And I suggest the authors explicitly list their main contributions in the Introduction section.

The seven domains in FlexMix were chosen to represent common, realistic scenarios where data access is restricted. For example, "Creative Writing" serves as a proxy for copyrighted materials like fiction and books, and "News" data represents copyrighted news data, which correspond to real-world lawsuits where major publishers like The New York Times [2] have recently restricted the use of their raw news corpus for pretraining LMs and Anthropic [3] face legal charges for using pirated books for pretraining. "Reddit" in FlexMix represents online discussion forums, which currently face tangible access restrictions; notably, this data is not downloadable following Reddit's 2023 access policy change [1]. These domains exemplify cases where data is not openly available due to its proprietary nature or user privacy concerns. While semantic overlap can exist between domains (e.g., textbooks containing math text), this complexity reflects real-world data.

FlexMix provides a valuable testbed for advancing research in privacy-conscious and data-flexible model pretraining. On this benchmark, we show that FlexOlmo can effectively leverage these restricted sets to produce a better LM. Finally, thank you for the suggestion regarding the contribution list; we will explicitly list our main contributions in the Introduction section of the final paper.

We hope our responses have addressed the reviewer's concerns. Please let us know if you may have any questions.

[1] https://the-decoder.com/reddit-ends-its-role-as-a-free-ai-training-data-goldmine/

[2] https://www.popsci.com/technology/nyt-generative-ai/

[3] https://fortune.com/2025/07/28/a-copyright-lawsuit-over-pirated-books-could-result-in-business-ending-damages-for-anthropic/

Thank you for your feedback! We will ensure all the revisions are incorporated into the final paper.

Regarding weakness 1, here is our clarification:

For the 2nd step of local training, FlexOlmo uses a Mixture-of-Experts model with two experts: a trainable local expert ($M_i$) and a frozen public expert ($M_{pub}$). Given a loss function $\mathcal{L}$ calculated on the local dataset $D_i$, the parameter updates via gradient descent are as follows:

Trainable Parameters

The local expert's parameters ($\theta_i$) and its router embedding ($r_i$) are updated.

$\theta_i \leftarrow \theta_i - \eta \frac{\partial \mathcal{L}}{\partial \theta_i}$

$r_i \leftarrow r_i - \eta \frac{\partial \mathcal{L}}{\partial r_i}$

Frozen Parameters

The public expert's parameters ($\theta_{pub}$), its router embedding ($r_{pub}$), and shared attention layers remain unchanged. The public expert and its router embedding are shared across different local training, serving as a common anchor that ensures all local experts learn to coordinate with the same reference model.

We are happy to answer any further questions the reviewer may have and respectfully ask that our revisions be considered in your final evaluation.

The paper does not report error bars for experimental results due to "prohibitive expense". While a large suite of benchmarks is used to mitigate this, the absence of statistical significance measures is a noted limitation.

We thank the reviewer for noting this. The costs associated with training the models unfortunately make it impossible to perform multiple runs for each experiment, which is required for reporting error bars. We acknowledge this is a limitation and will include this discussion in the paper.

While compute resources are reported, a formal analysis of how FLEXOLMO's computational complexity scales with the number of models (N) and objectives (K) is not provided, which could be useful for practitioners. How inference-time FLOPs/latency scale as the total number of experts (n) increases and as the number of active experts (k) varies (e.g., does activating only 4 out of 8 experts significantly reduce compute compared to a dense model of equivalent capacity?). Clarify how the overall training cost scales when adding new experts? While experts are trained independently, there might be implications for the shared anchor or router components.

• Training Cost: The cost of adding a new expert is independent of the total number of existing experts and only depends on the local expert size. Our model can scale efficiently in training as there is no joint training of all experts.

• Inference Cost: Similar to standard MoE, Inference costs are dependent on the number of active experts (k), not the total number of experts in the model. This means that activating only 4 out of 8 experts significantly reduces compute compared to a dense model of equivalent capacity.

We believe the FlexOLMo approach is very appealing in this regard, and we will include a detailed discussion regarding computational complexity scaling in the final version of the paper.

We hope our responses have addressed the reviewer's concerns. Please let us know if you may have any questions.

We thank the reviewer for their time and constructive feedback. We appreciate the acknowledgment that our method is both practical and efficient for real-world applications. We address each of your comments below.

The main limitation of the proposed work is its exclusive focus on large-scale datasets and continued pre-training. The scope of FlexOlmo would be more comprehensive if it could demonstrate SFT and RL based training for individual experts across different domains. The motivation for relying solely on large-scale datasets, specifically creating FlexMix and continued pre-training remains unclear.

We agree that extending the FlexOlmo framework to SFT and RL is a valuable direction for future research. Our choice to focus on contined pre-training is motivated by pressing, real-world cases where data owners possess large amounts of restricted, unstructured data rather than instruction-tuning or preference datasets. For example, the Reddit data [1] in FlexMix represents real-world unstructured data that now comes with access restrictions. Similarly, major publishers like The New York Times [2] have recently restricted the use of their raw news corpus for pretraining LMs. Anthropic [3] faces legal charges for using pirated books like libgen for pretraining. These access restrictions create a critical bottleneck: valuable data for improving LM is locked away. FlexOlmo addresses this problem by giving data owners more control over their data in pretraining. While we believe FlexOlmo's modular training approach would generalize effectively to SFT and RL, exploring that is a natural next step that remains beyond the scope of our current paper and opens up technical considerations that would require extensive experimentation. We will include a discussion on extending our work to SFT and RL in the future work section. Thanks for the suggestions.

Extending to Other Open-Source Architectures

We note that we cannot use Llama or Qwen models directly, as our research requires pre-training a model from scratch to ensure no overlap between the pre-training data and high-risk data added later. While it would be possible to pre-train models based on Llama- or Qwen-based architectures, it is computationally costly, requiring 512 H100 GPUs for three days, which exceeded our budget. We leave the exploration of FlexOlmo under alternative architectures for future work.

If experts are trained in complete independence, how do they learn to coordinate effectively

This is related to the core technical innovation of FlexOlmo. Our training method is explicitly designed to enable coordination of experts without joint training. We achieve this through two key mechanisms:

1. Coordinated Training with an Anchor Model: Instead of joint training, each data owner trains their expert alongside a frozen copy of the public model. This public model acts as a shared "anchor," ensuring that all independently trained experts learn to operate within a common functional space and can coordinate effectively when merged.

2. Domain-Informed Router Initialization: The router, which directs tokens to the appropriate experts, is initialized using domain-specific embeddings derived from each expert's training data. These embeddings encode useful information on how to differentiate between the experts, providing a strong starting point for the router.

Could the authors elaborate on what they mean by "fine-grained control of data access during inference"? How exactly does this mechanism work, and what are the specific benefits? Is there any quantitative analysis as support during evaluations that demonstrates the advantages of this approach?

“Fine-grained control” means that a single FlexOlmo model can deploy different combinations of experts (and their associated data) to different users based on specific access criteria such as permissions, access policies, or the user’s own preferences. Here is a concrete example – suppose a major newspaper contributed an expert. They could restrict access so that only users with an active subscription to that paper have the option to activate that expert, while others cannot. We quantitatively demonstrate the effect of data opt-out in section 5, showing that deactivating the News expert successfully removes its influence on in-domain tasks while having minimal impact on unrelated tasks.

Why the math expert improves coding performance

The improvement in coding performance from the Math expert is likely due to cross-domain generalization of reasoning abilities. The structured step-by-step problem-solving inherent in math text appears to enhance the model's underlying reasoning capabilities. This phenomenon has been observed in prior work; for example, [4] showed that training on mathematical data can significantly boost performance on coding benchmarks and vice versa.

What is the specific accuracy of the router? How can the combined model outperform individual expert X that was specifically trained on domain X

A traditional accuracy metric is not applicable because our router operates at the token level, not the sequence level. There is no single ground-truth expert label for a token. This token-level routing can dynamically select the most appropriate expert for each token in the input sequence. For example, consider an example that contains both math concepts and coding elements. With token-level routing, our model can route the math tokens to the math expert while directing the coding-related tokens to the coding expert. This token-level routing allows the model to leverage each expert’s expertise precisely where it's needed within a single example, rather than being constrained to use one expert for the entire example.

We hope our responses have addressed the reviewer's concerns. Please let us know if you may have any questions.

[1] https://the-decoder.com/reddit-ends-its-role-as-a-free-ai-training-data-goldmine/

[2] https://www.popsci.com/technology/nyt-generative-ai/

[3] https://fortune.com/2025/07/28/a-copyright-lawsuit-over-pirated-books-could-result-in-business-ending-damages-for-anthropic/

[4] Magistral

Thank you for your constructive feedback! We will incorporate a detailed discussion of SFT, RL, and CPT in the limitations section.

We agree that CPT, SFT, and RL each serve different purposes in the model development lifecycle. The choice between these approaches depends primarily on the format and characteristics of the restricted data:

• Supervised Fine-Tuning (SFT) is effective when the restricted data itself is in instruction format. For example, if a data owner possesses extensive user-model conversation histories, SFT can directly leverage this structured interaction data.

• Continued Pre-Training (CPT) is designed for absorbing domain-specific knowledge and is well-suited for unstructured data. This approach aligns with our current setup, where the restricted data consists of plain text such as news and books. Your suggestion to convert unstructured data into instruction format is a promising direction; it introduces technical considerations that would merit extensive experimentation.

Fundamentally, CPT and SFT are similar in nature, sharing the same next-token prediction objective, differing mainly in their input data format. Furthermore, prior research [1, 2] suggests that core model capabilities are primarily from from pretraining, making CPT the more effective method for integrating broad domain-specific knowledge. To make the model interact using chat format, a small, high-quality SFT dataset is often sufficient to align the model.

We acknowledge that a CPT-only model is not production-ready for end users. To enable instruction-following capabilities, people can easily perform SFT with a small number of examples on top of FlexOlmo (which is a base model) or even use in-context learning [3] to replace SFT/alignment. To validate the feasibility of SFT, we conducted a preliminary experiment. We performed SFT on our FlexOlmo model and compared it with the OLMo-2 instruction-following version. As shown in the table below, our initial results demonstrate that our model outperforms the baseline OLMo-2-7B on several key benchmarks. Through testing the fine-tuned models, we confirm that they possess instruction-following capabilities:

Although further hyperparameter optimizations may require additional time, this remains an active area of exploration. We will clarify the complete method and its preliminary validation in the final version of the paper.

Thank you again for this suggestion. We are happy to answer any further questions the reviewer may have and respectfully ask that our revisions be considered in your final evaluation.

[1] LIMA: Less Is More for Alignment [2] The False Promise of Imitating Proprietary LLMs [3] The Unlocking Spell on Base LLMs: Rethinking Alignment via In-Context Learning

Thank you for your continued feedback. We appreciate your thorough review and would like to address your specific concerns:

We would like to clarify FlexOlmo presents a new pretraining recipe for building a pretrained base LM, developed from scratch on unstructured data in a manner similar to OLMo [3], Llama [2] and DeepSeekMoE [4]. Our main experiments are designed to demonstrate the effectiveness of this new pretraining approach. This pretrained model can then be aligned for instruction-following through standard post-training procedures using a small, high-quality, general-domain instruction tuning dataset (as shown in our previous response). As shown in the FLOPs-controlled setup in Table 1, FlexOlmo proves to be a more effective base model compared to one trained with a standard MoE pretraining recipe [4]. With that context, we would like to address your specific points:

Regarding SFT Experimental Details:

For our Supervised Fine-Tuning experiments in the previous response, we used a general-domain instruction-tuning dataset from Tulu 3 [1]. The primary purpose of this experiment was to validate that FlexOlmo can be successfully aligned with user instructions through standard post-training procedures, thereby demonstrating the model's compatibility with existing alignment workflows.

Is the primary motivation for CPT-only approach for FlexOlmo comes from the difficulty of converting large and unstructured data to instruction format, or are there other reasons I miss?

The decision was not driven by difficulty of data conversion, but our research objective: developing a better pretraining recipe with new functionalities (decentralized training, data opt-in and out) for base LMs. We conduct controlled comparison against other pretraining approaches on the same data in the main experiments. Instruction-following capabilities can be added later by applying standard post-training techniques using curated instruction datasets from general domains like Tulu 3[1].

The evaluation details/discussions need to be more clear in that sense; especially in these "evaluation crisis" days with non-reproducable results.

We made all our models, training code, and evaluation scripts publicly available. Both our FlexOlmo models and baseline results can be fully reproduced using our released codebase. Our experimental setup ensures baseline pretraining approaches and our method are evaluated under identical resource constraints and using the same datasets.

[1] Tulu 3: Pushing Frontiers in Open Language Model Post-Training

[2] LLaMA: Open and Efficient Foundation Language Models

[3] OLMo: Accelerating the Science of Language Models

[4] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

Thank you for the thoughtful review and for recognizing our work as a meaningful approach to the critical challenge of data constraints in language model development. We address each of your comments below and have incorporated new experiments and analyses:

Why was federated learning not included as a baseline, given that it is a direct solution to the same data-sharing constraint? Additionally, have the authors considered comparisons to model routing methods such as: [1] Muqeeth, Mohammed, et al. "Learning to route among specialized experts for zero-shot generalization." ICML 2024.

Comparison to Muqeeth et al. (ICML 2024): Thanks for the suggestion. We have adapted the SMEAR method from Muqeeth et al. to our experimental setup. As shown in the table below, our proposed method outperforms SMEAR across all evaluation categories. We will include this result in the final version of the paper:

On Federated Learning (FL): We experimented extensively with existing FL algorithms (e.g., FedAvg, FedProx) using the NVFlare framework (https://github.com/NVIDIA/NVFlare). However, we encountered substantial engineering challenges in scaling these methods to our setting (training models with up to 37 billion parameters), primarily due to the synchronization required in FL. Existing FL methods are not yet supported for training at this scale. Additionally, FL methods do not support data opt-out at inference time, as they are designed to produce a single, monolithic model. We will add discussion of these challenges to the paper.

Our paper includes comparisons to four strong model merging baselines, including BTM and BTX. We believe these, along with the new SMEAR results above, provide a comprehensive comparison. We are happy to include additional baselines in the final version if the reviewer has additional specific suggestions.

The technical contributions over prior MoE-based approaches appear incremental rather than fundamentally novel.

We respectively disagree with this assessment. While Flexolmo borrows the MoE architecture, its core contribution is a fundamentally new training method that enables capabilities not supported by standard MoEs. Standard MoE models are trained end-to-end with joint access to all data. In contrast, our method introduces a new, coordination-based training algorithm that enables two key features: (1) distributed training without data sharing and (2) opting in and out of subsets of data during inference. These features are important for leveraging restricted datasets where centralized data pooling is not feasible.

Additionally, we chose to keep the standard MoE architecture in Flexolmo to ensure our method can be easily integrated into existing large-scale pretraining LM pipelines.

What is the hypothesized reason for FlexOlmo underperforming on Code4?

We attribute this to suboptimal routing. Ideally, coding tasks should be mostly handled by the code expert, but current routing sometimes directs them to other experts. To support this, we analyzed the router’s behavior on coding benchmarks and found that ~23% of tokens are routed to the Textbook and Math experts. We will include this analysis in the final version of the paper to accelerate future improvements.

Can the proposed framework accommodate settings where different data providers use different model architectures for their local experts?

The overall architecture of the local experts needs to be the same as the anchor model (e.g., same positional encoding). Varying the FFN’s hidden dimension is feasible for different local experts as long as the output vectors of the experts have the same dimension. The standard FFN expert has two linear layers: an up-projection to an intermediate dimension and a down-projection back to the output dimension. This intermediate dimension, and thus the FFN expert size, can vary based on a provider's local data size and compute constraints.

Limitation sections

We thank the reviewer for this suggestion. Below is a draft of a limitations section that we will include in the final version of the paper:

Data privacy and extraction risks. A primary concern with sharing model weights is the potential for data leakage. As we quantify in our data extraction analysis (Section 5.3), while the rate of extractable training data from an expert's weights is low, it may still be too high for high-stakes domains like healthcare or finance. A direction for future work is to enable data owners to apply differential privacy (DP) training during their local expert training. Since DP is orthogonal to the FlexOlmo architecture, each contributor could independently decide whether to adopt it, providing formal privacy guarantees.

Scaling and model size. The FlexOlmo architecture scales by adding new experts, which linearly increases the total number of model parameters. While the sparse activation could keep the number of active parameters constant during inference, the total model size can become substantial as more experts are integrated. Future work could explore techniques like expert distillation or pruning to manage model size.

Bias and ethics concerns of FlexOlmo. While the modularity of FlexOlmo provides powerful data control, it also introduces risks that require careful consideration. One concern is the potential for malicious use, where an adversary could intentionally contribute a low-quality expert trained on misinformation. Integrating this low-quality expert could compromise the integrity and safety of the entire model. Similarly, low-quality experts could compromise the model even without any ill intent. Addressing such scenarios would be valuable future work. Furthermore, the framework's guarantee of data owner control raises governance and ethical questions. For instance, an unethical data owner could engage in discriminatory gatekeeping: deciding to withhold a valuable expert (e.g., in medicine or finance) from certain users.

We hope our response and the additional results have addressed the reviewer's concerns. Please let us know if you may have any questions.

Thank you for the feedback and suggestion. Please let us know if you have any further questions. We would be happy to address them.