MGD$^3$ : Mode-Guided Dataset Distillation using Diffusion Models

MGD$^3$ : Mode-Guided Dataset Distillation using Diffusion Models

Speakers: Jeffrey A. Chan-Santiago, praveen tirupattur, Gaurav Kumar Nayak, Gaowen Liu, Mubarak Shah

Conference: ICML 2025

YouTube: https://slideslive.com/39044000

Overview

This article delves into MGD$^3$, a novel approach to Dataset Distillation (DD) that leverages the power of diffusion models with a unique mode-guided sampling strategy. Presented by Jeffrey Chan-Santiago and collaborators from UCF, IIT, and Cisco Research, the talk addresses critical challenges in generating compact, high-quality, and diverse synthetic datasets that can effectively substitute much larger original datasets for model training. The core innovation lies in its ability to sample a small dataset that is both diverse and representative, without the need for computationally expensive fine-tuning or the risk of mode collapse often associated with generative models.

Dataset distillation is a crucial area in machine learning, aiming to reduce the computational burden and storage requirements of large datasets while preserving downstream model performance. Traditional optimization-based DD methods often struggle with scalability and produce unrealistic synthetic data, while existing generative approaches can suffer from redundancy and a lack of diversity. MGD$^3$ offers a compelling solution by introducing Mode Discovery, Mode Guidance, and Stop Guidance to intelligently steer the sampling process of pretrained diffusion models. This results in state-of-the-art performance across various benchmarks, including ImageNet, and demonstrates remarkable generalizability, even working with models not originally trained on the target domain.

The significance of MGD$^3$ extends to practitioners and researchers alike. For those grappling with the sheer scale of modern datasets, MGD$^3$ provides a path to more efficient experimentation, faster model prototyping, and reduced training costs. By ensuring that distilled datasets are not only compact but also rich in information and diversity, the method promises to unlock new possibilities for democratizing access to large-scale machine learning, making advanced research and development more accessible without compromising on performance or data quality.

Background

▶ Watch: Introduction to Dataset Distillation and challenges (0:00)

The concept of Dataset Distillation (DD) emerged from the need to compress the knowledge contained within massive training datasets into significantly smaller, synthetic counterparts. The ultimate goal is that a model trained exclusively on this small, distilled dataset should achieve performance comparable to one trained on the full, original dataset. Typically, these distilled datasets are remarkably small, often comprising just 10 to 50 images per class, making the task of capturing the essence of the original data incredibly challenging.

Historically, DD approaches can be broadly categorized into two main types: optimization-based methods and generative methods.

Optimization-based approaches distill a dataset by formulating an optimization problem where a matching objective is minimized between the original and the distilled dataset. This often involves matching model statistics, such as gradients, feature distributions, or even full training trajectories. A common workflow involves taking a large dataset, using a model to compute certain learnable statistics (e.g., gradients), and then backpropagating towards the pixel space of a small, randomly initialized synthetic dataset to make its statistics match those of the original. While conceptually sound, these methods face several critical limitations:

• Time-consuming: They require a vast number of iterations for optimization.
• Scalability issues: They struggle to scale to larger datasets and high-resolution images due to the direct optimization in pixel space.
• Unrealistic outputs: The resulting synthetic images, even with some latent space optimization, often exhibit low-level features that are somewhat realistic but an overall image structure that remains "out-of-distribution" or visually unrealistic, which can hurt downstream model performance.

In contrast, generative dataset distillation aims to compress the information of a large dataset into a generative model itself. Once trained, this generative model can then be used to sample a distilled dataset. The rise of powerful generative models, particularly diffusion models, has made this approach increasingly attractive due to their exceptional ability to capture complex data distributions and generate high-fidelity, realistic images. Diffusion models work by iteratively denoising a random noise input, guiding it towards high-density regions of the data distribution until a realistic image is formed.

However, even with the prowess of diffusion models, challenges persist in generative DD:

• Redundancies: Simply sampling from a pretrained diffusion model can lead to redundancies, especially if the model tends to favor high-density regions, resulting in generated datasets with low diversity.
• Mode Collapse: Similar to issues observed in GANs, there's a risk of mode collapse, where the generative model fails to capture the full spectrum of data variations, concentrating on only a few dominant modes. This leads to a distilled dataset that is not truly representative.
• Computational Expense: Prior attempts to address diversity, such as Min-max diffusion, explicitly fine-tune the diffusion model with a representative diversity loss. While effective, this fine-tuning step is computationally expensive and hinders scalability, especially for larger datasets.

The problem, therefore, is how to effectively leverage pretrained diffusion models to sample a small, diverse, and representative dataset in just a few samples, without redundancies, mode collapse, or the need for costly fine-tuning. This is precisely the gap that MGD$^3$ aims to fill, by introducing a novel guidance mechanism that steers the diffusion process towards desired modes of the data distribution while maintaining image quality.

Key Findings

▶ Watch: MGD³'s core idea: mode discovery and guidance (2:00)

MGD$^3$ presents a significant leap forward in dataset distillation by addressing the critical challenges of diversity, realism, and computational efficiency. The talk highlighted several key findings that underscore its effectiveness:

• Diverse and Representative Dataset Generation without Fine-tuning: The most prominent finding is MGD$^3$'s ability to generate highly diverse and representative distilled datasets using a pretrained diffusion model, crucially, without any additional fine-tuning. This directly overcomes a major computational bottleneck present in prior generative DD methods like Min-max diffusion.
• Novel Guidance Mechanisms: The introduction of Mode Discovery, Mode Guidance, and Stop Guidance forms the algorithmic backbone of MGD$^3$. Mode Discovery identifies distinct data modes within each class, Mode Guidance steers the diffusion process towards these target modes, and Stop Guidance intelligently halts the guidance to preserve image quality, preventing degradation observed with continuous guidance.
• Superior Performance on Standard Benchmarks: MGD$^3$ consistently outperforms existing state-of-the-art dataset distillation methods, including both optimization-based approaches and fine-tuned generative methods like Min-max diffusion. This superiority was demonstrated across various ImageNet subsets (10, 20, 50 images per class), ImageNet-IDC, and even the full ImageNet-1K benchmark. While specific numerical improvements over baselines were not always explicitly stated, the presented charts clearly showed MGD$^3$'s method achieving higher validation accuracy.
• Component-wise Contribution: Ablation studies rigorously confirmed that each proposed component—Mode Discovery, Mode Guidance, and Stop Guidance—individually contributes to the overall performance improvement, validating the design choices and the synergistic effect of their combination.
• Optimal Guidance Timestep: Experimental analysis revealed an optimal timestep for applying Stop Guidance, specifically around timestep 20 in the diffusion trajectory. This finding is crucial for balancing diversity generation with image quality preservation.
• Enhanced Diversity in Latent Space: Qualitative evaluations, such as t-SNE plots of the generated datasets, visually confirmed that MGD$^3$ samples are more spread out in the latent space compared to baselines, directly indicating a higher degree of diversity within the distilled dataset. Visual examples further corroborated this, showing variations in backgrounds, poses, viewpoints, and even the number of subjects (e.g., dogs, individuals) within a class.
• Robustness to Untargeted Pretrained Models: MGD$^3$ demonstrated remarkable generalizability by achieving significant performance improvements even when applied to a Stable Diffusion model (a text-to-image model) for distillation on ImageNet benchmarks, despite Stable Diffusion not being specifically trained on the target dataset for this task. This highlights the method's potential for wide applicability with various pretrained diffusion models.

These findings collectively establish MGD$^3$ as a highly effective, efficient, and versatile solution for dataset distillation, pushing the boundaries of what is achievable in generating compact yet powerful synthetic training data.

Technical Deep Dive

▶ Watch: Detailed MGD³ pipeline: discovery, guidance, stop (3:30)

MGD$^3$ innovates by intelligently guiding the sampling process of latent diffusion models to create diverse and representative distilled datasets. The core idea is to leverage the implicit knowledge of data distribution within a pretrained diffusion model and explicitly guide its generation towards distinct "modes" within each class, while carefully managing the trade-off between diversity and image quality.

The technical framework of MGD$^3$ comprises three interconnected components: Mode Discovery, Mode Guidance, and Stop Guidance.

1. Motivation and the Challenge:

Diffusion models, when sampling from random noise, progressively denoise the input by following a gradient field towards high-density regions of the learned data distribution. While this produces realistic images, unguided sampling often leads to redundancies because the model naturally converges to the most prominent or "high-density" modes, especially when generating multiple samples for the same class. This issue of generating similar samples (e.g., same background, same faces) was clearly illustrated with examples of unguided generation. Prior work like Min-max diffusion attempted to solve this by fine-tuning the diffusion model with a representative diversity loss, but this is computationally prohibitive for large datasets. MGD$^3$ aims to achieve this diversity without fine-tuning.

1. MGD$^3$ Pipeline Overview:

The overall process begins with a training dataset. Since MGD$^3$ utilizes latent diffusion, the first step involves encoding the original images into their latent representations using a Variational Autoencoder (VAE) encoder. This allows the diffusion process to operate in a more compact and semantically rich latent space.

The pipeline then proceeds class by class:

• For a specific class, the generative feature distribution in the latent space is computed.
• Mode Discovery is applied to this feature distribution to identify N distinct modes per class, where N is the desired number of images per class in the distilled dataset.
• For each identified target mode, the pretrained latent diffusion model is used to generate an image. This generation starts from random noise.
• During the denoising trajectory, Mode Guidance is applied to steer the process towards the target mode.
• Crucially, Stop Guidance is then activated to stop the explicit guidance once the generation is sufficiently close to the mode, allowing the remaining denoising steps to proceed unguided.
• Finally, the generated latent representation is decoded back into pixel space using the VAE decoder. This entire process is repeated for all modes and all classes to form the complete distilled dataset.

1. Mode Discovery:

This component is responsible for identifying N distinct regions or "modes" within the latent feature distribution of a given class. The talk mentioned that several methods for mode discovery were evaluated, and Gaussian Mixture Model (GMM) was found to be the most effective. GMMs are probabilistic models that assume data points are generated from a mixture of several Gaussian distributions. By fitting a GMM to the latent features of a class, MGD$^3$ can identify the centroids or means of these Gaussian components as representative points for each mode MI. This provides concrete targets for the subsequent guidance mechanism.

1. Mode Guidance (G):

The core of MGD$^3$'s sampling strategy lies in its mode guidance mechanism. During the iterative denoising process of a diffusion model, at each timestep T, the model estimates the noise to be removed. MGD$^3$ modifies this estimated noise to steer the generation towards a specific target mode MI.

Let ε_theta(z_T, t) be the noise estimated by the pretrained diffusion model at timestep t for a noisy latent z_T. The mode guidance G is applied by updating this estimated noise:

ε'_theta(z_T, t) = ε_theta(z_T, t) - scale_factor * G

Here, G represents the direction towards the target mode MI. Specifically, G is defined as the direction towards the mode MI for the estimate of the denoised image at time T. This direction is scaled by a scale_factor and a scheduled variance at time T. The scale_factor is a hyperparameter that controls the strength of the guidance. By subtracting G from the estimated noise, the model is subtly nudged to generate features that align with the characteristics of the desired mode.

1. Stop Guidance:

While Mode Guidance is essential for diversity, continuous application throughout the entire denoising trajectory can degrade image quality. The talk showed that applying full mode guidance resulted in more diverse images but with noticeable quality degradation. This happens because the strong, continuous guidance can pull the latent representation away from the natural data manifold, leading to artifacts or less realistic details.

Stop Guidance is introduced to mitigate this. The idea is simple yet powerful: we don't need to apply mode guidance if the generation is already sufficiently close to the target mode. The experiments revealed that the optimal point to stop guidance occurs around the half-way mark of the denoising trajectory, specifically at timestep 20.

If guidance is stopped too early (e.g., at timestep 0, implying full guidance throughout), quality suffers. If guidance is stopped too late (e.g., at timestep 50, implying almost no guidance), diversity suffers because the unguided portion of the trajectory dominates, leading back to high-density regions. By stopping guidance at an optimal intermediate timestep, MGD$^3$ recovers image quality while retaining the diversity achieved in the initial guided steps. The remaining steps of the denoising process then proceed as unguided generation, allowing the diffusion model to refine the image naturally onto the data manifold, ensuring high fidelity.

This sophisticated interplay of mode identification, directed generation, and intelligent guidance cessation allows MGD$^3$ to create distilled datasets that are not only compact and diverse but also maintain the high visual quality characteristic of diffusion models.

Experimental Setup & Results

▶ Watch: Qualitative results: increased diversity and image quality (4:00)

The evaluation of MGD$^3$ was comprehensive, demonstrating its efficacy across various settings and comparing it against strong baselines. The experimental setup focused on widely accepted benchmarks in dataset distillation, primarily using subsets of ImageNet.

Datasets:

• ImageNet-Subset: The primary evaluation was conducted on subsets of the ImageNet dataset, a common practice in dataset distillation. Distilled datasets were generated with 10 images per class (IPC), 20 IPC, and 50 IPC.
• ImageNet-IDC: Another ImageNet-based subset, likely with different characteristics or sampling, was used for evaluation, referred to as ImageNet-IDC.
• ImageNet-1K: For a more challenging and large-scale assessment, MGD$^3$ was also evaluated on the full ImageNet-1K dataset.

Baselines:

The performance of MGD$^3$ was benchmarked against several key baselines, representing different paradigms of dataset distillation:

• Optimization-based DD methods: While not explicitly named in the transcript, the speaker referred to "previous methods" that optimization-based approaches struggled with scalability and realism, implying MGD$^3$ was compared against such techniques.
• Min-max diffusion: This was a crucial generative baseline. Min-max diffusion is a transformer-based UNet diffusion model that is explicitly fine-tuned with a representative diversity loss. This comparison highlighted MGD$^3$'s advantage of achieving diversity without fine-tuning.
• Transformer-based UNet diffusion without fine-tuning: This served as a baseline to show the performance of a standard diffusion model used for generative DD without any specific diversity-enhancing modifications.

Metrics:

The primary metric for evaluation was validation accuracy. This metric measures the performance of a downstream classifier model that has been trained exclusively on the distilled dataset, with the goal of matching or exceeding the accuracy achieved by a model trained on the full original dataset.

Headline Numbers and Improvements:

• MGD$^3$ consistently outperformed Min-max diffusion on evaluations across 10, 20, and 50 images per class on ImageNet-Subset and ImageNet-IDC. This is a significant result, as Min-max diffusion requires computationally expensive fine-tuning, whereas MGD$^3$ does not.
• The method also achieved state-of-the-art performance when scaled up to the full ImageNet-1K dataset, demonstrating its scalability beyond smaller subsets.
• A particularly compelling result was its performance with a Stable Diffusion model (a text-to-image model). MGD$^3$ significantly improved Stable Diffusion's evaluation performance on ImageNet-IDC and ImageNet-1K, even though Stable Diffusion was not specifically trained on these target datasets for distillation. This highlights the method's strong generalizability and ability to leverage existing powerful, albeit untargeted, generative models.

Ablation Studies:

Detailed ablation studies were conducted to understand the contribution of each proposed component:

• Component Contribution: The results showed that applying Mode Discovery, Mode Guidance, and Stop Guidance incrementally improved performance. Each component, when added to the pipeline, led to measurable gains in validation accuracy, confirming their synergistic effect and individual importance.
• Optimal Stop Guidance Timestep: A critical ablation explored the impact of the timestep at which guidance is stopped. The Y-axis represented validation accuracy, and the X-axis represented the timestep at which guidance was ceased.
• Stopping at timestep 0 (full guidance) resulted in a drop in performance due to quality degradation.
• Stopping at timestep 50 (almost no guidance) also resulted in lower performance, indicating insufficient diversity.
• The best performance was achieved by stopping guidance at timestep 20. This demonstrates the crucial balance achieved by Stop Guidance between ensuring diversity and preserving image quality.
• Mode Discovery Method: Various methods for discovering modes were evaluated. Gaussian Mixture Model (GMM) was identified as the most effective approach for estimating the N modes per class.

Qualitative Results:

• t-SNE Plots: t-SNE plots of the generated datasets for a particular class (e.g., 10 IPC) visually demonstrated that the blue dots representing MGD$^3$'s samples were more spread out in the latent space compared to baselines. This visually confirmed the enhanced diversity achieved by the method.
• Visual Examples: The presentation included visual examples that clearly illustrated the impact of MGD$^3$. Unguided generation produced samples that were very similar, often sharing backgrounds, poses, or even specific facial features. In contrast, samples generated with Mode Guidance showed more variation in backgrounds, poses, viewpoints, and even the number of individuals or objects (e.g., dogs) within the same class. When Stop Guidance was applied, this diversity was maintained, but the overall image quality was significantly recovered, appearing much more realistic than samples generated with continuous Mode Guidance.

These experimental results provide robust evidence for MGD$^3$'s superior performance, efficiency, and versatility in generating high-quality, diverse, and representative distilled datasets.

Practical Implications

▶ Watch: Ablation study: contribution of each MGD³ component (5:00)

MGD$^3$ brings several profound practical implications for practitioners, infrastructure teams, model builders, and deployers in the AI/ML landscape. By addressing core limitations of existing dataset distillation methods, it opens up new avenues for efficiency, scalability, and accessibility in machine learning development.

For Practitioners and Model Builders:

• Accelerated Experimentation and Prototyping: The ability to train models on significantly smaller datasets without a substantial drop in performance means that iterative development cycles can be drastically shortened. Model builders can quickly test new architectures, hyperparameter configurations, or training strategies, leading to faster innovation and product development.
• Reduced Training Costs: Training on large datasets is computationally expensive, incurring significant costs in terms of GPU hours and energy consumption. MGD$^3$'s distilled datasets allow for substantial reductions in these costs, making advanced ML accessible to a broader range of researchers and smaller organizations.
• Enhanced Model Accessibility: The sheer size of datasets like ImageNet can be a barrier to entry for many. By providing high-quality distilled versions, MGD$^3$ democratizes access to training resources, allowing more individuals and institutions to participate in cutting-edge ML research and application development.
• Improved Data Efficiency: MGD$^3$ ensures that the distilled data is not only compact but also rich in information and diversity. This means models trained on these datasets are exposed to a broader range of patterns and variations, potentially leading to more robust and generalized models, even with less data.
• Realistic Synthetic Data: Unlike optimization-based methods that often yield unrealistic or "out-of-distribution" images, MGD$^3$ leverages diffusion models to produce visually realistic and high-fidelity synthetic data. This is crucial for tasks where human interpretability or data fidelity is important, such as in data augmentation or synthetic data generation for privacy-preserving applications.

For Infrastructure Teams and Deployers:

• Reduced Storage Requirements: Distilled datasets are orders of magnitude smaller than their original counterparts. This directly translates to reduced storage costs and simpler data management, especially in cloud environments where data transfer and storage can be costly.
• Streamlined Data Pipelines: Smaller datasets are easier to move, process, and manage within MLOps pipelines. This can simplify data versioning, distribution, and deployment, leading to more agile and efficient infrastructure.
• Scalability to Large Datasets: MGD$^3$ explicitly addresses the scalability limitations of optimization-based methods, which struggle with high-resolution and large datasets. Its generative nature, coupled with efficient guidance, allows it to scale effectively, making it suitable for modern, massive datasets.
• Leveraging Pretrained Models: The method's ability to work effectively with pretrained diffusion models, even those not specifically trained for the target domain (e.g., Stable Diffusion on ImageNet), is a significant advantage. This means infrastructure teams can leverage existing powerful models without the need for costly and time-consuming domain-specific fine-tuning, maximizing the utility of available resources.

Tradeoffs and Limitations:

While MGD$^3$ offers substantial benefits, it's important to consider potential tradeoffs and limitations:

• Computational Cost of Mode Discovery: While the overall method avoids fine-tuning, the Mode Discovery step (e.g., fitting a GMM to latent features) does incur some computational cost, especially for very high-dimensional latent spaces or an extremely large number of classes. However, this is a one-time cost per dataset and significantly less than iterative fine-tuning.
• Hyperparameter Tuning: The scale_factor for mode guidance and the optimal timestep for stop guidance are hyperparameters that need to be tuned. While the talk provided an optimal timestep (20) for their experiments, this might vary slightly across different diffusion models or target datasets.
• Mode Collapse Risk with Extreme Scale: The speaker acknowledged that while MGD$^3$ helps prevent mode collapse on limited datasets (10-50 IPC), the issue could potentially reappear with significantly larger distilled datasets or if the mode discovery process isn't robust enough to capture subtle variations in very complex distributions. The method guarantees diversity for the limited set it samples, but it's not an infinite generator of unique modes.
• Dependence on Pretrained Diffusion Model Quality: The quality of the distilled dataset is inherently linked to the quality and capabilities of the underlying pretrained diffusion model. If the base diffusion model is poor or lacks sufficient knowledge of the data distribution, MGD$^3$ cannot magically compensate for those deficiencies, though it can still improve diversity.

In summary, MGD$^3$ provides a robust, efficient, and high-quality solution for dataset distillation that has the potential to significantly impact how machine learning models are developed and deployed. Its ability to generate diverse and realistic synthetic data without extensive fine-tuning makes it a valuable tool for addressing the challenges posed by the ever-growing scale of modern datasets.

Key Takeaways

• MGD$^3$ is a novel Dataset Distillation method that leverages pretrained latent diffusion models to generate compact, diverse, and representative synthetic training datasets.
• It introduces Mode Discovery, Mode Guidance, and Stop Guidance to intelligently steer the diffusion sampling process, preventing redundancies and mode collapse.
• No fine-tuning is required for the diffusion model, making MGD$^3$ significantly more computationally efficient than prior generative DD methods like Min-max diffusion.
• MGD$^3$ achieves state-of-the-art performance on ImageNet subsets (10, 20, 50 images/class) and the full ImageNet-1K, outperforming both optimization-based and fine-tuned generative baselines.
• Stop Guidance is crucial for image quality, halting explicit mode guidance at an optimal timestep (around 20) to balance diversity with visual fidelity, preventing degradation seen with continuous guidance.
• The method demonstrates strong generalizability, successfully improving performance even when applied to text-to-image models like Stable Diffusion on ImageNet benchmarks, highlighting its versatility across different pretrained generative models.

About the Speaker(s)

The talk was presented by Jeffrey A. Chan-Santiago, who collaborated with praveen tirupattur, Gaurav Kumar Nayak, Gaowen Liu, and Mubarak Shah. The research is a joint effort between institutions including UCF (University of Central Florida), IIT (likely Indian Institute of Technology), and Cisco Research. While specific titles and affiliations for all speakers were not detailed in the transcript, the collaboration between academic institutions and industrial research laboratories indicates a strong interdisciplinary approach to tackling advanced machine learning problems. Jeffrey Chan-Santiago was the primary presenter, articulating the technical details and experimental findings of MGD$^3$.