ICLR 2024 | NUS尤洋团队联合MIT/上海AI Lab等提出首个无损数据集蒸馏方法DATM

数据集蒸馏旨在以生成的方式将一个大的数据集压缩成一个小的数据集。得益于生成数据的高信息密度，在压缩率较高时，数据集蒸馏的效果远超数据集剪枝等其他数据压缩方法。然而，随着压缩率的逐渐降低，现有数据集蒸馏方法的效果逐渐变差，甚至弱于从原数据集中随机选取等量数据。这导致现有数据集蒸馏方法局限于高压缩率场景，无法实现对原数据集的无损压缩。

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

论文题目: Towards Lossless Dataset Distillation via Difficulty-Aligned Trajectory Matching 论文链接: https://arxiv.org/abs/2310.05773 代码链接: https://github.com/NUS-HPC-AI-Lab/DATM

一、动机

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

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

图一、匹配不同阶段的轨迹的蒸馏效果

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

二、方法

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

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

1. 
   

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

   
   

2. 
   

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

   
   

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

三、实验

3.1 效果比较

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

和多个蒸馏方法在不同数据集和不同IPC(image per class)设定下的比较

在高IPC情境下数据集蒸馏和数据集剪枝的表现比较

3.2 分析

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

匹配处于不同训练阶段的轨迹所得的蒸馏样本

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

不同IPC设定下蒸馏所得样本

匹配前期轨迹过程中图片的变化

匹配后期轨迹过程中图片的变化

四、总结

在这项工作中，作者发现应该根据压缩比率控制生成样本的难度。通过应用这一思想，作者提出了第一个可以在高/低压缩率皆能保持有效的数据集蒸馏算法，并首次实现了无损数据集蒸馏，将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.

Illustration From IconScout By Delesign Graphics

-The End-

本周上新！

扫码观看！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线500+期talk视频，3000+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈