Source-linked AI summary
Generalizing Dataset Distillation via Deep Generative Prior
George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A. Efros, Jun-Yan Zhu
TL;DR
Dataset Distillation lacks robust transfer to new architectures and effective scaling to high-resolution datasets. GLaD addresses this by optimizing distilled data in the intermediate latent space of pretrained generative models, improving generalization across settings while adding substantial computational overhead.
Problem
Existing dataset distillation methods struggle to generalize beyond the architecture used during optimization and often fail on high-resolution datasets.
Method
GLaD uses a deep generative prior to parameterize synthetic datasets in intermediate generative-model feature spaces and serves as an add-on to existing distillation methods.
Results
GLaD consistently improves performance across multiple datasets and distillation algorithms, with improvements in generalization extending to images at and beyond 512 × 512 resolution.
Takeaways & Limitations
Generative priors provide a practical path for making dataset distillation more architecture-general and applicable to more realistic, higher-resolution datasets.
Takeaways & Limitations
Using StyleGAN-XL introduces substantial memory and time costs because checkpointing requires a second generator pass, with additional batches increasing overhead.
Abstract
from arXiv · showhide
Dataset Distillation aims to distill an entire dataset's knowledge into a few synthetic images. The idea is to synthesize a small number of synthetic data points that, when given to a learning algorithm as training data, result in a model approximating one trained on the original data. Despite recent progress in the field, existing dataset distillation methods fail to generalize to new architectures and scale to high-resolution datasets. To overcome the above issues, we propose to use the learned prior from pre-trained deep generative models to synthesize the distilled data. To achieve this, we present a new optimization algorithm that distills a large number of images into a few intermediate feature vectors in the generative model's latent space. Our method augments existing techniques, significantly improving cross-architecture generalization in all settings.
1. Introduction
Dataset Distillation seeks compact synthetic datasets that preserve the training value of large datasets, but existing methods struggle with architecture transfer and high-resolution data. GLaD addresses these challenges by parameterizing distilled data in a deep generative model’s latent space and improves generalization across architectures, datasets, resolutions, and distillation algorithms.
- Dataset Distillation compresses a large dataset into a small synthetic dataset intended to produce comparable model performance when used for training.
- Existing methods often overfit a specific network architecture and perform poorly on datasets at resolutions of 128 × 128 or higher.
- Pixel-space parameterization can learn architecture-specific high-frequency patterns, while overly strong restriction to the real-data manifold reduces expressiveness and approaches coreset selection.
- GLaD parameterizes synthetic data in the intermediate feature space of generative models, providing a generative prior that improves generalizability without prohibitively restricting expressiveness.
- GLaD consistently improves distillation across multiple datasets and algorithms, including when the generator is trained on the target dataset without additional information.
- Experiments on CIFAR-10 and ImageNet subsets up to 512 × 512 report significant cross-architecture generalization improvements and reduced high-frequency noise in high-resolution distilled images.
- GLaD is designed as a plug-and-play addition to existing and future distillation methods, supporting more realistic datasets and better transfer across architectures.
2. Related Work
Prior work established dataset distillation and explored alternative parameterizations, while GLaD uses generative-model latent spaces to improve distillation performance. The paper analyzes latent-space trade-offs and presents an optimization procedure that updates latent variables through a distillation loss.
- Dataset Distillation was introduced to update a model toward desired behaviors using very little data, with later work adding learned labels, augmentation, and gradient or trajectory matching.
- Concurrent approaches reparameterize distilled data as basis sets or collections of lower-resolution images, whereas GLaD uses generative-model parameterizations for distillation performance.
- The choice of latent space trades off expressiveness against reconstruction quality, and the paper studies this trade-off for StyleGAN-family models.
- GLaD initializes latent variables, generates synthetic images, computes a distillation loss against the real training set, and updates the latents with stochastic gradient descent.
3. Generative Latent Distillation (GLaD)
Generative Latent Distillation (GLaD) replaces pixel-space optimization with latent-space parameterization using a deep generative prior, aiming to improve cross-architecture generalization while retaining expressivity. It integrates with existing distillation methods and benefits from intermediate latent spaces that balance realism and flexibility.
- GLaD method: GLaD optimizes latent codes of a differentiable generative model instead of raw pixels, allowing existing distillation objectives to parameterize synthetic datasets.The synthetic dataset is generated from latent vectors, which can be optimized with gradient matching, distribution matching, or trajectory matching losses.
- Motivation: Dataset distillation methods traditionally optimize a synthetic set against a backbone architecture, limiting generalization to unseen architectures.The reviewed methods seek synthetic data that trains a model comparably to the full dataset, but their pixel-space parameterization can overfit the backbone.
- Choosing a latent space: Earlier generative-model layers impose stronger priors and produce more realistic images, whereas later layers provide greater flexibility for fitting the distillation objective.The latent-space choice therefore trades reconstruction quality and realism against expressiveness.
- Results: GLaD also improves higher-resolution cross-architecture distillation, including 256×256 datasets using randomly initialized or domain-trained generators.The method is presented as an add-on to existing and future distillation methods, while generator and latent-space choices affect the resulting prior.
- Results: On ImageNet at 128×128, distilling into F-space rather than pixel space significantly improves cross-architecture generalization across methods and sampled datasets.Evaluation trains AlexNet, VGG11, ResNet18, and a Vision Transformer on synthetic sets distilled with a ConvNet.
4. Experiments
Experiments evaluate GLaD across datasets, architectures, latent spaces, generator choices, and resolutions. The method consistently improves cross-architecture generalization and produces more visually coherent high-resolution distilled images.
- Experimental setup: Experiments evaluate GLaD on CIFAR-10 and ImageNet subsets using existing distillation methods and unseen network architectures.The evaluation trains networks from scratch on distilled datasets and reports validation performance across architectures.
- Latent-space analysis: Earlier latent spaces impose stronger priors, whereas later spaces provide more optimization flexibility; F12, F16, and F20 are optimal for MTT, DC, and DM respectively.These choices were identified through ImageNet-subset experiments using the flamingo class.
- Cross-architecture generalization: GLaD slightly or significantly improves cross-architecture generalization for all three tested distillation methods across every dataset.The comparison covers MTT, DC, and DM with and without the deep generative prior.
- Cross-architecture generalization: On CIFAR-10, GLaD significantly improves original MTT and DM, while providing only marginal gains for DC.This shows that the benefit also appears on lower-resolution data, though its magnitude varies by distillation method.
- High-resolution synthesis: GLaD produces high-resolution artistic images, whereas pixel-space distillation at 512×512 degenerates into high-frequency patterns.Random-noise initialization also yields generator-dependent artistic properties and supports larger stylized images.
- Generator choices: Generators trained on FFHQ, Pokémon, or no data still improve cross-architecture generalization over pixel-space distillation, often matching or surpassing ImageNet generators.These results use DC for higher-resolution datasets, with MTT and DM results reported in the appendix.
5. Discussion and Limitations
The discussion presents GLaD as a plug-and-play generative-prior addition that improves generalization and extends distilled-image synthesis to at least 512×512 resolution. Its main limitation is increased memory and computation from generator passes.
- Discussion: GLaD significantly improves distilled-image generalization and produces high-quality images at resolutions up to and likely beyond 512×512.The method is presented as a plug-and-play addition to dataset distillation methods.
- Limitations: StyleGAN-XL introduces a large memory cost, while checkpointing mitigates memory usage at the expense of a second generator forward pass.Because one generator pass is time-consuming, the additional pass doubles that overhead.
- Limitations: Large synthetic sets require multiple generator batches, further increasing the method’s runtime cost.The paper notes that more efficient differentiable generators could reduce this cost in future.
A. More Visualizations
The section directs readers to the project webpage for additional visualizations.
- More Visualizations: Additional visualizations are available on the project webpage.The cited passage provides the project URL.
B. StyleGAN Latent Spaces
StyleGAN generates images through mapped style codes and a synthesis network, while GLaD optimizes intermediate feature inputs and later style codes to parameterize synthetic data.
- StyleGAN generation: StyleGAN maps a Gaussian vector z through an MLP to a style code W, which modulates convolutional layers in the synthesis network.The synthesis network combines these modulated convolutions with a learned constant to generate the image.
- GLaD latent spaces: In W+ space, GLaD optimizes a separate W code for each style block and uses the synthesis network to generate each distilled sample.The mapping network is used only to initialize W in this parameterization.
- GLaD latent spaces: In F_n space, GLaD directly optimizes the feed-forward input to the nth style block together with W codes for subsequent blocks.Earlier style blocks and their W codes are ignored in this parameterization.
C. Dataset Specifications
The high-resolution experiments use ImageNet 1k data and expert trajectories prepared from selected ImageNet classes with standardized resizing and center cropping.
- High-resolution data comes directly from the ImageNet 1k dataset using PyTorch’s built-in ImageNet loader.
- Expert trajectories use ImageNet training-set data, with classes selected from each given subset.
- Images are resized by their short side and center-cropped to the target resolution, matching the MTT preprocessing procedure.
- The ImageNet class composition of each dataset is listed in Table 5.
D. More Experimental Results
Additional experiments compare GLaD with pixel-space distillation on the backbone architecture, with larger synthetic sets, and against whole-dataset training baselines.
- On the backbone architecture, DC and DM with GLaD perform as well as or better than pixel-based versions, while MTT performs somewhat worse.The backbone is the architecture used for distillation, and GLaD is not expected to outperform pixel-based distillation there because it reduces overfitting to that architecture.
- With 10 distilled images per class, GLaD still tends to outperform pixel-space distillation.
- Whole-dataset training results provide baseline performance for each evaluated architecture.The experiments did not tune hyperparameters, which may help explain why ConvNet often performs best.
E. Hyper-Parameters and Experimental Details
The experiments combine existing open-source distillation implementations with shared augmentations, specified optimization settings, latent-space variants, and resolution-specific backbones.
- Experiments use open-source implementations of DC+DM, MTT, and TESLA for MTT and GLaD evaluations.
- Distilled images and latents use the original methods’ optimizers and hyperparameters, while W+ latents use one-tenth the learning rate.
- Optimization spaces in StyleGAN directly optimize only the latent variables boxed in red, excluding grayed-out variables and the unused mapping network.
- MTT uses N=10 synthetic steps per iteration, M=2 matched real epochs, a maximum starting epoch of 2, and 5k optimization iterations.
- The backbone changes with resolution: ConvNetD3 for 32 × 32, ConvNetD5 for 128 × 128, ConvNetD6 for 256 × 256, and ConvNetD7 for 512 × 512.
- All experiments apply the same differentiable augmentations: color, crop, cutout, flip, scale, and rotate.The augmentations use default parameters from the DSA codebase.