RegMix-D: Dynamic Data Mixing via Proxy Training Trajectories
For practitioners of large language model pretraining, RegMix-D offers a more efficient and effective dynamic data mixing method that outperforms static approaches with reduced proxy compute.
RegMix-D extends static data mixing to dynamic mixing by leveraging full loss trajectories from proxy runs, achieving consistent improvements over RegMix and DoReMi across 13 downstream tasks on 25B tokens of the Pile dataset with a 1B parameter model, even with only 128 proxy models (25% of RegMix's proxy compute budget).
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).