ICLR 2024｜NUS尤洋团队同MIT、上海AI Lab等单位提出首个无损数据集蒸馏方法

论文题目:
Towards Lossless Dataset Distillation via Difficulty-Aligned Trajectory Matching

论文地址:
https://arxiv.org/abs/2310.05773

代码地址:
https://gzyaftermath.github.io/DATM/

一、背景

数据集蒸馏旨在以生成的方式将一个大的数据集压缩成一个小的数据集。得益于生成数据的高信息密度，在压缩率较高时，数据集蒸馏的效果远超数据集剪枝等其他数据压缩方法。

然而，随着压缩率的逐渐降低，现有数据集蒸馏方法的效果逐渐变差，甚至弱于从原数据集中随机选取等量数据。这导致现有数据集蒸馏方法局限于高压缩率场景，无法实现对原数据集的无损压缩。

在本文中，作者分析了生成数据的难易以及其对压缩效果的影响，发现应该根据压缩比率控制生成数据的难易。通过应用这一策略，作者提出了第一个可以在低压缩率保持有效的数据集蒸馏方法，并首次实现了无损数据集蒸馏。

二、动机

数据蒸馏可以看作是一个从目标数据集提取特征并将之添加到生成数据上的过程。

目前，主流的数据集蒸馏方法大多通过匹配模型在目标数据集和生成数据集上的某种指标来完成特征的提取与压缩。其中，基于匹配模型训练轨迹[1](Matching Training Trajectory)的方法在多个基准下达到了最优。由于该方法通过匹配模型在生成数据集和目标数据集上的训练轨迹(匹配模型参数)来进行蒸馏，因此其效果受模型在目标数据集上的训练轨迹影响。

根据文献[2]可知，对于分类任务而言，模型倾向于在训练前期学习目标物体的简单特征而在后期学习难特征。因此我们推测匹配模型的前期轨迹就会倾向于在生成数据上添加更多的简单特征而匹配后期轨迹则会倾向于添加难特征。意识到这一点，我们开始探索匹配不同阶段路径对于数据集蒸馏的影响。

实验结果如图一所示，当IPC较低(即压缩率较高)时，匹配前期轨迹(即添加简单特征)是最有效的。而随着压缩率的降低，匹配后期轨迹则变成了最优选择，此时匹配前期路径甚至变得有害。基于这些实验现象，我们提出根据压缩比率控制生成特征的难易，并基于此提出了我们的方法。

三、方法

控制生成特征的难易对于基于轨迹匹配的方法来说较为简单，可以通过控制用于匹配的训练轨迹来实现。为了便于理解，我们用 来表示模型在目标数据集上的训练轨迹，并用 表示第 个epoch时模型的参数。为了能够灵活控制用于匹配的训练轨迹，我们为匹配范围设置了两个超参：匹配范围的下限 和上限 ，如下所示。

同时，为了使蒸馏数据集所含的信息更加丰富，我们将蒸馏数据的硬标签(one-hot label)替换为软标签并在蒸馏过程中不断优化它。然而这种做法面临着一些挑战，例如软标签的初始化可能会包含错误信息，以及优化软标签会让蒸馏变得不稳定。为了解决这些问题，我们提出了以下策略：

1.用可以被预训练模型正确分类的数据初始化蒸馏数据。

2.在蒸馏的早期阶段只生成简单特征(匹配前期轨迹)，以此来提高蒸馏稳定性。

为了实现第二点，我们为匹配范围新增一个可浮动的上限 ，让它在蒸馏初期较小并随着蒸馏的进行逐渐增加直至其上限 。最终，匹配的范围可以表示为：

四、实验结果

3.1 效果比较

我们在CIFAR10，CIFAR100和TinyImageNet数据集下和多个有代表性的数据集蒸馏方法进行了对比，并进一步探索了数据集蒸馏在高IPC情境下的表现。主要的实验结果如下，更多的结果请参考我们的论文。

3.2 分析

简单特征和难特征有什么特点？为什么压缩率较低时添加简单特征有害？我们对这些问题进行了探索。如下图所示，匹配前期轨迹(添加简单特征)会更大幅度的改变图片内容，将目标物体融入背景，使图片有更多关于目标物体的基本信息。这使得添加简单特征在压缩率较高时有效，因为此时蒸馏数据集太小以至于模型无法很好的学习到各个类别的基本特征。

此外，我们可以观察到添加简单特征会让图片损失关于目标物体的细节信息，而添加难特征(匹配后期轨迹)则会让目标物体具有更丰富的纹理细节。因此在压缩率较低时添加难特征的表现更好，因为此时蒸馏数据集本身就含有足够多的数据让模型学会基本的类别特征，添加更多的细节有利于模型识别离群样本。下图是蒸馏图片的可视化，更多分析请见论文。

五、总结

在这项工作中，作者发现应该根据压缩比率控制生成样本的难度。通过应用这一思想，作者提出了第一个可以在高/低压缩率皆能保持有效的数据集蒸馏算法，并首次实现了无损数据集蒸馏，将CIFAR-10，CIFAR-100压缩到了原数据集的1/5大小而不对训练ConvNet造成性能损失。

参考文献

[1] Dataset distillation by matching training trajectories. In CVPR, 2022.
[2] A closer look at memorization in deep networks. In ICML, 2017.

点击阅读原文进入CV社区

收获更多技术干货