Diffusion-Augmented Coreset Expansion for Scalable Dataset Distillation

TL;DR

This work tackles the scalability bottlenecks of dataset distillation by proposing a two-stage pipeline: (i) extract a compact coreset of informative image patches via patch-level cross-entropy using a trained teacher under a memory budget, and (ii) dynamically expand and super-resolve this coreset on-the-fly with Latent Diffusion Models to inject realism and controlled diversity. It further enhances diversity through latent-space mixup and anchors augmentations to real data, enabling high-quality, on-the-fly data generation during training. The method achieves significant gains on large-scale benchmarks and demonstrates strong cross-architecture transfer, addressing realism and diversity gaps that hinder prior approaches. Overall, diffusion-augmented coreset expansion offers a scalable, practical path toward high-fidelity, compact training data for diverse architectures with reduced storage and communication demands.

Abstract

With the rapid scaling of neural networks, data storage and communication demands have intensified. Dataset distillation has emerged as a promising solution, condensing information from extensive datasets into a compact set of synthetic samples by solving a bilevel optimization problem. However, current methods face challenges in computational efficiency, particularly with high-resolution data and complex architectures. Recently, knowledge-distillation-based dataset condensation approaches have made this process more computationally feasible. Yet, with the recent developments of generative foundation models, there is now an opportunity to achieve even greater compression, enhance the quality of distilled data, and introduce valuable diversity into the data representation. In this work, we propose a two-stage solution. First, we compress the dataset by selecting only the most informative patches to form a coreset. Next, we leverage a generative foundation model to dynamically expand this compressed set in real-time, enhancing the resolution of these patches and introducing controlled variability to the coreset. Our extensive experiments demonstrate the robustness and efficiency of our approach across a range of dataset distillation benchmarks. We demonstrate a significant improvement of over 10% compared to the state-of-the-art on several large-scale dataset distillation benchmarks. The code will be released soon.

Figures (5)

• Figure 1: Proposed framework illustration: starting with an image dataset $\mathcal{D}$, a teacher model is trained on the image-label pairs. Leveraging the uncertainty signal from the teacher's logits, following sun2024diversity, we identify the most important patch from each image to form a coreset. These patches are then upsampled, noise-corrupted using fixed random seeds, and processed through a multi-step diffusion model to achieve simultaneous super-resolution and introduce variations to the coreset. For each random seed and generated high-resolution image, the teacher's soft label is obtained. The student then uses these important patches and random seeds to recreate the high-resolution images and regress over the teacher's corresponding soft labels. Note that, similar to traditional geometric augmentation techniques, this super-resolution and augmentation process is performed on the fly and discarded once the student's gradient is computed.
• Figure 2: The extracted coreset for IPC=1 from ImageNette.
• Figure 3: Performing mixup in the latent space of the LDM's autoencoder.
• Figure 4: Qualitative illustration of the generated samples. By adjusting $\rho$, we can control the level of augmentation, allowing us to effectively distinguish between the contributions of super-resolution and augmentation.
• Figure 5: Study on the impact of patch size on student performance. In the RDED case, there is a trade-off between realism and diversity: reducing patch dimensions allows for more patches to fit in memory but significantly lowers the quality of downsampled patches without super-resolution. In contrast, our method benefits from increased performance by adding more patches, albeit at the cost of additional computation due to diffusion model calls. The x-axis represents varying $r$ values, with the numbers in parentheses indicating the corresponding patch sizes ($\frac{224}{r}$).