RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

Problem Formulation and Prior Approaches

The selection of domain mixture ratios is a vital factor in LLM pretraining, with substantial downstream performance impact. Existing automated mixture optimization methods, such as DoReMi (group DRO) and RegMix (regression-based proxy modeling), identify an ostensibly optimal mixture based on static proxy signals. These methods assume a single, immutable mixture suffices throughout pretraining. However, both theoretical work and empirical evidence indicate that the optimal mixture may shift over the course of training, invalidating the static mixture assumption. Dynamic mixing methods (e.g., Aioli, TiKMiX, AC-ODM) attempt to remedy this, but generally require additional model instrumentation or optimization machinery during target training.

REGMIX-D: Methodological Innovation

REGMIX-D proposes a regression-based dynamic mixture scheduling framework, extending RegMix while maintaining proxy efficiency and minimal additional computational overhead at deployment. The approach leverages full loss trajectories from proxy runs, not merely endpoint losses, to construct localized regression models that predict next-step losses conditioned on current step, mixture, and observed loss. This enables REGMIX-D to estimate optimal mixtures on a per-segment basis.

Two deployment modes are supported:

• Offline Schedule Generation: Prior to target training, recursively generate a mixture schedule based solely on proxy trajectories by feeding predicted losses into subsequent optimization decisions.
• Online Mixture Adaptation: During target model training, dynamically adjust mixtures at switch points using current observed validation losses, correcting for cross-scale loss discrepancies via power-law scaling.

Both modes require only a single regression model trained offline and Dirichlet-sampled mixture candidates; online regression querying incurs negligible overhead.

Experimental Evaluation

REGMIX-D was evaluated using the 1B TinyLlama on 25B tokens from the Pile (17 domains), employing a 1M proxy TinyLlama in main experiments. Comparisons were drawn against Human (Pile prior), DoReMi, and RegMix baselines. REGMIX-D was run at two proxy budgets: 128 and 512 proxy models (512 is RegMix's standard).

Main Results

• Across 13 downstream tasks from lm-eval-harness, both REGMIX-D offline and online variants consistently outperformed RegMix and DoReMi.
• REGMIX-D (128 proxy models) surpassed RegMix (512 proxies) in average downstream accuracy, demonstrating strong proxy-efficiency: only 25% of proxy compute was required.
• Online variants outperform offline in downstream metrics, attributable to better grounding in observed target dynamics.
• Gains are distributed broadly across task types rather than concentrated in specific categories; REGMIX-D achieved top-1 or top-2 results on 11 of 13 tasks.
• Increasing proxy count yields marginal improvements, indicating dynamic mixture scheduling (as opposed to proxy budget size) is the critical driver.
• Dynamic mixture trajectories diverge substantially from RegMix’s static mixture, with marked changes in pile-cc domain weighting.

Ablation Studies and Robustness

Ablations on switch point count, cross-scale loss correction, and proxy size revealed:

• The number of switch points (N) moderately affects performance, with N=5 optimal in main settings but N=3,7,9 also outperforming RegMix.
• Loss correction factor (β) for online mode has limited influence (maximum 0.25 Avg variation).
• Proxy model size scaling (1M–120M) yields virtually identical downstream scores, confirming proxy scale sufficiency.
• All studied REGMIX-D configurations exceeded RegMix baseline performance.

Minimal computational overhead was observed for REGMIX-D offline (identical to RegMix) and only +0.37% for the online mode.

Theoretical and Practical Implications

REGMIX-D demonstrates that localized regression modeling of proxy loss trajectories enables time-varying mixture schedules for LLM pretraining, substantially overcoming the suboptimality of static mixtures. Its proxy efficiency and low-overhead design make it amenable to practical deployment in large-scale settings. The method’s robust performance across hyperparameters and proxy scales implies adaptability for broader pretraining contexts without careful tuning.

Theoretically, REGMIX-D closes gaps in automated mixture selection by providing a flexible, regression-based scheduling paradigm grounded in actual training dynamics—eschewing reliance on additional optimization machinery or complex RL-based adaptation.

Limitations and Future Directions

REGMIX-D currently relies on a single target-domain validation loss (pile-cc) as the optimization signal, inherited from RegMix’s setup. Extension to multi-target or domain-agnostic objectives will require significant modification to the regression framework. Fair, quantitative comparison with other dynamic mixing methods (e.g., Aioli, TiKMiX) is hindered by differences in data, domain partitioning, and evaluation suites.

Potential future developments include:

• Generalizing REGMIX-D to support multi-domain or multi-objective schedules.
• Integrating richer proxy signals including gradient- or representation-level dynamics.
• Applying REGMIX-D framework to vision or multimodal model pretraining.

Conclusion

REGMIX-D offers a regression-based, proxy-efficient approach to dynamic data mixture scheduling for LLM pretraining, improving upon static and prior dynamic methods in downstream task performance, compute efficiency, and deployment simplicity. The approach’s robustness and minimal overhead make it a promising candidate for automated mixture optimization in large-scale LLM pipelines (2606.18663).