Dataset Distillation via the Wasserstein Metric

What came before

dm-2023-style mean matching and linear-kernel MMD are argued to be "not characteristic" and to neglect higher-order moments; at ImageNet scale, distribution-matching accuracy has lagged far behind decoupled BN-statistics-inversion methods such as sre2l-2023 and g-vbsm-2024.

The problem

Mean/MMD-based matching produces what the paper calls a "rigid mix-up" of the real feature distribution rather than a natural interpolation of it, and existing distribution-matching methods trail decoupled methods badly at ImageNet-1K scale.

The idea

Compute the Wasserstein barycenter of real features from a pretrained classifier -- which optimally interpolates a distribution rather than crudely averaging it -- and optimize synthetic images to match that barycenter in feature space, restoring the intra-class variation a single barycenter would erase by matching BatchNorm statistics separately per class.

How it works

Extracts penultimate-layer features of a classifier (ResNet-18/-50) pretrained once on the full real dataset (a "squeeze stage" akin to sre2l-2023's teacher); computes the Wasserstein barycenter of each class's real features via alternating optimal transport (about 10 iterations: subgradient descent on transport weights, a Newton step on barycenter positions); optimizes synthetic pixel images with a single-level loss combining an L2 distance to the barycenter features and a per-class BatchNorm-statistics term (BN running mean/variance computed separately per class rather than globally, so images keep the class's spread instead of collapsing to the barycenter's single point). No diffusion or other generative model is used, and there is no bi-level loop.

Evidence

ImageNet-1K, ResNet-18 eval: IPC=10, 38.2% vs. SRe2L 21.3%, G-VBSM 31.4% (Table 1); IPC=50, 57.6% vs. SRe2L 46.8%, G-VBSM 51.8%; IPC=100, 60.7% vs. the full-dataset ceiling of 63.1%. ImageNette, IPC=10: 64.8% vs. SRe2L 54.2%, DataDAM 59.4%. Cross-architecture (ImageNet-1K, IPC=50): ResNet-18 57.83%, ResNet-50 61.22%, ViT-Tiny 34.25% (markedly weaker on ViT). Cost (Table 3, ImageNette, IPC=1, RTX 3090): WMDD 0.013s/iter, 1.22GB, 207.5s total vs. SRe2L 0.015s/1.14GB/194.9s (comparable, since both use a single pretrained teacher rather than sampling many random networks per step) vs. DM 1.965s/9.93GB/4018s and DC 2.154s/11.90GB/6348s (18-30x more expensive). Ablation (Table 4, IPC=10): Wasserstein loss + per-class BN is best on every dataset (ImageNette 64.7%, TinyImageNet 41.8%, ImageNet-1K 38.1%); dropping per-class BN costs 4-11pp; replacing the Wasserstein term with vanilla MMD gives "near-random performance" (Figure 3).

Limitations

Requires a classifier pretrained on the full real dataset, tying the method's accuracy ceiling to that classifier's feature quality -- a real, if amortizable, cost not present in random-network distribution matching. Performance on Vision Transformers is notably weaker ("probably due to their data-hungry property," per the paper). Not tested outside vision. The paper acknowledges the method "inherently reflects existing biases" in the pretrained classifier's training data. The label regime used at evaluation is not fully specified in the sections read; given the SRe2L-style pipeline it is compared against, soft/relabeled labels are plausible but not confirmed here.