Generalizing Dataset Distillation via Deep Generative Prior

What came before

Pixel-space matching methods DC (gradient matching), DM (distribution matching) and MTT (trajectory matching) all parameterize the synthetic set as raw pixels; they overfit to the backbone architecture used during distillation and degrade into high-frequency noise above roughly 128x128, so cross-architecture generalization is their biggest weakness. IT-GAN (concurrent) had shown latent optimization can beat pixel optimization on CIFAR with a target-trained BigGAN, but only for one matching objective.

The problem

Directly optimizing pixels lets the synthetic set fit high-frequency, architecture-specific patterns of the one backbone network used during distillation, which do not transfer to other architectures, and this failure gets worse as resolution grows.

The idea

Regularize any existing matching-based distillation objective by parameterizing the synthetic images as the output of an intermediate layer of a frozen, pretrained deep generative model (GLaD) rather than raw pixels; an intermediate layer balances realism (too early = restrictive) against expressivity (too late = no regularization), and, surprisingly, the generator does not need to be trained on the target dataset at all for the regularization to help.

How it works

A synthetic sample is $G(z)$ for a frozen StyleGAN-XL generator $G$ and a learnable latent $z$ that concatenates an intermediate "F_n" feature map (the n-th layer of the synthesis network) with the W+ modulation codes of all subsequent layers; $z$ is optimized by backpropagating whichever matching loss is being used ($\mathcal{L}_{DC}$, $\mathcal{L}_{DM}$, or $\mathcal{L}_{MTT}$, computed exactly as in the underlying method) through $G$, which is otherwise treated as a plug-in add-on (Algorithm 1). $z$ is initialized either by a partial feed-forward pass through $G$ using the empirical latent distribution of the target class (when $G$ is trained on/near the target data) or by Gaussian noise matched in mean/variance (needed to make out-of-distribution or randomly-initialized generators usable). A gradient-checkpointing trick recomputes the generator's forward pass a second time only when computing $\partial\mathcal{S}/\partial\mathcal{Z}$, trading one extra forward pass for large memory savings (about 22GB per synthetic image at 128x128 in F0 space). The generator itself is never fine-tuned.

Evidence

ImageNet subsets at 128x128, IPC=1, cross-architecture average over AlexNet/ResNet-18/VGG-11/ViT-B/16 trained on sets distilled with a Depth-5 ConvNet backbone (Table 1): GLaD improves every one of MTT/DC/DM on every subset, e.g. ImageNet-A: MTT 33.4%->39.9%, DC 38.7%->41.8%, DM 27.2%->31.6%. At IPC=10 (Table 2) gains persist but shrink, e.g. DC ImageNet-B 45.1%->50.1%, DM ImageNet-C 47.5%->49.7%. On CIFAR-10 (Table 3, IPC=1) GLaD with a StyleGAN-XL trained on CIFAR-10 raises the 4-architecture average for MTT from 24.1% to 28.0% and for DM from 22.6% to 23.8%, but only marginally for DC (26.0%->26.3%). Labels are hard throughout (all three base methods are unmodified apart from the parameterization); training recipe is SGD with momentum, 500 epochs linear warmup + 500 epochs cosine decay, 5 seeds. Key ablation isolating the mechanism (Figure 3, Table 4): (a) sweeping which layer $F_n$ is used shows an inverted-U in cross-architecture accuracy, so an intermediate space, not the earliest or latest, is what helps; (b) at 256x256 with DC, generators trained on ImageNet, on unrelated FFHQ faces, on Pokémon, or left at random initialization all give comparable cross-architecture improvements over pixel space (e.g. ImageNet-B: pixel 32.8% vs ImageNet-G 41.5% vs Random-G 40.3% vs Pokemon-G 39.4% vs FFHQ-G 40.2%), showing the gain comes mostly from the generator's architectural/manifold prior, not from it having seen the target distribution.

Limitations

A second forward+backward pass through StyleGAN-XL is required per iteration even with checkpointing, which the authors state doubles the time overhead of an already slow generator pass, and large synthetic sets need multiple sub-batches, adding further overhead; only IPC=1/10 are tested, not larger budgets; only hard labels are used (the paper explicitly declines to use TESLA's soft-label assignment to isolate the generative-prior effect); the largest dataset tested is 10-class ImageNet subsets up to 512x512, not full ImageNet-1K; gains on CIFAR-10/DC are only marginal, so the benefit is resolution- and method-dependent, not universal.