在炼丹过程中,为了减少训练所需资源,MLer 有时会将大型复杂的大模型“蒸馏”为较小的模型,同时还要保证与压缩前相当的结果。这就是知识蒸馏,一种模型压缩 / 训练方法。
不过随着技术发展,大家也逐渐将蒸馏的对象扩展到了数据集上。这不,谷歌最近就提出了两种新的数据集蒸馏方法,在推特上引起了不小反响,热度超过 600:
像这样,将 50000 张标注图像的 CIFAR-10 数据集“蒸馏”缩小至 1/5000 大小,只基于 10 张合成数据点进行训练,模型的准确率仍可近似 51%:
△上:原始数据集 下:蒸馏后
而如果“蒸馏数据集”由 500 张图像组成(占原数据集 1% 大小),其准确率可以达到 80%。两种数据集蒸馏方法分别来自于 ICLR 2021 和 NeurIPS 2021 上的两篇论文。
通过两阶段循环进行优化
那么要如何才能“蒸馏”一个数据集呢?其实,这相当于一个两阶段的优化过程:
“内部循环”,用于在学习数据上训练模型
“外部循环”,用于优化学习数据在自然数据上的性能
通过内部循环可以得到一个核脊回归(KRR)函数,然后再外部循环中计算原始图像标注与核脊回归函数预测标注之间的均方误差(MSE)。
这时,谷歌提出的两种方法就分别有了不同的处理路线:
一、标注解释(LS)
这种方法直接解释最小化 KRR 损失函数的支持标注集(support labels),并为每个支持图像生成一个独特的密集标注向量。
△蓝:原始独热标注 橙:LS 生成的密集标注
二、核归纳点(KIP)
这种方法通过基于梯度的方法将 KRR 损失函数最小化,以此来优化图像和可能生成的数据。以 MNIST 为例,下图中的上、中、下三张图分别为原始的 MNIST 数据集、固定标注的 KIP 蒸馏图像、优化标注的 KIP 蒸馏图像。可以看出,在于对数据集进行蒸馏时,优化标注的效果最好:
对比已有的 DC(Dataset Condensation)方法和 DSP(Dataset Condensation with Differentiable Siamese Augmentation)方法可以看到:
如果使用每类别只有一张图像,也就是最后只有 10 张图像的蒸馏数据集,KIP 方法的测试集准确率整体高于 DC 和 DSP 方法。在 CIFAR-10 分类任务中,LS 也优于先前的方法,KIP 甚至可以达到翻倍的效果。
对此,谷歌表示:
这证明了在某些情况下,我们的缩小 100 倍的“蒸馏数据集”要比原始数据集更好。
两位华人作者
整个项目由萧乐超(Lechao Xiao)、Zhourong Chen、Roman Novak 三人合作完成。
其中萧乐超为 LS 方法的论文作者之一,本科毕业于浙江大学的应用数学系,在美国伊利诺大学厄巴纳-香槟分校(UIUC)取得博士学位,现在是谷歌大脑团队的一名科学家。他的主要研究方向是数学、机器学习和深度学习。
另一位华人科学家 Zhourong Chen 则是 KIP 方法的论文作者之一,本科毕业于中山大学,并在中国香港科技大学取得了计算机科学与工程系的博士学位,现是 Google Research 的一名软件工程师。
论文:
[1]https://openreview.net/forum?id=l-PrrQrK0QR
[2]https://openreview.net/forum?id=hXWPpJedrVP
开源地址:
https://github.com/google-research/google-research/tree/master/kip
参考链接:
https://ai.googleblog.com/2021/12/training-machine-learning-models-more.html