Joint Selection for Large-Scale Pre-Training Data via Policy Gradient-based Mask Learning

TL;DR

The paper addresses the challenge of selecting high-quality, diverse pre-training data for trillion-token scale LLMs. It introduces DATAMASK, a mask-learning framework that optimizes a joint objective over quality and diversity metrics via policy-gradient updates, enabling scalable data selection with orders-of-magnitude speedups over greedy methods. Empirically, DATAMASK yields a FineWeb-Mask subset (~1.5T tokens) that improves performance across 12 tasks for both 1.5B dense and 7B MoE models, confirming the value of joint metric optimization. The work demonstrates practical impact by delivering a scalable, efficient data selection approach suitable for trillion-scale pretraining, and provides insights into metric design, initialization, and training dynamics relevant to large language model development.

Abstract

A fine-grained data recipe is crucial for pre-training large language models, as it can significantly enhance training efficiency and model performance. One important ingredient in the recipe is to select samples based on scores produced by defined rules, LLM judgment, or statistical information in embeddings, which can be roughly categorized into quality and diversity metrics. Due to the high computational cost when applied to trillion-scale token pre-training datasets such as FineWeb and DCLM, these two or more types of metrics are rarely considered jointly in a single selection process. However, in our empirical study, selecting samples based on quality metrics exhibit severe diminishing returns during long-term pre-training, while selecting on diversity metrics removes too many valuable high-quality samples, both of which limit pre-trained LLMs' capabilities. Therefore, we introduce DATAMASK, a novel and efficient joint learning framework designed for large-scale pre-training data selection that can simultaneously optimize multiple types of metrics in a unified process, with this study focusing specifically on quality and diversity metrics. DATAMASK approaches the selection process as a mask learning problem, involving iterative sampling of data masks, computation of policy gradients based on predefined objectives with sampled masks, and updating of mask sampling logits. Through policy gradient-based optimization and various acceleration enhancements, it significantly reduces selection time by 98.9% compared to greedy algorithm, enabling our study to explore joint learning within trillion-scale tokens. With DATAMASK, we select a subset of about 10% from the 15 trillion-token FineWeb dataset, termed FineWeb-Mask. Evaluated across 12 diverse tasks, we achieves significant improvements of 3.2% on a 1.5B dense model and 1.9% on a 7B MoE model.

Figures (16)

• Figure 1: Average performance on 12 tasks across pre-training tokens for a 1.5B dense model trained on FineWeb, FineWeb-Semdedup, FineWeb-Edu, UltraFineWeb, FineWeb-DCLM, and our FineWeb-Mask, respectively.
• Figure 2: Visualization of text embeddings via t-SNE tsne on random subsets of FineWeb. White, light blue, and dark blue points correspond to samples that are top diverse, top high-quality, and samples selected by algorithms that exhibit both high diversity and quality. Light blue points show tighter clustering. Dark blue points are sparse in algorithms except for ours.
• Figure 3: Illustration of the quality and diversity scores when optimizing under DATAMASK with two initialization and two optimization strategies. Initialization strategies: same initialization (green and red lines) that all initial sampling probabilities are the same and quality initialization (black and blue lines) that initial sampling probabilities are proportional to the samples' quality score. Optimization strategies: diversity optimization (green and black lines) that only optimizes the diversity metric and joint optimization (blue and red lines) that optimizes both the diversity and quality metrics.
• Figure 4: Computation time of DiSF with greedy algorithm and DATAMASK when reaching the same optimized scores with varying data size. Selection ratio is 10%, and we specifically label the optimized score and required time.
• Figure 5: Heatmap of quality score and diversity score in 400 clusters. Deeper color means there are more clusters in a certain range of quality and diversity scores. Analysis are provided in Section \ref{['sec:insights']}.