IMS3: Breaking Distributional Aggregation in Diffusion-Based Dataset Distillation

What came before

Minimax fine-tunes a pretrained DiT with representativeness/diversity criteria; D4M, D3HR and VLCP (named DDVLCP here) all improve prototype construction, latent statistics, or added text conditioning, but all still sample from a diffusion model trained to maximize likelihood over the target class.

The problem

Diffusion models are trained to maximize generative likelihood, which concentrates probability mass (and therefore sampling) in high-density regions of the data manifold; but the samples a classifier most needs are near decision boundaries, which are typically low-density. The paper names this mismatch "distributional aggregation" and shows it visually (Fig. 1 t-SNE): Minimax and DiT samples cluster tightly while real data spreads more broadly across feature space.

The idea

Two independent fixes for two independent failure modes. (1) DDIM inversion of real images is known to be numerically unstable and drift toward low-density regions rather than reconstructing exactly; instead of treating this as a bug, fine-tune the generator to pull its denoising trajectory toward its own inversion trajectory at matching timesteps, deliberately injecting that low-density drift to widen coverage. (2) Standard sampling generates each class independently with no notion of inter-class structure, so explicitly select, from multiple candidate generated subgroups, the one whose feature centroid is both close to the real class centroid and far from other classes' subgroup centroids.

How it works

Backbone is DiT-XL/2, adapted with parameter-efficient Difffit adapters (as in Minimax) rather than full fine-tuning. Inversion-Matching (IM): compute the DDIM-inversion latent $z_t^{inv}$ of a real image at a sampled timestep $t$, and the standard forward-noised latent $z_t$; minimize $\mathcal{L}_{IM}=1-\cos(z_t^{inv},z_t)$ jointly with the standard denoising loss $\mathcal{L}=\mathcal{L}_{Diff}+\lambda_{IM}\mathcal{L}_{IM}$ ($\lambda_{IM}=0.002$), so the trajectory is pulled toward inversion-drifted (low-density) states without discarding fidelity. A theoretical section (the PF-ODE geometric-mean instability coefficient) argues this instability is near-inevitable in high dimensions, motivating exploiting rather than correcting it. Selective Subgroup Sampling (S3, training-free, applied after fine-tuning): draw $G$ candidate IPC-sized subgroups per class from the fine-tuned generator, embed all images with a frozen encoder $\phi$, compute each subgroup's L2-normalized centroid and the real class centroid $r_i$; select, via greedy search, the subgroup index per class minimizing $\mathcal{L}_{S^3}=\alpha\sum_i\log(1-\cos(c_{i,g_i},r_i)) - \frac{\beta}{(C-1)G}\sum_{i}\sum_{j\neq i}\sum_g\log(1-\cos(c_{i,g_i},c_{j,g}))$, i.e. pulling each class's chosen centroid toward its own real centroid while pushing it away from other classes' candidate centroids.

Evidence

ImageWoof (Table 1), IPC10/50 on ResNet-18: 41.3±1.1 / 60.1±1.1, vs Minimax 37.6±0.9/53.9±0.6, D4M 32.3±1.2/53.7±2.2, D3HR 39.6±1.0/57.6±0.4, VLCP(DDVLCP) 39.9±2.6/58.9±1.5 -- all under one shared hard/soft-label-whichever-is-higher protocol (paper reports the higher of the two per Minimax's and RDED's conventions). ImageNette (Table 2), ResNetAP-10, IPC10/50: 62.9±1.2/84.2±1.0 vs Minimax 58.6±0.4/83.7±0.4, D4M 60.9±1.7/77.7±1.1. ImageIDC (Table 3), ResNet-18, IPC1/10/50: 28.5±2.0/56.2±0.6/ 78.5±1.5 vs Minimax 22.4±0.6/51.9±1.4/78.3±0.2, DiT baseline 26.7±1.4/54.1±0.4/76.4±0.3. Ablation (Table 4, ImageWoof/ImageNette IPC10/50): DiT baseline 34.7/49.3 (Woof), 58.9/82.9 (Nette); +IM alone 37.3/53.5 and 60.0/81.5; +S3 alone 40.9/54.9 and 62.2/81.3; +ImS3 combined 41.8/61.0 and 62.9/84.2 -- isolating that IM raises coverage (accuracy) mostly at low-to-mid IPC while S3 adds inter-class separation, and combining them is superadditive at IPC50 (54.9 or 53.5 alone vs 61.0 combined on Woof). A real-data-accessibility check (Table 13, appendix) shows S3 still works using only diffusion-generated images (no real centroids) as the reference, at a moderate cost (IPC10 Woof: 39.6±0.4 gen-centroid vs 41.8±0.3 real-centroid).

Limitations

Stated: none explicitly flagged as a "Limitations" section in the main text (only future-facing Conclusion). Observed: Inversion-Matching requires per-dataset fine-tuning of an already target-trained DiT-XL/2 backbone (adapter-based, but still gradient updates and DDIM inversion during training), so it does not remove the target-trained-generator dependency CoDA later criticizes -- it deepens it, since inversion of real images is required at fine-tuning time. Selective Subgroup Sampling still needs access to real per-class images to compute reference centroids $r_i$ in its main configuration (a generated-centroid fallback is shown but underperforms, e.g. 39.6 vs 41.8 at IPC10 on ImageWoof). Evaluation reports "the higher of hard-label or soft-label protocol" per setting rather than a single fixed protocol, which makes its numbers harder to compare cell-by-cell against papers that fix one protocol throughout. The subgroup count $G$ and weights $\alpha,\beta$ are tuned per dataset, and no wall-clock/GPU-hour cost for either IM fine-tuning or the S3 candidate-generation-and-search step is given in the main text.