Prioritize Alignment in Dataset Distillation

📄 arXiv: 2408.03360v3 📥 PDF

作者: Zekai Li, Ziyao Guo, Wangbo Zhao, Tianle Zhang, Zhi-Qi Cheng, Samir Khaki, Kaipeng Zhang, Ahmad Sajedi, Konstantinos N Plataniotis, Kai Wang, Yang You

分类: cs.LG, cs.AI

发布日期: 2024-08-06 (更新: 2024-10-13)

备注: 19 pages, 9 figures

💡 一句话要点

提出PAD：通过对齐信息优先级，显著提升数据集蒸馏性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 数据集蒸馏 信息对齐 模型压缩 深度学习 轨迹匹配

📋 核心要点

  1. 现有数据集蒸馏方法在信息提取和嵌入过程中引入了未对齐的信息，导致蒸馏数据集质量下降。
  2. PAD通过修剪目标数据集和仅使用代理模型的深层来进行蒸馏，从而优先对齐信息，减少信息偏差。
  3. 实验表明，PAD在各种基准测试中显著提升了数据集蒸馏的性能，达到了当前最优水平。

📝 摘要（中文）

数据集蒸馏旨在将大型数据集压缩成一个显著更紧凑的合成数据集，同时不影响训练模型的性能。为了实现这一目标，现有方法使用代理模型从目标数据集中提取信息，并将其嵌入到蒸馏数据集中。因此，提取和嵌入信息的质量决定了蒸馏数据集的质量。本文发现，现有方法在信息提取和嵌入阶段都引入了未对齐的信息。为了缓解这个问题，我们提出了数据集蒸馏中的优先级对齐（PAD），它从以下两个角度对齐信息：1）我们根据压缩比修剪目标数据集，以过滤掉可以被代理模型提取的信息。2）我们仅使用代理模型的深层来进行蒸馏，以避免过度引入低级信息。这种简单的策略有效地过滤掉了未对齐的信息，并为主流的基于匹配的蒸馏算法带来了显著的改进。此外，基于轨迹匹配，PAD在各种基准测试中取得了显著的改进，实现了最先进的性能。

🔬 方法详解

问题定义：数据集蒸馏旨在将大规模数据集压缩为小规模的合成数据集，同时尽可能保留原始数据集的信息，使得在合成数据集上训练的模型性能接近于在原始数据集上训练的模型。现有方法在信息提取和嵌入过程中存在信息不对齐的问题，导致蒸馏后的数据集质量不高，模型性能下降。

核心思路：PAD的核心思路是通过优先考虑信息的对齐性来提高数据集蒸馏的性能。具体来说，它通过两个关键步骤来减少信息不对齐：一是通过修剪原始数据集，只保留代理模型能够有效提取的信息；二是只使用代理模型的深层特征进行蒸馏，避免引入过多的低级噪声信息。

技术框架：PAD方法主要包含两个阶段：1) 数据集修剪阶段：根据压缩比例，对原始数据集进行筛选，保留最具代表性的样本。这一步旨在过滤掉代理模型难以有效提取的信息。2) 蒸馏阶段：利用代理模型，但只使用其深层特征，将筛选后的数据集的信息嵌入到合成数据集中。这一步旨在避免引入过多的低级噪声信息。整个过程的目标是生成一个既小巧又具有代表性的合成数据集。

关键创新：PAD的关键创新在于其对信息对齐性的重视。与以往方法不同，PAD不是简单地将原始数据集的信息压缩到合成数据集中，而是首先对原始数据集进行筛选，然后只使用代理模型的深层特征进行蒸馏。这种方法可以有效地减少信息不对齐，提高蒸馏数据集的质量。

关键设计：在数据集修剪阶段，可以使用不同的策略来选择最具代表性的样本，例如基于梯度的选择、基于聚类的选择等。在蒸馏阶段，可以使用不同的损失函数来衡量代理模型在原始数据集和合成数据集上的性能差异，例如轨迹匹配损失、梯度匹配损失等。具体参数设置需要根据具体任务和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

PAD方法在多个基准数据集上取得了显著的性能提升。例如，在CIFAR-10数据集上，基于轨迹匹配的PAD方法相比于现有最佳方法，性能提升了超过5%。实验结果表明，PAD方法能够有效地减少信息不对齐，提高蒸馏数据集的质量，从而提升模型的泛化能力。

🎯 应用场景

PAD方法可应用于各种需要压缩数据集的场景，例如移动设备上的模型训练、资源受限环境下的机器学习等。通过使用PAD方法，可以在保证模型性能的前提下，显著减少数据集的存储空间和计算资源需求，从而提高机器学习算法的效率和可部署性。未来，该方法有望在边缘计算、联邦学习等领域发挥重要作用。

📄 摘要（原文）

Dataset Distillation aims to compress a large dataset into a significantly more compact, synthetic one without compromising the performance of the trained models. To achieve this, existing methods use the agent model to extract information from the target dataset and embed it into the distilled dataset. Consequently, the quality of extracted and embedded information determines the quality of the distilled dataset. In this work, we find that existing methods introduce misaligned information in both information extraction and embedding stages. To alleviate this, we propose Prioritize Alignment in Dataset Distillation (PAD), which aligns information from the following two perspectives. 1) We prune the target dataset according to the compressing ratio to filter the information that can be extracted by the agent model. 2) We use only deep layers of the agent model to perform the distillation to avoid excessively introducing low-level information. This simple strategy effectively filters out misaligned information and brings non-trivial improvement for mainstream matching-based distillation algorithms. Furthermore, built on trajectory matching, \textbf{PAD} achieves remarkable improvements on various benchmarks, achieving state-of-the-art performance.