Dataset Distillation for Pre-Trained Self-Supervised Vision Models

AI summary

Overview

Research area: Dataset distillation / efficient machine learning for computer vision foundation models.

Technical level: Intermediate. Readers should be comfortable with neural network training, gradient-based optimization, and the concept of self-supervised pre-training, but the core idea is intuitive and the paper explains it clearly.

Scope: The paper defines a new dataset distillation setting — synthesizing tiny image sets to train linear classifiers on top of frozen self-supervised vision backbones — and introduces a gradient-matching method called Linear Gradient Matching that beats every real-image baseline using just one synthetic image per class.

What This Paper Is About

Standard dataset distillation asks: can we replace a huge training set with a handful of synthetic images that teach a model the same thing? Almost all prior work answers this for models trained from scratch. But modern vision practice rarely trains from scratch — it fine-tunes or probes large pre-trained self-supervised models. This paper addresses the mismatch by distilling synthetic images specifically designed to train linear probes on frozen self-supervised backbones like CLIP, DINO-v2, EVA-02, and MoCo-v3. The goal is to reach near-full-dataset linear probe accuracy while only ever showing the classifier one labeled image per class.

Key Contributions

1. A new problem formulation. The authors reframe dataset distillation around pre-trained self-supervised feature extractors and linear probing rather than random-initialized training from scratch, separating it from the memory and stability bottlenecks that limit prior distillation methods on large models.

2. Linear Gradient Matching (LGM). A distillation objective that samples random linear classifiers, computes cross-entropy losses on real and synthetic images, and minimizes the cosine distance between the resulting gradients with respect to that classifier — a lightweight meta-loss that scales to large backbones.

3. Two regularization techniques that unlock cross-model transfer. A multi-scale pyramid image parameterization (from 1×1 up to 256×256, progressively grown during optimization) and channel color decorrelation, plus heavy differentiable augmentation, which together prevent synthetic images from overfitting to the distillation backbone.

4. Demonstrated interpretability uses. The distilled images serve as a diagnostic: they predict pairwise model alignment under the Platonic Representation Hypothesis, expose spurious-correlation sensitivity in adversarial datasets, and reveal out-of-distribution generalization capabilities of a backbone.

Main Findings

• One image per class beats every real-image baseline. On ImageNet-100 and ImageNet-1k, distilled images outperform "Neighbors" (real images closest to the synthetic embeddings), "Centroids" (real images closest to class means), and random selection, across all four backbones. On ImageNet-1k, a DINO-v2 linear probe trained on distilled data reaches 75.0% test accuracy, versus 83.0% for a probe trained on the full 1.3 million-image dataset.

• Distilled datasets transfer across architectures. A dataset distilled with one backbone trains competitive linear probes on other, unseen backbones. DINO-v2 produces the best cross-model distilled data (63.7% average on ImageNet-1k), and the strongest model on the full dataset also yields the most transferable synthetic data.

• The CLIP–MoCo pair is the notable exception. Cross-model transfer between CLIP and MoCo-v3 degrades sharply (46.6% and 57.2% normalized k-NN accuracy), which the authors attribute to poor alignment between those two embedding spaces.

• Augmentation is the single most important ingredient. Removing differentiable augmentations collapses same-model average accuracy from 87.2% to 68.6% and cross-model accuracy from 77.8% to 33.3% on ImageNet-100. The pyramid representation is the second most important, and matters far more for cross-model transfer than for same-model performance. Color decorrelation gives the smallest gain.

• Fine-grained classification shows an even larger margin. On Stanford Dogs and CUB-200-2011, distilled images outperform real-image baselines by a wider gap than on ImageNet, suggesting a single synthetic image can encode more class-discriminative information than any single real image.

• Adversarial datasets become interpretable. On Spawrious, where class labels are spuriously correlated with backgrounds, distilled images from DINO-v2 show recognizable dog breeds while MoCo-v3's show mostly backgrounds — matching MoCo-v3's catastrophic test accuracy of roughly 36% versus DINO-v2's 78%.

• Cross-model performance tracks representation alignment. The accuracy of a distilled dataset when evaluated on a different model correlates strongly with the mutual k-nearest-neighbor alignment of the two models' embedding spaces.

• Backbones can distill out-of-distribution data. DINO-v1, trained only on ImageNet, successfully distills ArtBench art styles into one image per class, and the synthetic images differ starkly from their nearest real neighbors — evidence the method synthesizes rather than retrieves.

Methodology in Plain English

The idea is to make synthetic images whose effect on a classifier looks like the effect of real data. At each optimization step, the method draws a fresh random linear classifier matrix. Real images and synthetic images are both pushed through the frozen feature extractor, then through that classifier, producing two cross-entropy losses. The gradients of those two losses with respect to the classifier weights capture "which direction the classifier wants to move." The authors compute the cosine distance between these two gradient vectors and back-propagate that distance all the way into the synthetic pixels, nudging them so that their gradient direction matches the real data's.

To stop the images from turning into adversarial-looking noise, three stabilizers are used. First, images are not stored as pixels but as a stack of resolutions from 1×1 up to 256×256 that are upsampled and summed before a sigmoid clamp; the optimization starts with only the coarsest level and adds finer levels over time, which keeps low frequencies dominant. Second, the color channels are learned in a decorrelated space and transformed back, discouraging the model's color biases from leaking into the images. Third, each synthetic image is augmented several times per step — flips, random resized crops, Gaussian noise — and the copies are concatenated into the training batch, so the optimization tries to make all augmented variants collectively useful rather than overfitting one view.

Evaluation is simple: distilling runs for 5,000 iterations at 224×224 with a ViT-B backbone, then a freshly initialized linear probe is trained to convergence on the synthetic set and tested on the real test set.

Why This Matters

Dataset distillation has traditionally been a research curiosity because it assumed training from scratch, which almost nobody does anymore. This paper repositions distillation as a tool for the dominant modern workflow: adapting frozen foundation models. The result that a single image per class can capture most of a linear probe's accuracy has direct practical value for anyone working with limited labels, and the interpretability findings give the method a second life as a diagnostic instrument rather than just a data compressor.

Real-world applications:

• Low-label domain adaptation: Quickly adapting medical imaging, satellite, or industrial inspection foundation models where expert annotation is expensive or impossible at scale.
• Model auditing and bias detection: Using distilled images as a visual probe to check whether a backbone relies on spurious cues before it is deployed in a high-stakes setting.
• Backbone selection and compatibility profiling: Predicting which pre-trained encoder will transfer well to a downstream task, and which pairs of encoders are interchangeable within a system.
• Privacy-preserving data sharing: Distributing tiny synthetic datasets that carry task-relevant signal without exposing the underlying real images.

Industry relevance: The cost of fine-tuning or probing foundation models is often dominated by the labels needed, not the compute. A method that reduces labeled data to one image per class and transfers across backbone vendors (CLIP, DINO, EVA, MoCo) lowers the barrier to adopting new foundation models and provides a cheap precedent to large-scale evaluation and compliance checks.

Future Directions

• Extending beyond linear probes. The current method optimizes only for a linear classifier on frozen features. Whether the same gradient-matching idea can distill data for adapter tuning, LoRA, or partial fine-tuning is an open question.
• Diagnosing and closing alignment gaps. The CLIP–MoCo transfer failure is observed but not explained theoretically. A deeper account of when distilled data transfers would sharpen the use of distillation as an alignment measurement.
• Mitigating, not just revealing, spurious correlations. Distilled images expose a backbone's reliance on backgrounds, but the paper does not attempt to correct the underlying data or model. Turning the diagnostic into a debiasing procedure is a natural next step.
• Scaling and compute efficiency. ImageNet-1k distillation required reducing augmentation rounds from 10 to 3 for compute reasons. Better optimization or amortized distillation could make large-scale distillation routine.
• Theory of the one-image-per-class regime. Why a single synthetic image suffices to capture a class's discriminative content — and when this breaks down — remains largely empirical.

Target Audience

Researchers in dataset distillation and dataset compression, self-supervised learning practitioners, and interpretability researchers interested in what frozen foundation models actually encode. It is also valuable for applied machine learning engineers who routinely fine-tune or probe pre-trained backbones under tight labeling budgets, and for anyone evaluating foundation models for bias or domain-transfer behavior.

Authors’ abstract

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.