Dataset Distillation for Pre-Trained Self-Supervised Vision Models

What came before

Essentially all prior dataset distillation (DC's gradient matching, MTT's trajectory matching, kernel-closed-form methods) targets training a randomly initialized model from scratch, and scaling those bi-level objectives to large models runs into memory limits and optimization instability (per the paper's reading of MTT and its extensions); meanwhile vision practice has shifted toward frozen, pretrained self-supervised backbones (CLIP, DINO-v2, EVA-02, MoCo-v3) with only a lightweight linear head trained on top, a regime existing distillation methods do not target.

The problem

No prior method distills a dataset specifically for training a linear probe on top of a frozen pretrained feature extractor, and a naive port of gradient matching to this setting (matching gradients of the full model or of a randomly-initialized model) does not obviously transfer, since the object of interest is now only the linear classifier's gradient in a fixed embedding space, not a whole network's parameters.

The idea

Restrict gradient matching to just the linear classifier: sample a fresh random linear layer at each distillation step, pass real and synthetic images through the frozen pretrained feature extractor and that linear layer, and optimize the synthetic images so the classification-loss gradient with respect to the linear layer matches between real and synthetic data — reducing dataset distillation, in this regime, to matching gradients of a single linear layer rather than an entire network.

How it works

Given a frozen self-supervised feature extractor $\phi$, at each step sample $W\sim\mathcal{N}(0,1)^ {c\times f}$, compute $\ell_{real}=CE(W\phi(X_{real}),Y_{real})$ and $\ell_{syn}=CE(W\phi(X_{syn}), Y_{syn})$, and minimize the meta loss $\mathcal{L}_{meta}=1-\cos(\text{vec}(\partial\ell_{real}/ \partial W),\text{vec}(\partial\ell_{syn}/\partial W))$, backpropagating through the inner gradient computation to update $X_{syn}$ (Eq. 1-2, Fig. 2). To prevent the resulting images from overfitting to the specific backbone used for distillation — which the paper finds happens with naive pixel optimization — three additions are used: (1) a multi-scale pyramid image parameterization (resolutions 1x1 to 256x256, rendered by bilinear upsampling and summing, optimized progressively from coarse to fine) rather than raw pixels; (2) a fixed decorrelated-color-space linear transform (as in feature- visualization work) to avoid color biases from the distilling model; (3) differentiable augmentations (flip, random-resized-crop, Gaussian noise) applied as multiple augmented copies per synthetic image per step, concatenated together in the loss. Stored artifact: synthetic images plus their (fixed) class labels; no generator or teacher network is kept after distillation. Feature extractor $\phi$ stays frozen throughout; only $X_{syn}$ (and the ephemeral random $W$) are optimized.

Evidence

Table 1 (ImageNet-100 and ImageNet-1k, 1 image/class, four backbones — CLIP, DINO-v2, EVA-02, MoCo-v3): Linear Gradient Matching beats Random, Centroid, and Nearest-Neighbor real-image baselines on every backbone/dataset; e.g. ImageNet-1k with DINO-v2: 75.0% vs. best real-image baseline (Centroids) 69.5%, and vs. full-dataset (1.28M images) 83.0% — reaching within 8 points of full-data accuracy from a single image per class. Table 2 (cross-model transfer): images distilled with one backbone train competitive linear probes on a different backbone's embeddings in most cells (e.g. DINO-v2-distilled ImageNet-1k evaluated on EVA-02: 65.4%), with one notable failure pair (CLIP ↔ MoCo-v3, attributed to poor model alignment, Sec. 4.6). Table 3 / Fig. 3, 5 (ablation, ImageNet-100): removing the pyramid representation causes the largest cross-model performance drop (high-frequency overfitting artifacts visibly appear in Fig. 5); removing color decorrelation has the smallest quantitative effect but the largest visual effect (oversaturated, incorrect colors); increasing the number of differentiable- augmentation rounds per step monotonically improves both same-model and cross-model performance (Fig. 3). Cost (Compute Budget, Appendix): distilling ImageNet-100 takes ~3 hours on 1 H200 GPU; ImageNet-1k takes ~12 hours on 4 H200 GPUs (limited to 3 augmentation rounds instead of 10 due to a data loading bottleneck); the paper notes trained linear probes converge in minutes vs. up to a full day training on the full dataset.

Limitations

Explicitly memory- and data-loading-bound — ImageNet-1k experiments had to use fewer augmentation rounds (3 vs. 10) than smaller experiments due to a data-loading bottleneck; the bi-level optimization forces use of PyTorch's DistributedParallel rather than DistributedDataParallel, causing slowdowns at the large batch sizes used; only linear-probe evaluation is studied, not full fine-tuning of the backbone; one backbone pair (CLIP and MoCo-v3) transfers poorly to each other, an unexplained outlier the paper attributes to representational misalignment rather than a mechanism failure; all experiments use ViT-B-sized backbones at 224x224, leaving larger backbone scales untested.