Taming Diffusion for Dataset Distillation with High Representativeness

What came before

D4M clusters real-image latents with K-means and Minimax fine-tunes a diffusion model with a cosine-similarity minimax objective to obtain representative latents; both still generate each distilled sample somewhat independently and rely on approximations (spherical K-means clusters, or direction-only cosine similarity) of the true latent distribution, which the paper's Lemma 3.1 argues is actually a complex multi-component Gaussian mixture, not something either approximation captures well. SRe2L/DWA/RDED need a teacher-architecture-specific synthesis step (BN statistics or patch selection), coupling the distilled set to one architecture.

The problem

The paper diagnoses three specific flaws in prior diffusion-based distillation: (i) inaccurate distribution matching — the VAE latent space per class is a genuinely multi-component Gaussian mixture (proved in Appendix A.1), so K-means' spherical-cluster assumption or cosine-similarity's direction-only matching cannot describe it well; (ii) distribution deviation with random noise — generating from freshly sampled random noise (rather than from the real data's own noised trajectory) injects unpredictable randomness that can shift the denoised output away from the intended distribution; (iii) separate sampling — each of the n distilled latents per class is matched to the target distribution individually, so the *set* of n latents as a whole is not guaranteed to match the target distribution even if each element does reasonably well.

The idea

Use DDIM inversion (not the stochastic DDPM forward process) to deterministically and bijectively map each real image's VAE latent into a noise-space domain that becomes provably closer to a single Gaussian as the number of inversion steps grows (Lemma 4.1); fit a per-class Gaussian to the inverted latents; then, instead of independently sampling n latents from that Gaussian, sample many candidate n-latent subsets and pick the one whose empirical mean/variance/skewness best match the fitted Gaussian ("group sampling"), before DDIM-sampling each selected latent back into an image.

How it works

A pretrained DiT + VAE (Peebles & Xie 2023, originally trained on ImageNet-1K) supplies the diffusion backbone; for CIFAR-10/100 and Tiny-ImageNet, this ImageNet-pretrained DiT is further fine-tuned for 400 epochs per dataset to adapt its generative capacity to the new label space and resolution (for ImageNet-1K itself, no fine-tuning is needed). For each class $\mathcal{C}$: (1) Domain Mapping — every real image's VAE latent $z_0$ is deterministically inverted via DDIM ($T{=}31$ steps in the main experiments) to a noise-space latent $z_T$ (Eq. 3), forming $\mathcal{Z}_{T,\mathcal{C}}$; (2) Distribution Matching — the empirical mean $\mu_{T,\mathcal{C}}$ and variance $\sigma^2_{T,\mathcal{C}}$ of $\mathcal{Z}_{T,\mathcal{C}}$ define a fitted Gaussian $\hat{\mathcal{Z}}_{T,\mathcal{C}}$; (3) Group Sampling — $m$ candidate subsets of $n{=}\text{IPC}$ i.i.d. latents are drawn from $\hat{\mathcal{Z}}_{T,\mathcal{C}}$ via the Ziggurat algorithm, each subset scored by a weighted sum $\mathcal{L}_{T,\mathcal{C}}=\lambda_\mu\mathcal{L}_\mu+\lambda_\sigma\mathcal{L}_\sigma+ \lambda_{\gamma_1}\mathcal{L}_{\gamma_1}$ of its mean/std/skewness deviation from the fitted Gaussian's parameters (Eqs. 5-9), and the lowest-scoring subset is kept; (4) each latent in the winning subset is DDIM-sampled back to $z_0$ and decoded to an image, paired with the class label. No architecture-specific matching or teacher model is used during synthesis at all — this is entirely training-free at the generation step. At evaluation, a separately pretrained teacher (ResNet-18, following SRe2L's convention) supplies soft labels to train the student, matching the label regime of SRe2L/DWA/D4M/RDED for the main comparison table; a separate section (6.4) additionally reports hard-label results matched to Minimax's own evaluation protocol on ImageWoof.

Evidence

Table 1 (four datasets, ResNet-18/ResNet-101, soft labels from a ResNet-18 teacher, following SRe2L): ImageNet-1K IPC=10 D3HR 44.3% vs. RDED 42.0% vs. D4M 27.9% vs. SRe2L 21.3% (ResNet-18); IPC=50 D3HR 59.4% vs. RDED 56.5%; CIFAR-10 IPC=50 D3HR 70.8% vs. RDED 62.1% (+8.7, the largest gap in the table, cited as a "12.5% relative improvement" for ResNet-18 and "17.4%" for ResNet-101); Tiny-ImageNet IPC=10 D3HR 44.4% vs. RDED 41.9%. Cross-architecture on ImageNet-1K IPC=10 (Table 2, train-on-one/test-on- another grid across ResNet-18/MobileNet-V2/VGG-11): D3HR beats RDED in every one of 9 cells, e.g. trained on MobileNet-V2, tested on MobileNet-V2: D3HR 46.4% vs. RDED 33.8% vs. SRe2L 10.2%. Hard-label comparison with Minimax (Table 5, ImageWoof, matching Minimax's own protocol): D3HR beats Minimax at every IPC/architecture, e.g. IPC=50/ResNetAP-10: D3HR 59.3% vs. Minimax 56.3%. Cost: sampling (group sampling + DDIM decode) takes 2.6s/class on a single RTX A6000 for ImageNet-1K IPC=10 with $m{=}10^6$ candidate subsets; storage can be reduced to just the per-class Gaussian's mean/variance plus the shared DiT weights, smaller than storing all distilled latents/images. Key ablations (Table 3, ImageNet- 1K IPC=10, ResNet-18): domain mapping via DDIM inversion (Base-RS, 41.6%) beats DDPM-style random-noise domain mapping (Base-DDPM, 37.3%) by 4.3 points (paper's text calls this an "11.5%" relative gain), isolating the deterministic-inversion fix; adding group sampling with all three statistics ($\mathcal{L}_\mu+\mathcal{L}_\sigma+\mathcal{L}_{\gamma_1}$, 44.3%) beats Base-RS (individual Gaussian sampling, 41.6%) by 2.7 points and beats any single statistic alone (42.3-42.6%), showing the subset-level matching is a distinct, additive contribution on top of domain mapping; Figure 4 shows group sampling also reduces run-to-run variance versus individual random sampling.

Limitations

The DiT backbone is either already ImageNet-1K-pretrained (used as-is for the ImageNet-1K experiments) or explicitly fine-tuned for 400 epochs per dataset (CIFAR-10/100, Tiny-ImageNet) — every reported number therefore uses a generator matched (by pretraining or fine-tuning) to the target distribution, with no off-target-generator ablation; the main results (Table 1) rely on teacher soft-label supervision at evaluation like its SRe2L-family baselines, so the paper does not fully separate generation-quality gains from label-regime effects except in the separate hard-label section (6.4), which is scoped only to ImageWoof/Minimax's setting; the single-Gaussian approximation (Lemma 4.1) is asymptotic in the number of inversion steps $T$, and the paper's own Fig. 5 shows accuracy peaks at $T{=}31$ and *declines* for larger $T$ because more inversion steps destroy structural information even as they improve the Gaussian fit — a stated trade-off, not a solved problem; excludes Minimax from the main Table 1 comparison (addressed separately) because Minimax "requires training multiple diffusion models" at large scale, a methodological choice the paper flags but that also avoids a fully matched head-to-head in the primary table.