Group-Level Data Selection for Efficient Pretraining
This work addresses the challenge of improving pretraining efficiency for language models, offering a novel method that significantly enhances speed and quality, though it is incremental as it builds on existing data selection techniques.
The paper tackles the problem of inefficient data selection for language model pretraining by introducing Group-MATES, a group-level approach that uses a relational data influence model and clustering to optimize selection, achieving 3.5%-9.4% relative performance gains over random selection and reducing token requirements by up to 1.75x.
In this paper, we introduce Group-MATES, an efficient group-level data selection approach to optimize the speed-quality frontier of language model pretraining. Specifically, Group-MATES parameterizes costly group-level selection with a relational data influence model. To train this model, we sample training trajectories of the language model and collect oracle data influences alongside. The relational data influence model approximates the oracle data influence by weighting individual influence with relationships among training data. To enable efficient selection with our relational data influence model, we partition the dataset into small clusters using relationship weights and select data within each cluster independently. Experiments on DCLM 400M-4x, 1B-1x, and 3B-1x show that Group-MATES achieves 3.5%-9.4% relative performance gains over random selection across 22 downstream tasks, nearly doubling the improvements achieved by state-of-the-art individual data selection baselines. Furthermore, Group-MATES reduces the number of tokens required to reach a certain downstream performance by up to 1.75x, substantially elevating the speed-quality frontier. Further analyses highlight the critical role of relationship weights in the relational data influence model and the effectiveness of our cluster-based inference. Our code is open-sourced at https://github.com/facebookresearch/Group-MATES.