Self-Supervised Dataset Distillation for Transfer Learning

What came before

Supervised dataset distillation methods (DSA, DM, MTT, and kernel-ridge-regression methods KIP and FRePo) optimize synthetic images against a labeled classification objective and are not designed to produce a dataset useful for self-supervised (label-free) pre-training; a naive fix would substitute a standard SSL loss (e.g. Barlow Twins, SimCLR) for the classification loss inside the same bi-level optimization used by those methods.

The problem

The naive substitution is unstable in practice, and the paper proves why: SSL losses are stochastic functions of a randomly sampled data augmentation or input mask, and Theorem 1 shows that the gradient of the outer SSL loss with respect to the synthetic images, when the inner loop uses only a finite number of stochastic augmentation samples, is a biased estimator of the true gradient in bi-level optimization (unlike in standard non-bilevel SSL training, where the same stochastic gradient is unbiased) — because bias/covariance terms coupling the augmentation randomness to the inner optimization's implicit dependence on $X_s$ do not cancel.

The idea

Remove the randomness from both the inner and outer objectives by replacing the SSL loss with a deterministic mean-squared-error regression onto learnable target feature vectors: distill both synthetic examples $X_s$ and paired learnable target representations $Y_s$, train the inner model to regress $Y_s$ from $X_s$ (no augmentation, no randomness), and match the resulting model's representations on the real data to those of a self-supervised teacher model pretrained on the full dataset with a real SSL objective.

How it works

Pretrain a target model $g_\phi$ (ResNet-18) on the full unlabeled dataset $X_t$ with a real SSL objective (Barlow Twins) — the only place actual self-supervised training with augmentations happens. Then jointly optimize $(X_s, Y_s)$ by inner objective $\mathcal{L}_{inner}=\frac{1}{2}\|Y_s-\hat{g}_ \theta(X_s)\|_F^2$ and outer objective $\min_{X_s,Y_s}\frac{1}{2}\|g_\phi(X_t)-\hat{g}_{\theta^*(X_s, Y_s)}(X_t)\|_F^2$. Following FRePo, $\hat{g}_\theta$ is split into a feature extractor $f_\omega$ and linear head $h_W$; a pool of $l{=}10$ feature extractors is maintained and periodically reset, and for each meta-update the linear head is solved in closed form via kernel ridge regression, $h_{W^*}(v) = v^\top f_\omega(X_s)^\top(K_{X_s,X_s}+\lambda I_m)^{-1}Y_s$, avoiding the need to unroll the inner loop. Stored artifact: the synthetic images $X_s$ and their learned target representations $Y_s$ (no labels, no generator). At deployment, a fresh feature extractor and head are pretrained on $(X_s,Y_s)$ by MSE regression, then the head is discarded and the feature extractor is fine-tuned with a fresh task-specific head on the labeled target dataset.

Evidence

Table 1 (CIFAR100→{CIFAR10, Aircraft, Cars, CUB2011, Dogs, Flowers}, ConvNet3, 1,000 synthetic images = 2% compression): KRR-ST beats all 8 baselines (Random, Kmeans, DSA, DM, MTT, KIP, FRePo, w/o pre) on every target, e.g. Cars 28.68% vs. best baseline KIP 23.12%, Flowers 67.88% vs. KIP 63.04%. Table 2 (TinyImageNet→6 targets, ConvNet4, 2,000 images = 2%): larger margins, e.g. Aircraft 58.83% vs. MTT 48.21%, Cars 49.26% vs. MTT 30.35%. Table 3 (ImageNet→7 targets, ConvNet4, 1,000 images ≈0.08%, FRePo the only feasible supervised baseline due to memory): Aircraft 57.17% vs. FRePo 39.03%, Cars 46.95% vs. FRePo 20.00%. Architecture-generalization (Fig. 3, TinyImageNet-distilled, evaluated on VGG11/AlexNet/MobileNet/ResNet10): KRR-ST wins in nearly every architecture x target cell (one exception: MobileNet on Aircraft). Target-data-free KD (Table 4, TinyImageNet-distilled surrogate, CIFAR10 teacher): KRR-ST beats all baselines across 5 student architectures (e.g. ConvNet4 58.31% vs. best baseline Kmeans 52.00%), and notably the supervised DD baselines (DSA/DM/MTT/FRePo) underperform even Random selection in this setting.

Limitations

Only classification is used as the downstream fine-tuning task, despite the method being framed as general transfer learning; relies on a single SSL objective (Barlow Twins) for the target model, so it is unclear how sensitive results are to the choice of SSL algorithm; the theoretical bias result (Theorem 1) motivates the method but is not itself an empirical ablation of how much bias the naive bi-level SSL baseline actually suffers; distillation cost (160,000 meta-update iterations per dataset) is not reported in wall-clock or GPU-hours; largest scale tested is ImageNet at 64x64 resolution condensed to ~0.08%, well below full ImageNet-1K resolution/scale work in the image-classification literature.