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

arXiv: https://arxiv.org/abs/2512.24265 · HTML: https://mlsanity.com/paper/2512.24265/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 43 references that are in the map.
- First posted on arXiv 30 Dec 2025.

## Builds on (most cited first, 29 of 29)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://mlsanity.com/paper/2212.03533/) (2022)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [RACE: Large-scale ReAding Comprehension Dataset From Examinations](https://mlsanity.com/paper/1704.04683/) (2017)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [SemDeDup: Data-efficient learning at web-scale through semantic deduplication](https://mlsanity.com/paper/2303.09540/) (2023)
- [Seed1.5-Thinking: Advancing Superb Reasoning Models with Reinforcement Learning](https://mlsanity.com/paper/2504.13914/) (2025)
- [QuRating: Selecting High-Quality Data for Training Language Models](https://mlsanity.com/paper/2402.09739/) (2024)
- [A Proximal Stochastic Gradient Method with Progressive Variance Reduction](https://mlsanity.com/paper/1403.4699/) (2014)
- [Organize the Web: Constructing Domains Enhances Pre-Training Data Curation](https://mlsanity.com/paper/2502.10341/) (2025)
- [Arctic-Embed: Scalable, Efficient, and Accurate Text Embedding Models](https://mlsanity.com/paper/2405.05374/) (2024)
- [Ultra-FineWeb: Efficient Data Filtering and Verification for High-Quality LLM Training Data](https://mlsanity.com/paper/2505.05427/) (2025)
- [Programming Every Example: Lifting Pre-training Data Quality Like Experts at Scale](https://mlsanity.com/paper/2409.17115/) (2024)
- [Combinatorial Optimization for Panoptic Segmentation: A Fully Differentiable Approach](https://mlsanity.com/paper/2106.03188/) (2021)
- [GneissWeb: Preparing High Quality Data for LLMs at Scale](https://mlsanity.com/paper/2502.14907/) (2025)
- [Combatting Dimensional Collapse in LLM Pre-Training Data via Diversified File Selection](https://mlsanity.com/paper/2504.20644/) (2025)

## Built on by (most cited first, 5 of 5)

- [One Algorithm, Two Goals: Dual Scoring for Parameter and Data Selection in LLM Fine-Tuning](https://mlsanity.com/paper/2605.06166/) (2026)
- [SimReg: Achieving Higher Performance in the Pretraining via Embedding Similarity Regularization](https://mlsanity.com/paper/2605.08809/) (2026)
- [Spokes: Optimizing for Diverse Pretraining Data Selection](https://mlsanity.com/paper/2606.15216/) (2026)
- [Towards Scalable Data Diversification for Language Model Pretraining via Leverage Score Sampling](https://mlsanity.com/paper/2609.32484/) (2026)
- [Probabilistic Multi-Source Data Selection via Learnable Categorical Priors: Variance Reduction and Guaranteed Derandomization](https://mlsanity.com/paper/2610.14942/) (2026)

## Abstract (arXiv)

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.