Optimizing Distributional Geometry Alignment with Optimal Transport for Generative Dataset Distillation

2025-11-28

AI Summary

Dataset distillation aims to replace large-scale real datasets with compact synthetic data for training high-performance models, yet existing methods merely match global statistics while neglecting instance-level features and intra-class structure, limiting generalization. This paper proposes a generative distillation framework grounded in optimal transport (OT), formulating distillation as minimizing the OT distance to achieve fine-grained alignment of both global distributions and local geometric structures. Key innovations include: (i) OT-guided diffusion sampling, (ii) soft label recalibration enforcing label-image consistency, and (iii) OT-based logits matching. Evaluated on ImageNet-1K with 10 images per class (IPC=10), our method achieves an average accuracy improvement of ≥4% across diverse architectures over state-of-the-art methods, with significantly improved training efficiency.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies

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.

Problem

• Optimizes dataset distillation using Optimal Transport for geometric alignment
• Addresses instance-level and intra-class variation oversight in current methods
• Enhances generalization by preserving local modes and fine-grained patterns

Innovation

• Uses Optimal Transport for fine-grained distribution alignment
• Implements OT-guided diffusion sampling for latent space matching
• Applies OT-based logit matching with soft-label distributions