# DATAMASK **Repository Path**: ByteDance-Seed/DATAMASK ## Basic Information - **Project Name**: DATAMASK - **Description**: Joint Selection for Large-Scale Pre-Training Data via Policy Gradient-based Mask Learning - **Primary Language**: Unknown - **License**: Apache-2.0 - **Default Branch**: main - **Homepage**: None - **GVP Project**: No ## Statistics - **Stars**: 0 - **Forks**: 0 - **Created**: 2026-09-15 - **Last Updated**: 2026-09-17 ## Categories & Tags **Categories**: Uncategorized **Tags**: None ## README # DATAMASK 中文 | [English README](ReadMe.md)

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

| 📑 Paper | 🐱 FineWeb-Mask Data |

**Ziqing Fan**1,2 , Yuqiao Xian1,\* , Yan Sun3, Li Shen4 _1 ByteDance Seed, 2 Shanghai Jiao Tong University, 3 University of Sydney, 4 Sun Yat-sen University Shenzhen Campus._
## 📖 Overview
DATAMASK PIPELINE
**Motivation** 在本研究中，我们重新审视了基于指标的数据筛选方法，并观察到：基于质量指标（如 FineWeb-Edu、Ultra-FineWeb 和 FineWeb-DCLM）进行样本筛选，在长期预训练中表现出严重的收益递减；而基于多样性指标（如 FineWeb-Semdedup）的筛选则剔除了过多有价值的高质量样本。这两种情况均限制了预训练大语言模型（LLMs）的能力。 **Method: DATAMASK** 为了解决这一问题，如上图流程所示，我们提出了一种用于大规模预训练数据筛选的新颖且高效的优化框架，该框架能够在一个统一的流程中同时优化多种类型的指标。该方法将筛选过程建模为掩码学习（Mask Learning）问题，具体步骤包括：数据掩码的迭代采样、基于采样掩码与预定义目标计算策略梯度，以及更新掩码采样的概率参数 (logits)。 **Results: FineWeb-Mask** 通过基于策略梯度的优化及多种加速改进策略，DATAMASK 与贪婪算法（基于 DiSF 算法估算）相比，将筛选时间显著缩短了 98.9%，使我们能够在万亿级 Token 规模下探索联合学习。利用 DATAMASK，我们从包含 15 万亿 Token 的 FineWeb 数据集中筛选出了约 10% 的子集，并将其命名为 FineWeb-Mask。该子集在数千亿 Token 的预训练后取得了显著的性能提升，从而证明了该方法的有效性。 ## 🚀 Quick Start 0. 克隆仓库: ```bash git clone https://github.com/ByteDance-Seed/DATAMASK.git cd DATAMASK ``` 1. 准备数据 基于我们的源代码，您的数据应为 .parquet 格式。在该 parquet 文件中，您需要提供 idx、scores 和 text 字段。idx（例如 chunk_id）用于在大规模分布式筛选过程中标识数据集中被选中的样本。在我们的论文中，数据结构如下所示： | chunk_id | quality_score | feature_arr | |--------|--------|--------| | 0 | 8 | size(768) | | 1 | 4 | size(768) | | ... | ... | ... | 2. 确定您的评分指标 在代码中，我们提供了三种多样性评分的实现，并将其与质量评分结合以执行优化。您可以通过修改 utils/utils.py 中的函数，来自定义指标的数量及其组合方式。 3. 超参数 在 DATAMASK 优化中，我们引入了多个超参数。我们在论文中针对这些参数提供了详尽的消融实验研究。更多细节请参考论文。 * n_epochs: 更新步数。 * partial: 用于批量训练的样本占总样本的比例。 * algorithm: 多样性评分类型，选项包括 ["DiSF", "Facility", "Pair_simi"]。 * max_lr 和 min_lr: 初始学习率，以及达到 n_epochs 时的最终学习率。我们使用的是线性调度器。 * select_ratio: 筛选比例。 * lamb: 用于平衡质量与多样性的 lambda 系数。 * n_rollout: 每个 epoch 的 rollout（推演/采样）次数。 * init: logit 初始化策略，选项包括 same 和 quality。 4. 快速上手 在准备好数据、优化指标和超参数后，您可以尝试运行以下代码进行快速测试： ```bash python3 train_mask.py \ --input your_input_path.parquet \ --output your_output_path.parquet \ --device cuda:0 \ --n_epochs 5000 \ --algorithm DiSF \ --partial 0.1 \ --max_lr 10 \ --min_lr 1 \ --n_rollout 128 \ --select_ratio 0.3 \ --init quality \ --lamb 0.5 ; ``` ## 🔎 Text Feature Visualizations
Feature
为了直观呈现这一困境，我们在 FineWeb 的随机子集上利用 t-SNE 对文本嵌入进行了可视化。其中，白色、浅蓝色和深蓝色数据点分别对应多样性最高的样本、质量最高的样本，以及由兼顾高多样性与高质量的算法所筛选出的样本。浅蓝色点呈现出更紧密的聚类状态。除我们的算法外，深蓝色点在其他算法的结果中均较为稀疏。这表明，基于质量评分筛选样本（即深蓝色和紫色数据点）会导致比原始数据分布更紧密的聚类，从而意味着更高的语义冗余度以及信息多样性的降低。 ## 🌟 Optimization Curves and One Optimization Ablation 在此，我们展示了优化曲线以及一项针对 rollout 次数 G 的优化消融实验。如图所示，我们在调整 G 值的过程中，记录了优化后的 Facility Location 值随计算时间的变化情况。结果表明，G 值过小会导致训练发散，而 G 值过大则会带来过高的计算开销。基于三种多样性指标进行调优后，我们推荐 G = 128 或 256，这是在所有测试用例中既能保持稳定又能取得近乎最优结果的最小值。关于更多的消融实验细节，请参阅论文。
Optimization curves and one optimization ablation in terms of rollout number G
## 📈 Detailed Performance 接下来，我们展示了 1.5B Dense 模型和 7B MoE 模型在预训练期间各任务的详细性能表现。上半部分基于稠密模型，而下半部分基于 MoE 模型。 | **1.5B Dense** | **RACE-H** | **RACE-M** | **HellaSwag** | **NQ** | **OBQA** | **KQAPro** | **MMLU** | **TrivalQA** | **ARC-Challenge** | **SIQA** | **PIQA** | **WinoGrande** | | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | | FineWeb | 40.8 | 51.2 | 58.9 | 11.1 | 48.6 | 45.1 | 33.9 | 35.2 | 31.9 | 48.9 | 75.1 | 58.5 | | FineWeb-Semdedup | 40.5 | 51.7 | 57.2 | 11.0 | 45.0 | 43.5 | 33.4 | 30.2 | 31.6 | 48.7 | 74.9 | 57.9 | | FineWeb-Edu | 42.3 | 51.4 | 57.6 | 12.1 | 53.6 | 42.5 | **37.8** | 37.8 | **44.5** | 49.2 | 74.2 | 59.0 | | UltraFineWeb-en | 41.8 | 53.0 | 57.8 | 10.9 | 50.6 | 42.2 | 37.2 | 30.6 | 44.2 | 48.9 | 75.9 | 57.1 | | FineWebPro | 43.1 | 52.2 | 61.3 | 12.4 | **51.5** | 42.0 | 36.2 | 38.6 | 43.0 | 50.1 | 75.2 | 61.0 | | FineWeb-DCLM | 43.6 | 52.9 | **61.4** | 11.1 | 48.2 | 43.4 | 34.8 | 37.8 | 40.7 | 50.4 | **76.2** | **61.4** | | **FineWeb-Mask (Ours)** | **43.8** | **53.7** | 56.4 | **14.1** | 51.4 | **47.0** | 36.5 | **47.3** | 40.9 | **51.4** | 74.4 | 59.8 | | | | | | | | | | | | | | | | **7B MoE** | **RACE-H** | **RACE-M** | **HellaSwag** | **NQ** | **OBQA** | **KQAPro** | **MMLU** | **TrivalQA** | **ARC-C** | **SIQA** | **PIQA** | **Wino** | | FineWeb | 42.1 | 54.1 | **69.9** | 17.9 | 53.6 | 49.8 | 35.8 | 53.8 | 40.7 | 52.8 | **78.6** | 63.7 | | FineWeb-Semdedup | 41.4 | 55.3 | 67.7 | 17.0 | 53.0 | 49.3 | 35.6 | 49.1 | 38.4 | 49.4 | 77.4 | 65.1 | | FineWeb-Edu | 42.8 | **56.5** | 65.8 | 16.6 | 53.8 | 48.5 | 41.3 | 50.4 | **50.0** | 50.4 | 77.1 | 64.5 | | UltraFineWeb-en | 42.4 | 54.0 | 64.9 | 13.9 | **55.4** | 42.0 | **41.4** | 38.6 | 49.3 | 49.7 | 78.3 | 60.9 | | FineWebPro | 42.8 | 55.8 | 68.9 | 15.5 | 55.2 | 46.3 | 40.0 | 49.8 | 47.9 | **52.8** | 77.9 | 64.3 | | FineWeb-DCLM | **43.2** | 54.8 | 69.3 | 18.7 | 54.0 | 49.8 | 40.4 | 54.6 | 47.0 | 52.1 | 78.5 | 64.0 | | **FineWeb-Mask (Ours)** | 42.8 | 55.4 | 66.1 | **19.5** | 55.0 | **51.4** | 39.5 | **61.7** | 45.9 | 51.2 | 77.5 | **65.1** | ## **📚** Citation ```bibtex @misc{fan2025jointselectionlargescalepretraining, title={Joint Selection for Large-Scale Pre-Training Data via Policy Gradient-based Mask Learning}, author={Ziqing Fan and Yuqiao Xian and Yan Sun and Li Shen}, year={2025}, eprint={2512.24265}, archivePrefix={arXiv}, primaryClass={cs.CL}, url={https://arxiv.org/abs/2512.24265}, } ``` ## **📧** Contact - Ziqing Fan: fanziqing.knight@bytedance.com ; zqfan_knight@sjtu.edu.cn - Yuqiao Xian (corresponding author): xianyuqiao.eric@bytedance.com ## **📄** License [Apache 2.0](./LICENSE.md)