Source-linked AI summary
Improved Distribution Matching for Dataset Condensation
Ganlong Zhao, Guanbin Li, Yipeng Qin, Yizhou Yu
TL;DR
Dataset condensation seeks smaller training sets, but optimization-oriented methods are computationally expensive and scale poorly. IDM improves distribution matching by addressing feature imbalance and unvalidated embeddings with three techniques. The resulting method improves over DM and most optimization-oriented methods while requiring fewer resources and supporting larger datasets.
Problem
Optimization-oriented condensation methods are computationally expensive and scale poorly, while naive distribution matching suffers from imbalanced feature numbers and unvalidated embeddings.
Method
IDM uses partitioning and expansion augmentation, efficient and enriched model sampling, and class-aware distribution regularization to improve distribution matching.
Results
IDM significantly improves over DM, surpasses most optimization-oriented methods, and applies to large datasets such as ImageNet Subset with fewer computational resources.
Takeaways & Limitations
IDM makes distribution matching competitive with optimization-oriented condensation while retaining high efficiency and scalability.
Takeaways & Limitations
The method assumes models can be trained for at most K × Nmax times before replacement, requiring K × Nmax to be large enough to capture informative real-data features.
Abstract
from arXiv · showhide
Dataset Condensation aims to condense a large dataset into a smaller one while maintaining its ability to train a well-performing model, thus reducing the storage cost and training effort in deep learning applications. However, conventional dataset condensation methods are optimization-oriented and condense the dataset by performing gradient or parameter matching during model optimization, which is computationally intensive even on small datasets and models. In this paper, we propose a novel dataset condensation method based on distribution matching, which is more efficient and promising. Specifically, we identify two important shortcomings of naive distribution matching (i.e., imbalanced feature numbers and unvalidated embeddings for distance computation) and address them with three novel techniques (i.e., partitioning and expansion augmentation, efficient and enriched model sampling, and class-aware distribution regularization). Our simple yet effective method outperforms most previous optimization-oriented methods with much fewer computational resources, thereby scaling data condensation to larger datasets and models. Extensive experiments demonstrate the effectiveness of our method. Codes are available at https://github.com/uitrbn/IDM
1. Introduction
Dataset condensation reduces large datasets to smaller training sets, but optimization-oriented methods are computationally expensive and scale poorly. IDM addresses distribution matching’s feature imbalance and embedding shortcomings with three techniques, improving performance and scalability.
- Motivation: Dataset condensation creates smaller datasets that retain model-training validity and reduce storage costs.The condensed datasets can also benefit downstream tasks.
- Limitations of prior methods: Optimization-oriented methods match gradients or intermediate parameters because directly matching converged model parameters is difficult.They impose similar optimization paths between models trained on real and condensed data.
- Limitations of prior methods: Nested optimization makes prior methods computationally expensive and poorly scalable to large datasets and models.They optimize condensed data and model parameters alternately, often requiring pretrained models and additional computation or storage.
- Distribution matching: Distribution Matching avoids model optimization by matching feature distributions extracted from real and condensed datasets with randomly initialized models.This removes the expensive nested loops and improves efficiency and scalability.
- IDM: IDM addresses class misalignment caused by imbalanced feature numbers and unvalidated embeddings through three techniques.The techniques are partitioning and expansion augmentation, efficient and enriched model sampling, and class-aware distribution regularization.
- IDM: IDM significantly improves over DM, surpasses most optimization-oriented methods, and scales to large datasets such as ImageNet Subset.The method is described as highly efficient and scalable.
2. Related Works
Dataset condensation methods include meta-learning, optimization-based approaches, distribution matching, generative modeling, and coreset selection. Distribution matching improves scalability by avoiding bi-level optimization, but its efficiency has historically been accompanied by inferior performance.
- Dataset Condensation: Dataset condensation preserves information needed to train models while supporting continual learning, neural architecture search, federated learning, and privacy-preserving machine learning.Data Distillation pioneered the topic with meta-learning, while later methods introduced techniques such as soft labels.
- Distribution Matching: Distribution Matching discards bi-level optimization and matches feature distributions, reducing memory and computation for large-dataset condensation.Its high efficiency has come with inferior performance that limits further application.
- Generative modeling: Generative modeling methods are not compared because they use different experimental settings.
- Coreset Selection: Coreset selection identifies representative samples using criteria such as compactness, diversity, or forgetfulness.These heuristic criteria are unrelated to target tasks, cannot guarantee an optimal solution, and depend on original-image quality.
3. Problem Definition
Dataset condensation synthesizes a much smaller labeled set intended to retain the training power of a large dataset. Distribution Matching replaces nested optimization with MMD-based matching of real and condensed feature distributions using randomly initialized networks.
- Dataset Condensation: Given a labeled training set T, condensation synthesizes a much smaller labeled set S with similar model-training power.The target size satisfies |S| ≪ |T|.
- Dataset Condensation: The original condensation objective compares the losses of models trained on the real and condensed datasets.The models share the same architecture and use cross-entropy loss.
- Distribution Matching: Distribution Matching matches feature distributions ϕθ(x_i) and ϕθ(s_j) for real and condensed samples using Maximum Mean Discrepancy.This converts the condensation objective into a feature-distribution matching problem.
- Distribution Matching: DM samples network parameters from Pθ0 and avoids training those networks, reducing computational cost.Pθ0 denotes the distribution of randomly initialized network parameters.
4. Methodology
IDM improves distribution matching by addressing imbalanced feature counts and insufficient embeddings through augmentation, enriched model sampling, and class-aware regularization.
- Shortcomings of Distribution Matching: Naive distribution matching compares thousands of real-image features with only a few synthetic-image features, making direct distribution matching inherently difficult.This imbalance is especially problematic when the real dataset is not highly redundant.
- Partitioning and Expansion Augmentation: Partitioning and expansion augmentation splits each synthetic image into l × l pieces and expands them back to the original size before feature extraction.The procedure increases synthetic features from |S| to l^2|S| without increasing the synthetic dataset size.
- Efficient and Enriched Model Sampling: IDM enriches random model embeddings by sampling models initialized from Pθ0 and trained for different numbers of iterations, forming Pθ(T).This adds training iterations as an additional sampling dimension so more informative features can be extracted.
- Efficient and Enriched Model Sampling: A memory-efficient model queue estimates Pθ(T) through Sample, Train, Push, and Pop operations while controlling computation and model storage.The queue maintains model diversity across initialization states and training iterations, with Pop limiting its size to Nmax.
- Class-aware Distribution Regularization: The enriched embeddings still produce scattered features, so IDM adds a class-aware classification-loss regularizer to make synthetic features more distinguishable.The authors associate the scattering with matching means without adequately matching higher-order distribution moments.
- Overall Objective: IDM combines enriched distribution matching with class-aware regularization in an overall loss weighted by λ_reg.The distribution-matching term samples model parameters from Pθ(T), while the regularizer is Acc_ϕLCE.
5. Experiment
Experiments evaluate IDM across multiple datasets, condensation ratios, ablations, architectures, and continual-learning settings. IDM improves substantially over DM and often surpasses optimization-oriented methods while remaining effective across challenging settings.
- Experimental Setup: IDM is evaluated on CIFAR-10, CIFAR-100, TinyImageNet, and ImageNet Subset using 1, 10, and 50 images per class.Image classification accuracy is used as the evaluation proxy for models trained on condensed sets.
- Comparison with Previous Methods: IDM surpasses DM by 9.7% on CIFAR-10 with 10 Img/Cls and 15.4% on CIFAR-100 with 10 Img/Cls.It also outperforms all compared optimization-oriented methods in the reported comparison.
- Comparison with Previous Methods: Pre-trained models outperform random models for DM across six CIFAR-100 training durations, while IDM’s mixed model sampling performs best.The proposed sampling avoids storing many pre-trained models and tuning the number of training epochs.
- Ablation Study: On CIFAR-100, model sampling, distribution regularization, and augmentation improve DM by 2.3%, 2.3%, and 5.4%, respectively.On ImageNet Subset, model sampling decreases performance by 1.1%, whereas distribution regularization improves it by 3.7%; regularization requires model sampling.
- Ablation Study: A 2 × 2 partition performs best on ImageNet Subset with 1 Img/Cls, while 3 × 3 is slightly worse because it discards more image detail.High partition counts are reported as unsuitable for low-resolution datasets such as CIFAR-10 and CIFAR-100.
- Architectural Generalization: IDM with ConvNet performs best across ConvNet, AlexNet, VGG11, and ResNet18 evaluation architectures and significantly outperforms DM.More complex architectures generally perform worse because of optimization difficulty and noisier extracted features.
- Continual Learning: In continual learning, IDM reaches 39.3% for both 5-step and 10-step settings, exceeding Random, DSA, and DM.The corresponding performances are 24.8%/25.2% for Random, 31.0%/29.8% for DSA, and 33.8%/33.7% for DM.
6. Conclusion
The paper introduces IDM, an improved distribution-matching method for dataset condensation that addresses shortcomings of naive distribution matching. It improves performance while using fewer computational resources and scaling to larger datasets with more categories.
- IDM addresses imbalanced feature numbers and unvalidated embeddings through three techniques: partitioning and expansion augmentation, enriched model sampling, and class-aware distribution regularization.
- IDM achieves significant improvements over previous methods while requiring fewer computational resources.
- IDM can be applied to larger datasets with more categories at minimal extra cost.
- Extensive experiments demonstrate IDM’s effectiveness and the competitiveness of distribution-matching methods.
1. Implementation Details
The implementation details specify a fixed learning rate for Table 1 experiments and an interval between model-queue operations to reduce condensation training cost.
- Table 1 experiments use a learning rate of 0.2.
- An extra interval between Push and Train model-queue operations reduces condensation training cost.
- For CIFAR-10 and CIFAR-100, the model-queue interval is set to 30.
2. Cross-architectural Experiments
Across three additional architectures, the proposed method outperforms distribution matching in all 16 evaluated settings.
- 2. Cross-architectural Experiments: The method is compared with distribution matching across three additional architectures in Table 5.These experiments complement Table 4 in Section 5.5 of the main paper.
- 2. Cross-architectural Experiments: 16 settings are evaluated across the four-by-four comparison grid.The passage describes these as 4×4 different settings.
- 2. Cross-architectural Experiments: The proposed method outperforms distribution matching in all 16 settings.The full comparison is reported in Table 5.
3. Start with Pre-trained Models
For larger and more difficult datasets, the model queue can require substantial training to produce meaningful embeddings, so IDM modifies it to start from pretrained models and preserve model diversity.
- A large K and Nmax may be needed for the model queue to contain sufficiently trained models on larger, more difficult datasets.Models trained for K×Nmax iterations must be sufficiently effective to extract meaningful embeddings.
- Table 5 evaluates IDM against DM on CIFAR-10 with 10 Img/Cls using three additional architectures.
- IDM enables the model queue to start with pretrained networks to reduce the computational burden of training from scratch.
- A small pretrained model set supplies randomly selected models to the model queue after t training iterations.
- Learning-rate perturbations and class-subset training diversify pretrained models and reduce training effort on ImageNet Subset.
4. Visualization of Condensed Synthetic Sets
The figures visualize condensed synthetic sets on CIFAR-100 and CIFAR-10, comparing results with and without partition and expansion augmentation. Without the augmentation, CIFAR-100 images show repetitive textures and potentially lower pixel utilization, while original DM exhibits similar texture patterns.
- On CIFAR-100 with 1 image per class, the method’s synthetic sets are visualized with and without partition and expansion augmentation.Figures 9 and 10 provide the with-augmentation and without-augmentation comparisons.
- Without partition and expansion augmentation, CIFAR-100 images contain repetitive textures that might indicate lower pixel utilization as fine-grained details are discarded during forward propagation.
- Original DM results on CIFAR-100 share texture patterns similar to those produced by the proposed method.
- On CIFAR-10 with 10 images per class, the synthetic sets are visualized with and without the proposed partition and expansion augmentation.Figures 12 and 13 show the two augmentation settings.