Distributional Dataset Distillation with Subtask Decomposition
作者: Tian Qin, Zhiwei Deng, David Alvarez-Melis
分类: cs.LG
发布日期: 2024-03-01
💡 一句话要点
提出分布式数据集蒸馏方法以提高存储效率
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 数据集蒸馏 分布式学习 内存效率 计算机视觉 深度学习
📋 核心要点
- 现有的蒸馏方法在将数据集压缩为显式原型时，常常导致意外的存储成本，表现不够理想。
- 本文提出的分布式数据集蒸馏（D3）通过编码每类的最小统计信息，结合解码器实现更高效的内存使用。
- 在实验中，D3方法在TinyImageNet和ImageNet-1K上表现优异，特别是在ImageNet-1K上提升了6.9%的性能。
📝 摘要（中文）
本文探讨了神经网络在特定任务数据集训练中学习到的知识，提出了数据集蒸馏的概念。现有方法在蒸馏显式原型时常常表现不佳，导致蒸馏标签的存储成本意外增加。为此，本文提出了分布式数据集蒸馏（D3），通过最小化每类统计信息来编码数据，并结合解码器，将数据集蒸馏为更为紧凑的分布表示，显著提高了内存效率。此外，提出的联邦蒸馏方法将数据集分解为子集，利用子任务专家并行蒸馏，最终重新聚合。实验结果表明，该方法在TinyImageNet和ImageNet-1K上达到了最先进的效果，尤其是在ImageNet-1K上，在每类2张图像的存储预算下，性能提升了6.9%。
🔬 方法详解
问题定义：本文旨在解决现有数据集蒸馏方法在压缩显式原型时的存储成本过高的问题。现有方法往往未能有效利用每类数据的统计信息，导致性能和存储效率不佳。
核心思路：提出的分布式数据集蒸馏（D3）方法通过编码每类的最小统计信息，结合解码器生成紧凑的分布表示，从而提高了内存效率，避免了显式原型带来的冗余存储。
技术框架：D3方法的整体架构包括数据编码模块和解码器模块。首先，数据集被分解为多个子集，利用子任务专家并行进行蒸馏，最后将各子集的结果重新聚合，形成最终的分布表示。
关键创新：D3的核心创新在于通过最小化每类的统计信息来替代传统的原型蒸馏方法，从而显著降低了存储需求，并提高了蒸馏效率。与现有方法相比，D3在存储和计算效率上具有明显优势。
关键设计：在实现过程中，D3采用了特定的损失函数来优化每类统计信息的提取，并设计了高效的网络结构以支持并行蒸馏。此外，联邦蒸馏的策略使得子任务专家能够独立工作，从而加速了整体蒸馏过程。
🖼️ 关键图片
📊 实验亮点
实验结果显示，D3方法在TinyImageNet和ImageNet-1K数据集上均取得了最先进的性能，尤其是在ImageNet-1K上，在每类仅使用2张图像的存储预算下，性能提升达6.9%。这一显著的提升证明了D3在数据集蒸馏领域的有效性和创新性。
🎯 应用场景
该研究的潜在应用领域包括计算机视觉、自然语言处理等需要处理大规模数据集的任务。通过提高数据集的存储和处理效率，D3方法能够帮助研究人员和工程师在资源受限的环境中更有效地训练模型，推动相关领域的发展。未来，该方法可能会在更多实际应用中得到推广，尤其是在需要快速迭代和高效存储的场景中。
📄 摘要（原文）
What does a neural network learn when training from a task-specific dataset? Synthesizing this knowledge is the central idea behind Dataset Distillation, which recent work has shown can be used to compress large datasets into a small set of input-label pairs ($\textit{prototypes}$) that capture essential aspects of the original dataset. In this paper, we make the key observation that existing methods distilling into explicit prototypes are very often suboptimal, incurring in unexpected storage cost from distilled labels. In response, we propose $\textit{Distributional Dataset Distillation}$ (D3), which encodes the data using minimal sufficient per-class statistics and paired with a decoder, we distill dataset into a compact distributional representation that is more memory-efficient compared to prototype-based methods. To scale up the process of learning these representations, we propose $\textit{Federated distillation}$, which decomposes the dataset into subsets, distills them in parallel using sub-task experts and then re-aggregates them. We thoroughly evaluate our algorithm on a three-dimensional metric and show that our method achieves state-of-the-art results on TinyImageNet and ImageNet-1K. Specifically, we outperform the prior art by $6.9\%$ on ImageNet-1K under the storage budget of 2 images per class.