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

• URL: http://arxiv.org/abs/2512.24265v1
• Date: Tue, 30 Dec 2025 14:38:34 GMT
• Title: Joint Selection for Large-Scale Pre-Training Data via Policy Gradient-based Mask Learning
• Authors: Ziqing Fan, Yuqiao Xian, Yan Sun, Li Shen,
• Abstract summary: A fine-grained data recipe is crucial for pre-training large language models.<n>One important ingredient in the recipe is to select samples based on scores produced by defined rules, LLM judgment, or statistical information in embeddings.<n>We introduce DATAMASK, a novel and efficient joint learning framework designed for large-scale pre-training data selection.
• Score: 17.10988637170463
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Exploring Instruction Data Quality for Explainable Image Quality Assessment [58.345719195248314]
  We investigate the role of data quality of instruction tuning dataset for explainable IQA.<n>We find that selecting a subset of the data set randomly can even lead to better results than training with the entire instruction tuning dataset.<n>We propose a clustering-based data selection framework with three stages: clustering feature extraction, cluster quota allocation, and cluster sampling strategy.
  arXiv  Detail & Related papers  (2025-10-04T17:12:54Z)
• Oversampling and Downsampling with Core-Boundary Awareness: A Data Quality-Driven Approach [2.334306891078381]
  We propose a method to systematically identify and differentiate between two types of data.<n>By prioritizing high-quality, decision-relevant data, our approach can be extended to text, multimodal, and self-supervised learning scenarios.<n>This work paves the way for future research in data-efficient learning, where intelligent sampling replaces brute-force expansion.
  arXiv  Detail & Related papers  (2025-09-24T07:55:07Z)
• Large Language Models are Demonstration Pre-Selectors for Themselves [57.101804269100185]
  In-context learning (ICL) with large language models (LLMs) delivers strong few-shot performance by choosing few-shot demonstrations from the entire training data.<n>FEw yet Essential Demonstration prE-selectoR is a novel pre-selection framework that identifies a representative subset of demonstrations.<n>FEw yet Essential Demonstration prE-selectoR can reduce training data size by over 20% while maintaining performance.
  arXiv  Detail & Related papers  (2025-06-06T12:29:03Z)
• T-SHIRT: Token-Selective Hierarchical Data Selection for Instruction Tuning [5.963754140027611]
  Token-Selective HIeRarchical Data Selection for Instruction Tuning (T-SHIRT) is a novel data selection framework.<n>We demonstrate that models instruction-tuned on a curated dataset can outperform those trained on the entire large-scale dataset.
  arXiv  Detail & Related papers  (2025-06-02T04:59:17Z)
• Add-One-In: Incremental Sample Selection for Large Language Models via a Choice-Based Greedy Paradigm [50.492124556982674]
  This paper introduces a novel choice-based sample selection framework.<n>It shifts the focus from evaluating individual sample quality to comparing the contribution value of different samples.<n>We validate our approach on a larger medical dataset, highlighting its practical applicability in real-world applications.
  arXiv  Detail & Related papers  (2025-03-04T07:32:41Z)
• Star-Agents: Automatic Data Optimization with LLM Agents for Instruction Tuning [71.2981957820888]
  We propose a novel Star-Agents framework, which automates the enhancement of data quality across datasets. The framework initially generates diverse instruction data with multiple LLM agents through a bespoke sampling method. The generated data undergo a rigorous evaluation using a dual-model method that assesses both difficulty and quality.
  arXiv  Detail & Related papers  (2024-11-21T02:30:53Z)
• A CLIP-Powered Framework for Robust and Generalizable Data Selection [51.46695086779598]
  Real-world datasets often contain redundant and noisy data, imposing a negative impact on training efficiency and model performance.<n>Data selection has shown promise in identifying the most representative samples from the entire dataset.<n>We propose a novel CLIP-powered data selection framework that leverages multimodal information for more robust and generalizable sample selection.
  arXiv  Detail & Related papers  (2024-10-15T03:00:58Z)
• Data curation via joint example selection further accelerates multimodal learning [3.329535792151987]
  We show that jointly selecting batches of data is more effective for learning than selecting examples independently. We derive a simple and tractable algorithm for selecting such batches, which significantly accelerate training beyond individually-prioritized data points.
  arXiv  Detail & Related papers  (2024-06-25T16:52:37Z)
• How to Train Data-Efficient LLMs [56.41105687693619]
  We study data-efficient approaches for pre-training language models (LLMs) We find that Ask-LLM and Density sampling are the best methods in their respective categories. In our comparison of 19 samplers, involving hundreds of evaluation tasks and pre-training runs, we find that Ask-LLM and Density are the best methods in their respective categories.
  arXiv  Detail & Related papers  (2024-02-15T02:27:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.