← 返回 2026-02-11

OPUS:面向大语言模型预训练中每轮迭代高效且有原则的数据选择 OPUS: Towards Efficient and Principled Data Selection in Large Language Model Pre-training in Every Iteration

Shaobo Wang, Xuan Ouyang, Tianyi Xu, Yuzheng Hu, Jialin Liu, Guo Chen, Tianyu Zhang, Junhao Zheng, Kexin Yang, Xingzhang Ren, Dayiheng Liu, Linfeng Zhang 📅 2026-02-05 👍 355 2026-07-13 08:35
LLM训练 优化器感知 动态采样 数据选择 预训练优化

在优化器诱导的更新空间中定义数据效用,实现高效动态选择

前置知识

动态数据选择(Dynamic Data Selection)

在大语言模型预训练过程中,根据模型当前状态和训练动态,实时选择最有价值的训练样本。与静态过滤(训练前一次性筛选)不同,动态方法在每个训练步骤根据模型的演化需求调整选择策略。核心思想是:模型在不同训练阶段需要不同类型的数据,早期可能需要基础语法数据,后期需要复杂推理数据。动态选择通过评估每个候选样本对模型改进的贡献来实现这一目标。

本文提出的OPUS框架正是动态数据选择方法,理解这一概念是理解论文核心贡献的基础。

优化器预条件(Optimizer Preconditioning)

现代优化器(如AdamW、Muon)通过维护梯度的移动平均和自适应学习率来变换原始梯度方向。AdamW使用二阶矩估计对梯度进行逐元素缩放,Muon则通过Newton-Schulz正交化对矩阵参数进行变换。这种预条件操作改变了参数更新的实际方向,使得优化器诱导的更新空间与原始梯度空间存在本质差异。

本文的核心创新在于将数据选择从原始梯度空间转移到优化器诱导的更新空间,理解预条件机制是理解这一创新的关键。

Ghost技术(Ghost Technique)

一种高效计算逐样本梯度的技术,利用线性层梯度的秩-1外积结构。对于线性层权重W,样本z的梯度可以分解为输入激活向量a与输出梯度向量b的外积:∇W L(z) = a ⊗ b。这允许在标准前向/反向传播过程中收集激活和梯度统计量,无需显式构造高维梯度矩阵,大幅降低内存和计算开销。

OPUS采用Ghost技术实现高效的数据效用评分,是方法可扩展性的关键技术支撑。

CountSketch投影

一种随机投影技术,将高维向量压缩到低维空间同时保持内积的无偏估计。通过哈希函数和符号函数将d维向量映射到m维(m远小于d),满足E[⟨Π(g₁), Π(g₂)⟩] = ⟨g₁, g₂⟩。这使得在压缩空间中计算的相似度与原始空间一致,实现高效的梯度相似度计算。

OPUS使用CountSketch将效用计算投影到低维空间,是实现4.7%低开销的核心技术。

Boltzmann采样

一种软采样策略,根据效用分数的概率分布进行选择,而非贪心地选择top-K。候选样本z被选中的概率与其效用分数U_z成指数关系:p(z) ∝ exp(U_z/τ),其中τ是温度参数。这确保高效用样本被优先选择,同时保持互补样本的非零选择概率,避免过度拟合到代理信号的噪声。

OPUS采用Boltzmann采样替代传统贪心选择,是保证数据多样性和训练稳定性的关键设计。

研究动机

随着高质量公开文本接近枯竭(即'数据墙'现象),大语言模型预训练正从'更多token'转向'更好token'。然而现有数据选择方法存在根本性缺陷:静态过滤方法(如FineWeb-Edu分类器、DCLM质量分类器)依赖固定的、与训练无关的启发式规则,假设样本效用在训练过程中保持恒定,无法适应模型演化的需求。更关键的是,现有动态选择方法(如GREATS)在原始梯度空间中评估候选样本,隐式假设SGD优化动态,但现代LLM训练依赖AdamW和Muon等自适应优化器,它们通过动量和自适应预条件变换梯度方向。这种假设与实际优化器几何的错配导致选择的样本并非真正最优,如论文Figure 2所示,原始梯度方法与优化器实际更新路径存在显著偏差。

本文的目标是本文旨在设计一个动态数据选择框架,能够:(1)在优化器诱导的更新空间中定义数据效用,使评分与实际训练轨迹对齐;(2)保持计算效率,避免逐样本梯度物化的高昂开销;(3)支持多样性和稳定性,防止选择退化。具体目标是在30B token预算内达到或超越200B token训练的性能,实现约8倍的计算效率提升。

与已有工作不同的是,本文的独特切入点在于认识到数据选择本质上是优化过程的一部分,而非独立的预处理步骤。已有工作忽视了一个关键事实:优化器的预条件操作从根本上改变了更新方向,因此在原始梯度空间中定义的效用与实际训练目标存在几何错配。OPUS通过推导AdamW和Muon的线性化预条件算子,将效用定义在优化器实际使用的参数更新空间中,首次实现了评分空间与训练空间的严格对齐。这种'优化器感知'的视角是本文区别于所有先前工作的本质创新。

核心方法

OPUS的直觉可以用一个比喻来理解:想象你在指导一个学生学习,传统方法根据学生的'原始反应'(梯度)来判断哪些题目有价值,但学生实际是通过'眼镜'(优化器)看世界的,不同眼镜会改变视野。OPUS的核心思想是透过学生实际使用的眼镜来评估题目的价值。技术路线上,OPUS在每个训练步骤执行以下流程:首先从优化器状态构造预条件算子P_t;然后通过检索构建与目标基准分布对齐的代理池D_proxy;接着使用Ghost技术和CountSketch高效计算每个候选样本在优化器诱导空间中的效用分数;最后通过Boltzmann采样选择多样化的更新批次。整个流程如论文Figure 3所示,形成一个完整的'代理池构建→原始梯度计算→高效梯度投影→迭代效用估计'流水线。

OPUS的核心创新在于首次将动态数据选择的效用函数定义在优化器诱导的更新空间中,而非原始梯度空间。对于AdamW优化器,通过冻结RMS几何并推导线性化形式,得到近似对角预条件算子P_t^{AdamW} ≈ C_t · Diag(1/√(v̂_{t-1} + ε));对于Muon优化器,通过冻结Newton-Schulz算子并线性化,得到密集层级预条件算子P_t^{Muon} = κ_t S_t。效用函数的核心形式为U_z ≈ η⟨u_z, g_{proxy}⟩ - η²⟨u_z, G⟩,其中第一项衡量样本与代理方向的对齐程度(alignment),第二项惩罚与已选样本的冗余(redundancy),u_z = P_t ∇L(z; θ_t)是优化器诱导的有效更新方向。这种公式化确保选择的样本不仅减少损失,而且特别与基准相关的优化景观子空间对齐。

方法步骤详情

OPUS的完整步骤如下:(1)批次采样:从训练流D_tr读取候选缓冲区B_t = {z_1, ..., z_N},大小为N=32个序列。(2)预条件算子构建:从优化器O在步骤t的状态构造P_t,对于AdamW使用对角近似,对于Muon使用线性化Newton-Schulz算子。(3)代理特征生成:从代理池D_proxy采样K_proxy个样本,通过Ghost因子计算每层代理草图ψ_{t,proxy}^{(r)}。(4)候选特征生成:对每个候选z,隐式计算预条件草图特征ϕ_t^{(r)}(z) = Π_r(P_t^{(r)} · a_z^{(r)} ⊗ b_z^{(r)})。(5)软采样循环:目标批次大小K=⌊ρN⌋(ρ=0.5),初始化已选集和历史Φ。对每个选择位置,计算所有未选候选的效用U_z,通过Softmax概率p(z*) ∝ exp(U_z/τ)采样,更新历史Φ以实现冗余惩罚。(6)模型更新:使用选定批次B̂_t训练模型。

技术新颖性

OPUS的技术新颖性体现在多个层面:首先,它是首个将数据选择效用与优化器几何严格对齐的框架,通过推导AdamW和Muon的闭式线性化预条件算子,实现了评分空间与训练空间的一致性。其次,BENCH-PROXY构建方法创新性地使用语义检索从预训练语料中选取与目标基准分布对齐的代理样本,既保持了梯度估计的稳定性又避免了分布偏移。第三,Ghost技术与CountSketch的结合实现了O(d_in + d_out)的投影成本(对于AdamW),而非朴素实现的O(d_in · d_out),使方法可扩展到LLM规模。最后,Boltzmann软采样替代贪心top-K,在保持高效用样本优先选择的同时维护了数据多样性,避免了估计噪声下的过拟合。

OPUS流水线概述
Figure 3: OPUS流水线概述

实验结果

OPUS在多样化的语料库、质量层级、优化器和模型规模上展现了强大且一致的性能提升。在FineWeb数据集上的从头预训练实验中,GPT-2 XL使用Muon优化器,OPUS在30B token预算内达到41.75的平均分数,超越Random选择(40.29)1.46个百分点,甚至接近60B token Random训练的41.29。在GPT-2 Large上,OPUS在AdamW设置下达到41.43,显著超越所有基线。最关键的是,OPUS实现了8倍的计算效率提升,仅需30B token即可达到200B token训练的效果。在FineWeb-Edu的严格评估中,尽管OPUS仅从低质量子集(score=3)动态选择,而基线使用高质量分区(score 4+5),GPT-2 XL仍达到44.99的最优平均分,超越所有在更优数据上训练的基线。在域外泛化评估中,OPUS在BBH、RACE等不在代理集中的基准上同样表现最佳(平均40.07 vs 39.42 for GREATS),证明其反映的是更通用的训练信号质量而非窄化的过拟合。领域困惑度分析显示,OPUS在10个领域的平均困惑度最低(GPT-2 XL: 3.26 vs 3.28 for PPL),表明其选择策略广泛提升了知识压缩能力。在持续预训练场景,Qwen3-8B-Base在SciencePedia上仅用0.5B token即超越3B token全量训练的性能,实现6倍数据效率增益。

参数与优化器分配
Table 1: 参数与优化器分配
基准评估配置
Table 2: 基准评估配置
FineWeb数据集30B token训练后的评估结果
Table 3: FineWeb数据集30B token训练后的评估结果
FineWeb-Edu数据集30B token训练评估
Table 4: FineWeb-Edu数据集30B token训练评估
领域特定困惑度分析
Table 6: 领域特定困惑度分析
采样和验证策略消融研究
Table 7: 采样和验证策略消融研究
超参数敏感性分析
Table 8: 超参数敏感性分析
OPUS在10个基准上平均超越随机选择2.2%准确率,在GPT-XL上实现8倍计算减少
Figure 1: OPUS在10个基准上平均超越随机选择2.2%准确率,在GPT-XL上实现8倍计算减少
FineWeb-Edu数据集上的验证损失曲线
Figure 4: FineWeb-Edu数据集上的验证损失曲线
持续预训练的领域分解结果
Figure 5: 持续预训练的领域分解结果
SciencePedia持续预训练结果
Figure 6: SciencePedia持续预训练结果
效率和计算成本分析
Figure 7: 效率和计算成本分析
查看结构化数据
任务指标本文基线提升
FineWeb预训练(GPT-2 XL + Muon) 10基准平均准确率 41.75 40.29 (Random) +1.46
FineWeb预训练(GPT-2 Large + AdamW) 10基准平均准确率 41.43 40.38 (UltraFineweb) +1.05
FineWeb-Edu(GPT-2 XL,仅score=3数据) 10基准平均准确率 44.99 42.59 (Random 60B) +2.40
域外泛化(GPT-2 XL) 6基准平均准确率 40.07 39.42 (GREATS) +0.65
持续预训练(Qwen3-8B) SciAssess平均 0.5B token超越基线 3B token全量训练 6×效率

局限与改进

尽管OPUS取得了显著成果,仍存在若干局限性需要考虑。首先,BENCH-PROXY构建依赖于目标基准的验证集,这意味着需要预先知道下游评估分布,在开放式预训练场景中可能不适用。其次,当前实验主要在GPT-2规模(774M-1.5B参数)和Qwen3-8B上验证,在更大规模模型(如70B+)上的效果尚待验证,虽然方法设计上具有可扩展性。第三,代理池的语义检索使用Arctic-Embed-L v2模型,其质量可能影响代理方向的准确性,在多语言或特定领域语料中可能需要调整。第四,超参数敏感性分析显示,温度τ=0.9、缓冲区大小64、投影维度8192是最优配置,但这些设置可能因数据分布和模型规模而异,需要额外调优。最后,OPUS的4.7%开销虽然是轻量级的,但在极大规模训练中(数万亿token),累积的额外计算仍然可观。

独立分析的弱点

通过独立分析,我识别出OPUS的几个潜在弱点及改进方向:(1)代理方向的单峰性:当前BENCH-PROXY从单一基准验证集构建代理池,可能无法捕获多维度能力需求。改进方向是构建多代理池或使用混合代理,同时对齐知识、推理、代码等多个目标分布。(2)线性化近似的精度损失:冻结优化器状态的线性化假设在训练早期(状态变化快)可能不够准确。可以探索自适应调整冻结频率或使用二阶修正项。(3)序列级选择粒度:OPUS在序列级别进行选择,但序列内token的价值可能高度异质。未来可以探索更细粒度的token级或chunk级选择。(4)缺乏理论收敛保证:虽然论文提供了启发式推导,但缺乏严格的收敛性分析,特别是在非凸优化设置下的理论保证。(5)对初始化的依赖:Ghost技术和CountSketch的投影质量依赖于初始化的哈希函数和种子,可能存在方差。

未来方向

基于本文成果,未来研究方向包括:(1)扩展到数据混合优化:当前OPUS针对单一数据源的选择,自然延伸是将其优化器感知思想应用于多源数据混合比例优化。(2)与课程学习结合:利用OPUS的效用分数设计自适应课程,根据模型能力变化调整数据难度分布。(3)异构硬件适配:在不同硬件配置(如GPU/TPU混合集群)上优化投影和采样的并行策略。(4)长期训练稳定性:研究OPUS在数百B乃至数T token训练中的稳定性,特别是代理方向漂移的累积效应。(5)多模态扩展:将优化器感知选择思想扩展到图文、视频等多模态预训练场景。(6)元学习框架:将选择策略本身作为可学习组件,通过元学习自动适应不同训练阶段。

复现评估

论文在复现性方面提供了较好的支持。代码层面,虽然论文未明确说明开源,但详细描述了所有技术细节(Algorithm 1)、超参数设置(Section 6.1)和实现细节(如CountSketch种子42、投影维度8192)。数据方面,使用公开可获取的FineWeb、FineWeb-Edu和SciencePedia数据集,代理池构建流程有完整描述。算力方面,实验在8GPU上运行,使用GPT-2规模模型,对于学术实验室是可复现的。但需要注意:(1)Muon优化器的具体实现细节可能因版本而异;(2)Ghost技术的实现需要修改标准训练循环以收集激活/梯度统计量;(3)BENCH-PROXY的语义检索需要Arctic-Embed-L v2模型。总体复现难度中等,预计需要2-3周的工程努力。