Dataset Distillation as Pushforward Optimal Quantization
作者: Hong Ye Tan, Emma Slade
分类: cs.LG, cs.CV, math.OC, stat.ML
发布日期: 2025-01-13 (更新: 2025-09-03)
备注: Modified abstract, additional experiments based on diffusion transformers
💡 一句话要点
将数据集蒸馏重构为推前最优量化问题，提升ImageNet性能。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 数据集蒸馏 最优量化 深度学习 图像分类 ImageNet 扩散模型 生成模型 聚类算法
📋 核心要点
- 数据集蒸馏旨在减少训练数据量，但现有方法要么计算成本高昂，要么性能受限。
- 论文将解耦数据集蒸馏方法重构为最优量化问题，通过最小化投影距离寻找代表性数据。
- 实验表明，该方法在ImageNet-1K上优于现有方法，并在高图像/类别的设置下达到SOTA性能。
📝 摘要（中文）
数据集蒸馏旨在寻找一个合成训练集，使得在该合成数据上训练的模型能够达到与在真实数据上训练的模型相似的性能，同时显著降低计算需求。现有方法大致可分为两类：具有神经网络训练启发式的双层优化问题，以及通过匹配数据分布来绕过双层优化的解耦方法。后一种方法在训练和蒸馏数据集的大小方面具有速度和可扩展性的主要优势。我们证明，当配备编码器-解码器结构时，经验上成功的解耦方法可以被重新表述为最优量化问题，其中通过最小化期望投影距离来找到一组有限的点来近似底层概率测度。特别是，我们将现有的解耦数据集蒸馏方法与经典的最优量化和Wasserstein重心问题联系起来，证明了基于扩散的生成先验的蒸馏数据集的一致性。我们提出了一种基于潜在空间聚类的最优量化数据集蒸馏方法。与之前的SOTA方法D extsuperscript{4}M相比，我们在ImageNet-1K数据集上以微不足道的额外计算量实现了更好的性能和模型间泛化，并在更高的每类图像设置中实现了SOTA性能。在更强的扩散Transformer模型中使用蒸馏噪声初始化，我们在ImageNet-1K及其子集上获得了SOTA蒸馏性能，优于扩散引导方法。
🔬 方法详解
问题定义：数据集蒸馏旨在用远小于原始数据集的合成数据集训练模型，使其性能接近在原始数据集上训练的模型。现有方法，特别是双层优化方法，计算成本高昂，限制了其在大规模数据集上的应用。解耦方法虽然速度快，但其理论基础和与最优量化的联系尚不明确。
核心思路：论文的核心思路是将解耦数据集蒸馏方法视为一个最优量化问题。具体来说，就是寻找一组有限的点（即蒸馏数据集）来近似原始数据的概率分布，目标是最小化期望投影距离。通过将蒸馏问题转化为最优量化，可以利用已有的量化理论和算法来解决数据集蒸馏问题。
技术框架：该方法首先使用一个编码器将原始数据映射到潜在空间。然后在潜在空间中，使用聚类算法（如k-means）找到一组代表性的聚类中心。这些聚类中心就构成了蒸馏数据集。最后，使用一个解码器将潜在空间的聚类中心映射回原始图像空间。整个框架可以看作是一个编码器-解码器结构，其中编码器负责提取特征，聚类算法负责量化，解码器负责重建图像。
关键创新：论文的关键创新在于将解耦数据集蒸馏方法与最优量化理论联系起来，并证明了现有方法可以被视为最优量化的特例。此外，论文还提出了一种基于潜在空间聚类的具体实现方法，该方法在实验中取得了良好的效果。这种联系为理解和改进数据集蒸馏方法提供了一个新的视角。
关键设计：论文使用编码器-解码器结构来学习数据的潜在表示。编码器和解码器可以是任何神经网络结构，例如卷积神经网络或Transformer。聚类算法可以使用k-means或其他聚类算法。损失函数通常是重建误差，即原始数据与重建数据之间的距离。此外，还可以使用正则化项来防止过拟合。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在ImageNet-1K数据集上取得了显著的性能提升，优于之前的SOTA方法D extsuperscript{4}M。具体来说，该方法在模型间泛化方面表现更好，并且在高图像/类别的设置下达到了SOTA性能。此外，通过使用蒸馏噪声初始化，该方法在更强的扩散Transformer模型上获得了SOTA蒸馏性能，优于扩散引导方法。
🎯 应用场景
该研究成果可应用于资源受限的场景，例如移动设备或嵌入式系统，在这些场景中，存储和计算资源有限，无法使用完整的数据集进行训练。此外，该方法还可以用于加速模型训练，提高训练效率，并为联邦学习等隐私保护场景提供支持。
📄 摘要（原文）
Dataset distillation aims to find a synthetic training set such that training on the synthetic data achieves similar performance to training on real data, with orders of magnitude less computational requirements. Existing methods can be broadly categorized as either bi-level optimization problems that have neural network training heuristics as the lower level problem, or disentangled methods that bypass the bi-level optimization by matching distributions of data. The latter method has the major advantages of speed and scalability in terms of size of both training and distilled datasets. We demonstrate that when equipped with an encoder-decoder structure, the empirically successful disentangled methods can be reformulated as an optimal quantization problem, where a finite set of points is found to approximate the underlying probability measure by minimizing the expected projection distance. In particular, we link existing disentangled dataset distillation methods to the classical optimal quantization and Wasserstein barycenter problems, demonstrating consistency of distilled datasets for diffusion-based generative priors. We propose Dataset Distillation by Optimal Quantization, based on clustering in a latent space. Compared to the previous SOTA method D\textsuperscript{4}M, we achieve better performance and inter-model generalization on the ImageNet-1K dataset with trivial additional computation, and SOTA performance in higher image-per-class settings. Using the distilled noise initializations in a stronger diffusion transformer model, we obtain SOTA distillation performance on ImageNet-1K and its subsets, outperforming diffusion guidance methods.