Efficient Dataset Distillation via Minimax Diffusion

What came before

Pixel-level (IDC-1) and embedding-level (GLaD) sample-wise iterative optimization methods scale their parameter space with IPC and resolution: IDC-1 takes over 90 GPU-hours to distill IPC=100 on ImageWoof (longer than training on ImageWoof itself), and its per-pixel modifications shrink as IPC grows, so gains over random images vanish or reverse at large IPC; GLaD needs prohibitive GPU memory as IPC grows. Separately, sampling directly from a pretrained diffusion model (DiT) or naively fine-tuning it (Difffit) had not been evaluated as a distillation method.

The problem

The paper's own diagnostic (t-SNE, Fig. 3) shows raw DiT samples are representative of high-density regions of the real distribution but fail to cover its sparser regions (low diversity), while random real-image selection has the opposite problem (diverse but not representative) — so neither raw diffusion sampling nor coreset selection alone gives an effective surrogate set, and this gap becomes the paper's target instead of a matching loss to optimize against.

The idea

Fine-tune a pretrained class-conditional diffusion model (DiT) with two extra minimax loss terms, one pulling each generated sample's embedding toward the least-similar real sample in a memory bank (representativeness) and one pushing it away from the most-similar previously-generated sample in another memory bank (diversity), so a single short fine-tuning run yields a generator whose *sampling* (no further per-sample optimization) gives both properties at once, at any IPC.

How it works

DiT (ImageNet-pretrained, latent diffusion with encoder $E$/decoder $D$) is fine-tuned with Difffit-style parameter-efficient fine-tuning using $\mathcal{L}=\mathcal{L}_{simple}+\lambda_r \mathcal{L}_r+\lambda_d\mathcal{L}_d$, where $\mathcal{L}_{simple}$ is the standard diffusion noise- prediction loss, $\mathcal{L}_r=\arg\max_\theta \min_{m}\sigma(\hat{\mathbf{z}}_\theta(\mathbf{z}_t, \mathbf{c}),\mathbf{z}_m)$ pulls the predicted clean embedding toward the *least* similar real sample stored in a real-embedding memory bank $\mathcal{M}$ (cosine similarity $\sigma$), and $\mathcal{L}_d= \arg\min_\theta \max_{d}\sigma(\hat{\mathbf{z}}_\theta(\mathbf{z}_t,\mathbf{c}),\mathbf{z}_d)$ pushes it away from the *most* similar previously-predicted embedding stored in a second memory bank $\mathcal{D}$ of adjacent-iteration generations. Fine-tuning uses $\lambda_r{=}0.002$, $\lambda_d{=}0.008$, batch size 8, 8 epochs, AdamW lr 1e-3, image size 256x256, single RTX 4090; a theoretical section frames the combined objective as an idealized tri-level optimization (Eq. 7) approximating a stochastic-control problem, arguing representativeness and diversity can both be satisfied without a strict trade-off. At deployment, the fine-tuned model is sampled (50 DDIM-style denoising steps) for the desired number of images per class, conditioned only on the class label — no test-time guidance and no further per-sample optimization.

Evidence

ImageWoof, hard labels, ResNet-18/ResNetAP-10/ConvNet-6 (Table 1, all methods reproduced under the same protocol): at IPC=100, Ours 65.7%/64.5%/61.1% vs. second-best DiT baseline 58.9%/56.4%/53.4% (+5.5-8.1 points, the headline gain the abstract cites); at IPC=10, Ours 39.2%/39.2%/37.0% vs. DiT 34.7%/34.7%/34.2%; full-dataset upper bound ~86-89%. ImageNette/ImageIDC (Table 3, hard labels, ResNet-10): consistent smaller but positive margins over DiT and DM at all IPC 10/20/50. ImageNet-1K (Table 4, RDED's validation protocol, 224x224): IPC=10 Ours 44.3% vs. RDED 42.0% vs. SRe2L 21.3% vs. DiT baseline 39.6%; IPC=50 Ours 58.6% vs. RDED 56.5% vs. SRe2L 46.8% — a decoupled-synthesis (SRe2L/RDED) head-to-head under a protocol the paper adopts but does not fully restate (RDED's own protocol typically trains with teacher soft-label relabeling, so this comparison may not be a hard-label-only one, unlike the paper's own ImageWoof/Nette/IDC tables). Distribution-matching quality (Table 2, MMD to real features): Ours lowest (4.0 at IPC=50) vs. DM 4.8, DiT 5.4, IDC-1 6.7. Cost: under 1 hour to fine-tune+generate IPC=100 for a 10-class subset on a single RTX 4090, vs. IDC-1's >90 hours (the "under one-twentieth" claim), with GPU memory constant across IPC (unlike IDC-1/GLaD). Key ablation (Table 5, ResNet-10, ImageWoof/ImageIDC): naive representativeness alignment (Eq. 2, pulling to the distribution center rather than the minimax version) *degrades* accuracy at every IPC relative to plain Difffit fine-tuning (e.g. ImageWoof IPC=50: 47.1% vs. 51.0% baseline); the minimax diversity term alone gives the largest single gain (54.9%/68.4% at IPC=50 on Woof/IDC); combining minimax representativeness + minimax diversity (the full method) gives the best result (56.3%/69.6%), confirming the minimax formulation (not naive centroid-matching) is what makes the representativeness term useful at all.

Limitations

The diffusion model (DiT) is pretrained on ImageNet-1K, i.e. matched to the target distribution for every experiment in the paper (ImageWoof/Nette/IDC/100/1K are all ImageNet subsets or ImageNet-1K itself); no experiment tests an off-target or non-ImageNet-pretrained generator, so none of the reported gain isolates generator-domain-mismatch robustness the way GLaD's ablation did for GANs; no CIFAR/MNIST or small-image experiments (the method targets exactly the ImageNet-scale, high-resolution regime where pixel-space methods struggle); the paper's own stated limitation is that it addresses only the classification setting and leaves other data domains to future work; fine-tuning still requires a class- conditional generator pretrained at the target's scale/resolution, so the "under 1 hour" distillation cost excludes the (much larger) upfront DiT-pretraining cost; the ImageNet-1K comparison follows RDED's protocol by reference without the paper independently confirming its own label regime there.