Efficient Dataset Distillation for Pre-Trained Self-Supervised Models via Statistical Flow Matching

TL;DR

This paper introduces a new statistical flow matching method for dataset distillation that significantly reduces computational costs while maintaining or improving performance on pre-trained self-supervised models.

Contribution

It proposes a stable, efficient framework that aligns statistical flows for synthetic data generation, reducing memory and runtime compared to prior gradient matching methods.

Findings

Achieves comparable or better performance with 10x lower GPU memory usage.

Reduces runtime by 4x compared to state-of-the-art methods.

Introduces classifier inheritance for efficient inference with minimal additional storage.

Abstract

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.