GeoDM: Geometry-aware Distribution Matching for Dataset Distillation

TL;DR

GeoDM tackles the limitation of Euclidean-only distribution matching in dataset distillation by embedding real and synthetic data into a learnable product manifold that combines Euclidean, hyperbolic, and spherical geometries. The framework introduces learnable curvature and geometry weights, plus a geometry-aware optimal-transport loss, to align distributions while preserving manifold structure. Theoretical results show tighter generalization bounds for product-space matching, and extensive experiments across MNIST, CIFAR, and high-resolution datasets demonstrate consistent accuracy gains over state-of-the-art methods. GeoDM provides a principled approach to manifold-aware distillation with robust cross-architecture performance and ablations confirming the contribution of each component. The work highlights practical benefits for reducing data footprints without sacrificing/domain-relevant geometric structure, suggesting broader applicability of mixed-curvature representations in data synthesis tasks.

Abstract

Dataset distillation aims to synthesize a compact subset of the original data, enabling models trained on it to achieve performance comparable to those trained on the original large dataset. Existing distribution-matching methods are confined to Euclidean spaces, making them only capture linear structures and overlook the intrinsic geometry of real data, e.g., curvature. However, high-dimensional data often lie on low-dimensional manifolds, suggesting that dataset distillation should have the distilled data manifold aligned with the original data manifold. In this work, we propose a geometry-aware distribution-matching framework, called \textbf{GeoDM}, which operates in the Cartesian product of Euclidean, hyperbolic, and spherical manifolds, with flat, hierarchical, and cyclical structures all captured by a unified representation. To adapt to the underlying data geometry, we introduce learnable curvature and weight parameters for three kinds of geometries. At the same time, we design an optimal transport loss to enhance the distribution fidelity. Our theoretical analysis shows that the geometry-aware distribution matching in a product space yields a smaller generalization error bound than the Euclidean counterparts. Extensive experiments conducted on standard benchmarks demonstrate that our algorithm outperforms state-of-the-art data distillation methods and remains effective across various distribution-matching strategies for the single geometries.

Figures (4)

• Figure 1: 3D embeddings of CIFAR10 data in Euclidean (left), hyperbolic (middle), and spherical (right) spaces with fitted manifolds: in hyperbolic space the points exhibit a hierarchical pattern and align well with the hyperbolic surface, while in spherical space they concentrate on a sphere, revealing inherent spherical geometry; the Euclidean space struggles to capture such geometric structure.
• Figure 2: Overall pipeline of our GeoDM: Real and synthetic data are first processed by a curved feature network, where hyperbolic and spherical branches incorporate learnable curvature and dedicated projection heads. The resulting embeddings are mapped to the Euclidean–hyperbolic–spherical product space. DM aligns real and synthetic features, while a geometry-aware optimal transport (OT) loss further couples the three geometries and preserves class-conditional mass.
• Figure 3: Comparison between product spaces and single spaces on CIFAR-10 with IPC=10.
• Figure 4: Condensed images

Theorems & Definitions (19)

• Theorem 4.1: Geometry-driven risk decomposition
• Theorem 4.2: Product spaces yield tighter bounds
• Theorem A.1: Geometry-driven risk decomposition repeat
• proof
• Lemma A.2: Hyperbolic volume growth and large separated subsets
• proof
• Lemma A.3: Trees embed into hyperbolic space with constant distortion
• proof
• Lemma A.4: Lower bounds for embedding trees into Euclidean/Hilbert
• proof