Olmix: A Framework for Data Mixing Throughout LM Development

Abstract: Data mixing -- determining the ratios of data from different domains -- is a first-order concern for training 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.

• The paper introduces a comprehensive framework for designing data mixtures in LM development, emphasizing empirical studies and optimal proxy model selection.
• It demonstrates that sparse and dense sampling strategies vary by topic and source levels, with swarm size scaling linearly in the number of domains.
• The reuse strategies (FullMixtureReuse and PartialMixtureReuse) achieve over 95% of full recomputation performance while significantly reducing computational cost.

Olmix: A Framework for Data Mixing Throughout LLM Development

Introduction and Problem Motivation

The construction of data mixtures—determining how to allocate training tokens among multiple domains or sources—is a critical aspect of large-scale LLM (LM) pretraining. Optimal mixing ratios significantly influence downstream task performance, especially as models are trained on increasingly heterogeneous and evolving corpora. However, two fundamental challenges impede the effectiveness and practicality of current mixing approaches:

1. Mixing Configuration Uncertainty: The literature exhibits little consensus or rigorous justification for key design choices in mixing methods, with most works assuming idealized or limited scenarios that do not account for real-world data constraints.
2. Evolving Domain Sets: In production-scale LM development, the set of available domains undergoes continual evolution—domains are added, removed, revised, or partitioned throughout the data curation process—which existing mixing algorithms do not address systematically.

   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?

The Olmix framework systematically addresses both challenges, providing rigorous empirical guidance for data mixture design and introducing principled methods for efficient mixture recomputation under evolving domain sets (2602.12237).

Offline Mixing Schema: Empirical Foundations

Olmix builds upon the prevailing offline mixing schema, widely adopted in prior works, wherein small-proxy “swarm” models are trained on multiple mixtures, a regression surrogate is fit to predict downstream performance from mixture weights, and the optimal mix is identified via surrogate optimization.

Figure 2: The offline mixing schema used by many existing methods (\S\ref{sec:schema}).

Systematic Study of Mixing Configuration

Olmix provides the first comprehensive empirical study of seven key design choices, spanning proxy model scale, swarm construction, regression target granularity, and data constraint handling:

1. Proxy Model Size: Reliable signal transfer to target scale requires proxies with at least 15M parameters. Proxy models smaller than this demonstrably yield degraded rank correlation with the performance of target LMs.

   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.

2. Swarm Size Scaling: Contrary to previous assumptions of quadratic or super-linear scaling, empirical results show sample complexity is linear in the number of domains ($\mathcal{O}(m)$), with strong performance achieved as $K \geq 3(m+1)$.

   Figure 4: Error vs. Swarm Size. Curves collapse across different $m$, indicating $\mathcal{O}(m)$ runs are needed.

3. Swarm Distribution: At the topic level, sparse Dirichlet sampling (where only subsets of domains are represented in each mixture) yields superior results; at the source level, dense sampling is preferred. The choice is inherently data-dependent.

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 data-dependent behavior.

1. Regression Model Family: The log-linear mixing law delivers the best average downstream Bits-Per-Byte (BPB) and regression fit, superseding alternatives such as LightGBM, BiMix, and Gaussian Process regressors—particularly when sample complexity recommendations are met.

Figure 6: Left: Fit versus swarm size across regression models. The log-linear model achieves the best regression fit. Right: Downstream performance of regression models for $K=128$. The log-linear model achieves the best average downstream task BPB.

1. Regression Granularity: Fitting a separate surrogate for each downstream task, as opposed to family-level or global regression, results in both improved regression fits and better final mixture performance.
2. Optimization under Data Constraints: When domain data is finite, directly incorporating repetition constraints (i.e., capping how often a domain can be repeated) at the optimization level, not in swarm construction, is empirically shown to maintain performance without violating practical data constraints. The proposed mixture weights shift non-trivially as the constraint is relaxed.

   Figure 7: Proposed Mixture Weights versus Repetition Factor. Enforcing varying caps on repetition $k$ leads to significantly different proposed mixes.

3. Optimization Solver Choice: Although the exact solver minimizes the surrogate loss, adding moderate KL-divergence regularization towards the natural domain distribution further improves final downstream LM performance, likely due to transfer noise between proxy and target scales.

Figure 8: Downstream performance (left) and 30M predicted performance (right) across optimization solvers. Moderate KL regularization helps achieve the best downstream performance, even if the exact solver achieves the lowest predicted BPB at proxy scale.

Collectively, these findings inform the OlmixBase method—a robust offline mixing pipeline that specifies practical and empirically justified choices for all aforementioned components.

Mixture Recomputation for Evolving Domain Sets

Mixing algorithms in the prior literature universally assume a static domain set, but in applied LM development the data composition is highly dynamic. Olmix formalizes the evolving domain problem and introduces mixture reuse strategies to amortize proxy training and optimize efficiency across iterative domain updates.

Figure 9: Overview of Mixture Recomputation Strategies. Full recomputation applies OlmixBase on the updated domain set, while mixture reuse methods freeze ratios for unaffected domains and only recompute over affected domains.

FullMixtureReuse and PartialMixtureReuse

• FullMixtureReuse: Upon a domain update, Olmix partitions the domain set into unaffected and affected subsets. Proxy computational effort is amortized by freezing ratios for the unaffected domains and only recomputing mixture weights for the affected part (plus the single new “virtual domain” representing the frozen subset). Theoretical analysis reveals that the performance gap to full recomputation is strictly controlled by the “reuse gap”—i.e., how much the optimal mixtures for the unaffected domains shift under the new downstream tasks and constraints—and the coupling between affected and unaffected domains.
• PartialMixtureReuse: To further reduce coupling effects (when some unaffected domains remain highly correlated with the new or updated domains), Olmix introduces selective recomputation of a subset of previously unaffected domains, providing a tunable spectrum between full recomputation cost and maximum reuse efficiency.

Experimental Results

Comprehensive experiments, simulating five successive domain updates yielding a 64-domain mixture, demonstrate that FullMixtureReuse and PartialMixtureReuse consistently retain at least 95% of the downstream performance gains of full recomputation—at a fraction (down to 26%) of the computational cost in terms of proxy runs required.

Figure 10: Performance improvement versus cost of mixing under evolving domains. FullMixtureReuse and PartialMixtureReuse achieve >95% of the improvement of full recomputation while using at least 67% fewer proxy runs.

PartialMixtureReuse can reach the same final BPB as the full recomputation baseline in $3.05\times$ fewer training steps, demonstrating substantial gains in data efficiency.

Figure 11: Downstream performance across training for PartialMixtureReuse (best seed) versus the natural distribution. PartialMixtureReuse reaches the natural distribution's final performance in $3.05\times$ fewer steps.

Qualitative analysis confirms the learned mixtures under PartialMixtureReuse closely follow those produced by full recomputation, deviating significantly from simple size-proportional (natural) mixing.

Figure 12: Proposed mixes for different mixing strategies. PartialMixtureReuse and full recomputation produce more similar mixtures to each other than to the natural distribution.

Theoretical Analysis and Empirical Validation

The Olmix framework justifies mixture reuse strategies with precise theoretical bounds, showing that under a log-linear performance model, the performance degradation from reuse is bounded by the product of the coupling constant and the reuse gap. Empirical validation confirms that as the reuse gap grows (for example, if newly introduced domains have high utility or are large in scale), the performance of reuse strategies degrades controllably.

Figure 13: Performance vs Reuse Gap. The optimality of $\tilde{p}$ after the domain update is correlated with the success of mixture reuse.

Figure 14: Performance gap vs reuse gap vs $1 - \rho^\star$ for significant domain-addition operations, validating the mixture reuse theory developed in the paper.

Targeted application of PartialMixtureReuse—selectively recomputing individual previously unaffected domains with high coupling—recovers most of the remaining performance gap, confirming the theoretical predictions.

Figure 15: Performance of PartialMixtureReuse (recompute DCLM:software_development) for $D_1=$DCLM, $D_2'=$Stack-Edu. Selective recomputation closes the reuse and performance gap introduced by coupling.

Figure 16: Largest $\kappa$ reduction occurs when selectively recomputing high-coupling domains, clarifying where PartialMixtureReuse is most effective.

Implications and Future Directions

The Olmix framework sets a new standard for empirical rigor in data mixing, codifying best practices across proxy selection, swarm design, surrogate fitting, and constraint handling. Its mixture reuse algorithms drastically enhance the cost-effectiveness and data-efficiency of iterative LM development—directly supporting workflows where model pretraining data is continually refined and updated.

On the theoretical front, Olmix provides tools for precisely quantifying when mixture reuse will be effective and where full recomputation is necessary, guiding both practitioners and future research on adaptive mixture learning.

Potential future work includes joint design with online mixing algorithms, automated coupling-detection for optimal PartialMixtureReuse, and large-scale validation in >30B parameter regimes.

Conclusion

Olmix provides a principled, empirically validated methodology for both configuring and maintaining high-quality data mixtures during the entire LM developmental lifecycle. Its integration of rigorous empirical design, automated reuse with theoretical guarantees, and scalability to evolving domains represents a substantial advance in the operational science of foundation model pretraining (2602.12237).