AutoData: 预训练数据选择的智能体搜索

# AutoData：面向预训练数据选择的智能体搜索
来源：https://arxiv.org/html/2609.19754  
Dhruv Srikanth 单位：Weco AI  
Bingchen Zhao 单位：Weco AI  
Zhengyao Jiang 单位：Weco AI  
Yuxiang Wu 通讯作者：y\.meng@uva\.nl, yuxiang@weco\.ai 单位：Weco AI

###### 摘要
近期，LLM智能体在通过执行反馈编辑模型和训练代码以自动化机器学习工程方面展现出潜力。然而，数据在很大程度上仍处于这种智能体优化循环之外。我们将预训练数据选择构建为基于文档特征（即词汇统计、类别标签和困惑度）的启发式工程问题。我们推出AutoData，一个直接搜索可执行选择算法的智能体。与先前优化固定领域集合权重的数据混合方法不同，AutoData在更丰富的程序空间中搜索评分、分层和随机选择规则，通过使用代理模型的验证反馈迭代优化算法，自动发现特征交互。在一夜的搜索中，AutoData发现了一种选择算法，其性能优于现有手工设计的精选流程。尽管仅在小型代理模型上进行搜索，所发现的方案能够迁移到更大规模，并改善下游指标CORE。这些结果表明，数据工程可以被视为一个智能体机器学习问题，将自主研究从模型和训练代码优化扩展到数据层面。

## 1 引言
大型语言模型的性能取决于模型架构和数据。在最近的NanoChat排行榜（Karpathy, 2025 (https://arxiv.org/html/2609.19754#bib.bib7)）中，训练方案已由自动研究智能体（Karpathy, 2026 (https://arxiv.org/html/2609.19754#bib.bib8)）发现，但这些智能体将训练数据视为静态的。这忽略了一个首要杠杆，即数据。最近NanoChat的一项结果显示，仅将训练数据从FineWeb-Edu（Penedo et al., 2024 (https://arxiv.org/html/2609.19754#bib.bib1)）替换为NVIDIA-ClimbMix（Diao et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib27)），就能将GPT-2的墙钟训练时间减少27%，在此规模上比大多数架构级改进带来的收益更大。本文中，我们通过将智能体搜索扩展到数据选择来弥合这一差距。现有的手工设计数据精选方法结合了去重、质量分类、基于困惑度的过滤、重要性采样和领域混合优化（Xie et al., 2023b (https://arxiv.org/html/2609.19754#bib.bib2); Li et al., 2024a (https://arxiv.org/html/2609.19754#bib.bib26); Penedo et al., 2024 (https://arxiv.org/html/2609.19754#bib.bib1); Ankner et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib4); Thrush et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib5)）。然而，手工设计的方法通常依赖于这些启发式方法的固定组合，从而限制了设计空间和迭代速度。我们推出AutoData，一个用于数据选择的智能体搜索框架。我们将数据选择构建为一个启发式工程问题，其中每个文档由一组特征表征，即词汇统计、类别标签、困惑度和LLM标注的质量信号。AutoData建立在AIDE风格（Jiang et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib6)）的LLM智能体之上，该智能体在这些特征上搜索选择算法。在每一步中，智能体提出一个可执行的选择策略，在所选子集上训练一个小的代理模型，观察验证反馈，并在下一次迭代中优化该策略。然后，AutoData返回具有最佳代理验证性能的选择算法。我们的结果表明，智能体搜索是预训练数据工程的一个有效且实用的范式。在一个小型GPT-2代理模型（125M）上进行200个搜索步骤后，AutoData发现的方案在两个搜索目标上（包括验证每字节比特数val-bpb和下游CORE准确率）均优于DCLM（Li et al., 2024a (https://arxiv.org/html/2609.19754#bib.bib26)）、困惑度过滤、RegMix（Liu et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib25)）和默认的ClimbMix排序（Diao et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib27)）。无需重新调优，这些发现的方案可以跨模型规模迁移。AutoData在125M到897M的范围内取得了最佳的val-bpb，并且在统计上显著优于所有基线。对所发现方案的分析表明，它们超越了单特征排序和阈值过滤。相反，AutoData发现了组合评分规则与保持多样性的选择机制相结合的策略，在保持数据分布广泛覆盖的同时，更倾向于高质量的文档。总体而言，AutoData表明数据工程是自主AI研究的自然下一个前沿。AutoData没有将精选视为固定的、手工设计的预处理流程，而是使数据选择变得可搜索：智能体如何使用直接验证反馈来发现评分文档、组合信号以及保持多样性的方法。这重新将预训练数据精选构建为一个可执行方案上的优化问题，将自主AI研究扩展到塑造模型学习的数据层面。

图1：AutoData概述。一个LLM智能体在特征标注的文档池上提出数据选择函数。每个选定的子集用于预训练一个代理语言模型，其产生的验证分数为后续搜索提供反馈。

## 2 背景
### 2.1 NanoChat任务
我们使用NanoChat速通任务作为语言模型预训练的实验平台（Karpathy, 2025 (https://arxiv.org/html/2609.19754#bib.bib7)）。NanoChat提供了一个轻量级的GPT风格预训练设置，这种小规模环境使得迭代训练实验变得可行。此外，这个速通环境已经产生了超越排行榜本身的见解。例如，源于nanoGPT风格速通实验的Muon（Jordan et al., 2024 (https://arxiv.org/html/2609.19754#bib.bib13)）最近挑战了AdamW（Loshchilov and Hutter, 2019 (https://arxiv.org/html/2609.19754#bib.bib14)）在大规模语言模型训练中的主导地位。排行榜已经表明数据质量是训练效率的一个首要因素。仅将训练语料库从FineWeb-Edu（Penedo et al., 2024 (https://arxiv.org/html/2609.19754#bib.bib1)）切换到ClimbMix（Diao et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib27)），就能将达到CORE阈值所需的时间减少大约27%。然而，当前的速通仅使用了ClimbMix的一小部分——大约是完整语料库的前2%。这引发了一个自然问题：在相同的计算预算下，我们是否可以从完整的ClimbMix池中选择一个更有效的2%子集。

### 2.2 基础语料库：NVIDIA ClimbMix
我们在ClimbMix（Diao et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib27)）上进行数据选择实验，这是NVIDIA发布的一个400B标记的英语预训练语料库。ClimbMix是通过一个手工设计的流程产生的，该流程在Common Crawl上使用文档聚类、参考模型评分和聚类级重新加权进行精选。我们使用完整的ClimbMix作为源语料库，研究LLM生成的选择策略能否识别出比NanoChat中使用的默认策略更优的2%紧凑子集。

## 3 AutoData
AutoData建立在AIDE框架（Jiang et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib6)）之上，这是一个由LLM驱动的代码优化代理，它使用执行反馈来改进候选程序。我们通过优化选择算法（以代码形式）来适应预训练数据选择任务。如图1（https://arxiv.org/html/2609.19754#S1.F1）所示，LLM智能体在给定的数据池上提出数据选择方法，执行选择函数以构建训练数据，并通过训练一个小型代理语言模型来评估选择方法。评估指标作为反馈提供给LLM智能体，引导后续搜索朝向更有效的数据选择方法。

### 3.1 问题表述
我们将预训练数据选择构建为在选择算法空间上的搜索。给定候选文档池$\mathcal{D}=\{d_{i}\}_{i=1}^N$，文档级特征$\{\mathbf{x}_{i}\}_{i=1}^N$，以及固定预算$B$，选择算法是一个函数$f:\big(\mathcal{D},\,\{\mathbf{x}_{i}\}_{i=1}^N,\,B\big)\rightarrow\mathcal{S},\quad \mathcal{S}\subset\mathcal{D},\ |\mathcal{S}|=B$，它将池、特征库和预算映射为一个预算约束子集。设$\mathcal{A}_{\theta}(\mathcal{S})$表示在$\mathcal{S}$上使用固定超参数$\theta$预训练得到的模型，$J(\cdot)$是一个评估函数（例如，验证每字节比特数或下游任务准确率）。AutoData在选择算法空间$\mathcal{F}$中搜索以最大化经验训练性能：$f^{*}=\arg\max_{f\in\mathcal{F}}\;J\!\left(\mathcal{A}_{\theta}\big(f(\mathcal{D},\{\mathbf{x}_{i}\},B)\big)\right)$。

### 3.2 特征库
为了支持组合式选择，我们沿着四个互补维度为智能体提供文档级特征：词汇统计、类别标签、参考模型困惑度和LLM生成的内容标注。这些特征从多个角度刻画文档，捕捉表面属性、主题多样性、由参考模型估计的学习难度以及内容的事实性。
##### 词汇特征。对于每个文档，我们计算其标记长度（通过nanochat分词器）以及$n \in \{1,\ldots,5\}$的不同n-gram比率。文档长度作为粗略的质量信号：非常短的文档通常是片段，而非常长的文档往往包含混合或离题内容。不同的n-gram比率捕捉重复性和模板化文本。低值通常表示样板文本或重复措辞，而短文本上的高值通常标记为损坏或随机字符串。
##### 类别特征。我们使用WebOrganizer分类器为每个文档分配一个主题标签和一个格式标签。在ClimbMix池中，主题和格式共同产生$24 \times 24=576$个联合类别，涵盖领域（例如，新闻、科学、金融）和文档类型（例如，文章、FAQ、列表文章、广告）。
##### 困惑度特征。我们在一个小型参考模型（Qwen2.5-0.5B-Base）下计算每个文档的每字节比特数。基于先前使用小型参考模型进行基于困惑度剪枝的工作（Ankner et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib4)），此信号捕捉相对于参考分布的文档难度：低困惑度标记易于预测的文本，而高困惑度标记噪声内容。
##### LLM标注特征。对于候选池的一个子集，我们使用来自Gemini-3-Flash的LLM生成标注。这些标注提取三个文档级计数：
- •n_factual：错误事实主张的数量；
- •n_rsteps：明确的推理步骤数量；
- •n_rerrors：无效推理步骤的数量。
我们旨在研究内容级信号是否在廉价的词汇、类别和困惑度特征之外提供补充信息。标注指南见附录A.3（https://arxiv.org/html/2609.19754#A1.SS3）。

### 3.3 自演化循环
在每次迭代中，LLM智能体根据先前评估候选及其分数的摘要，提出一个候选选择函数$f \in \mathcal{F}$。提出的程序分三个阶段执行：(i) 它从文档池构建训练子集$\mathcal{S}=f(\mathcal{D},\{\mathbf{x}_{i}\},B)$；(ii) 使用生成的子集在固定超参数$\theta$下预训练一个代理语言模型；(iii) 在保留的验证集上，使用评估函数$J(\cdot)$评估代理模型。所得分数作为反馈返回给智能体，连同先前轨迹的摘要记录，为下一次提议提供信息。通过这个自演化循环，AutoData使用经验的小规模训练性能作为目标信号来优化预训练数据选择。我们实验中使用的具体代理架构、验证集和评估指标在第4节（https://arxiv.org/html/2609.19754#S4）中描述。

## 4 实验
### 4.1 实验设置
##### 数据池。我们将AutoData应用于NVIDIA的完整ClimbMix语料库（Diao et al., 2025 (https://arxiv.org/html/2609.19754#bib.bib27)），这是一个400B标记的英语预训练语料库。该语料库包含6,542个训练分片，覆盖$N=553,155,584$篇文档。我们保留一个分片作为保留验证集，并将其从所有选择池中排除。
##### 选择预算。AutoData从ClimbMix池中选择$B=14,374,266$篇文档的子集，对应于语料库的约2.6%。此预算与标准NanoChat预训练任务中使用的数据规模相匹配。因此，所有方法都在相同的文档级选择预算下进行比较。
##### 代理模型。对于每个候选策略，在完整的$B$文档子集上训练代理模型计算成本高昂。我们改为在从其选定池中均匀采样的880,000文档子集上评估每个策略。我们训练一个深度为8的GPT-2模型，目标参数-数据比=10，对应于大约0.42B训练标记。每次运行在单个H100 GPU上大约需要10分钟。
(a) val-bpb上的优化信号。
(b) CORE上的优化信号。
图2：AutoData在基础LLM（GPT-5.5、Claude-Opus-4.7和Gemini-3-Pro-Preview）上在完整ClimbMix池上的搜索轨迹。每一步提出一个新的数据选择方案，通过训练一个代理GPT-2（深度=8）模型进行评估。
##### 子采样协议。为了减少单次小规模运行的高随机种子方差，我们通过训练四个代理模型来评估每个选择策略，使用相同的选择子样本，仅改变代理训练种子。这四个训练在4×H100 GPU上并行运行。固定子样本并仅改变训练种子可以隔离训练引起的变异，从而获得每个选择器性能的稳定估计。
##### 评估信号。AutoData使用代理模型性能作为搜索反馈。我们考虑两种反馈信号：验证每字节比特数（val-bpb）和下游CORE。对于每个信号，我们使用该信号作为优化目标运行一次单独的AutoData搜索。然后，我们报告每个发现方案的这两个指标：优化指标衡量直接改进，而非优化指标则测试该方案是否超越搜索期间使用的反馈信号进行泛化。val-bpb在保留的ClimbMix分片上测量，值越低表示语言建模性能越好。遵循Li et al.（2024a）(https://arxiv.org/html/2609.19754#bib.bib26)，CORE被用作下游评估指标。它定义为