解决GPT-4o耗电难题！DeepMind新算法训练效率提升13倍，能耗降低10倍！

news2024/9/20 5:52:23

01 有更好的解决方案吗？

02 从“超级batch”中筛选数据

03 技术介绍

04 实验结果

生成可学习batch

谷歌DeepMind推出的新算法JEST，将LLM训练的迭代次数减少了13倍，计算量降低了10倍，有望重塑AI未来。

GPT-4o早已成为耗能巨头：一天耗电超过50万度，相当于1.7万个美国家庭的用电量！

然而，大模型对能源的消耗不仅如此。

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

随着AI计算需求的增加，还需要用水来冷却计算系统。

研究表明，微软的用水量从2021年到2022年增加了34%，ChatGPT每处理5-50个提示就会消耗接近半升水。

没体验过OpenAI最新版GPT-4o？快戳最详细升级教程，几分钟搞定：
升级ChatGPT-4o Turbo步骤https://www.zhihu.com/pin/1768399982598909952

如何使用WildCard正确方式打开GPT-4o，目前 WildCard 支持的服务非常齐全，可以说是应有尽有！

官网有更详细介绍：WildCard

01 有更好的解决方案吗？

最近，谷歌DeepMind研究团队提出了一种加快AI训练的新方法——多模态对比学习与联合示例选择（JEST），大大减少了所需的计算资源和时间。

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

JEST以13倍更少的迭代次数和10倍更少的计算量，超越了最先进的模型！

预训练的参考模型已经学会了哪些数据是优质或有用的，然后通过模型引导选择这些精心筛选的小型数据集。

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

网友激动地表示，“我没想到这么快就会发生。

模型能够自主选择训练数据的能力是巨大的，因为它使训练变得显著更容易，你不再需要猜测什么是高质量的训练数据，你有一个能够‘理解’什么样的数据对自身学习最有价值的模型”。

02 从“超级batch”中筛选数据

无论是语言、视觉还是多模态模型，数据质量是预训练性能的重要驱动因素。

例如Phi-3、Gemma 2等模型的成功表明，更少、更高质量的数据可能实现更强大的性能。

要筛选出高质量的数据，数据管道的建立变得尤为重要。

现有的方法主要有两种：

1）手动管理

2）基于模型的数据管理，用正在训练模型的特征选择高质量数据。

前者成本高昂且难以扩展，后者则有望为多模态LLM实现Scaling Law。
然而，现有方法忽略了一个事实。

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

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

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

论文提出的JEST算法正是为了解决这个问题，原理很简单：直接从“超级batch”中筛选出“子batch”。

03 技术介绍

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

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

一种直观的启发式方法是直接选择损失值最高的batch，这种方法可称为“硬学习”（hard learner）。

这种方法具有丢弃琐碎数据的理想属性，适用于小型、干净的数据集；但对于较大、较少管理的数据集可能不太适用，因为它仍然会采样到噪声数据。

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

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

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

使用的顺序采样方法则受到了block Gibbs采样的启发。在第n次迭代、对第B_n个batch进行采样时，依据如下概率公式对块{X_k}进行无替换采样：

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

实验发现，在使用迭代数N=16且每次迭代独立采样b/N=2048个样本时，便可恢复出具有高度学习性的batch。

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

因此，论文采用了在线模型近似的方法和高效的FlexiViT架构，仅使用32×32分辨率的patch来评估“超级batch”，与全分辨率、16×16 patch的方法相比，减少了72%的FLOP和67%的挂钟时间（wall-clock time）。

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

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

04 实验结果

图1展示了使用JEST或FlexiJEST方法后效率的显著提升。

左图显示，JEST++在训练数据量减少13.1倍的情况下，仍能达到与原有SigLIP基线模型相同的准确率。

即使考虑到额外引入的评分成本，JEST++在浮点运算次数（FLOP）上也提升了近10倍（中图）。

右图对比了JEST++/FlexiJEST++（绿色）与之前的方法（灰色），相较于经典模型CLIP和EVA-CLIP，JEST++/FlexiJEST++在计算成本和性能上都实现了双重提升。