Optimizing Training Data via Metagradient Descent

Curating a training dataset from a mass of unfiltered data is a necessary and influential step in any large-scale machine learning pipeline. We frame dataset selection as an optimization problem over per-example data counts and solve it with metagradient descent (MGD). At a high level, our algorithm operates as follows: we start with a randomly chosen set of training data, then iteratively update the dataset selection using metagradients with respect to importance weights placed on each training datapoint. To do this at scale, we develop: (1) REPLAY, an algorithm and implementation for computing scalable metagradients using recursive checkpointing and efficient differentiation of individual training steps; (2) a diagnostic for *metasmoothness*, allowing us to modify the training routine to obtain useful gradients; (3) a continuous surrogate for optimizing data counts; and (4) a block-SGD-inspired algorithm for improving data optimization. Combining these four contributions, MGD reaches average scores of 22% on DataComp-small and 40.2% on DataComp-medium, and improves over LESS and full-data training for Gemma-2B instruction tuning on BBH and MMLU. These results show that repeated gradient-based dataset updates can improve both pre-training and instruction-tuning data.