RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories

2026-06-17

AI Summary

This work addresses the limitations of static data mixing strategies in large language model pretraining, which fail to adapt to the dynamic learning requirements across training stages. The authors propose a regression-based method that leverages the complete loss trajectory of a small-scale proxy model to predict optimal data mixture proportions, enabling dynamic scheduling. Notably, this approach is the first to utilize full-training loss trajectories for optimizing data mixing policies and supports both online and offline deployment. Evaluated on the Pile dataset (25B tokens) using a 1B-parameter model, the method consistently outperforms RegMix and DoReMi across 13 downstream tasks while incurring only 25% of the proxy model’s computational overhead, thereby achieving significant performance gains with reduced computational cost.

Abstract

Data mixture selection is critical for Large Language Model pretraining. Existing methods such as RegMix select a single static mixture by fitting a regression model on small-scale proxy runs. We propose RegMix-D, a simple extension of RegMix to dynamic mixing. Our key observation is that proxy runs produce not only endpoint losses, but also full loss trajectories, which can be used to further improve data mixture. By training regression model on these trajectories, we can predict optimal mixtures at multiple training stages. RegMix-D supports two deployment modes: an offline variant that generates a complete mixture schedule before target training, and an online variant that adapts the mixture during training using observed loss. Experiments on 25B tokens of the Pile dataset with a 1B parameter target model show that RegMix-D consistently improves over RegMix and DoReMi across 13 downstream tasks while remaining proxy-efficient: it surpasses RegMix even with only 128 proxy models (25% of RegMix's proxy compute budget).

Problem

• data mixture selection
• dynamic mixing
• large language model pretraining
• proxy training trajectories

Innovation

• dynamic data mixing
• proxy training trajectories
• regression-based mixture selection
• large language model pretraining
• efficient proxy modeling