Olmix: A Framework for Data Mixing Throughout LM Development

1. Problem Statement

Data mixing—determining the optimal ratios of data from different domains (e.g., code, web text, math) for training Large Language Models (LMs)—is critical for downstream performance. However, existing methods face two significant challenges when applied to real-world LM development:

1. The Configuration Problem: The design space for offline mixing methods is poorly understood. Existing literature lacks consensus on key hyperparameters (e.g., swarm size, regression model family, data constraints), often relying on unjustified choices or ignoring practical constraints like limited data availability.
2. The Evolving Domain Problem: Real-world development is iterative. Datasets are frequently added, removed, partitioned, or revised. Existing methods assume a fixed domain set, requiring a full recomputation of the mixing strategy from scratch after every update. This is computationally prohibitive as the number of updates and domains grows.

2. Methodology

The authors propose Olmix, a framework addressing both challenges through a comprehensive empirical study and a novel reuse mechanism.

A. Solving the Configuration Problem (OlmixBase)

The authors formalize the "offline mixing schema" (Swarm Construction $\to$ Regression $\to$ Optimization) and conduct a systematic empirical study to identify optimal design choices for pre-training Olmo 3. Key findings include:

• Swarm Size: The number of proxy runs ($K$) required scales linearly with the number of domains ($m$). Specifically, $K \approx 3(m+1)$ is sufficient for strong performance, contradicting prior assumptions of quadratic scaling.
• Proxy Model Size: Proxy models must be at least 15M parameters to achieve strong rank correlation with target models (1B scale). Smaller models (1M) are unreliable.
• Regression Model: Log-linear models (adapted from Data Mixing Laws) achieve the best overall downstream performance and regression fit, outperforming LightGBM, Gaussian Processes, and Power Laws, provided the swarm size is sufficient.
• Regression Granularity: Fitting per-task regression models yields better performance than per-family or aggregated models, despite higher noise.
• Data Constraints: Existing methods assume infinite data. Olmix explicitly incorporates repetition constraints (limiting how many times a domain can be sampled) directly into the optimization step. This prevents performance degradation caused by oversampling low-availability, high-utility domains.
• Optimization Solver: An exact solver with KL regularization (towards the natural distribution) outperforms pure exact solvers and search-based heuristics, balancing optimization accuracy with robustness against regression noise.

These findings culminate in OlmixBase, a concrete configuration used throughout the Olmo 3 development.

B. Solving the Evolving Domain Problem (Mixture Reuse)

To handle domain updates without full recomputation, Olmix introduces Mixture Reuse.

• Concept: When a domain set $D$ updates to $D'$, the method freezes the relative ratios of unaffected domains ($D_{fix}$) and only recomputes the weights for affected domains ($D_{comp}$) and the total weight of the frozen group.
• Mechanism:
  1. FullMixtureReuse: Aggregates all unaffected domains into a single "virtual domain." The optimization space is reduced from $m'$ dimensions to $1 + |D_{comp}|$.
  2. PartialMixtureReuse: A middle-ground approach that selectively recomputes a subset of "unaffected" domains if they are highly coupled with the new domains (e.g., recomputing "Software Development" topics when adding new code data).
• Theoretical Guarantee: The authors derive bounds showing that the performance gap between reuse and full recomputation depends on:
  1. The Reuse Gap: How much the optimal ratios for unaffected domains change after the update.
  2. Coupling: The degree to which unaffected and affected domains impact the same downstream tasks. If coupling is low and the reuse gap is small, reuse performs nearly identically to full recomputation.

3. Key Contributions

1. Comprehensive Configuration Study: The first large-scale empirical study of the offline mixing schema, providing concrete, data-driven guidelines for swarm size, model architecture, regression families, and constraint handling.
2. Mixture Reuse Framework: A novel mechanism to efficiently update data mixtures as domain sets evolve, significantly reducing computational costs.
3. Theoretical Analysis: Formal bounds characterizing when mixture reuse is optimal, linking performance gaps to the "reuse gap" and domain coupling.
4. OlmixBase: A robust, open-source mixing method used to train the state-of-the-art Olmo 3 family of models.

4. Experimental Results

The authors evaluated Olmix in a simulated real-world development scenario involving 5 domain updates (adding, removing, partitioning, and revising datasets), culminating in a 64-domain set.

• Performance vs. Cost:
  • FullMixtureReuse achieved 95% of the performance gains of full recomputation while requiring 74% fewer proxy runs (216 runs vs. 832 runs).
  • PartialMixtureReuse achieved 98% of full recomputation's gains (+12.0% over natural distribution) with 67% fewer runs.
• Data Efficiency: The best mixture found via reuse was 3.05$\times$ more data-efficient than the natural distribution (proportional to domain sizes), reaching the same final performance in one-third of the training steps.
• Theoretical Validation: Empirical measurements confirmed that the theoretical bounds (Reuse Gap and Coupling) accurately predict the performance gap between reuse strategies and full recomputation.
• Robustness: The method performed consistently across different update operators (Add, Remove, Partition, Revise) and data constraints ($R=1T$ and $R=6T$).

5. Significance

• Practical Impact: Olmix provides a scalable solution for the iterative nature of modern LM development. It allows practitioners to update data mixtures frequently without incurring the massive computational cost of retraining proxy swarms from scratch.
• Methodological Rigor: By moving away from heuristic tuning to empirically grounded design choices (e.g., linear scaling of swarm size, log-linear regression), the paper sets a new standard for data mixing research.
• Resource Efficiency: The framework enables high-performance model training with significantly reduced compute budgets, making advanced data mixing accessible to organizations with limited resources.
• Open Science: The framework, code, and data are released to support the broader community in developing efficient and effective language models.