Optimizing Distributional Geometry Alignment with Optimal Transport for Generative Dataset Distillation

AI summary

Overview

• Research area: Computer vision, specifically dataset distillation — compressing a large training set into a small synthetic one — combined with optimal transport (OT) theory and latent diffusion generative models.
• Technical level: Advanced. The paper assumes familiarity with Wasserstein distance, Sinkhorn iterations, entropy-regularized optimal transport, latent diffusion sampling and guidance, and soft-label distillation.
• Scope (one sentence): The paper reformulates generative dataset distillation as an optimal transport distance minimization problem, decomposes it into three tractable stages, and evaluates it on ImageNet-1K, ImageNet-100, ImageNette and CIFAR-100 across six architecture families.

What This Paper Is About

Large-scale dataset distillation methods typically compress a dataset by matching global distributional statistics such as mean and variance, or by matching global gradient statistics. The authors argue this overlooks instance-level characteristics and intra-class variation, so the distilled set loses fine-grained structure and generalizes poorly. Their goal is to align the geometry of the real and distilled distributions at both global and instance levels by minimizing an optimal transport distance throughout the distillation pipeline, from image sampling through label assignment to student model training.

Key Contributions

1. An OT-based reformulation of dataset distillation. The task is posed as minimizing the Wasserstein distance between the true data-label distribution and the distribution induced by a student model trained on the distilled set, and this total distance is decomposed into three tractable components.
2. A theoretical decomposition using two principles: a triangle-inequality partition of the discrepancy before and after distilled set construction, plus a multiplicative contraction term that models the benefit of soft labels over hard one-hot labels, producing the structured upper bound in Equation 8.
3. Three pipeline components targeting the terms of that bound: (1) OT-guided diffusion sampling (OTG), which aligns latent distributions of real and distilled images; (2) label-image-aligned soft relabeling (LIA), which adapts label distribution complexity to the number of images per class (IPC); and (3) OT-based logit matching (OTM), which aligns student model logits with soft-label distributions via a batch-wise Sinkhorn distance.
4. Broad empirical validation across ResNet, MobileNet, EfficientNet, Swin Transformer, ConvNet and ConvNeXt on ImageNet-1K, ImageNet-100, ImageNette and CIFAR-100, reporting at least a 4% accuracy improvement under IPC=10 settings for each architecture on ImageNet-1K.

Main Findings

• ImageNet-1K with ResNet-18 (IPC=10): the method reaches 52.9 ± 0.1 at 300 epochs and 58.6 ± 0.3 at 1000 epochs, compared with 45.5 ± 0.5 for DiT-IGD and 48.6 ± 0.3 for EDC. At IPC=50 it reaches 61.9 ± 0.5 (300 epochs) and 64.2 ± 0.4 (1000 epochs), versus 59.8 ± 0.3 for DiT-IGD and 58.0 ± 0.2 for EDC.
• Other ImageNet-1K architectures (IPC=10, 1000 epochs): MobileNet-V2 57.6 ± 0.1 (EDC 45.0 ± 0.2), EfficientNet-B0 62.4 ± 0.1 (EDC 51.1 ± 0.3), Swin Transformer 63.7 ± 0.2 (EDC 46.0 ± 0.5), ConvNeXt 67.0 ± 0.1 (EDC 54.4 ± 0.2). Gains are reported as larger at lower IPC settings.
• ImageNette with soft labels (IPC=10): ConvNet-6 74.5 ± 0.3, ResNetAP-10 77.8 ± 0.8, ResNet-18 79.0 ± 0.3, against DiT-IGD* at 69.6 ± 1.0, 73.6 ± 1.3 and 74.8 ± 0.7 respectively. Full-dataset reference accuracies are 94.3 ± 0.5, 94.6 ± 0.5 and 95.3 ± 0.6.
• ImageNet-100: ResNet-18 at IPC=10 reaches 47.7 ± 0.3 (RDED 36.0 ± 0.3, DELT 28.2 ± 1.5, SRe2L 9.5 ± 0.4) and 72.6 ± 0.1 at IPC=50; ResNet-101 at IPC=10 reaches 36.3 ± 0.5; MobileNet at IPC=10 reaches 43.2 ± 0.2.
• CIFAR-100 with ConvNet-3: 50.7 ± 0.2 at IPC=10, 57.5 ± 0.3 at IPC=50 and 58.7 ± 0.2 at IPC=100, versus RDED at 48.1 ± 0.3, 57.0 ± 0.1 and 58.1 ± 0.4.
• Ablation on ImageNette at IPC=10: on ResNet-18 the pipeline improves from 67.7 (hard label) and 69.1 (soft label) to 77.2 without OTG, 77.8 with OTG, 77.5 without OTM and 79.0 for the full method. The same trend appears for ConvNet-6 (61.9, 67.0, 72.5, 74.3, 73.2, 74.5) and ResNetAP-10 (66.5, 68.0, 74.2, 76.4, 75.9, 77.8). "w/o LIA" is defined as soft relabeling using the high-IPC teacher ensemble without adapting to the current IPC.
• Runtime overhead is small: on ImageNet-1K, mean sampling time per class is 97.1s without OTG versus 97.7s with OTG at IPC=10, and 537.4s versus 540.3s at IPC=50; mean matching time per epoch is 23.2s versus 23.3s at IPC=10 and 126.1s versus 126.6s at IPC=50.
• Generation cost: the paper reports for IPC=10 on ImageNet-1K on 8×4090 hardware that EDC requires 3h presample + 3h postsample + 5h recover + 0.4h relabel, while the proposed method requires 3.4h diffusion sampling + 0.3h relabel.
• Contraction factor: Table 10 on ImageNette reports α values of 0.906, 0.903 and 0.643 with corresponding average accuracies of 76.0, 76.2 and 77.1, indicating that IPC-adaptive relabeling lowers the contraction factor while accuracy rises.

Methodology in Plain English

The authors start from the idea that two distributions with the same mean or variance can still differ in shape (Figure 1a), so statistics-based matching misses geometric structure. They therefore measure the gap between the real distribution and the distribution induced by the distilled data with the Wasserstein distance, which is a true metric and satisfies the triangle inequality — unlike KL divergence or cosine similarity. Applying the triangle inequality and then substituting soft labels for hard labels yields an upper bound with two parts: the dataset discrepancy between real data, and the logit matching error between soft labels and student outputs. The dataset discrepancy is further factored into a class-wise transport term multiplied by a contraction factor that captures how much soft labels help, giving three terms to minimize in sequence.

• Stage one — OT-guided diffusion sampling (OTG): while generating the n-th latent for a class, the method draws a random batch of real images of that class, encodes them into latents, and computes the optimal transport distance between those real latents and the concatenation of already-sampled latents plus the current one. The transport plan is solved approximately with entropy-regularized Sinkhorn iterations using an ℓp-norm cost matrix. The gradient of this distance is added to the usual trajectory and diversity guidance terms in the diffusion update, so sampling is steered toward the geometry of the real latent distribution. Decoded outputs form the distilled image set.
• Stage two — label-image-aligned soft relabeling (LIA): the complexity of the soft labels is matched to the expressive capacity of the distilled set, which is governed by IPC. At low IPC, fewer teacher models produce simpler, lower-entropy soft labels to avoid mismatched supervision; at high IPC, a larger and more diverse teacher ensemble produces finer-grained soft labels. Each image's soft label is the average of the selected teachers' logit outputs, and the paper adopts the same region-level soft-label storage strategy as FKD for fair comparison.
• Stage three — OT-based logit matching (OTM): rather than matching each logit independently, the method computes a batch-wise Sinkhorn distance between the batch of soft labels and the batch of student logits using an ℓp-norm cost matrix, which captures relationships between samples. The final training loss combines cross-entropy on one-hot labels, MSE against soft labels, and the Sinkhorn distance loss, each with its own scalar weight.

Why This Matters

• Research impact: the paper offers a principled theoretical frame — a decomposable Wasserstein upper bound — for why soft labels and instance-level alignment help dataset distillation, and it shows that instance-level geometric alignment can be added to generative distillation at negligible runtime cost (for example, 97.1s to 97.7s per class at IPC=10). It also provides a direct argument against global-statistics matching in model-inversion methods, which the authors note discard fine-grained instance-level information.
• Knowledge transfer: distilled datasets can carry the behavior of large models to smaller ones, which the paper lists among the practical uses of dataset distillation.
• Federated learning: compact distilled sets reduce what must be communicated and stored across clients, which the paper cites as an application area.
• Continual learning: a small distilled set can stand in for past data when a model must retain earlier tasks, also cited by the authors.
• Storage, computation and energy: the paper frames dataset distillation as a response to prohibitive storage, computation and energy costs of data expansion, which matters for large-scale deployments such as ImageNet-1K.

Future Directions

• Automating the IPC-aware teacher selection. LIA depends on a discrete choice of teacher subset per IPC regime; a learned or continuous version of this choice could remove manual tuning and might further reduce the contraction factor.
• Tuning the OT machinery. The entropy regularization weights and loss weights (the paper reports α1 = 1, γ1 ∈ {1000, 3000}, α2 = 0.1, γ2 = 0.1 and p = 1 with ℓ1-norm in the main text) are fixed or chosen from a small set; how sensitive the results are to these, and whether the bound's components should be weighted differently, is not settled in the provided content.
• Extending beyond the evaluated settings. The provided content covers ImageNet-1K, ImageNet-100, ImageNette and CIFAR-100; whether the OT formulation transfers to other modalities, higher resolutions, or IPCs beyond those tested is not reported.
• Closing the remaining gap to full-data training. On ImageNette at IPC=10 the best reported result is 79.0 ± 0.3 versus full-dataset accuracies of 94.3 ± 0.5, 94.6 ± 0.5 and 95.3 ± 0.6 — a substantial remaining gap that motivates further work. Details of appendices C, D, E, F and G.2 (dataset descriptions, symbol list, pseudocode, hyperparameters) are not included in the provided content and are therefore not summarized here.

Target Audience

• Researchers working on dataset distillation, dataset condensation and data-efficient training.
• Practitioners applying optimal transport or Sinkhorn algorithms to generative modeling and distribution matching.
• Engineers who need to compress large training corpora for storage, compute or energy reasons in production pipelines.
• Graduate students with a background in diffusion models and probability metrics who want a worked example of turning a distribution-matching objective into a staged, implementable training recipe.

Authors’ abstract

Dataset distillation seeks to synthesize a compact distilled dataset, enabling models trained on it to achieve performance comparable to models trained on the full dataset. Recent methods for large-scale datasets focus on matching global distributional statistics (e.g., mean and variance), but overlook critical instance-level characteristics and intraclass variations, leading to suboptimal generalization. We address this limitation by reformulating dataset distillation as an Optimal Transport (OT) distance minimization problem, enabling fine-grained alignment at both global and instance levels throughout the pipeline. OT offers a geometrically faithful framework for distribution matching. It effectively preserves local modes, intra-class patterns, and fine-grained variations that characterize the geometry of complex, high-dimensional distributions. Our method comprises three components tailored for preserving distributional geometry: (1) OT-guided diffusion sampling, which aligns latent distributions of real and distilled images; (2) label-image-aligned soft relabeling, which adapts label distributions based on the complexity of distilled image distributions; and (3) OT-based logit matching, which aligns the output of student models with soft-label distributions. Extensive experiments across diverse architectures and large-scale datasets demonstrate that our method consistently outperforms state-of-the-art approaches in an efficient manner, achieving at least 4% accuracy improvement under IPC=10 settings for each architecture on ImageNet-1K.