Always Learning, Always Mixing: Efficient and Simple Data Mixing All The Time

ListenOne Data-Mixing Algorithm to Rule Them All10:25

Show notes

OP-Mix (On-Policy Mix) reframes data mixing as a single online decision problem spanning pretraining, continual midtraining, and continual instruction tuning — replacing separate proxy models with one LoRA adapter per domain trained from the current model, then simulating mixtures by interpolating those adapters. A single instantiation hits state-of-the-art across all three phases, notably matching on-policy self-distillation with 95% less compute. We unpack the mechanism, the load-bearing numbers, the ablations, and what "weight-merging as data-mix proxy" might mean for multimodal continual training. **Sources** - Always Learning, Always Mixing: Efficient and Simple Data Mixing All The Time

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

OP-Mix (On-Policy Mix) unifies pretraining, continual midtraining, and continual instruction tuning under a single data-mixing algorithm by replacing separate small proxy models with one LoRA adapter per data domain, trained directly from the current model, and then simulating candidate mixtures via linear interpolation between those adapters. The single instantiation is SOTA across all three phases: matches MergeMix with 14% less compute in pretraining (150M–530M OLMo, Chinchilla-optimal), nearly matches full retraining with 66% less compute in continual midtraining, and matches on-policy self-distillation (SDFT) with 95% less compute in continual instruction tuning on Qwen2.5-7B-Instruct. The reframing is: data mixing is an online decision problem, not a phase-specific one, and LoRA + linear-mode-connectivity is the cheapest known machinery for solving it.

Key claims

Prior data-mixing methods split cleanly into “no proxy models” (ADO/DoGE/GRAPE/MergeMix) and “adapts to new data” (OLMix), but no method has both properties — this is the gap OP-Mix fills [Table 1].
Off-policy proxy models (small models trained from scratch on candidate mixtures, RegMix / OLMix style) diverge from the target model’s learning dynamics at scale, and open-weight releases typically ship without a matching small-model proxy — making these methods “defunct after pretraining” [§2].
The core primitive is: train one LoRA per domain from the current model, sample interpolation points on the (K−1)-simplex, evaluate each interpolated adapter, fit a regression over (interpolation weight → performance), minimize to get the optimal old-vs-new tradeoff, then distribute weight across all datasets and do the final training run [§3, Algorithm 1].
In pretraining (150M/300M/530M OLMo, Chinchilla-optimal on Algebraic Stack + ArXiv + c4 + Reddit + StackExchange), OP-Mix beats ERM by 6.3% in average perplexity and matches MergeMix with 14% less compute; outperforms OLMix by 5–6% at every scale [§4.1 Fig. 4].
In continual midtraining (DataDecide 150M/300M/530M cyclically permuted over 5 domains), OP-Mix nearly matches full retraining (the skyline) while using up to 66% less compute [§4.1 Fig. 5, Fig. 9].
In continual instruction tuning on Qwen2.5-7B-Instruct across Tool Use / Science / Medical, OP-Mix on top of standard SFT hits 60.0% mean accuracy — matching SDFT’s 60.2% at 95% less compute — and combining OP-Mix + SDFT gives the best overall 61.9%, showing data mixing and objective modification are orthogonal axes [§4.1 Fig. 6].
Merging LoRAs closely tracks the true full-training loss surface as a function of mixture proportion: OP-Mix’s average loss increase from optimal is 0.9% vs 2.9% for a fixed 10%-data-replay baseline [§5.1 Fig. 7, Fig. 8].
Theoretical bound: OP-Mix’s suboptimality decomposes into (a) LoRA-approximates-full-FT error and (b) linear-interpolation-approximates-mixing error, both empirically small and non-increasing across continual-learning stages [§5.2, Remark on OP-Mix performance gap].
Ablation: naively merging the trained LoRAs into the base model at the optimized weight (LoRA-Merge) — i.e. skipping the final full-training step — is significantly worse than OP-Mix, indicating LoRA works as a proxy but not as the final policy [§4.1 Fig. 10].
The linear-mode-connectivity assumption OP-Mix leans on is empirically supported by prior work observing that linear interpolation between LoRA-fine-tunes of a shared base model does not spike loss [§5.2 Corollary B.2].

Method

Let D = {D_1, ..., D_m} be the current set of domains and θ the current model. When K new domains D_{m+1}, ..., D_{m+K} arrive:

Train one LoRA per new domain on the current model, giving θ^LoRA_{D_{m+i}} for each — a cheap on-policy approximation of what full fine-tuning on D_{m+i} would produce.
Train one “old” LoRA θ^LoRA_{D_old} on the existing mixture, so the search only adjusts the ratio between the old aggregate and the new domains rather than re-searching over the full simplex.
Sample interpolation points on the (K)-simplex over {D_old, D_{m+1}, ..., D_{m+K}} and evaluate each interpolated merge — no additional training, just forward passes.
Fit a regression on (interpolation weight, evaluation loss) pairs and minimize it to find w*, the optimal old-vs-new tradeoff.
Distribute w* across individual datasets and run the final full training pass with the resulting mixture.

For pretraining, OP-Mix warms up with 20% ERM (equal sampling) before reintroducing each dataset as a “new” domain to trigger the same adjustment loop.

The pattern generalizes MergeMix (weight-merging as data-mix proxy on full-parameter fine-tunes) by (a) replacing full FT with LoRA and (b) allowing new domains to be incorporated online. It generalizes the offline proxy-model pipeline of OLMix/RegMix/CLIMB by removing the separately-initialized proxy model entirely.

Results

Pretraining perplexity vs baselines at Chinchilla-optimal budgets: OP-Mix beats ERM by 6.3% average, matches MergeMix at 14% less compute, and beats OLMix by 5–6% at 150M, 300M, and 530M [Fig. 4].
Downstream evaluations (Appendix Table 2): OP-Mix is best or second-best across a battery of held-out tasks, mirroring the perplexity ranking.
Continual midtraining on the DataDecide model suite: OP-Mix nearly matches full-retraining (the skyline that trains from scratch on all accumulated data at each stage) with 66% less compute; the LoRA-Merge ablation (using merged LoRAs as the final policy, skipping full FT) is significantly worse [Fig. 9, Fig. 10].
Continual instruction tuning on Qwen2.5-7B-Instruct across Tool Use → Science → Medical: SFT 58.7% / SDFT 60.2% / OP-Mix+SFT 60.0% / OP-Mix+SDFT 61.9%; OP-Mix+SFT uses 95% less total compute than SDFT to hit essentially the same accuracy [Fig. 6].
Pareto frontier: across all three phases, no baseline achieves better performance at lower compute than OP-Mix [Fig. 2].
Loss-surface fidelity: the LoRA-merge proxy tracks the true full-training loss surface closely enough that OP-Mix picks mixtures within 0.9% average loss of the true optimum, vs 2.9% for a fixed 10%-replay baseline [§5.1 Fig. 7].
Scale limits stated by authors: experiments top out at 530M for pretraining/midtraining and 7B for instruction tuning; frontier-scale (70B+) behavior is untested and the paper does not characterize behavior as domain count grows to 10 or 100.

Why it’s interesting

OP-Mix is the third clean datapoint in a converging cluster around “weight-merging as data-mixing proxy”: Decouple Searching from Training: Scaling Data Mixing via Model Merging for Large Language Model Pre-training (DeMix) makes the bet with full-parameter component models at Qwen3-1.7B scale, MergeMix (arxiv 2601.17858, cited here as the direct antecedent) makes it in mid-training, and OP-Mix generalizes both by using LoRA proxies + explicit domain-expansion machinery so the same algorithm covers pretraining, midtraining, and continual instruction tuning. The three papers agree on the linearity hypothesis — that training on a mixture is approximately a linear combination in weight space of training on each component — but disagree on where to spend compute: DeMix on scale-matched component models, MergeMix on full FT, OP-Mix on cheap on-policy LoRAs. For Luma this closes a specific gap Kyle flagged when filing DeMix: the scale-transfer worry (“we can’t afford to train components at target scale”) is exactly what LoRA-as-proxy sidesteps, since the LoRA is fit from the target model itself.

The other reason this matters for the team’s reading list is the 95%-compute-cut vs on-policy distillation result. On-policy distillation (On-Policy Distillation) has been the recommended continual-learning recipe on the wiki; OP-Mix says most of that gain is recoverable via data mixing alone, and the two compose to give an additional lift. That reframes SDFT-style methods as an orthogonal-axis improvement rather than the default fix — a useful prior when designing multimodal continual-training pipelines where compute is already the binding constraint. It also complements Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training (CLIMB) — same iterative-predictor-loop skeleton, but CLIMB uses small proxy models trained from scratch and OP-Mix uses on-policy LoRAs, exactly the axis DeMix and OP-Mix argue is load-bearing.