[Paper] Olmix: A Framework for Data Mixing Throughout LM Development

Source: arXiv - 2602.12237v1

Overview

Training large language models (LLMs) often involves pulling data from many different sources—news articles, code repositories, scientific papers, and more. Deciding how much of each source to use (the “mixing ratio”) can dramatically affect model quality, yet most existing methods assume a static set of domains and provide little guidance on the myriad design choices involved. The paper Olmix: A Framework for Data Mixing Throughout LM Development tackles this gap by (1) systematically mapping the design space of mixing strategies and (2) introducing mixture reuse, a technique that lets developers update their data mixes efficiently as the pool of domains evolves during a model’s lifecycle.

Key Contributions

• Comprehensive empirical study of the mixing‑method design space, pinpointing which hyper‑parameters and heuristics actually matter for strong performance.
• Mixture reuse algorithm that re‑optimizes only the ratios of domains that changed, re‑using previously computed ratios for the rest.
• Real‑world simulation of five successive domain‑set updates (additions, deletions, splits) that mirrors how production teams iterate on data pipelines.
• Compute savings: mixture reuse achieves the same downstream performance as recomputing the mix from scratch while cutting the required compute by ~74 %.
• Performance boost: models trained with Olmix’s mixing strategy outperform a baseline that trains on the raw concatenated data by +11.6 % on downstream evaluation tasks.

Methodology

1. Define the mixing design space – The authors enumerate the knobs that existing mixing methods manipulate, such as:

   • Domain weighting heuristics (e.g., uniform, size‑based, loss‑based)
   • Optimization objective (e.g., minimizing validation loss, maximizing task‑specific metrics)
   • Update frequency (how often the mix is recomputed)
   • Constraints (max/min per‑domain data, total budget)

2. Empirical grid search – They run a large‑scale grid search across these knobs on a suite of public corpora (Wikipedia, Common Crawl, code, scientific text, etc.) to see which combinations consistently yield the best validation loss and downstream scores.

3. Mixture reuse mechanism – When the domain set changes (e.g., a new dataset is added), the algorithm:

   • Identifies affected domains (the new one, any removed or split domains).
   • Keeps the old ratios for unchanged domains.
   • Re‑optimizes only the ratios for the affected subset using the same objective as the original mix.
     This is essentially a warm‑start for the mixing optimizer, avoiding a full recompute.

4. Evaluation pipeline – The authors simulate a realistic development cycle: after each of five domain‑set updates they train a fresh LM using (a) the full recomputed mix, (b) mixture reuse, and (c) a naïve “no‑mix” baseline. They then fine‑tune each model on several downstream tasks (question answering, code completion, summarization) and report task‑specific metrics.

Results & Findings

• Design‑space insights: Loss‑based weighting (using a small validation set to gauge per‑domain difficulty) consistently outperformed simple size‑based or uniform mixes. Adding a minimum‑data constraint prevented catastrophic forgetting of low‑resource domains.
• Mixture reuse robustness: Even after multiple, non‑trivial domain changes (including splitting a large corpus into thematic sub‑domains), reuse maintained performance, confirming that the optimal ratios for unchanged domains are stable across updates.

Practical Implications

• Faster iteration cycles – Teams can now tweak their data pipelines (add a new domain, drop noisy data, or re‑segment a corpus) without paying the full cost of re‑optimizing the entire mix. This is especially valuable for large‑scale LLM projects where each full mix computation can cost thousands of GPU hours.
• Better resource allocation – By identifying the most impactful mixing heuristics, developers can focus engineering effort on loss‑based weighting and constraint handling rather than trial‑and‑error with arbitrary ratios.
• Continuous data‑drift handling – In production, data sources evolve (e.g., new APIs, updated documentation). Olmix’s reuse strategy provides a principled way to keep the model’s training distribution aligned with the latest data without destabilizing performance.
• Open‑source potential – The framework is modular; it can be dropped into existing training pipelines (e.g., Hugging Face 🤗 Transformers, DeepSpeed) as a pre‑processing step that outputs a weighted sampling schedule.

Limitations & Future Work

• Scope of domains – The empirical study focuses on a handful of public corpora; exotic or highly imbalanced domains (e.g., low‑resource languages) may behave differently.
• Optimization overhead – While mixture reuse cuts compute dramatically, the initial full‑mix optimization still requires a non‑trivial budget, which could be prohibitive for very large domain sets.
• Dynamic weighting during training – The current approach recomputes a static mix before each training run. Future work could explore online mixing where ratios adapt continuously as the model’s loss landscape evolves.
• Task‑specific mixing – The paper optimizes for a generic validation loss; extending the framework to directly target downstream task metrics (e.g., BLEU for translation) could yield further gains.

Olmix offers a pragmatic, data‑centric toolset that bridges the gap between academic mixing strategies and the messy realities of production LLM development. By demystifying the design space and providing a compute‑efficient reuse mechanism, it equips engineers to iterate faster, allocate data more intelligently, and ultimately ship higher‑quality language models.

Authors

• Mayee F. Chen
• Tyler Murray
• David Heineman
• Matt Jordan
• Hannaneh Hajishirzi
• Christopher Ré
• Luca Soldaini
• Kyle Lo

Paper Information

• arXiv ID: 2602.12237v1
• Categories: cs.LG, cs.AI, cs.CL
• Published: February 12, 2026
• PDF: Download PDF