用强化学习优化大规模数据筛选,兼顾质量与多样性。
Joint Selection for Large-Scale Pre-Training Data via Policy Gradient-based Mask Learning
- 将数据筛选建模为掩码学习,通过策略梯度联合优化质量与多样性。
- 在15万亿token的FineWeb上仅用1.1%时间完成筛选,选中约10%数据。
- 新数据集使小模型性能提升3.2%,大模型提升1.9%,适合训练高效LLM。
精细的数据选择对大语言模型预训练至关重要,能显著提升训练效率和模型表现。当前方法多基于质量或多样性指标筛选样本,但二者常难以兼顾:单纯追求质量导致收益递减,过度强调多样性则会丢弃优质样本。为此,我们提出DATAMASK,一种高效的联合学习框架,可统一优化多个指标。该框架将数据选择视为掩码学习问题,通过迭代采样、基于目标计算策略梯度并更新采样概率,在加速技术加持下,相比贪心算法减少98.9%的筛选时间,实现对15万亿token级数据集(FineWeb)的联合优化。最终筛选出约10%的子集(FineWeb-Mask),在12项任务上验证,1.5B参数稠密模型性能提升3.2%,7B参数MoE模型提升1.9%,显著优于基线。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。