转载自:量子位(QbitAI)
只要一块6年前的2080Ti,就能做大模型数据蒸馏?
来自上交大EPIC实验室等机构的一项最新研究,提出了一种新的数据集蒸馏方法——
NFCM
。
与前SOTA相比,新方法的
显存占用只有1/300,并且速度提升了20倍
,相关论文获得了
CVPR满分
。
NCFM引入了一个辅助的神经网络,将数据集蒸馏重新表述为一个极小化极大
(minmax)
优化问题。
在多个基准数据集上,NCFM都取得了显著的性能提升,并展现出可扩展性。
在CIFAR数据集上,NCFM只需2GB左右的GPU内存就能实现无损的数据集蒸馏,用2080Ti即可实现。
并且,NCFM在连续学习、神经架构搜索等下游任务上也展现了优异的性能。
将数据蒸馏转化为minmax优化
NCFM的核心是
引入了一个新的分布差异度量NCFD,并将数据集蒸馏问题转化为一个minmax优化问题
。
通过交替优化合成数据以最小化NCFD,以及优化采样网络以最大化NCFD,NCFM在提升合成数据质量的同时,不断增强分布差异度量的敏感性和有效性。
特征提取与频率参数采样
NCFM的第一步,是进行
特征提取
,也就是从真实数据集和合成数据集中分别采样一批数据,并将其输入到特征提取网络中。
特征提取网络将原始数据从像素空间映射到一个特征空间,得到对应的特征表示,目的是提取数据的高层语义特征,为后续的分布匹配做准备。
特征提取网络可以是一个预训练的模型,也可以是一个随机初始化的模型,这里NCFM采用了一种混合方式。
接下来,NCFM
引入了一个轻量级的神经网络
作为采样网络,它接受一个随机噪声作为输入,输出一组频率参数。
这些频率参数将用于对特征函数
(Characteristic Function,CF)
进行采样。
特征函数计算与分布差异度量
对于每一个频率参数,将其与特征表示进行内积运算,然后取复指数,就得到了对应的CF值。
这两个CF值都是复数,其中实部刻画了数据在该频率上的分布范围,捕捉分布的散度或多样性;虚部则反映了数据在该频率上的分布中心,捕捉分布的典型性或真实性。
通过
比较真实数据和合成数据的CF值
,就可以全面地度量它们在特征空间上的分布差异。
为了定量地度量真实数据和合成数据之间的分布差异,NCFM引入了一个称为神经特征函数差异
(Neural Characteristic Function Discrepancy,NCFD)
的度量。
NCFD综合考虑了所有采样频率上的CF差异,将其汇总为一个标量值。NCFD越小,说明两个分布越接近;NCFD越大,说明两个分布差异越大。
minmax优化
有了NCFD这个分布差异度量,NCFM的优化目标就很清晰了——
最小化NCFD,使得合成数据和真实数据的分布尽可能接近;同时,望最大化NCFD对合成数据的敏感度,使之能够准确反映合成数据的变化。
为了同时实现这两个目标,NCFM
引入了一个minmax优化框架
:
通过交替进行极小化阶段和极大化阶段的优化,NCFM不断改进合成数据的质量,同时也不断强化NCFD度量的敏感性和准确性。
模型微调与标签生成
为了进一步提升合成数据的质量,NCFM在优化过程中还引入了两个额外的步骤——模型微调和标签生成。