MGD$^3$: Mode-Guided Dataset Distillation using Diffusion Models

What came before

D4M clusters real-image latents but samples from the noised modes directly, which the paper argues limits diversity; Minimax Diffusion fixes diversity by fine-tuning the diffusion model itself with minimax representativeness/diversity losses, at real computational cost (Minimax needs ~10 GPU-hours to fine-tune for ImageNet-100 IPC=10, and the paper reports IDC-1 needs >100). GLaD/H-GLaD/LD3M optimize generative latents with a matching objective but do not scale past small IPC/resolution (20 IPC at 256x256 is called out as their practical ceiling).

The problem

Pretrained diffusion models oversample the densest modes of a class's data distribution, and when the number of dominant modes is smaller than the target IPC this produces redundant, near-duplicate samples — a diversity failure that persists even after Minimax's fine-tuning (the paper's t-SNE analysis, Fig. 5-6, shows Minimax reaches broader but still incomplete coverage) and that costs real GPU time to fine-tune around.

The idea

Fix the diversity problem entirely at sampling time, without ever touching the diffusion model's weights: discover a class's distinct modes by clustering real-image latents, steer each sample's denoising trajectory toward a specific assigned mode with a lightweight guidance term, and stop that guidance before the final refinement stage of denoising so the model can still finish each sample with unconstrained, artifact-free detail.

How it works

A frozen DiT-XL/2 (or, in a robustness check, LDM's U-Net, or a general-purpose text-to-image Stable Diffusion model) generates class-conditioned samples via classifier-free guidance (scale 4.0, 50 steps). Mode Discovery: K-means (k=IPC) clusters real images' VAE-encoder latents per class into $N$ modes $\mathbf{M_c}=\{m_1,...,m_N\}$ (any clustering algorithm works per the paper's own ablation). Mode Guidance: at each reverse step, a guidance vector $\mathbf{g}_t=(m_i-\hat{x}_0^t)$ (the direction from the model's current denoised-latent estimate toward the assigned target mode $m_i$) is added to the noise prediction, $\hat{\epsilon}_\theta(x_t,t,c)=\tilde\epsilon_\theta(x_t,t,c)+\lambda\cdot \mathbf{g}_t\cdot\sigma_t$ (Eq. 6, $\lambda{=}0.1$), steering that sample's trajectory toward its assigned mode. Stop Guidance: guidance is set to zero once the timestep falls below $t_{stop}{=}25$ (of 50 steps, i.e. roughly the boundary of the "refinement stage" per prior diffusion-stage analysis), letting the model finish denoising unconstrained to avoid the artifacts that full-trajectory guidance produces. No fine-tuning, no matching network, no expert trajectories, and (unlike D4M's TTM or SRe2L- family methods) no soft-label training-time matching is required by the core method, though soft labels are still used for the ImageNet-1K evaluation protocol.

Evidence

ImageNette/ImageIDC, hard-label protocol, ResNet-10-avg-pool, IPC=10/20/50 (Table 1): DiT+MGD3 66.4%/ 71.2%/79.5% (Nette) vs. MinMaxDiff 62.0%/66.8%/76.6% vs. raw DiT 59.1%/64.8%/73.3% — +4.4/+4.4/+2.9 points over the prior SOTA the abstract cites; on IDC, MGD3 55.9%/61.9%/72.1% vs. MinMax 53.1%/59.0%/ 69.6%. ImageNet-100, hard-label, IPC=20 (Table 2): MGD3 beats MinMaxDiff by 1.3-1.6 points across ConvNet-6/ResNetAP-10/ResNet-18 (e.g. ResNetAP-10: 33.9% vs. 32.3%); at IPC=10, IDC-1 (a pixel-optimization method) edges out MGD3 slightly but at far higher cost. ImageNet-1K, soft-label (RDED-style region soft labels, ResNet-18 teacher/student), IPC=10/50: MGD3 beats prior SOTA by 1.3/1.6 points (Fig. 4d); scaling to larger student backbones at IPC=50 (Table 3): MGD3 ResNet-101 67.7% vs. EDC 64.9% vs. D4M 63.4% vs. RDED 61.2%. Generative-prior comparison, ImageNet-A-E, IPC=10, hard-label (Table 4): MGD3 beats GLaD/H-GLaD/LD3M (labeled "LM3D" in this paper) by a wide margin, e.g. ImageNet-A 63.4% vs. LD3M 57.0% vs. H-GLaD 55.1% vs. GLaD 53.1%. Off-target generator test: a general-purpose text-to-image Stable Diffusion model (not trained on any target dataset) with MGD3 guidance beats the same Stable Diffusion model without guidance by 3.4/2.3 points on ImageNet-1K IPC=10/50 (Fig. 4c), demonstrating the guidance mechanism transfers to a genuinely off-target generator, though absolute accuracy trails the ImageNet-trained-DiT variant. Cost: 0.42 GPU-hours to generate ImageNet-100 IPC=10 vs. Minimax's ~10 hours (fine-tuning) and IDC-1's >100 hours. Key ablation (Table 5, ImageNette IPC=10): Mode Discovery alone gives 53.2%/57.1%/53.5% (ConvNet-6/ResNetAP-10/ResNet-18); +Mode Guidance raises it to 57.5%/63.8%/62.0%; +Stop Guidance further to 59.6%/66.4%/64.4% — each of the three stages adds measurable accuracy, with Mode Guidance contributing the largest single jump. Sampler-agnosticism (Table 6): the method improves both DDPM and DDIM samplers, and both LDM (U-Net) and DiT (Transformer) backbones, similarly.

Limitations

The headline numbers use a DiT pretrained on ImageNet-1K, i.e. matched to the target distribution for every ImageNet-derived experiment; the off-target Stable Diffusion test (Fig. 4c) is real but shows a smaller improvement and is reported only for ImageNet-1K, not for ImageNette/IDC/ImageNet-100, so its generality across datasets is untested; hard-label and soft-label protocols are used for different dataset groups within the same paper (ImageNette/IDC/ImageNet-100 hard, ImageNet-1K soft), so results are not directly comparable across the paper's own tables; mode discovery uses K-means with k fixed equal to IPC, which ties the number of distinct modes assumed to the budget rather than to the true number of underlying data clusters; the paper does not report results below IPC=10 or above IPC=100/ IPC=50 for ImageNet-1K's largest backbones.