PaperlayerView exact revision on arXiv

Source-linked AI summary

Olmix: A Framework for Data Mixing Throughout LM Development

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

arXiv:2602.12237v1cs.LGcs.AIcs.CL

TL;DR

Data mixing is important for LM performance, but existing methods leave configuration choices underexplained and generally assume fixed domain sets. Olmix studies configuration choices, incorporates data constraints, and introduces mixture reuse for evolving domains. Across five updates, mixture reuse matches full recomputation with 74% less compute and improves over the natural distribution by 11.6%.

  • Problem

    Existing mixing methods provide limited configuration guidance, often ignore data constraints, and generally assume fixed domain sets despite iterative LM development.

  • Method

    Olmix empirically studies seven offline-mixing design choices, incorporates repetition constraints, and reuses mixture ratios while recomputing weights affected by domain updates.

  • Results

    +11.6% improvement over the natural distribution with 74% fewer proxy runs; FullMixtureReuse captures 95% of full recomputation’s gains across five updates.

  • Takeaways & Limitations

    Mixture reuse provides a lower-cost alternative to recomputing mixtures from scratch as domain sets evolve, while retaining most of full recomputation’s gains.

  • Takeaways & Limitations

    The work focuses on offline mixing, its theoretical analysis assumes log-linear regression, and PartialMixtureReuse requires selecting unaffected domains to recompute.

Abstract

from arXiv · show

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.

1 Introduction

Olmix addresses two practical data-mixing challenges in LM development: configuring methods on initial domains and efficiently recomputing mixtures as domains evolve. It studies configuration choices empirically and introduces mixture reuse for updated domain sets.

  • Data mixing is crucial for downstream performance, but manual tuning or exhaustive search can require thousands of GPU hours.
  • Existing methods provide limited guidance for configuring mixing methods because design choices are often unjustified, inconsistent, or ignore data constraints.
  • Olmix empirically studies seven configuration choices and finds that proxy-run requirements scale linearly with the number of domains.
  • Different regression families perform best at different swarm sizes, while the adapted log-linear model achieves the best overall downstream performance.
  • Data constraints incorporated into mixture optimization control sample repetition while maintaining performance and significantly shape the proposed mix.
  • Mixture reuse freezes relative weights among unaffected domains and recomputes only their total weight alongside affected-domain weights.
  • Across five updates culminating in 64 domains, FullMixtureReuse improves over the natural distribution by +11.6% while requiring 74% fewer proxy runs.

2 Related Work

Related work includes offline, dynamic, and online mixing methods, alongside research on data constraints and iterative LM data development. Olmix differs by explicitly supporting varied domain updates and integrating repetition constraints into mixing.

  • Offline mixing methods fit performance from static mixture experiments, using either parametric regressors or nonparametric models such as LightGBM and Gaussian Processes.
  • DoReMi and DoGE dynamically explore mixture weights with one or two proxy runs rather than many static proxy runs.
  • Online mixing methods adjust mixture weights during the final training run, unlike Olmix’s domain-set dynamics before that run.
  • Existing mixture work largely assumes fixed domain sets, whereas Olmix handles adding, removing, partitioning, and revising domains.
  • Prior data-constrained work caps repetition or models diminishing returns, while Olmix integrates repetition constraints directly into data mixing.
  • LM data-development projects document iterative curation processes including quality filtering, deduplication, and domain selection.

3 The Mixing Configuration Problem

Olmix studies how to configure offline data-mixing methods for language-model development, addressing poorly justified design choices and practical constraints. Its findings define OlmixBase, including proxy-model, swarm, regression, granularity, and repetition-constraint choices.

  • Olmix approach: Olmix empirically studies the offline mixing schema and uses the findings to define OlmixBase for Olmo 3 development.The schema’s configuration includes seven key design choices.
  • Motivation: Existing methods provide limited guidance on configuration choices, disagree on regression models, and often overlook data constraints.These gaps motivate a systematic study of the offline mixing schema.
  • Swarm construction: Proxy models with ≥15M parameters achieve strong rank correlation with 1B target models, whereas 1M models are unreliable; Olmix uses 30M proxies.The 30M proxy setup uses 3B tokens and achieves a Spearman correlation of 89.6 with 1B target models.
  • Swarm construction: O(m) proxy runs are sufficient for strong performance as the number of domains m grows, and the authors recommend K ≥3(m + 1) runs with log-linear regression.Error curves collapse across domain counts when swarm size is scaled linearly with m.
  • Swarm construction: Sparse swarms perform better for topic-level domains, while dense swarms perform better for source-level domains, so swarm choice depends on the domain set.The result holds for both downstream BPB and regression fit.
  • Regression modeling: Log-linear regression achieves the best overall downstream performance, while per-task regression provides the strongest combination of regression fit and downstream performance.Different regression families excel at different swarm sizes, but log-linear models are recommended when K ≥3(m + 1).
  • Data constraints: In data-constrained training, constrained optimization satisfies repetition limits while maintaining performance, and the constraint changes allocations across domains.Tighter repetition limits increase allocations to domains such as literature, while relaxed limits increase allocations to high-utility domains such as software development.

4 The Evolving Domain Problem

Olmix addresses evolving domain sets by formalizing mixture recomputation after updates and reusing ratios from unaffected domains. Its reuse strategies reduce optimization cost while preserving performance when affected and unaffected domains are suitably decoupled.

  • Problem setup: Olmix studies domain sets that evolve through adding, removing, partitioning, or revising domains during language-model development.The framework separates unaffected domains D1 from affected domains D2 transformed into D′2.
  • Problem setup: Full recomputation reapplies OlmixBase after every update, requiring O(m′) proxy runs whose costs accumulate across updates.This provides the high-performance baseline against which reuse methods are compared.
  • Mixture reuse: Mixture reuse freezes relative ratios for unaffected domains, aggregates them into a virtual domain, and recomputes only the virtual-domain weight and affected-domain weights.The optimization dimensionality falls from m′ to 1 + |D′2|.
  • Mixture reuse: FullMixtureReuse applies OlmixBase in the collapsed space, requiring O(|Dcomp|) proxy runs before expanding the result back to the full domain set.The method enforces the reused relative ratios while optimizing the total unaffected-domain weight and affected-domain mix.
  • Performance conditions: Reuse matches full recomputation when the optimal mix changes little and coupling between fixed and recomputed domains is low.Empirical validation finds that the reuse gap tracks actual performance gaps and that success for additions correlates with small 1 − ρ⋆ and low coupling.
  • Performance conditions: PartialMixtureReuse selectively recomputes some nominally unaffected domains to reduce coupling while retaining lower cost than full recomputation.For example, recomputing a software-development web topic alongside added code data improves performance compared with reusing all web topics.

5 Experimental Results

Across five evolving-domain updates, mixture reuse maintained performance close to full recomputation while substantially reducing proxy-run costs. PartialMixtureReuse further narrowed the performance gap by selectively recomputing coupled domains.

  • Mixture reuse performance: FullMixtureReuse achieved 95% of full recomputation’s improvement (+11.6% versus +12.2%) using 74% fewer proxy runs (216 versus 832).It reuses ratios for unaffected domains while recomputing affected ones.
  • Mixture reuse performance: PartialMixtureReuse achieved 98% of full recomputation’s performance (+12.0%) with 272 runs, still 67% fewer than full recomputation.It selectively recomputes some unaffected domains to close the remaining performance gap.
  • Comparisons: At matched budgets of 216–272 proxy runs, FullMixtureReuse and PartialMixtureReuse improved over the natural distribution by +11.6% and +12.0%, respectively.At the same budget, swarm reuse achieved +11.4%, while full recomputation with c = 1 achieved +10.8%.
  • Data efficiency: PartialMixtureReuse reached the natural distribution’s final BPB in approximately 20,000 steps versus 61,000, a 3.05× data-efficiency speedup.This compares training-step efficiency rather than final performance alone.
  • Analysis: Mixture reuse performance tracked the reuse gap, while PartialMixtureReuse improved performance by reducing coupling between reused and recomputed domains.Recomputing DCLM:software_development when adding Stack-Edu reduced the coupling term κ more than recomputing other DCLM topics.

6 Discussion

Olmix addresses configuration and evolving-domain challenges in data mixing, but its scope and automation remain limited. The reported design-choice findings include linear swarm-size scaling and strong performance from log-linear regression, with evidence extending to larger proxy settings.

  • Contributions: Olmix combines an empirical study of mixing design choices with mixture reuse for updating mixtures as domain sets evolve.The framework targets both initial method configuration and repeated recomputation during LM development.
  • Limitations: The theoretical analysis of mixture reuse assumes log-linear regression models and has unclear extension to non-parametric, non-convex, non-differentiable objectives.The authors expect natural extensions to some other parametric models.
  • Limitations: PartialMixtureReuse requires domain expertise to determine which unaffected domains should be recomputed because no general automated selection method is provided.The paper gives DCLM:software_development as an intuitive example when Stack-Edu code data is added.
  • Scope: The study focuses on offline mixing, while online methods that adjust mixtures during training may require different design choices.The authors identify extending mixture reuse and the design-choice study to online mixing as future work.
  • Design-choice findings: When swarm size was scaled as K = c(m + 1), error curves collapsed across m, supporting O(m) runs for strong downstream performance at 30M and 1B scales.The results used m = 6 and m = 24 in the additional analysis.

A.4 RQ4: Regression model family

The regression-model study finds that model families operate best in different swarm-size regimes, helping explain prior disagreement. Across source-level comparisons, log-linear regression provided the strongest overall combination of fit and downstream performance.

  • Regression model families: Log-linear regression outperformed other approaches across downstream performance and regression-fit comparisons at the source level.The comparisons averaged results across three random train-test splits of the swarm.
  • Regression model families: Different regression model families excelled at different swarm sizes and domain counts, indicating distinct sample-complexity regimes.The authors identify these regimes as a potential explanation for disagreement across existing methods.
  • Model adaptations: The adapted regression models represent downstream task performance as a function of training-domain mixture weights rather than directly modeling proxy-model training steps.The BiMix adaptation simplifies the mixing law by absorbing constants into fitted parameters, while the AutoScale adaptation uses a task-level constant.
  • Optimization: For the evaluated solver comparisons, the study selected an exact solver with AutoScale and a search-based solver for BiMix.The choice followed the downstream-performance results summarized in the solver comparison table.
  • Regression granularity: At 30M scale, per-task regression performed best, followed by per-family and aggregated regression granularities.The performance ordering matched the corresponding regression-fit findings.

A.6 RQ6: Data repetition constraints

The appendix examines how repetition constraints and optimization choices affect proposed mixtures. It also documents mixture-reuse parameterizations for removing, partitioning, and revising domains.

  • Repetition constraints: Varying the repetition factor changes the proposed weights across all 24 DCLM topics.The full result is presented as an expanded version of the repetition-constraint figure.
  • Optimization: The solver comparison evaluates performance and predicted performance across optimization methods when mixing across sources.The exact optimizer minimized predicted performance, while Exact + KL (0.05) produced the strongest downstream BPB in the reported comparison.
  • Mixture reuse operators: Removing a domain can reuse the remaining domains’ relative ratios, producing an updated mixture without additional recomputation.The example renormalizes [0.25, 0.5] to [0.33, 0.67].
  • Mixture reuse operators: Partitioning a domain replaces full recomputation with optimization over a virtual unaffected-domain aggregate and the newly formed domains.The example recomputes over three domains instead of all four.
  • Mixture reuse operators: Revising a domain similarly optimizes over a virtual aggregate of unaffected domains plus the revised domain.The example reduces the recomputation problem from three domains to two.

C Proofs for Section 4.3

This section introduces notation for separating unaffected and changed domains and for expressing mixtures through their relative mass and conditional ratios.

  • Objective representation: The objective is written as a bivariate function of the mixture over changed domains and the fixed-domain conditional mixture.This representation supports analyzing how the optimal mixture changes when the fixed-domain component varies.
  • Mixture decomposition: A mixture can be decomposed into weights and conditional distributions over the fixed and complementary domain sets.The decomposition uses π for the mass on Dfix and 1 − π for the mass on D \ Dfix.
  • Sensitivity quantities: Feature vectors associated with fixed and complementary domains are separated, with αfix and αcomp recording their task-wise norms.These norm vectors are used to characterize sensitivity of the objective to mixture changes.

C.2 Proof for Theorem 1

Theorem 1 establishes exactness when the reused ratios on unaffected domains equal their post-update optimum, and bounds the performance gap when they differ under stated modeling and convexity assumptions.

  • Assumptions: The analysis assumes a log-linear relationship between mixture ratios and task performance, exact optimization, strong convexity in the optimized variables, and mutual feasibility.The exact-optimization assumption excludes proxy-transfer and parameter-estimation errors, which would contribute approximation terms.
  • Exactness condition: Zero performance gap occurs when the reused fixed-domain mixture equals the optimal post-update mixture on those domains.Under this equality condition, the mixture-reuse and standard mixing problems are equivalent.
  • Proof strategy: The proof reduces the performance-gap analysis to comparing objectives along changes in the fixed-domain mixture.First-order conditions, mean-value bounds, and strong convexity control how the optimizer and objective vary between two fixed-domain mixtures.
  • Performance-gap bound: The resulting bound depends linearly on the distance between the reused and optimal fixed-domain mixtures, with constants determined by gradient and feature-norm quantities.The bound incorporates quantities such as a_max, b_max, αfix, αcomp, and the reduced objective.

C.3 Proof for Theorem 2

This section analyzes domain-set updates and identifies conditions under which mixture reuse is exact or admits a bounded deviation from the fully recomputed optimum.

  • Approximate reuse: The non-exact update analysis relies on strong convexity, mutual feasibility, and bounds on feature norms to control optimizer and objective changes.These assumptions support argmin-stability and gradient-based performance-gap bounds.
  • Adding domains: For added domains, reuse is exact when the post-update optimum assigns no mass to the added domains.In that case, restricting the post-update problem to zero mass on added domains preserves the optimum.
  • Removing domains: For removed domains, reuse is exact when the pre-update optimum assigns zero mass to every removed domain.The pre-update and restricted post-update problems then have the same optimum on the remaining domains.
  • Partitioning domains: For partitioning, reuse preserves the unaffected-domain ratios when the original domain’s conditional distribution over subdomains equals the token-count distribution.The partition constraint does not change the optimum when the post-update solution satisfies this natural conditional distribution.
  • Revising domains: For revising a domain, the reused mixture is exact when the revised domain has the same task-specific feature vectors before and after revision.Identical objectives and feasible sets imply identical pre-update and post-update optimizers.

D.1 Experimental Setup

The experiments train 1B-parameter decoder-only transformer models and evaluate them on downstream tasks using specified optimization, hardware, and tokenization settings.

  • Evaluation: The downstream evaluation tasks are listed in Table 9, and domain token counts are reported in Table 8.These tables define the evaluation task set and provide the per-domain data breakdown.
  • Model: The experiments use 1B-parameter decoder-only transformers based on the OLMo 2 architecture.The models use 16 layers, 16 attention heads, d_model 2048, and head dimension 128.
  • Training: Training uses batch size 512, learning rate 0.0018, cosine scheduling with warmup and linear decay, and sequence length 4096.The Dolma 2 tokenizer is used for all models.
  • Compute: Full models are trained on 32 NVIDIA H100 GPUs, while proxy models use one NVIDIA H100 GPU.This setup distinguishes the computational resources used for target and proxy training.

D.2 Implementation Details

The implementation reuses prior mixture information across evolving domains and applies consistent proxy-model training and swarm-sizing procedures.

  • FullMixtureReuse and PartialMixtureReuse carry each newly recomputed mix forward as the next prior mixture.The initial mixture is produced by OlmixBase on the initial web corpus.
  • All evaluated swarm-based methods train 30M-parameter proxy models for 5x Chinchilla with batch size 64, learning rate 0.007, and sequence length 2048.
  • Proxy-run counts scale with the number of recomputed domains using K ≈ c(m + 1), with c values of 1, 2, or 3.For two domains, the method uses one-dimensional search instead of regression.
  • Swarm Reuse is presented as an alternative algorithm that reuses prior mixing information.

D.3 Additional Results

Additional results show that mixture reuse preserves performance across domain updates while reducing proxy-run requirements and remaining effective under low compute budgets.

  • 99% of full recomputation’s performance improvement (+6.94% versus +6.97%) requires 74% fewer proxy runs (216 versus 832).
  • At low budgets, swarm reuse performs similarly to FullMixtureReuse.
  • At low budgets, mixture reuse and swarm reuse require as few as 76 proxy runs, compared with 267 for full recomputation.The low-budget evaluation spans 76 to 272 total proxy runs.
  • The experiments cover Add, Remove, Partition, and Revise updates, with mixture weights and evaluation-suite details provided in the accompanying tables.Tables 10–13 list stage-wise swarm details for the evaluated recomputation and reuse strategies.
  • Figure 24 compares performance improvement with mixing cost, while Figure 25 reports results for R = 1T and R = 6T.Figure 26 compares performance gaps and total variation distances against full recomputation across update operators.
  • Mixture reuse achieves similar performance to full recomputation across Add, Remove, and Partition updates, while Revise yields a 0.21% total variation distance.Downstream performance was not evaluated for Revise because variation was expected to be dominated by noise.
Loading 2602.12237v1…