编辑：桃子 乔杨

ChatGPT早已成为世界耗能大户：一天用掉超50万度电，相当于1.7万个美国家庭的用电量！

然而，大模型对能源的吞噬，远不仅如此。

国际能源署（IEA）预测，从2022年到2026年，数据中心的用电量将翻一番。

论文地址：https://arxiv.org/pdf/2406.17711

预训练的参考模型，已经学习了什么样的数据是有「优质的」或「有用的」。然后通过模型，来引导数据选择那些精心筛选过的小型数据集。

这一发现揭示了，数据筛选水平可以作为评判Scaling Law的一个新维度。

网友激动表示，「我没想到这么快就会发生。模型能够自主选择训练数据的能力是巨大的，因为它使训练变得显著更容易，你不再需要猜测什么是高质量的训练数据，你有一个能够『理解』什么样的数据对自身学习最有价值的模型」。

前谷歌、苹果软件工程师称赞道，这项研究非常令人印象深刻。

前者成本高昂且难以扩展，后者则有望为多模态LLM实现Scaling Law。

然而，现有方法忽略了一个事实。

如果仅在单个数据点的层面进行筛选，就没有考虑到数据集以及batch的总体组成。毕竟，训练数据是以batch为单位，数据点之间的依赖性不可忽视。

许多计算机视觉的研究都曾表明，hard negatives（表达空间中相近但标签不同的样本）相比可被平凡解的数据簇，能提供更有效的学习信号。

那么如何让模型以batch为单位筛选数据呢？

论文提出的JEST算法正是要解决这个问题，原理很好理解：就是直接从「超级batch」中筛选出「子batch」。

技术介绍

用数学语言来描述这个问题，就是从大小为B的「超级batch」中提取出与学习最相关的子batch ℬ={,∈[1,…,]}⊂，过滤比率可以写作=1−/。

之前的优先采样（prioritized sampling）会使用基于模型的评分函数对每个数据点打分，再按比例采样。JEST则直接对整个子batch评分，再按照batch级别的分数采样。

一种最直观的启发式方法就是在现有模型参数 : hard⁢(ℬ|)=ℓ⁢(ℬ|) 中，直接选择损失值最高的batch，这种方法可被称之为「硬学习」（hard learner）。

这种方法具有丢弃琐碎数据的理想属性，已被证明适用于小型、干净的数据集；然而对于较大、较少管理的数据集往往弊大于利，因为它依旧会采样到噪声数据。

另一种方法常用于多模态，使用具有参数 ∗:^easy⁢(ℬ|∗)=−ℓ⁢(ℬ|∗) 的参考模型为预训练模型采样数据。但作者依旧否定了这个方案，因为它无法直接反映模型当前的状态，可能过度依赖参考模型的选择，而且不易于扩展。

最后，论文选择借鉴ICML 2022年的一篇论文中提到的方法，将上述两方面的评分结合起来：^learn⁢(ℬ|,∗)=hard⁢(ℬ|) ^easy⁢(ℬ|∗)=ℓ⁢(ℬ|)−ℓ⁢(ℬ|∗)，并将这种启发式方法称为「可学习性评分」（learnability score）。

其中，batch上的损失值ℓ⁢(ℬ|)是各数据点之和，使用sigmoid对比损失函数计算（sigmoid-contrastive loss），因为相比softmax对比损失而言，它的扩展性更强。

由于batch上的对比损失可以分解为每个样本的条件损失之和，因此可学习性评分可被分解为单个样本可学习性评分⁢(|,∗,ℬ)之和，写作：

将X_k块添加到B_n中来更新当前采样的batch，直至迭代数n=N时终止。算法的总体流程如下图所示：

实验中发现，使用迭代数N=16且每次迭代时独立采样b/N=2048个样本时，就足以恢复出学习性非常高的batch。

可学习性评分中涉及到使用参考模型为数据点打分，之前的方法惯常使用额外的小型模型，但这会增加每次迭代的计算成本，降低总体FLOP效率增益。

因此论文使用了在线模型近似的方法以及效率较高的FlexiViT架构，只使用降低分辨率的32×32的patch来评估「超级batch」，与全分辨率、patch大小为16×16的方法相比减少了72%的FLOP，以及67%的挂钟时间（wall-clock time）。

此外，论文还提出了进行多分辨率训练的技巧。将每个batch随机分成两半，使用不同分辨率编码后再拼接起来，提升了评分过程和训练的效率。

下图详细描述了全分辨率JEST和多分辨率Flexi-JEST方法的伪代码实现。

在最终性能方面，当过滤90%的数据时，JEST也带来了高达6%的显著提升（图3，中间，蓝色曲线）。

值得注意的是，这种scaling行为这种性能提升在独立样本选择方法中，并没有观察到。（图3，中间，橙色曲线）。

最后，研究者还评估JEST是否也改善了，除可学习性之外的其他优先标准。

图3右侧显示了使用easy-reference优先选择的模型在不同过滤比例下的性能。

与基于可学习性的优先选择一致，JEST仍优于独立样本选择，特别是在高过滤比例下（在这种情况下，独立样本选择导致性能下降）。

优先选择具有最高损失的数据产生了较小的收益，并且随着过滤更多数据而更快地退化（图10）。

由于基于可学习性的JEST产生了最佳的scaling行为，研究人员在后续实验中保留了这一标准。

随着数据batch中被过滤的比例增加，基于可学习性评分的JEST变得更加高效。

然而，评分的成本会带来显著的提升：过滤超级batch 80%的数据会导致每次迭代的浮点运算量是IID训练的4倍，或者在缓存参考模型得分时是2.3倍。

尽管JEST在训练迭代次数方面（以下简称「训练效率」）显著提高了效率，但额外的评分浮点运算降低了其相对于IID基准的计算效率（图1，左vs右）。

因此，作者还研究了一种计算效率更高的变体，称为Flexi-JEST，它使用多分辨率训练和低分辨率评分，将总开销降低到仅比基准高10%（图4，左）。

这些近似方法对性能有什么影响？

正如预期的那样，Flexi-JEST的每次迭代性能相对于JEST有所下降，但仍然比IID有显著的加速（图1，左；图4，中）。

然而，考虑到总浮点运算量的减少，每次迭代性能的下降是非常有利的：最好的Flexi-JEST模型与40B Siglip运行产生相同的平均性能，但浮点运算量减少了9.9倍，比全分辨率JEST少2倍（图1，右；图4，中）。

这些实验表明了多分辨率训练和联合示例选择之间的协同效应，前者为加速后者提供了高效和准确的评分能力。

实验结果，还指出了数据策划策略的帕累托前沿（pareto front）。

如果以计算为代价来最大化训练速度或训练效率，全分辨率JEST方法相对于可比的IID训练运行，可以产生高达13倍的加速。

扩展数据筛选

假设参考模型性能与JEST性能之间的普遍解耦，可能仅仅是由数据筛选所施加的数据集大小限制造成的。

为了理解这种效果，研究人员在WebLI-curated上训练了5个参考模型，同时改变所见的总样本数（从2.5亿到30亿）。

在这种情况下，图5（右）显示了改进的参考模型与更好的JEST预训练之间存在着显著的相关性。

这表明「解耦」现象主要可以归因于参考模型因筛选后数据集大小减少而导致的饱和。

此外，研究人员还注意到，当数据集达到饱和时，图5（右）中的相关性开始崩解，即在10个epoch或者看到10亿个样本之后。

这些结果表明，JEST可能会从进一步扩大参考数据集的数据筛选中获益。

鉴于使用WebLI-curated 对数据进行扩展整理能显著提高参考模型的性能，作者提出了是否有必要在原始WebLI数据集上进行预训练的问题。

然而，在评估参考模型在不同数据集上的性能时，却发现：虽然它在2个下游任务上的性能优于WebLI预训练，但在其他6个任务上的性能，以及平均性能都明显低于WebLI预训练（表 5）。

最后，论文应用JEST 在公开的LAION-2B数据集上进行预训练，删除了其中不安全的图像-文本对，但没有进行其他的预先过滤。

这个数据规模相比的SOTA方法DBP减少了4×，但JEST 依旧远远超过了所有之前的离线数据管理方法。

之前提到过，用于预训练的WebLI-curated是原始数据集WebLI过滤后得到的，以求筛选出高质量的图像-文本对齐的数据。

如表3所示，这种离线数据管理流程对IID（独立同分布）训练方法的性能至关重要，但JEST 则表现出了对预过滤流程的鲁棒性。即使没有过滤，JEST 的性能也没有出现明显下滑，降低了模型对基础数据集的要求。

结论和局限性

总体来说，JEST方法展现出了「数据质量引导」（data quality bootstrapping）方法的巨大潜力，即使用小规模精选数据集来指导对更大的、未经管理的数据集的学习。

最近的研究表明，在下游任务未知时，静态数据集的过滤会限制模型性能。这篇论文的结果则表明，相比单独选择样本的方法，在线构建batch能提高预训练的效率。

无论是使用JEST参考模型对数据集进行预评分，还是通过可学习性评分来根据模型需求进行动态调整，都可以成为通用基础数据集的更有效率的替代方案。

论文的最后，作者也提出了该方法的局限性。虽然JEST同时实现了性能增益和训练成本降低，但依旧依赖于小型、精心管理的参考数据集，它指定了未经管理的更大数据集中优先考虑的分布。

因此，未来的工作可以探索一种方法，从指定的下游任务中如何推断出参考数据集的组成和分布。

参考资料：

https://www.reddit.com/r/singularity/comments/1dw7xnf/google_deepminds_jest_method_can_reduce_ai/

https://decrypt.co/238730/new-ai-training-technique-is-drastically-faster-says-google