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.

翻译：精细化的数据配方对于大型语言模型的预训练至关重要，它能显著提升训练效率和模型性能。该配方中的一个重要环节是基于既定规则、LLM判断或嵌入中统计信息生成的分数来筛选样本，这些分数大致可分为质量指标和多样性指标。由于在应用于万亿规模词元的预训练数据集（如FineWeb和DCLM）时计算成本极高，这两类或更多类型的指标很少在单次选择过程中被联合考虑。然而，在我们的实证研究中发现，基于质量指标筛选样本在长期预训练中会出现严重的收益递减现象，而基于多样性指标筛选则会剔除过多有价值的高质量样本，这两者都限制了预训练LLM的能力。因此，我们提出了DATAMASK，一个新颖高效的大规模预训练数据联合学习框架，它能在统一过程中同时优化多种类型的指标，本研究特别聚焦于质量与多样性指标。DATAMASK将选择过程视为一个掩码学习问题，涉及迭代采样数据掩码、基于预定义目标与采样掩码计算策略梯度，以及更新掩码采样逻辑。通过基于策略梯度的优化及多种加速增强技术，其选择时间相比贪心算法显著减少了98.9%，使得本研究得以在万亿词元规模内探索联合学习。利用DATAMASK，我们从15万亿词元的FineWeb数据集中筛选出约10%的子集，命名为FineWeb-Mask。在12项多样化任务上的评估表明，我们在一个15亿参数的稠密模型上取得了3.2%的显著提升，在一个70亿参数的MoE模型上取得了1.9%的提升。