Diffusion-Augmented Coreset Expansion for Scalable Dataset Distillation
作者: Ali Abbasi, Shima Imani, Chenyang An, Gayathri Mahalingam, Harsh Shrivastava, Maurice Diesendruck, Hamed Pirsiavash, Pramod Sharma, Soheil Kolouri
分类: cs.CV, cs.AI
发布日期: 2024-12-05
💡 一句话要点
提出扩散增强型 Coreset 扩展方法，用于可扩展的数据集蒸馏
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 数据集蒸馏 Coreset选择 扩散模型 数据压缩 知识蒸馏
📋 核心要点
- 现有数据集蒸馏方法在高分辨率数据和复杂模型下计算效率低，限制了其应用。
- 该论文提出一种两阶段方法，先用 Coreset 压缩数据，再用扩散模型动态扩展和增强 Coreset。
- 实验表明，该方法在多个数据集蒸馏基准上显著优于现有技术，提升超过 10%。
📝 摘要（中文）
随着神经网络规模的快速增长，数据存储和通信需求日益增加。数据集蒸馏作为一种有前景的解决方案应运而生，它通过解决双层优化问题，将来自大型数据集的信息压缩成一组紧凑的合成样本。然而，当前的方法在计算效率方面面临挑战，尤其是在处理高分辨率数据和复杂架构时。最近，基于知识蒸馏的数据集精简方法使这一过程在计算上更加可行。然而，随着生成式基础模型的最新发展，现在有机会实现更大的压缩，提高蒸馏数据的质量，并为数据表示引入有价值的多样性。在这项工作中，我们提出了一种两阶段解决方案。首先，我们通过仅选择信息量最大的图像块来压缩数据集，形成一个 Coreset。接下来，我们利用生成式基础模型来动态扩展这个压缩集，实时提高这些图像块的分辨率，并为 Coreset 引入可控的变异性。我们广泛的实验证明了我们的方法在各种数据集蒸馏基准测试中的鲁棒性和效率。与最先进的方法相比，我们在几个大规模数据集蒸馏基准测试中取得了超过 10% 的显著改进。代码即将发布。
🔬 方法详解
问题定义：数据集蒸馏旨在将大型数据集压缩成一个小的合成数据集，同时保留原始数据集的关键信息，以便用这个小数据集训练的模型能够达到与用原始大数据集训练的模型相近的性能。现有方法，尤其是在处理高分辨率图像和复杂模型时，计算成本很高，限制了其在大规模数据集上的应用。
核心思路：该论文的核心思路是利用 Coreset 选择最具代表性的数据子集，然后利用生成式扩散模型来增强和扩展这个 Coreset。Coreset 降低了计算复杂度，而扩散模型则提高了蒸馏数据的质量和多样性。
技术框架：该方法包含两个主要阶段：1) Coreset 选择：从原始数据集中选择最具代表性的图像块，形成一个压缩的 Coreset。2) 扩散增强：利用预训练的生成式扩散模型，对 Coreset 中的图像块进行超分辨率重建和变异性增强，从而生成更丰富、更高质量的蒸馏数据集。
关键创新：该方法的关键创新在于将 Coreset 选择与扩散模型相结合，既降低了计算成本，又提高了蒸馏数据的质量和多样性。与传统的基于优化的数据集蒸馏方法相比，该方法更加高效且易于扩展。
关键设计：Coreset 选择可以使用各种采样策略，例如 k-中心点聚类。扩散模型可以使用预训练的图像生成模型，例如 Stable Diffusion。损失函数可以包括重构损失和对抗损失，以确保生成的数据与原始数据分布相似。具体的参数设置和网络结构需要根据具体的数据集和任务进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在多个大规模数据集蒸馏基准测试中取得了显著的性能提升，超过现有最先进方法 10% 以上。这表明该方法在提高数据集蒸馏的效率和质量方面具有显著优势，尤其是在处理高分辨率数据时。
🎯 应用场景
该研究成果可应用于各种需要数据压缩和加速模型训练的场景，例如移动设备上的模型部署、联邦学习中的数据传输、以及大规模数据集上的快速原型设计。通过减少数据存储和通信成本，该方法可以促进人工智能技术在资源受限环境中的应用。
📄 摘要（原文）
With the rapid scaling of neural networks, data storage and communication demands have intensified. Dataset distillation has emerged as a promising solution, condensing information from extensive datasets into a compact set of synthetic samples by solving a bilevel optimization problem. However, current methods face challenges in computational efficiency, particularly with high-resolution data and complex architectures. Recently, knowledge-distillation-based dataset condensation approaches have made this process more computationally feasible. Yet, with the recent developments of generative foundation models, there is now an opportunity to achieve even greater compression, enhance the quality of distilled data, and introduce valuable diversity into the data representation. In this work, we propose a two-stage solution. First, we compress the dataset by selecting only the most informative patches to form a coreset. Next, we leverage a generative foundation model to dynamically expand this compressed set in real-time, enhancing the resolution of these patches and introducing controlled variability to the coreset. Our extensive experiments demonstrate the robustness and efficiency of our approach across a range of dataset distillation benchmarks. We demonstrate a significant improvement of over 10% compared to the state-of-the-art on several large-scale dataset distillation benchmarks. The code will be released soon.