Dataset Distillation via the Wasserstein Metric

📄 arXiv: 2311.18531v3 📥 PDF

作者: Haoyang Liu, Yijiang Li, Tiancheng Xing, Peiran Wang, Vibhu Dalal, Luwei Li, Jingrui He, Haohan Wang

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

发布日期: 2023-11-30 (更新: 2025-07-02)

备注: Accepted to ICCV 2025. Project page at https://liu-hy.github.io/WMDD/ and code is available at https://github.com/Liu-Hy/WMDD

💡 一句话要点

提出WMDD方法以解决数据集蒸馏效率问题

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

关键词: 数据集蒸馏 Wasserstein度量 特征重心 分布匹配 机器学习 高分辨率数据集

📋 核心要点

  1. 现有的数据集蒸馏方法在生成合成数据集时，往往无法有效捕捉原始数据的分布特征，导致模型性能下降。
  2. WMDD方法通过引入Wasserstein度量，计算特征的Wasserstein重心，以优化合成数据与原始数据分布的对齐。
  3. 实验结果表明，WMDD在多个高分辨率数据集上表现出色，达到了当前最先进的性能，展示了其广泛的适应性。

📝 摘要（中文）

数据集蒸馏（DD）旨在生成一个紧凑的合成数据集，使模型在性能上能够与在完整大数据集上训练相媲美，从而显著降低计算成本。本文提出WMDD（基于Wasserstein度量的数据集蒸馏），该方法利用Wasserstein度量增强分布匹配。我们计算预训练分类器特征的Wasserstein重心，以捕捉原始数据分布的基本特征。通过优化合成数据以与该重心在特征空间中对齐，并利用每类的BatchNorm统计量来保留类内变异性，WMDD在保持分布匹配方法效率的同时，在各种高分辨率数据集上实现了最先进的结果。我们的广泛实验展示了WMDD的有效性和适应性，突显其在大规模机器学习应用中的潜力。

🔬 方法详解

问题定义：本文旨在解决现有数据集蒸馏方法在生成合成数据集时无法有效捕捉原始数据分布特征的问题。这导致模型在训练时性能下降，无法充分利用大数据集的优势。

核心思路：WMDD方法的核心思路是利用Wasserstein度量来增强合成数据与原始数据分布之间的匹配。通过计算特征的Wasserstein重心，WMDD能够更好地捕捉原始数据的基本特征，从而优化合成数据的生成过程。

技术框架：WMDD的整体架构包括几个主要模块：首先，利用预训练分类器提取特征；其次，计算这些特征的Wasserstein重心；最后，通过优化合成数据使其在特征空间中与重心对齐，同时保留类内变异性。

关键创新：WMDD的主要创新在于引入Wasserstein度量作为数据蒸馏的基础，显著提升了分布匹配的效率和效果。这一方法与传统的蒸馏方法相比，能够更准确地反映原始数据的分布特征。

关键设计：在WMDD中，关键的设计包括使用Wasserstein重心来捕捉特征分布，以及利用每类的BatchNorm统计量来保留类内的变异性。此外，损失函数的设计也确保了合成数据与重心的有效对齐。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，WMDD在多个高分辨率数据集上实现了最先进的性能，相较于传统方法，性能提升幅度达到XX%（具体数据未知）。这一成果表明WMDD在数据集蒸馏领域的有效性和适应性，具有重要的研究价值。

🎯 应用场景

WMDD方法在多个领域具有广泛的应用潜力，尤其是在需要高效训练大规模模型的场景中，如计算机视觉、自然语言处理等。通过生成高质量的合成数据集，WMDD能够降低计算成本，提高模型训练的效率，推动机器学习技术的进一步发展。

📄 摘要（原文）

Dataset Distillation (DD) aims to generate a compact synthetic dataset that enables models to achieve performance comparable to training on the full large dataset, significantly reducing computational costs. Drawing from optimal transport theory, we introduce WMDD (Wasserstein Metric-based Dataset Distillation), a straightforward yet powerful method that employs the Wasserstein metric to enhance distribution matching. We compute the Wasserstein barycenter of features from a pretrained classifier to capture essential characteristics of the original data distribution. By optimizing synthetic data to align with this barycenter in feature space and leveraging per-class BatchNorm statistics to preserve intra-class variations, WMDD maintains the efficiency of distribution matching approaches while achieving state-of-the-art results across various high-resolution datasets. Our extensive experiments demonstrate WMDD's effectiveness and adaptability, highlighting its potential for advancing machine learning applications at scale.