Olmix: Principled Data Mixing for LLMs

• Olmix is a principled framework for data mixing in LLM training, integrating offline configuration, optimization, and mixture reuse for robust performance.
• It systematically explores key design axes—such as proxy model size, swarm size, and regression granularity—to achieve state-of-the-art bits-per-byte metrics.
• Olmix employs efficient mixture reuse algorithms that adapt to evolving domains, reducing proxy runs by up to 74% while maintaining high model utility.

Olmix is a principled framework for data mixing in the training and continual development of large LMs under evolving domain sets and practical data constraints. It encompasses both a systematic study of the offline data mixing configuration space and a suite of algorithms for efficient mixture recomputation as datasets are added, removed, or revised throughout real-world LM development cycles (Chen et al., 12 Feb 2026).

1. Formal Definition and Problem Setup

Let the domain set be $D = \{D_1, \dots, D_m\}$, where each domain $D_i$ contains $N_i$ available tokens. A mixture is a probability vector

$$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$

specifying that $p_i R$ tokens from $D_i$ are used when training for $R$ tokens total. An LM of size $S$ parameters trained for $R$ tokens on mixture $p$ is denoted $D_i$0. For $D_i$1 downstream tasks, the $D_i$2-th task is evaluated via the bits-per-byte (BPB) metric

$D_i$3

with lower values indicating better performance. The data mixing objective is to find

$D_i$4

In data-constrained settings, a repetition bound $D_i$5 is imposed such that

$D_i$6

This constraint prevents excessive resampling from low-resource domains, addressing overfitting and maintaining data diversity.

2. Offline Mixing Schema and Configuration Space

Olmix builds on a three-step offline mixing schema: Swarm $D_i$7 Regression $D_i$8 Optimization. The configuration space is parameterized along seven key design axes, each empirically studied:

1. Proxy model size ($D_i$9): Proxy LMs with $N_i$0 parameters are required to achieve Spearman $N_i$1 correlation with full-scale (1B) models. $N_i$2-parameter proxies yield only $N_i$3.
2. Swarm size $N_i$4 vs. domain count $N_i$5: Error in the mixing objective collapses when $N_i$6; empirically, $N_i$7 suffices.
3. Swarm distribution $N_i$8: Sparse swarms (allowing $N_i$9) benefit topic-level domains by uncovering exclusion, while dense swarms suit source-level domains. Dirichlet priors centered on the natural distribution $$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$0 proved robust; uninformative priors degraded performance.
4. Regression family $$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$1: Tested surrogates include LightGBM, Gaussian Process, BiMix, AutoScale, and log-linear. The per-task log-linear model

$$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$2

yields the highest Pearson correlation ($$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$3 at $$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$4) and best downstream BPB.

1. Regression granularity: Fitting per-task surrogates (one model per downstream task) achieves the best fit ($$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$5) and BPB (0.765 vs. 0.777 per-family, 0.774 aggregated).
2. Data repetition constraints: Constraints $$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$6 are enforced solely in the optimization step ($$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$7 is typical), rather than during swarm construction. Constraining the swarm degrades performance, while unconstrained swarms plus constrained optimization maintain efficacy.
3. Optimization solver: Mixture optimization minimizes

$$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$8

subject to the repetition bounds, using CVXPY with a Kullback-Leibler (KL) penalty ($$p \in \Delta^{m-1}, \qquad \Delta^{m-1} = \left\{ p \in \mathbb{R}^m_{\ge 0} : \sum_{i=1}^m p_i = 1 \right\},$$9). This regularization mitigates overfitting to the regression surrogate.

These choices form the OlmixBase method, which achieves state-of-the-art results on rigorous BPB benchmarks.

3. Key Empirical Insights

Empirical studies underpin Olmix's recommendations for data mixing in LM training:

• Proxy LMs with $p_i R$0 parameters at $p_i R$1 Chinchilla scaling reliably transfer mixing strategies to 1B-parameter models.
• Swarm sample complexity is linear in the number of domains, with $p_i R$2 sufficient for error minimization.
• Domain granularity (topic or source) informs whether to use sparse or dense swarm priors.
• Per-task log-linear regression achieves average BPB = 0.765; LightGBM and BiMix trail at 0.775 and 0.778, respectively.
• Enforcing repetition bounds in the optimization only allows for higher utility mixes and steadier performance.
• Mixes optimized with a small KL penalty ($p_i R$3) outperform pure convex solvers by reducing over-optimization risk.

4. Mixture Reuse Under Evolving Domains

Olmix introduces algorithms for efficient mixture recomputation when domains change. Rather than re-running the entire offline schema ($p_i R$4 proxy runs for updated $p_i R$5), mixture reuse freezes weights for unaffected domains and optimizes for new/affected domains only. Partition $p_i R$6, prescribing

$p_i R$7

with a collapsed mixture $p_i R$8, $p_i R$9. Sampling $D_i$0 runs in this reduced space and expanding to $D_i$1 via $D_i$2, one fits per-task log-linear surrogates and optimizes $D_i$3 with repetition and KL constraints. PartialMixtureReuse recomputes a superset of old domains to further control coupling, with corresponding complexity.

Reuse Method	Proxy Runs per Update	Performance (% BPB improvement vs. natural)
Full Recompute	832	+12.2
FullMixtureReuse	216	+11.6
PartialMixtureReuse	272	+12.0
Swarm Reuse	268	+11.4

Mixture reuse provides 95% of the full recompute's gains while requiring 74% fewer proxy runs. At 6T tokens, FullMixtureReuse achieves 99% of the full recompute’s improvement (Chen et al., 12 Feb 2026).

5. Experimental Evaluation on Realistic Development Pipelines

A five-stage simulation, mirroring OLMo 1→3 development, demonstrates Olmix's utility:

1. Initial: 24 topic domains (DCLM).
2. Add: +15 programming-language domains.
3. Add: +6 STEM corpora (ArXiv, FineMath-3+, PDFs, Wikipedia, AlgebraicStack, Pes2o).
4. Revise: Improved PDF processing replaces previous pipeline.
5. Remove/Partition: Drop AlgebraicStack; split PDFs into 21 finer topics (final $D_i$4).

1B-parameter OLMo2 LMs are trained for $D_i$5B tokens per stage, evaluated on $D_i$6 downstream tasks (math, codegen, QA, commonsense), and measured in BPB. The baseline is sampling in proportion to token count ($D_i$7), and methods compared include OlmixBase, FullMixtureReuse, PartialMixtureReuse, and swarm reuse.

Results reveal that Olmix achieves a $D_i$8 speedup in convergence: the best mix reaches the OLMo natural BPB in $D_i$9K steps, compared to $R$0K for naive mixing. All findings are statistically robust across seeds (min–max intervals $R$1 of mean).

6. Practical Recommendations

Derived from empirical and theoretical analysis, Olmix’s primary recommendations are:

• Use $R$2M-parameter proxies at $R$3 Chinchilla scaling for reliable mixture transfer.
• Set swarm size $R$4 and sample from Dirichlet$R$5; select sparse/dense by domain type.
• Fit per-task log-linear surrogates of the form $R$6.
• Enforce repetition limits $R$7 in optimization only ($R$8).
• Solve for the optimal mixture with CVXPY and a KL regularization term ($R$9).
• When domains shift, employ FullMixtureReuse to reduce proxy runs from $S$0 to $S$1; with high coupling, use PartialMixtureReuse.

These procedures supply both guidance for initial mixture design and tools for efficient, robust updating as the domain landscape changes and expands.

7. Significance and Context within LLM Training

Olmix empirically establishes principled, statistically robust protocols for data mixing under evolving training corpora—a persistent, under-addressed challenge in LM pipeline management. It systematizes surrogate model choice, regression granularity, sample complexity scaling, domain constraint enforcement, and mixture optimization for optimal downstream task performance. The introduction of mixture reuse methods allows practical, compute-efficient adaptation to new datasets without sacrificing performance. These advances directly address gaps in real-world data curation and model improvement, providing a concrete toolkit for practitioners and researchers training large LMs on dynamic, multi-domain corpora (Chen et al., 12 Feb 2026).