SelMatch: Effectively Scaling Up Dataset Distillation via Selection-Based Initialization and Partial Updates by Trajectory Matching

📄 arXiv: 2406.18561v1 📥 PDF

作者: Yongmin Lee, Hye Won Chung

分类: cs.CV, cs.LG

发布日期: 2024-05-28

备注: ICML 2024

💡 一句话要点

SelMatch：通过选择初始化和轨迹匹配的部分更新，有效扩展数据集蒸馏规模

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

关键词: 数据集蒸馏 轨迹匹配 选择初始化 部分更新 合成数据 小样本学习 模型压缩

📋 核心要点

  1. 现有数据集蒸馏方法在高IPC下难以捕捉复杂特征，导致性能下降，甚至不如随机选择。
  2. SelMatch通过选择初始化和轨迹匹配的部分更新，自适应地调整合成数据集的难度，从而提升性能。
  3. 在CIFAR和TinyImageNet上的实验表明，SelMatch在不同子集比例下均优于现有方法。

📝 摘要（中文）

数据集蒸馏旨在从大型数据集中为每个类别合成少量图像（IPC），以近似完整数据集训练，同时最大限度地减少性能损失。虽然在非常小的IPC范围内有效，但许多蒸馏方法的效果会降低，甚至在IPC增加时表现不如随机样本选择。我们对各种IPC规模下基于轨迹匹配的最新蒸馏方法的研究表明，即使IPC增加，这些方法也难以将较难样本的复杂、稀有特征纳入合成数据集，导致简单和困难测试样本之间持续存在覆盖差距。受此观察的启发，我们引入了SelMatch，一种新型蒸馏方法，可以有效地随IPC扩展。SelMatch使用基于选择的初始化和通过轨迹匹配的部分更新来管理合成数据集的所需难度级别，以适应IPC规模。在CIFAR-10/100和TinyImageNet上进行测试时，SelMatch在5%到30%的子集比例下始终优于领先的仅选择和仅蒸馏方法。

🔬 方法详解

问题定义：数据集蒸馏旨在用少量合成数据近似完整数据集的训练效果。现有基于轨迹匹配的蒸馏方法在高IPC（每个类别图像数）下，难以有效学习并保留困难样本中的复杂和稀有特征，导致模型在困难样本上的泛化能力不足，性能提升受限。现有方法无法很好地平衡简单样本和困难样本的表示学习。

核心思路：SelMatch的核心思想是根据IPC的大小，自适应地调整合成数据集的难度。通过选择初始化，优先选择具有代表性的样本作为初始合成数据集；然后，通过轨迹匹配的部分更新，逐步引入更多信息，避免过早陷入局部最优，从而更好地覆盖整个数据集的特征空间。

技术框架：SelMatch主要包含两个阶段：1) 选择初始化：从原始数据集中选择一部分样本作为合成数据集的初始值。选择策略旨在选取最具代表性的样本，例如选择训练过程中梯度变化最大的样本。2) 轨迹匹配的部分更新：使用轨迹匹配方法更新合成数据集，但不是每次迭代都更新所有合成样本，而是根据一定的策略选择部分样本进行更新。这种部分更新的策略有助于避免过拟合简单样本，并鼓励模型学习困难样本的特征。

关键创新：SelMatch的关键创新在于结合了选择初始化和部分更新的轨迹匹配。选择初始化能够快速构建一个具有代表性的合成数据集，而部分更新的轨迹匹配能够逐步引入更多信息，避免过拟合，从而在高IPC下也能保持良好的性能。与现有方法相比，SelMatch能够更好地平衡简单样本和困难样本的表示学习。

关键设计：SelMatch的关键设计包括：1) 选择策略：可以使用不同的选择策略来初始化合成数据集，例如基于梯度幅度或多样性的选择。2) 部分更新策略：可以根据样本的梯度或损失值来选择需要更新的样本。3) 轨迹匹配损失函数：使用标准的轨迹匹配损失函数来优化合成数据集，例如最小化真实数据和合成数据训练的模型参数轨迹之间的距离。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SelMatch在CIFAR-10/100和TinyImageNet数据集上，在5%到30%的子集比例下，始终优于领先的仅选择和仅蒸馏方法。例如，在CIFAR-10上，SelMatch在30%的子集比例下，相比于基线方法取得了显著的性能提升，验证了其在高IPC下的有效性。

🎯 应用场景

SelMatch可应用于数据隐私保护、模型压缩和加速等领域。在数据隐私保护方面，可以使用SelMatch生成合成数据集，用于训练模型，避免直接使用敏感数据。在模型压缩和加速方面，可以使用SelMatch生成小规模数据集，用于训练轻量级模型，从而降低计算成本和存储空间。该方法在资源受限的边缘设备上具有潜在的应用价值。

📄 摘要（原文）

Dataset distillation aims to synthesize a small number of images per class (IPC) from a large dataset to approximate full dataset training with minimal performance loss. While effective in very small IPC ranges, many distillation methods become less effective, even underperforming random sample selection, as IPC increases. Our examination of state-of-the-art trajectory-matching based distillation methods across various IPC scales reveals that these methods struggle to incorporate the complex, rare features of harder samples into the synthetic dataset even with the increased IPC, resulting in a persistent coverage gap between easy and hard test samples. Motivated by such observations, we introduce SelMatch, a novel distillation method that effectively scales with IPC. SelMatch uses selection-based initialization and partial updates through trajectory matching to manage the synthetic dataset's desired difficulty level tailored to IPC scales. When tested on CIFAR-10/100 and TinyImageNet, SelMatch consistently outperforms leading selection-only and distillation-only methods across subset ratios from 5% to 30%.