Olmix: A Framework for Data Mixing Throughout LM Development

TL;DR

Olmix, a framework that addresses two challenges of data mixing: the configuration space for developing a mixing method is not well understood and in practice, the domain set evolves throughout LM development as datasets are added, removed, partitioned, and revised.

Abstract

Data mixing -- determining the ratios of data from different domains -- is a first-order concern for training language models (LMs). While existing mixing methods show promise, they fall short when applied during real-world LM development. We present Olmix, a framework that addresses two such challenges. 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, in practice, the domain set evolves throughout LM development as datasets are added, removed, partitioned, and revised -- a problem setting largely unaddressed by existing works, which assume fixed domains. We study how to efficiently recompute the mixture after the domain set is updated, leveraging information from past mixtures. We introduce mixture reuse, a mechanism that reuses existing ratios and recomputes ratios only for domains affected by the 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.

Figures (27)

• Figure 1: Two problems with data mixing encountered during LM development: (1) How to best configure your mixing method? (2) How to efficiently mix under evolving domain sets?
• Figure 2: The offline mixing schema used by many existing methods (§\ref{['sec:schema']}). We study the design choices needed to configure the schema to develop a strong mixing method (§\ref{['sec:design_choices']}-§\ref{['sec:study']}).
• Figure 3: Correlation in performances between pairs of proxy models and 1B target models trained on the same mixture. Proxy models with over 15M parameters achieve strong rank correlation.
• Figure 4: Error vs. Swarm Size. Curves collapse across different $m$, indicating $\mathcal{O}(m)$ runs are needed. Results are averaged across 3 random seeds, and the intervals indicate min and max results.
• Figure 5: Performance (left) and regression fit (right) of dense versus sparse swarms. While sparse swarms perform better at the topic level, dense swarms perform better at the source level, suggesting that behavior is data-dependent.

Theorems & Definitions (38)

• Theorem 1: Performance gap bound
• Theorem 2
• Lemma 1: Reparametrized mixture reuse problem
• proof
• Lemma 2
• proof
• Lemma 3: FOC Inequality
• proof
• Lemma 4: Mean value bound for $\triangledown_r F$ along a segment
• proof