Dataset Distillation via Relative Distribution Matching and Cognitive Heritage
This work addresses efficiency bottlenecks in dataset distillation for researchers and practitioners, offering a more scalable solution with incremental improvements in speed and resource usage.
The paper tackles the computational and memory inefficiency of dataset distillation for classification tasks by introducing statistical flow matching, which reduces GPU memory usage by 10x and runtime by 4x while maintaining or improving performance compared to state-of-the-art methods.
Dataset distillation seeks to synthesize a highly compact dataset that achieves performance comparable to the original dataset on downstream tasks. For the classification task that use pre-trained self-supervised models as backbones, previous linear gradient matching optimizes synthetic images by encouraging them to mimic the gradient updates induced by real images on the linear classifier. However, this batch-level formulation requires loading thousands of real images and applying multiple rounds of differentiable augmentations to synthetic images at each distillation step, leading to substantial computational and memory overhead. In this paper, we introduce statistical flow matching , a stable and efficient supervised learning framework that optimizes synthetic images by aligning constant statistical flows from target class centers to non-target class centers in the original data. Our approach loads raw statistics only once and performs a single augmentation pass on the synthetic data, achieving performance comparable to or better than the state-of-the-art methods with 10x lower GPU memory usage and 4x shorter runtime. Furthermore, we propose a classifier inheritance strategy that reuses the classifier trained on the original dataset for inference, requiring only an extremely lightweight linear projector and marginal storage while achieving substantial performance gains.