Unlocking Dataset Distillation with Diffusion Models

📄 arXiv: 2403.03881v4 📥 PDF

作者: Brian B. Moser, Federico Raue, Sebastian Palacio, Stanislav Frolov, Andreas Dengel

分类: cs.CV, cs.AI, cs.LG

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

🔗 代码/项目: GITHUB

💡 一句话要点

提出LD3M方法以解决数据集蒸馏中的梯度消失问题

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

关键词: 数据集蒸馏 扩散模型 梯度消失 合成样本 深度学习

📋 核心要点

  1. 现有的数据集蒸馏方法在利用扩散模型时面临梯度消失的问题，导致合成样本优化效果不佳。
  2. 本文提出的LD3M方法通过预训练的潜在扩散模型实现了端到端的梯度蒸馏，克服了传统方法的局限性。
  3. 在多个ImageNet子集上，LD3M相较于现有最先进的方法，准确率提升最高可达4.8个百分点，显示出其优越性。

📝 摘要（中文）

数据集蒸馏旨在将数据集浓缩为更小但高度代表性的合成样本。尽管扩散模型在生成基准中表现优异，但现有蒸馏方法主要依赖于GAN或自编码器，未能有效利用扩散模型。本文提出了潜在数据集蒸馏与扩散模型（LD3M），这是首个通过预训练的潜在扩散模型端到端学习梯度蒸馏潜在和类别嵌入的方法。通过在每个反向步骤中注入线性衰减的跳跃连接，LD3M在多个ImageNet子集上实现了最高4.8个百分点的准确率提升，展示了其在数据集蒸馏中的有效性。

🔬 方法详解

问题定义：本文解决的是数据集蒸馏中的梯度消失问题。现有方法在长去噪链中进行简单的反向传播，导致梯度信号减弱，影响合成样本的优化效果。

核心思路：LD3M方法通过引入线性衰减的跳跃连接，从初始噪声状态注入梯度信号，确保在多个反向步骤中保持有效的梯度传递，从而实现更好的样本优化。

技术框架：该方法基于预训练的潜在扩散模型，整体流程包括数据集的潜在表示学习、类别嵌入的优化以及通过扩散模型生成合成样本。主要模块包括梯度蒸馏和反向传播机制。

关键创新：LD3M是首个能够通过扩散模型进行端到端学习的蒸馏方法，其创新之处在于有效解决了梯度消失问题，显著提升了合成样本的质量和代表性。

关键设计：在设计中，采用了线性衰减的跳跃连接，确保梯度信号在多个时间步长中得以保留。此外，损失函数的选择和网络结构的设计也经过精心调整，以优化蒸馏效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LD3M在多个ImageNet子集上实现了显著的性能提升，准确率提升最高可达4.8个百分点（1 IPC）和4.2个百分点（10 IPC），相较于现有最先进的方法，展示了其在数据集蒸馏中的有效性和优势。

🎯 应用场景

该研究在数据集蒸馏领域具有广泛的应用潜力，尤其是在需要高效数据处理和存储的场景中，如计算机视觉、自然语言处理等领域。通过生成高质量的合成样本，LD3M能够帮助提升模型训练效率，降低计算资源消耗，推动智能系统的快速发展。

📄 摘要（原文）

Dataset distillation seeks to condense datasets into smaller but highly representative synthetic samples. While diffusion models now lead all generative benchmarks, current distillation methods avoid them and rely instead on GANs or autoencoders, or, at best, sampling from a fixed diffusion prior. This trend arises because naive backpropagation through the long denoising chain leads to vanishing gradients, which prevents effective synthetic sample optimization. To address this limitation, we introduce Latent Dataset Distillation with Diffusion Models (LD3M), the first method to learn gradient-based distilled latents and class embeddings end-to-end through a pre-trained latent diffusion model. A linearly decaying skip connection, injected from the initial noisy state into every reverse step, preserves the gradient signal across dozens of timesteps without requiring diffusion weight fine-tuning. Across multiple ImageNet subsets at 128x128 and 256x256, LD3M improves downstream accuracy by up to 4.8 percentage points (1 IPC) and 4.2 points (10 IPC) over the prior state-of-the-art. The code for LD3M is provided at https://github.com/Brian-Moser/prune_and_distill.