GeoDM: Geometry-aware Distribution Matching for Dataset Distillation

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 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.

Dataset distillation aims to create tiny, high-quality versions of giant datasets to speed up AI training. Most existing tools only work in Euclidean spaces, which fails to capture the intrinsic, curved geometry—like hierarchies or symmetries—found in complex real-world data. Our research introduces GeoDM, a framework that maps data into a combination of flat, spherical, and hyperbolic spaces. By using learnable "curvature" parameters, GeoDM automatically adapts to the unique geometric structure of the data it is trying to condense. We also used a alignment tool, optimal transport, to ensure the distilled data remains faithful to the original. Our experiments prove that by respecting the underlying geometry, we can achieve better performance with fewer images than ever before.