Source-linked AI summary
Dataset Condensation with Distribution Matching
Bo Zhao, Hakan Bilen
TL;DR
The paper addresses the high synthesis cost of dataset condensation, where bilevel optimization and second-order derivatives undermine the goal of reducing training computation. It matches real and synthetic feature distributions across randomly sampled embedding spaces, achieving substantially faster synthesis while retaining comparable or better performance and scaling to larger datasets and applications.
Problem
Dataset condensation reduces training-set size with learned synthetic data, but prior synthesis remains computationally expensive because of bilevel optimization and second-order derivatives.
Method
The method learns synthetic images by matching real and synthetic data distributions with maximum mean discrepancy across embedding spaces from randomly initialized deep networks.
Results
45× faster synthesis than the state-of-the-art is reported for learning 50 images/class on CIFAR10, alongside applications to TinyImageNet, ImageNet-1K, continual learning, and neural architecture search.
Takeaways & Limitations
The method enables condensation of larger synthetic sets and more challenging datasets while providing synthetic memory and proxy sets for downstream use.
Takeaways & Limitations
Coreset selection remains constrained by greedy selection and by the information contained in selected original samples, while the method relies on randomly initialized network embeddings.
Abstract
from arXiv · showhide
Computational cost of training state-of-the-art deep models in many learning problems is rapidly increasing due to more sophisticated models and larger datasets. A recent promising direction for reducing training cost is dataset condensation that aims to replace the original large training set with a significantly smaller learned synthetic set while preserving the original information. While training deep models on the small set of condensed images can be extremely fast, their synthesis remains computationally expensive due to the complex bi-level optimization and second-order derivative computation. In this work, we propose a simple yet effective method that synthesizes condensed images by matching feature distributions of the synthetic and original training images in many sampled embedding spaces. Our method significantly reduces the synthesis cost while achieving comparable or better performance. Thanks to its efficiency, we apply our method to more realistic and larger datasets with sophisticated neural architectures and obtain a significant performance boost. We also show promising practical benefits of our method in continual learning and neural architecture search.
1. Introduction
Rising model-training costs motivate smaller training sets, but coreset selection is information-limited and dataset condensation remains expensive to synthesize. The paper proposes distribution matching in sampled embedding spaces to efficiently learn synthetic data and extend condensation to larger settings and downstream applications.
- Motivation: Training costs rise rapidly with larger models and datasets, while repeated training for architecture and task design multiplies computation.An exhaustive CIFAR10 neural architecture search used 100 TPU years, whereas training the best architectures took dozens of TPU minutes.
- Limitations of Existing Solutions: Coreset selection is efficient but greedily selects original samples, limiting both its optimization horizon and the information available in the selected set.The cited criteria include dataset-center distance, diversity, clustering, misclassification frequency, and negative implicit gradients.
- Limitations of Existing Solutions: Dataset condensation synthesizes a small training set for comparable testing accuracy, but prior methods require costly bilevel optimization or graph unrolling.Gradient matching avoids some unrolling, and closed-form optimization offers another efficiency improvement, but synthesis remains expensive.
- Proposed Approach: The proposed method minimizes maximum mean discrepancy between real and synthetic data across embedding spaces generated by randomly initialized deep networks.Real and synthetic samples are embedded with randomly sampled networks, and the synthetic set is optimized to reduce their distribution discrepancy.
- Contributions: 45× faster synthesis for 500 CIFAR10 images is reported, while the method also supports larger synthetic sets, larger datasets, continual learning, and neural architecture search.The method synthesizes 1250 images per CIFAR10 class and is applied to TinyImageNet and ImageNet-1K, with downstream use as memory and proxy data.
2. Methodology
The method replaces expensive bilevel dataset-condensation optimization with distribution matching between real and synthetic data in sampled embedding spaces. It optimizes only synthetic data, using class-wise matching and differentiable augmentation.
- Dataset condensation seeks a small synthetic set whose trained models achieve comparable performance to models trained on the original set.
- Existing methods optimize synthetic data through bilevel training of network parameters, requiring additional outer-loop updates and second-order derivatives.
- The proposed method embeds images into lower-dimensional spaces and matches real and synthetic feature distributions across randomly sampled embedding functions.
- Maximum mean discrepancy provides the distribution-distance objective, estimated empirically because the ground-truth distributions are unavailable.
- The method applies the same randomly sampled differentiable augmentation to real and synthetic batches, enabling synthetic data to benefit from semantic-preserving transformations.
- Synthetic data are optimized class-wise with stochastic gradient descent while embedding-network parameters are sampled rather than optimized.
3. Experiments
Experiments evaluate distribution matching across datasets, synthetic-set sizes, normalization choices, architectures, network distributions, continual learning, and neural architecture search. The method scales to larger settings, matches or exceeds baselines in several comparisons, and substantially reduces synthesis time.
- Datasets and setup: Experiments cover MNIST, CIFAR10, CIFAR100, TinyImageNet, and ImageNet-1K, with synthetic sets trained for 20,000 or 10,000 iterations depending on dataset scale.The larger datasets contain 100K images across 200 classes and 1.3M images across 1K classes, respectively.
- Comparison to baselines: Training-set synthesis methods generally outperform coreset selection; DM outperforms DSA at 50 images/class on CIFAR10 and CIFAR100, while DSA is stronger at 1/10 image(s) per class.The comparison attributes this pattern partly to DSA's inner-loop model optimization under limited synthetic-set sizes.
- Larger datasets: 27 hours on one Tesla V100 condenses TinyImageNet into three sets, while 28 hours on ten GTX 1080 GPUs condenses ImageNet-1K into the same settings.Prior bilevel-optimization methods typically do not scale to these larger, higher-resolution datasets.
- Larger datasets: 24.1% testing accuracy with 50 images/class on TinyImageNet exceeds Herding's 16.7%, while DM reaches 11.4% on ImageNet-1K versus random selection's 7.54%.On TinyImageNet, DM recovers 60% of whole-dataset baseline performance using 10% of the data.
- Batch normalization: 65.2%, 48.0%, and 28.2% testing accuracy on CIFAR10, CIFAR100, and TinyImageNet result from using Batch Normalization with 50 images/class.These improve over DM with default instance normalization by 2.2%, 4.4%, and 4.1%, respectively.
- Cross-architecture generalization: DM learned on ConvNet reaches 65.2% in cross-architecture testing and exceeds DSA by 8.0%, 8.9%, and 9.7% when evaluated with AlexNet, VGG, and ResNet.The study trains synthetic data on one architecture and evaluates it by training another architecture from scratch on the synthetic set.
- Network distribution ablation: DM performs consistently across nine network distributions, with performance variance below 1 standard deviation when networks are grouped by validation accuracy.The networks are pretrained on the original training set before being sampled for synthetic-data learning.
- Continual learning: DM produces the best condensed memory in continual learning, reaching 34.4% and 34.6% in 5-step and 10-step learning, respectively.The corresponding DSA results are 31.7% and 30.5%, while random selection reaches 24.8% in both settings.
4. Conclusion
The paper presents distribution matching as an efficient dataset condensation method that avoids bi-level optimization and second-order derivatives. This efficiency supports larger datasets, larger synthetic sets, and practical applications in continual learning and neural architecture search.
- Distribution matching avoids bi-level optimization and second-order derivative computation, allowing synthetic data from different classes to be learned independently and in parallel.
- The method extends dataset condensation to TinyImageNet and ImageNet-1K and to 1250 images per class on CIFAR10.
- 45 times faster than the state-of-the-art for learning 50 images/class synthetic data on CIFAR10.
- The method produces more informative continual-learning memory and better proxy sets for accelerating neural architecture search.
A.1. Dataset Condensation
The appendix reports implementation choices for reproducing DSA comparisons and adapting training procedures to new classes and images. These choices include coarse hyperparameter searches, batch normalization settings, and class-specific training data.
- DSA results on CIFAR100 were obtained by running released code and coarsely searching outer- and inner-loop steps, both set to 10.
- Batch normalization was enabled in ConvNet experiments when obtaining the reported DSA results in Tables 2 and 3.
- DSA training used only new training classes and images at each learning step rather than incorporating the old model and memory.
A.3. Neural Architecture Search
The NAS evaluation compares proxy-set methods and early stopping under controlled training budgets, using rank correlation to assess whether proxy training identifies strong architectures. Distribution matching provides the highest reported ranking correlation while using a small proxy set.
- The Spearman rank correlation measures agreement between proxy-set and whole-dataset performance for the top 5% of architectures.
B.1. Comparison to Generative Models
The paper compares its condensed samples with VAE, GAN, and GMMN generative models on CIFAR10. Its samples are more data-efficient for training ConvNets, while traditional generative-model samples are no more informative than randomly selected real images.
- The comparison includes DC-VAE, BigGAN with differentiable augmentation, and GMMN, alongside the proposed method.
- Unlike VAE, GAN, and GMMN image-generation objectives, the proposed method synthesizes informative training samples for efficiently training deep networks.
- The proposed method outperforms traditional generative models by large margins when ConvNets train on CIFAR10 synthetic images and are evaluated on real test images.
- Samples from traditional generative models are not more informative than randomly selected real images under the reported comparison.
B.2. Comparison to MMD Baseline
The pixel-space MMD baseline performs best with a linear kernel, but remains inferior to the proposed method while outperforming generative models.
- Only pixel-space MMD with a linear kernel produces better synthetic images than randomly selected real images.Linear, polynomial, RBF, and Laplacian kernels with various hyperparameters were tested.
- Figure F8 compares performance rank correlation between proxy-set and whole-dataset training.
- The pixel-space linear-kernel MMD baseline outperforms all generative models but remains inferior to the proposed method.
B.3. Comparison to GTN and KIP Methods
The paper compares its method with GTN and KIP, reporting comparable performance using fewer synthetic images in one MNIST comparison and better results in some settings than KIP.
- On MNIST, GTN reports 4,096 synthetic images with performance comparable to the paper’s 50 images per class, or 500 images total.
- KIP replaces neural-network optimization in bilevel condensation with closed-form kernel ridge regression and applies ZCA preprocessing.
- 98.6 ± 0.1% versus 98.3 ± 0.1% is reported when learning 50 images per class in one comparison with KIP.
C. Extended Visualization and Analysis
Synthetic images learned with lower-validation-accuracy networks look blurrier and show checkerboard patterns, whereas those from higher-accuracy networks are more colorful, yet their generalization performance is similar.
- Synthetic images learned with lower-validation-accuracy networks look blurrier and contain obvious checkerboard patterns.
- Synthetic images learned with higher-validation-accuracy networks look colorful and contain some twisty patterns.
- Despite different visual appearances, synthetic images learned from different network parameter distributions have similar generalization performance.The paper attributes the visual differences mainly to background patterns while describing the semantics as similar.
D. Connection to Gradient Matching
The paper connects last-layer gradient matching to weighted feature matching. With randomly initialized networks, this becomes mean feature matching with a constant weight.
- Setup: The analysis embeds each real and synthetic sample with a neural network and focuses on the weights and gradients of the final linear classification layer.
- Gradient-feature relation: Eq. 9 expresses each last-layer gradient vector as a feature vector weighted by the classification probability error, and vice versa.
- Gradient-feature relation: The weight is larger when the predicted probability differs more from the ground-truth one-hot label.
- Batch matching: Averaging corresponding gradient components over same-class real and synthetic batches yields weighted mean-feature matching.
- Random initialization: With randomly initialized networks and uniform predicted probabilities, last-layer mean gradient matching is equivalent to mean feature matching multiplied by a constant.