D$^4$M: Dataset Distillation via Disentangled Diffusion Model

What came before

Bi-level (STM, synthesis-time matching: DC, DSA, CAFE, MTT, TESLA) and SRe2L's dual-time-matching (DTM: squeeze-recover-relabel) methods both couple image synthesis to a specific network architecture — STM matches gradients/features/trajectories against one architecture during synthesis, DTM recovers images from a fixed teacher's BN statistics — so both must regenerate a dataset (or at least verify it) for each new evaluation architecture, and STM in particular becomes computationally infeasible or inaccurate at ImageNet-1K scale. GLaD constrains the output to a GAN's manifold but the inner loop still matches through a specific network.

The problem

Every synthesis-time matching approach biases the synthetic image space toward whatever architecture supplies the loss, which the paper argues both caps cross-architecture generalization and produces images with degraded human-legible semantics (Fig. 2); STM also scales matching cost/GPU-memory with the chosen architecture's size (MTT needs 79.9GB on ImageNet-1K at 128x128 over one V100-class setup), making ImageNet-1K distillation impractical for the STM paradigm.

The idea

Remove architecture matching from the synthesis step entirely: extract per-category prototypes (cluster centroids) from a frozen, off-the-shelf latent diffusion model's own latent space, condition the frozen diffusion model on the prototype + a text-encoded label to generate images, and push all learning into Training-Time Matching (TTM) — i.e. soft-label relabeling during downstream network training — so no architecture ever touches the synthesis process.

How it works

A frozen, publicly released Stable Diffusion V1-5 checkpoint (a general text-to-image LDM pretrained on LAION-scale web image-text pairs, not on the target classification dataset) supplies encoder $\mathcal{E}$, decoder $\mathcal{D}$, text encoder $\tau_\theta$, and denoising U-Net $\mathcal{U}_t$, all frozen. Real images per category are encoded to latents $z=\mathcal{E}(\mathcal{T})$, then Mini-Batch k-Means clusters them into $C$ prototypes $z^c$ per class (for IPC<100, $C$ is set equal to IPC; for larger IPC, e.g. 10-50 prototypes with multiple samples generated per prototype). Each prototype is noised via the forward process to a chosen strength (0.7, an img2img-style partial-noising parameter) and denoised while cross-attending to the label's text embedding $\tau_\theta(L)$ with classifier-free guidance scale 8, then decoded: $\mathcal{S}=\mathcal{D}(\tilde{Z}^c)$. No gradients flow through the diffusion model at any point (synthesis is training-free). At evaluation, a teacher network (ResNet-18/50/101 pretrained on the real dataset) generates soft labels for the distilled images at every training epoch (an FKD-style scheme, Eq. 8) and the student is trained with KL- divergence to these soft predictions rather than hard one-hot labels — this soft-relabeling (TTM) is presented as itself a distribution-matching mechanism, distinct from the synthesis step.

Evidence

ImageNet-1K, IPC=100, ResNet-18 teacher/student, soft-relabel TTM (Table 2): D4M 66.5% vs. SRe2L 52.8% vs. full-dataset ceiling 69.8%; IPC=10: D4M 27.9% vs. SRe2L 21.3% vs. TESLA (STM, ConvNetD4 teacher) 7.7%. Tiny-ImageNet IPC=100: D4M 51.4% vs. SRe2L 50.8% (R18); a "D4M-G" variant, built by simply extracting 200 of the already-distilled ImageNet-1K classes rather than distilling Tiny-ImageNet from scratch, nearly matches or exceeds dedicated D4M (53.3% R18), evidencing the claimed architecture/ dataset-versatility. CIFAR-10/100 (Table 1, small-scale, STM baselines KIP/FRePO/DSA/CAFE/TESLA copied from originals): D4M beats FRePo and TESLA on CIFAR-100 IPC=10 by 2.5/3.3 points (45.0% vs. 42.5%/41.7%) but underperforms TESLA on CIFAR-10 IPC=10 (56.2% vs. 66.4%), and the paper itself notes STM still wins at small category counts/IPC, TTM only becoming favorable as classes grow. Cost (Table 6): on ImageNet- 1K, D4M synthesis is 2.7s/image at 6.1GB GPU memory (fixed, architecture-free) vs. SRe2L 5.2s/34.8GB and MTT 45.0s/79.9GB (both architecture-dependent and rising with matching-network size) — D4M is 3.82x faster than SRe2L end-to-end per the paper's own count. Key ablations: (1) Table 3 — applying an additional STM step (BN-distribution matching) on top of D4M's already-generated images *reduces* accuracy at every teacher (e.g. R18 teacher: 27.9%->23.6%, a 4.3-point drop), because it disrupts the fused image+text information already encoded by the diffusion conditioning — architecture matching is actively harmful once TTM is in place, not merely unnecessary; (2) Table 4 — replacing prototype initialization with random noise drops ImageNet-1K IPC=10 accuracy from 27.9% to 15.6% (a 12.3-point gap), confirming prototypes (not just "any diffusion sample") are the mechanism.

Limitations

The diffusion model (Stable Diffusion V1-5) is a general, off-the-shelf text-to-image model trained on web-scale LAION data, never fine-tuned or trained on CIFAR/ImageNet/Tiny-ImageNet — this is the clearest case in the group of a generator with zero target-dataset training, but also means all reported gains are entangled with soft-label relabeling (TTM) at evaluation, so the paper does not isolate how much of its advantage over SRe2L comes from image quality/diversity versus from label regime (both differ simultaneously in every ImageNet-1K/Tiny-ImageNet comparison). The paper's own stated limitation is a "significant performance degradation" at extreme low budgets (IPC=1/10), where CIFAR-10 IPC=10 D4M (56.2%) trails TESLA (66.4%) by 10 points despite winning at larger IPC/more classes; text prompts are only class-name-level (no richer per-image captions, unlike later VLM-prompted successors); the strength/guidance-scale hyperparameters (0.7/8) are tuned on ImageNet-1K/ResNet and not shown to transfer automatically to other domains.