Olmix: A Framework for Data Mixing Throughout LM Development

Data mixing---determining the ratios of data from different domains---is a first-order concern for training language models (LMs), but existing mixing methods have poorly understood design choices and assume that the set of domains remain fixed throughout development. We present Olmix, a framework that addresses two challenges encountered during LM development. First, the configuration space for developing a mixing method is not well understood---design choices across existing methods lack justification or consensus and overlook practical issues like data constraints. We conduct a comprehensive empirical study of this space, identifying which design choices lead to a strong mixing method. Second, the domain set evolves throughout LM development as datasets are revised and expanded---a problem setting largely unaddressed by existing works. We study how to efficiently recompute the mixture after the domain set is updated, given an existing mix from before the update. We introduce mixture reuse, a mechanism that reuses existing relative ratios and recomputes ratios only for domains affected by an update. Over a sequence of five domain-set updates mirroring real-world LM development, mixture reuse matches the performance of fully recomputing the mix after each update with 74% less compute and improves over training without mixing by 11.6% on downstream tasks.

When you train a large language model, one key decision is the data mixture: how much of each kind of text (web pages, code, books, etc.) to include in training. This mixture strongly affects how good the model is across a variety of tasks. We introduce Olmix, a framework for determining a better data mixture in practice. We first study the many design choices behind existing data mixing methods to understand which ones lead to better models. We then address a practical problem: training datasets often change over time as sources are added, removed, cleaned, or updated. Rather than recomputing the entire mixture after every change, Olmix reuses the parts of the old mixture that are still relevant and only updates the parts affected by the change. In experiments that mimic a realistic sequence of dataset updates, this reuse strategy performs as well as recomputing the mixture from scratch while using 74% less compute. It also improves final model performance by 11.6% compared with not optimizing the mixture at all.