💻 computer science

Towards Consistent and Efficient Dataset Distillation via Diffusion-Driven Selection

This paper proposes a novel dataset distillation framework that leverages pre-trained diffusion models to select distinctive image patches based on predicted noise loss and intra-class clustering, enabling a streamlined one-step process that outperforms existing generation-based methods on large-scale datasets.

Original authors: Xinhao Zhong, Shuoyang Sun, Zhaoyang Xu, Xulin Gu, Bin Chen, Min Zhang, Yaowei Wang

Published 2026-06-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xinhao Zhong, Shuoyang Sun, Zhaoyang Xu, Xulin Gu, Bin Chen, Min Zhang, Yaowei Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a chef trying to teach a new apprentice how to cook a massive, complex banquet (like a full ImageNet dataset with thousands of dishes). Usually, you'd have to show the apprentice every single dish, every ingredient, and every step. This takes forever, requires a huge kitchen (memory), and burns a lot of gas (computing power).

Dataset Distillation is the idea of creating a tiny "cheat sheet" or a "mini-menu" that is so perfect, the apprentice can learn the whole banquet just by studying these few examples.

The Problem with Previous "Cheat Sheets"

In the past, scientists tried to make these mini-menus in two ways:

Pixel-by-Pixel Optimization: Trying to mathematically tweak every single pixel of a few images to be perfect. This is like trying to sculpt a masterpiece out of clay by hand; it's incredibly slow and only works for small, simple dishes (like CIFAR-10). When you try it on a huge banquet (ImageNet), the kitchen explodes with complexity.
AI Generation: Using a powerful AI (a Diffusion Model) to invent new images from scratch that look like the real dishes. The problem here is that the AI was trained on a different menu (the internet) than the one you are teaching. It might invent a "Golden Retriever" that looks like a fluffy cloud or a "Shark" that looks like a fish with wings. It creates a distribution shift—the fake food doesn't quite taste like the real thing.

The New Solution: "Diffusion-Driven Selection"

The authors of this paper propose a clever new way to make the cheat sheet. Instead of asking the AI to create new food, they ask the AI to act as a critic to find the best parts of the real food already in the kitchen.

Here is how their method works, using a simple analogy:

Step 1: The "Blind Taste Test" (Stage I)

Imagine you have a giant photo of a real Golden Retriever. You want to know: "Which part of this photo is the most Golden Retriever-like?"

You take a pre-trained AI (Stable Diffusion) and ask it to guess the noise (the static) needed to turn a blank screen into a "Golden Retriever."
Then, you ask it to guess the noise needed to turn a blank screen into just "a dog" (without the specific breed).
The Magic: By comparing these two guesses, the AI highlights the specific pixels that make the dog look like a Golden Retriever specifically. It's like the AI pointing a laser pointer at the dog's floppy ears and golden fur, saying, "This is the important part!"
They crop out these "laser-pointed" patches from the real images. They aren't making new images; they are just cutting out the best pieces of the real ones.

Step 2: The "Grouping Party" (Stage II)

Now you have thousands of these "best pieces." Some are ears, some are tails, some are noses.

If you just throw them all in a bag, it's chaotic.
So, the authors use a clustering technique (like sorting socks by color and pattern) to group similar patches together.
They pick the most representative "sock" from each group to ensure the final mini-menu has a good variety (diversity) but still captures the essence of every class.

Why This is a Game-Changer

The paper claims three main victories:

No More "Fake Food": Because they are cutting pieces from real images rather than generating fake ones, the data distribution matches the original perfectly. There is no "hallucination" or weird artifacts.
One-Step Process: Old methods often had to re-run the whole process every time you wanted to change the number of images or the types of classes. This new method is like a "one-click" button. You run it once, and you get a dataset that works for many different settings.
Speed and Scale: It works on massive datasets (like ImageNet-1K) and complex models (like ResNet-101) where previous methods failed or took too long.

The Result

In their experiments, this "cut-and-paste" method, guided by the AI's ability to recognize features, consistently beat the state-of-the-art methods. It produced smaller datasets that trained models just as well (or better) than the massive original ones, without the heavy computational cost of trying to optimize pixels or generate new images from scratch.

In short: Instead of asking an AI to paint a new picture of a dog, this method asks the AI to point to the most dog-like parts of a real photo, cuts them out, and arranges them into a perfect, compact study guide.

Technical Summary: Towards Consistent and Efficient Dataset Distillation via Diffusion-Driven Selection

Problem Statement

Dataset distillation aims to synthesize a compact dataset that enables training deep models to achieve performance comparable to that of the full original dataset, thereby reducing memory and computational costs. However, existing optimization-based methods (e.g., gradient matching) face a vast pixel-level optimization space, limiting their feasibility to small-scale datasets like CIFAR-10/100. While recent approaches leverage pre-trained diffusion models to generate synthetic images directly, they suffer from significant limitations:

Distribution Shifts: Methods that fine-tune diffusion models (either the model itself or image prototypes in latent space) often struggle with the distributional mismatch between the pre-training data (e.g., LAION, ImageNet-1K) and the target dataset, leading to class-irrelevant or noisy outputs.
Inefficiency and Rigidity: Current diffusion-based methods often require model fine-tuning for specific datasets or multiple distillation steps when handling varying class combinations or compression ratios (IPC settings), hindering practical scalability.

Methodology

The authors propose DDDS (Diffusion-Driven Dataset Selection), a framework that diverges from generative approaches by utilizing pre-trained diffusion models for patch selection rather than image generation. The method operates in a one-step process without fine-tuning the diffusion model or training a proxy classifier.

Stage I: Optimal Class-Relevant Feature Selection

Instead of generating new images, the method identifies the most informative regions within the original real images.

Differential Noise Prediction: For each input image $z$ and a target class label $c$ , the pre-trained Stable Diffusion (SD) model predicts noise maps under two conditions: with the class text prompt $c$ (e.g., "An image of golden retriever") and without the label (classifier-free guidance, $\phi$ ).
Representativeness Scoring: The difference between these two noise predictions, $M(z|c) = \mathbb{E}_{\epsilon, t}[N_t(z, \epsilon, c) - N_t(z, \epsilon, \phi)]$ , quantifies the class-relevant information. A higher differential loss indicates that the region is crucial for distinguishing the class.
Patch Extraction: The noise map is upsampled to the original image resolution. The image is divided into patches, and a score $R(p|c)$ is computed for each patch by summing the differential noise values. The patches with the highest scores are selected as candidates.

Stage II: Visual Elements Aggregation

To mitigate randomness and ensure diversity among selected patches, a clustering-based selection strategy is applied.

Feature Extraction: The selected patches are processed to extract DIFT (Diffusion Feature) representations. This involves adding noise at a specific time step to the patch latent and passing it through the U-Net with text conditioning to obtain intermediate layer activations.
Clustering and Ranking: K-means clustering is performed on the DIFT features. The selection follows a dual-ranking rule:
1. Intra-cluster: Patches are ranked by their distance to the cluster centroid (favoring central, representative samples).
2. Inter-cluster: Clusters are ranked by the median representativeness score ( $R(p|c)$ ) of their members.
Final Selection: The top-ranked clusters (e.g., top 10) are prioritized to ensure both representativeness and diversity. If the Images Per Class (IPC) is low, multiple patches are concatenated to form a single synthetic image.

Key Contributions

Orthogonal Paradigm: The paper introduces a dataset distillation framework that leverages pre-trained diffusion models for localization (selection) rather than generation. This approach directly addresses distributional discrepancies between the diffusion prior and the target dataset.
Efficient One-Step Process: The method eliminates the need for model fine-tuning, proxy model training, or repetitive distillation across different settings (class combinations or IPCs). It achieves a fully unrestricted, single-pass distillation process.
State-of-the-Art Performance: Extensive experiments demonstrate that the proposed method consistently outperforms existing training-free and optimization-based methods across various datasets (CIFAR, ImageNet subsets, ImageNet-1K) and model architectures (ResNet, ViT, EfficientNet).

Experimental Results

High-Resolution Datasets: On ImageNet-1K and its subsets (ImageNette, ImageWoof, ImageNet-100), DDDS achieves significant improvements over baselines like SRe2L, RDED, Minimax, and D4M. For instance, on ImageNette with IPC=100, it achieves 92.0% accuracy compared to 91.0% for the next best method (MGD3).
Low-Resolution Datasets: On CIFAR-10/100 and TinyImageNet, the method matches or exceeds the performance of information-matching methods (e.g., DSA, IDM) and other diffusion-based approaches.
Cross-Architecture Generalization: The distilled datasets perform robustly across diverse teacher-student architectures, including CNNs and Vision Transformers (ViTs).
Efficiency: The one-step process results in a constant time cost regardless of the number of class combinations or IPC settings, offering a substantial speedup (e.g., 4x speedup when using a subset of 300 real images for patch extraction) compared to methods requiring repeated distillation.

Significance and Claims

The authors claim that this work provides a fresh perspective on leveraging diffusion models for dataset distillation. By shifting from generation to selection, the method effectively restores the original data distribution without the biases introduced by generative priors. The framework is highlighted as a solution to the scalability and consistency challenges in large-scale dataset distillation, offering a practical, training-free, and highly efficient alternative to current state-of-the-art techniques. The source code is made available to facilitate further research.