Dataset Distillation for Pre-Trained Self-Supervised Vision Models
This work addresses the challenge of efficiently training linear probes on large pre-trained models, offering a tool for fine-grained classification and model interpretability, though it is incremental as it builds on existing distillation methods for a new context.
The paper tackles the problem of dataset distillation for pre-trained self-supervised vision models by introducing Linear Gradient Matching, which optimizes synthetic images to match gradients from real data in linear classifiers, resulting in synthetic data that outperforms real-image baselines and generalizes across models, such as training a CLIP probe competitively using DINO-distilled data.
The task of dataset distillation aims to find a small set of synthetic images such that training a model on them reproduces the performance of the same model trained on a much larger dataset of real samples. Existing distillation methods focus on synthesizing datasets that enable training randomly initialized models. In contrast, state-of-the-art vision approaches are increasingly building on large, pre-trained self-supervised models rather than training from scratch. In this paper, we investigate the problem of distilling datasets that enable us to optimally train linear probes on top of such large, pre-trained vision models. We introduce a method of dataset distillation for this task called Linear Gradient Matching that optimizes the synthetic images such that, when passed through a pre-trained feature extractor, they induce gradients in the linear classifier similar to those produced by the real data. Our method yields synthetic data that outperform all real-image baselines and, remarkably, generalize across pre-trained vision models, enabling us, for instance, to train a linear CLIP probe that performs competitively using a dataset distilled via a DINO backbone. Further, we show that our distilled datasets are exceptionally effective for fine-grained classification and provide a valuable tool for model interpretability, predicting, among other things, how similar two models' embedding spaces are under the platonic representation hypothesis or whether a model is sensitive to spurious correlations in adversarial datasets.