CoDA: From Text-to-Image Diffusion Models to Training-Free Dataset Distillation

What came before

Minimax, LD3M, IGD and MGD3 fine-tune or otherwise rely on a diffusion model (often a DiT) that was itself pretrained on the full target dataset, e.g. Minimax's DiT backbone needed roughly 3,000 TPU-days of ImageNet pretraining before distillation could start. D4M avoids that by using a general text-to-image latent diffusion model but only initializes sampling from noised K-means prototypes and then lets the model denoise unguided.

The problem

Two failures, argued from Figure 1: (1) target-trained generators reintroduce the very cost DD is meant to avoid ("putting the cart before the horse"); (2) a truly off-the-shelf text-to-image model (SDXL) suffers distributional mismatch -- its web-scale prior samples cluster in a narrow region that does not cover the target class manifold (52% vs 60% for a target-pretrained LDM on ImageNette IPC10) -- and D4M's initialize-then-drift approach lets the general prior override the informative prototype during unguided denoising.

The idea

Identify the dataset's "intrinsic core distribution" -- its high-density, most representative real samples -- with a density-based (not K-means) discovery pipeline in VAE latent space, then continuously steer an off-the-shelf diffusion model's denoising trajectory toward one assigned representative sample per output image for the entire sampling run, rather than only at initialization.

How it works

Distribution Discovery: VAE-encode each class's images (no target-trained feature extractor), apply UMAP dimensionality reduction (tuned to increase local density contrast) then HDBSCAN (robust to non-convex clusters and outliers, unlike K-means); take each cluster's highest-membership point as a representative sample. Since HDBSCAN's cluster count rarely equals IPC, a 3-strategy post-processing scheme (SplitCluster via recursive HDBSCAN+K-means seeding, K-means on the leftover outlier pool, and ForcedSplit with a relaxed min_cluster_size) forces the representative set $S_r$ to exactly IPC samples, deliberately keeping outliers usable rather than discarding them (unlike LOF-based purification, e.g. in VLCP/D4M). Distribution Alignment: at every denoising step, from the CFG text-conditional prediction $\epsilon_\theta(z_t,t,c)$ the clean-latent estimate $\hat z_0(z_t)$ is computed in closed form, a guidance vector $g(z_t)=s_j-\hat z_0(z_t)$ pulls it toward the $j$-th representative sample $s_j\in S_r$ (one fixed real sample assigned per generated image for the whole trajectory), scaled by strength $\gamma$ and algebraically converted into an equivalent additive correction $\Delta\epsilon_\theta$ in noise space so it composes with standard classifier-free guidance; the last "Prior Injection Steps" (PIS) revert to plain CFG so the model's own prior can finish photorealistic detail. Base generator: SDXL, DPM++ Karras sampler, 50 steps, CFG scale 5.0, class-name-only text prompt; nothing is fine-tuned.

Evidence

ImageIDC/ImageNette, ResNet10-AP, 256x256 (Table 1): CoDA-on-SDXL beats prior SOTA by 2.6/2.7/5.5 points on ImageIDC and 2.4/0.9/3.6 on ImageNette at IPC 10/20/50 (e.g. ImageIDC IPC50: 77.6±0.6 vs SDXL+D4M 56.6±0.2 vs SDXL+MGD3 61.4±1.4 vs SDXL raw 45.1±1.8); MGD3 transplanted onto SDXL collapses relative to its ImageNet-trained-LDM version, evidencing that MGD3's success depended on the target-trained prior, not its guidance algorithm alone. ImageWoof (Table 2): CoDA beats the prior SOTA (Minimax/MGD3 on a target-trained model) at IPC>=50 (e.g. IPC50 ResNet-18: 61.2±0.9 vs Minimax 57.1±0.6) but trails slightly at IPC<=20 (IPC10 ResNetAP-10: 39.2±0.7 vs MGD3 40.4±1.9). ImageNet-1K soft-label protocol, ResNet-18 teacher/student (Table 4, shared RDED-style KD protocol from Sun et al. 2024): IPC10 44.3±0.1 (on par with Minimax 44.3±0.5, below IGD 46.2±0.6, +16.4 over D4M 27.9±0.2); IPC50 60.4±0.2, a new SOTA exceeding IGD 60.3±0.4, Minimax 58.6±0.3 and D4M 55.2±0.1. Ablation of components (Table 7, "R" vs "G"): the CPU-only Distribution Discovery output alone (real, selected/resized images, no diffusion sampling at all) already beats the prior SOTA on Places365 IPC10 (used to show zero-shot domain transfer: 47.0±1.6 SDXL+Ours vs 41.7±1.6 random vs 34.4±1.6 SDXL+K-means); adding Distribution Alignment (G) improves further (e.g. ImageIDC IPC50: R 75.7±1.9 to G 77.6±0.6). Step-by-step ablation on MGD3-as-base (Fig. 3) shows UMAP preprocessing alone gives a partial gain over raw K-means on SDXL, but only the full HDBSCAN discovery pipeline surpasses the ImageNet-trained MGD3 baseline. Cost: Distribution Discovery takes 2.2 GPU-hours for all 1,000 ImageNet-1K classes (CPU-bound after VAE encoding); Distribution Alignment adds negligible per-step overhead, generating a 224x224 image in 2.7s (25-step sampling can even outperform 50-step, at 1.3s per 256x256 image).

Limitations

Stated (Appendix A.3): the two key Distribution Discovery hyperparameters (UMAP n_neighbors, HDBSCAN min_cluster_size) drift slightly across datasets since each has its own core-distribution density, so some per-dataset grid search is needed even though it is CPU-cheap; the authors flag automatic hyperparameter selection as future work. Observed: still trails target-trained diffusion methods (IGD, Minimax) at low IPC on hard fine-grained datasets (ImageWoof IPC<=20) and at ImageNet-1K IPC10, so the training-free/off-the-shelf claim does not yet dominate uniformly across budgets; representativeness is defined purely by density in VAE-latent space, inheriting whatever class-mixing the VAE's Gaussian prior induces (Fig. 1b-c show poor separability there relative to a target-trained classifier's feature space); the guidance is a single real-sample target per generated image for the entire trajectory, so diversity depends entirely on how many distinct representative samples Distribution Discovery finds, not on any diversity-specific guidance term (contrast with MGD3's explicit diversity stage or Minimax's minimax loss).