Dataset Distillation:把整个数据集压缩成几张图

今天读了一篇关于数据集蒸馏(Dataset Distillation)的文章,核心思路是:对数据集进行蒸馏,得到一组合成图片和一个学习率,用这个学习率在蒸馏得到的图片上只需要很少的迭代步数,就能达到接近在原始数据集上完整训练的精度。两个网络采用相同的初始化方式(fixed network initialization),文章后面也有随机初始化、预训练初始化等情形下的实验。

蒸馏得到的数据集并不一定具有正确的像素分布,极端情况下可以压缩到每个类别只保留一张图片。这样做的好处是:可以快速地重新加载出一个具备相同能力的网络,而不必每次都重新跑完整的训练流程。

基本方法

具体做法是:将网络的权重视为合成训练图片的可微函数,然后对这些合成图片进行优化。但这种做法需要预先指定网络的初始权重。为此,文章后续开发了适用于随机初始化网络的方法,以及支持训练多个 epoch 的可迭代版本,最后还用线性模型分析了达到与原始数据集同等性能所需的蒸馏数据数量下界。

分阶段探索

第一阶段:固定初始化 + 单步梯度下降

在得到学习率和蒸馏数据之后,只需一步梯度下降就可以在真实测试集上得到不错的结果。

θ1=θ0η~θ0(x~,θ0)\theta_1 = \theta_0 - \tilde{\eta}\,\nabla_{\theta_0}\,\ell(\tilde{\mathbf{x}}, \theta_0)

其中学习率和蒸馏数据集通过如下方式得到:

x~,η~=argminx~,η~L(x~,η~;θ0)=argminx~,η~(x,θ1)=argminx~,η~(x,θ0η~θ0(x~,θ0)),\begin{aligned} \tilde{\mathbf{x}}^*, \tilde{\eta}^* &= \arg\min_{\tilde{\mathbf{x}}, \tilde{\eta}} \mathcal{L}(\tilde{\mathbf{x}}, \tilde{\eta}; \theta_0) = \arg\min_{\tilde{\mathbf{x}}, \tilde{\eta}} \ell(\mathbf{x}, \theta_1) \\ &= \arg\min_{\tilde{\mathbf{x}}, \tilde{\eta}} \ell\big(\mathbf{x},\, \theta_0 - \tilde{\eta}\,\nabla_{\theta_0}\,\ell(\tilde{\mathbf{x}}, \theta_0)\big), \end{aligned}

第二阶段:随机初始化

上面得到的学习率和蒸馏图片都是在特定权重下优化出来的,换其他初始化方式可能并不实用。蒸馏图片实际上同时编码了原始数据的信息和用来蒸馏的网络初始化权重信息,换一套初始权重就可能失效。

为了应对随机初始化的情形,将蒸馏图片和学习率的优化方式改为以下形式——即在随机初始化的情况下,对权值分布的期望权重处进行优化:

x~,η~=argminx~,η~Eθ0p(θ0)L(x~,η~;θ0),\tilde{\mathbf{x}}^*, \tilde{\eta}^* = \arg\min_{\tilde{\mathbf{x}}, \tilde{\eta}} \mathbb{E}_{\theta_0 \sim p(\theta_0)}\,\mathcal{L}(\tilde{\mathbf{x}}, \tilde{\eta}; \theta_0),

具体算法如下:

算法 1  Dataset Distillation

输入:p(θ₀):初始权重分布;M:蒸馏数据的数量
输入:α:步长;n:batch 大小;T:优化迭代次数;η̃₀:η̃ 的初始值
 1: 随机初始化 x̃ = {x̃ᵢ}ᵢ₌₁ᴹ,η̃ ← η̃₀
 2: for 每个训练步 t = 1 到 T do
 3:     取一个 minibatch 的真实训练数据 xₜ = {x_{t,j}}ⱼ₌₁ⁿ
 4:     采样一个 batch 的初始权重 θ₀⁽ʲ⁾ ~ p(θ₀)
 5:     for 每个采样得到的 θ₀⁽ʲ⁾ do
 6:         用 GD 计算更新后的参数:θ₁⁽ʲ⁾ = θ₀⁽ʲ⁾ − η̃ ∇_{θ₀⁽ʲ⁾} ℓ(x̃, θ₀⁽ʲ⁾)
 7:         在真实训练数据上评估目标函数:L⁽ʲ⁾ = ℓ(xₜ, θ₁⁽ʲ⁾)
 8:     end for
 9:     更新 x̃ ← x̃ − α ∇_{x̃} Σⱼ L⁽ʲ⁾,以及 η̃ ← η̃ − α ∇_{η̃} Σⱼ L⁽ʲ⁾
10: end for
输出:蒸馏数据 x̃ 与优化后的学习率 η̃

第三阶段:线性模型分析

用一种线性模型来说明该方法的属性与局限性,并讨论了初始权重分布对蒸馏效果的影响。

第四阶段:多步梯度下降与多 epoch

在上述算法基础上做了改进,类似于正常神经网络训练的过程,允许蒸馏数据集被使用多个 epoch。

第五阶段:不同初始化分布与不同学习目标下获取蒸馏图片

探索了在多种初始化分布和学习目标下如何获取蒸馏图片。蒸馏得到的数据可以在预训练好的模型上进行很快的 fine-tune。

分类任务实验结果

在分类任务上的实验结果如下:

Ours: Fixed init.Ours: Random init.Baselines (GD steps): Random realBaselines (GD steps): Optimized realBaselines (GD steps): k-meansBaselines (GD steps): Average realBaselines (K-NN): Random realBaselines (K-NN): k-means
MNIST96.679.5 ± 8.168.6 ± 9.873.0 ± 7.676.4 ± 9.577.1 ± 2.771.5 ± 2.192.2 ± 0.1
CIFAR1054.036.8 ± 1.221.3 ± 1.523.4 ± 1.322.5 ± 3.122.3 ± 0.718.8 ± 1.329.4 ± 0.3

文章还有一些其他实验,这里暂时先忽略。