Unlocking Dataset Distillation with Diffusion Models

What came before

GLaD backpropagates gradient/distribution/trajectory-matching losses through a frozen StyleGAN-XL's latents, but needs costly GAN inversion for initialization and its W+ multi-level latent space is expensive to optimize; sampling-based diffusion methods (Minimax Diffusion, D4M) avoid backpropagating through the denoising chain entirely by selecting or fine-tuning toward representative samples; an autoencoder-only approach (Duan et al., 2023) optimizes latents decoded by an LDM's autoencoder but ignores the diffusion process itself.

The problem

Naive backpropagation of a matching loss through the T-step reverse diffusion chain to the initial latent $\mathcal{Z}$ suffers vanishing gradients: the paper measures gradient norms for $\mathcal{Z}$ falling roughly tenfold as the number of diffusion steps $T$ increases from 10 to 90 (Table 1), which the paper attributes to the product of T bounded Jacobians in the chain rule (Eq. 1) and which has forced every prior diffusion-based distillation method to avoid true end-to-end latent optimization.

The idea

Inject a residual/skip connection from the fixed initial noisy state $\mathbf{z}_T$ into every reverse step, with a linearly decaying weight from 1 (at $t=T$) to ~0 (as $t\to0$), so gradients from the distillation loss have a short path back to the learnable latent $\mathcal{Z}$ that bypasses most of the Jacobian chain, without ever fine-tuning the diffusion model's own weights.

How it works

A pretrained ImageNet-trained LDM (2x-compression autoencoder $\mathcal{E}/\mathcal{D}$, frozen U-Net denoiser $\mu_\theta$) is used without fine-tuning. Learnable latent codes $\mathcal{Z}$ (initialized by encoding a small set of randomly-selected real images via $\mathcal{E}$, not by GAN-style inversion) and learnable class-conditioning codes $\mathbf{c}$ (initialized from the pretrained class-embedding network) are optimized. The modified reverse step $\mathbf{z}_{t-1}\leftarrow((1-t/T)\mu_\theta(\mathbf{c},\mathbf{z}_t,\gamma_t)+(t/T)\mathbf{z}_T)+\sigma_t^2\varepsilon_t$ (Eq. 7) replaces the standard predicted mean with a decaying blend of the model's prediction and the fixed $\mathbf{z}_T$; the final $\mathbf{z}_0$ is decoded to $\mathcal{S}=\mathcal{D}(\mathbf{z}_0)$ and plugged into an existing distillation objective (DC, DM, or MTT) exactly as GLaD does for StyleGAN-XL, with the resulting loss backpropagated through the whole modified chain into $\mathcal{Z}$ and $\mathbf{c}$. Gradient checkpointing (as in GLaD) manages VRAM by recomputing forward segments during the backward pass rather than storing the full graph. Default diffusion steps are $T{=}10$ (128x128) or $T{=}20$ (256x256).

Evidence

CIFAR-10 IPC=1, cross-arch avg. over AlexNet/ResNet18/VGG11/ViT, hard labels (Table 2a): LD3M(DC) 28.6% vs. GLaD(trained G) 26.3% vs. pixel-space DC 26.0%; LD3M(DM) 24.9% vs. GLaD(DM) 23.8%; sampling-based D4M reported ~10% under this same protocol. CIFAR-10 IPC=50 ConvNet (Table 2b, note differing optimization/IPC across rows): LD3M 73.2% vs. D4M 72.8% vs. decoupled SRe2L 60.2% (SRe2L's number is from its own IPC=1K report). ImageNet 10-class subsets, 128x128, IPC=1, avg. over 10 subsets and 4 archs (Table 3): LD3M improves over GLaD by +3.76% (DC), and over GLaD DM by roughly +16.3 points relative to pixel-space DM; at IPC=10 (Table 4, ImageNet A-E) LD3M beats GLaD by +2.52% (DC) and +3.46% (DM), reaching 47.1-47.3% average. At 256x256, IPC=1 (Table 5) LD3M beats GLaD by ~+2pp averaged across three generator-pretraining sources (ImageNet/FFHQ/Random) and by ~+7pp over pixel space. Cost: at T=20 on an A100-40GB, LD3M uses slightly less peak memory than GLaD (29.4GB vs. 31.2GB) and is faster (574 vs. 693 min total distillation). Key ablations: (1) Table 6 isolates the diffusion process itself — removing it (autoencoder-only, akin to Duan et al.) drops DC/IPC=1 ImageNet-A-E average from 36.5% to 35.3%, statistically similar to GLaD's 35.4%, showing the diffusion prior alone (without the gradient fix) adds nothing; (2) Table 7 isolates the gradient-flow fix — learnable conditioning alone gives 15.8% avg (ImageNette, MTT, IPC=1), adding learnable $\mathcal{Z}$ raises it to 22.3% (still below GLaD's 26.6%), and only adding the decaying residual (Eq. 7) reaches 28.1%, confirming the skip connection, not latent learnability alone, is what makes diffusion-based end-to-end optimization competitive; (3) Table 8 shows real-image latent initialization vastly outperforms Gaussian-noise initialization for DC/DM (e.g. DC ImageNette 32.9% vs. 13.1%), while MTT is comparatively robust to initialization.

Limitations

The LDM used is trained on ImageNet, i.e. matched to the target distribution for the ImageNet-subset experiments (only the 256x256 table separately tests FFHQ- and randomly-initialized generators, and even there the gap between ImageNet/FFHQ/Random pretraining is small, similarly to GLaD's own out-of-distribution-generator finding); the paper's own stated limitation (Appendix E) is that the linear decay schedule for the residual may not fully solve vanishing gradients at larger T, and alternative (non-linear) decay schedules are left untested; evaluated only on CIFAR-10 and 10-class ImageNet subsets up to 256x256, not full ImageNet-1K (explicitly named as future work, "contingent on computational resources"); hard labels only, single RTX A6000/A100 hardware; the SRe2L/D4M numbers in Table 2b are copied from their original papers under different IPC/optimization protocols, flagged by the paper itself as not a like-for-like comparison.