CaO$_2$: Rectifying Inconsistencies in Diffusion-Based Dataset Distillation

What came before

Minimax Diffusion fine-tunes a diffusion model for representativeness/diversity, and D4M clusters latents into category prototypes; both generate images purely to match the target distribution, with no reference to how a classifier will actually use them. RDED instead selects and stitches real image crops using an evaluation-style pipeline. Matching-based methods (DC/DM/MTT/TESLA/DATM) do align distillation with the classification objective by construction, but do not scale efficiently to ImageNet.

The problem

The paper names two specific failure modes never previously isolated: (1) Objective Inconsistency — a conditional diffusion model's images are optimized purely for generative fidelity/likelihood, so some generated images, when passed through a classifier, are predicted to belong to a different class than the one they were conditioned on; (2) Condition Inconsistency — because no real diffusion model achieves perfect conditional likelihood (only an ELBO is optimized, with nonzero training loss), a generated image's latent is never *exclusively* associated with its conditioning label, weakening the image-label correlation that supervised training on the distilled set depends on.

The idea

Treat evaluation-time correctness as a two-stage post-hoc filter and refinement applied to an already-pretrained, untouched diffusion model: first over-generate a pool of candidate images per class and keep only the ones a lightweight external classifier confidently assigns to the correct class (fixing Objective Inconsistency), then nudge each kept image's latent, via a short diffusion-loss- minimizing optimization under its class condition, toward a region where the pretrained model's own conditional likelihood is higher (fixing Condition Inconsistency) — all without ever fine-tuning the generator.

How it works

A pretrained DiT (256x256, ImageNet-trained, untouched/frozen) is sampled (50 DDIM-style steps, fixed seed) to produce an image pool of size $mN$ per class ($m{=}2$ or $4$, $N{=}$IPC). Stage 1 (Objective- guided Sample Selection): a lightweight pretrained classifier (ResNet-18) scores the pool via softmax, and only images predicted as the conditioned class are kept, choosing the top-$N$ most confident (easy tasks/higher IPC per the task-oriented variation, Sec. 3.4) or least confident (harder tasks/lower IPC) correctly-classified samples; classes with insufficient correct samples backfill with random pool images. Stage 2 (Condition-aware Latent Optimization): each selected image's latent $\mathbf{x}$ is perturbed by Gaussian noise to a randomly sampled early timestep $t\in[1,\hat T]$, $\hat T \ll T$, and optimized to minimize the diffusion denoising loss $\mathbb{E}_{t,\varepsilon}[\|\epsilon_\theta(\mathbf{x}_t,\hat{\mathbf{c}},t)-\varepsilon\|_2^2+ \lambda\|\epsilon_\theta(\mathbf{x}_t,\hat{\mathbf{c}},t)-\varepsilon\|_\infty]$ (Eq. 4, $L_\infty$ regularization to bound the latent shift; 100 iterations, Adam, lr 0.0006, $\lambda{=}10$), i.e. moving the latent toward a region of higher conditional likelihood under the frozen diffusion model itself, not toward a classifier's decision boundary. The condition $\hat{\mathbf{c}}$ used for this optimization is the true class label for easier tasks and an unconditional (classifier-free-guidance null) embedding for harder tasks (Eq. 5), chosen per-dataset from validation accuracy. Both stages generalize to a Masked Autoregressive (MAR) generation backbone and can be layered on top of an already Minimax-fine- tuned DiT as a plug-in.

Evidence

ImageWoof/ImageNette, best-of-hard-or-soft-label protocol (paper reports whichever evaluation regime scores higher for each method, per method, per cell), ResNet-18/50/101 (Table 1): ImageNette IPC=10 ResNet-18 CaO2 65.0% vs. RDED 61.4%/Minimax 61.4% vs. SRe2L 29.4%; ImageWoof IPC=10 ResNet-18 CaO2 45.6% vs. RDED 38.5%/Minimax 40.1%. ImageNet-100/ImageNet-1K (Table 2, same mixed-best-protocol convention): ImageNet-1K IPC=10 ResNet-18 CaO2 46.1% vs. RDED 42.0% vs. Minimax 44.3% vs. D4M 27.9%; IPC=50 CaO2 60.0% vs. Minimax 58.6% vs. RDED 56.5%. Reported average gains over best baselines: +1.6% (ImageWoof), +4.3% (ImageNette), +1.8% (ImageNet-100), +1.5% (ImageNet-1K) — the abstract's headline "2.3% average." Key ablation (Table 3, ImageWoof/ImageNette, ResNet-18): OSS (sample selection) alone and CLO (latent optimization) alone each give comparable, roughly equal-sized gains over the no-component baseline (e.g. ImageWoof IPC=10: baseline 38.7%, +OSS 42.6%, +CLO 42.1%, both combined 45.6%) — the paper reads this as evidence the two inconsistencies are separate, similarly-sized problems, not one problem in two guises. Regularization ablation (Table 4): $L_\infty$ norm beats $L_1$, $L_2$, and no regularization. Selection-pool-size ablation (Fig. 3b): accuracy peaks at pool size 2x-4x IPC and does *not* keep improving with a larger pool, because a larger pool, while more diverse, can be less representative once filtered. Backbone generalization (Table 5): CaO2 improves a MAR backbone by +3.3%/+2.0% (ImageWoof/ImageNette averages) and, layered on top of an already Minimax-fine-tuned DiT, improves it further by +4.3%/+2.0% — evidence the method is complementary to, not merely a substitute for, fine- tuning-based diversity/representativeness objectives.

Limitations

The DiT backbone is ImageNet-trained, i.e. matched to the target distribution for every reported number; the paper's own Limitations appendix (Sec. 11) states this explicitly as a constraint — "DiTs can only deal with ImageNet classes" — and notes that switching to a general text-to-image model (e.g. Stable Diffusion) to handle arbitrary classes is not tested here and would trade this limitation for larger model size and a loss of classification-aware constraints; the paper's own framing acknowledges the two inconsistencies stem from "the fundamental difference between generation and discrimination," which its post-hoc filter-and-nudge approach only partially closes rather than resolving with a unified objective; the headline comparison table reports, per method, whichever of two different label-regime evaluation protocols (hard-label vs. soft-label via KD) scores higher, which makes Table 1/2 numbers not a single controlled label-regime comparison across methods, even though it is applied evenhandedly to CaO2 and all baselines alike.