Towards Consistent and Efficient Dataset Distillation via Diffusion-Driven Selection

TL;DR

This paper tackles the challenge of distilling large-scale datasets by leveraging a diffusion-prior to localize informative image patches rather than generate synthetic images. It introduces a two-stage framework: Stage I uses diffusion-noise prediction differences to identify class-relevant regions, and Stage II clusters and ranks these patches to ensure diversity, enabling a one-step distillation without model training. The approach achieves state-of-the-art results on ImageNet-1K and derivatives across various IPC settings, while significantly improving training efficiency and generalization across architectures. Overall, the diffusion-driven patch selection provides a practical, scalable alternative to pixel-space optimization and generation-based distillation, reducing distributional bias and repetition across configurations.

Abstract

Dataset distillation provides an effective approach to reduce memory and computational costs by optimizing a compact dataset that achieves performance comparable to the full original. However, for large-scale datasets and complex deep networks (e.g., ImageNet-1K with ResNet-101), the vast optimization space hinders distillation effectiveness, limiting practical applications. Recent methods leverage pre-trained diffusion models to directly generate informative images, thereby bypassing pixel-level optimization and achieving promising results. Nonetheless, these approaches often suffer from distribution shifts between the pre-trained diffusion prior and target datasets, as well as the need for multiple distillation steps under varying settings. To overcome these challenges, we propose a novel framework that is orthogonal to existing diffusion-based distillation techniques by utilizing the diffusion prior for patch selection rather than generation. Our method predicts noise from the diffusion model conditioned on input images and optional text prompts (with or without label information), and computes the associated loss for each image-patch pair. Based on the loss differences, we identify distinctive regions within the original images. Furthermore, we apply intra-class clustering and ranking on the selected patches to enforce diversity constraints. This streamlined pipeline enables a one-step distillation process. Extensive experiments demonstrate that our approach consistently outperforms state-of-the-art methods across various metrics and settings.

Figures (14)

• Figure 1: Comparison of diffusion-based distillation methods. (a) Fine-tuning a diffusion model pre-trained on ImageNet-1K to directly generate images. (b) Fine-tuning the prototypes of the original images in latent space to generate images. (c) Identifying the most representative regions in the real images. The images generated by our method more effectively represent the feature distribution of the original dataset.
• Figure 2: (a): Stage $\textup{I}$ leverages image and corresponding differential class text prompts as inputs to the diffusion model to select the regions that most represent class-relevant features. (b): Stage $\textup{II}$ first computes DIFT of the patches and performs clustering to aggregate the most representative visual elements for each class. The evaluation process utilizes soft labels provided by a pre-trained teacher model. Rather than existing methods employing diffusion models, our proposed method maximizes the restoration of the original data distribution.
• Figure 3: Visualization comparison of images from the top ten classes of the ImageNet-1k dataset is presented (not cherry picked). From top to bottom, each row corresponds to the real dataset, RDED, Minimax, D$^4$M, and ours under IPC=50. Compared to existing methods, our method fully adheres to the distribution of the real dataset and demonstrates a sharper focus on class-specific features.
• Figure 4: Comparison of t-SNE visualizations of features extracted by a pre-trained ResNet-18 on ImageWoof under IPC=100.
• Figure 5: The separate time cost of individual tasks and total time cost of finishing all the tasks. Our method achieves a constant time cost due to one-step distillation process, which eliminates the need for repeated distillation.