Optimizing Distributional Geometry Alignment with Optimal Transport for Generative Dataset Distillation

TL;DR

The paper recasts dataset distillation as an Optimal Transport distance minimization problem to preserve fine-grained distributional geometry across both latent image and label spaces. It introduces a three-fold OT framework: OT-guided diffusion sampling to align latent distributions, label-image-aligned soft relabeling to match label complexity, and OT-based logit matching to transfer geometry to student models. The approach is decomposed into tractable components with theoretical motivation and validated on ImageNet-1K and CIFAR-100, showing consistent improvements over a wide range of architectures, especially under low IPC. The results demonstrate improved generalization, efficiency, and robustness, signaling practical impact for scalable, high-fidelity dataset distillation and downstream learning in resource-constrained settings.

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.

Figures (11)

• Figure 1: (a) Distributions with identical mean or variance may differ geometrically, causing biases in global-statistics-based optimization. (b) We decompose the OT objective into three stages: OT-guided diffusion sampling, label-image-aligned soft relabeling, and OT-based logit matching. (c) Our method consistently outperforms prior approaches across architectures on ImageNet-1K (IPC = 10).
• Figure 2: Comparison of generated images from different methods on ImageNet-100 (IPC = 10).
• Figure 3: Training accuracy curves of ResNet-18 (left) and MobileNet (right). IGD: DiT-IGD.
• Figure 4: Effect of $\beta_1$ (OT sampling weight) on ImageNette nette (IPC=10).
• Figure 5: Effect of $\lambda_1$ (entropy regularization weight for sampling) on ImageNette nette (IPC=10).