GeoDM: Geometry-aware Distribution Matching for Dataset Distillation

AI summary

Overview

Research area: Dataset distillation (dataset condensation) for computer vision, specifically distribution-matching methods, using Riemannian geometry and manifold learning.

Technical level: Advanced. The paper assumes familiarity with Riemannian manifolds (hyperbolic and spherical geometry, exponential/logarithmic maps, tangent spaces), optimal transport, and the dataset-distillation literature.

Scope: This paper proposes GeoDM, a distribution-matching framework that distills datasets inside a Cartesian product of Euclidean, hyperbolic, and spherical spaces with learnable curvature and weights, and supports it with a generalization-error analysis and experiments on MNIST, CIFAR-10, and CIFAR-100.

What This Paper Is About

Dataset distillation tries to synthesize a tiny set of images that lets a model train to roughly the same accuracy as on the full dataset. Nearly all existing distribution-matching methods compare real and synthetic data in a flat Euclidean embedding space, which can only express linear structure and ignores curvature. GeoDM instead matches distributions in a product of curved manifolds so the distilled data can inherit hierarchical (hyperbolic) and cyclical/directional (spherical) structure present in real data.

Key Contributions

1. Reframing dataset distillation as product-manifold matching. The authors build on the manifold hypothesis and treat distillation as aligning distributions between product manifolds, where each factor carries its own curvature — Euclidean (flat), hyperbolic (negative curvature), or spherical (positive curvature).

2. The GeoDM framework. A distribution-matching pipeline operating in the Cartesian product of Euclidean, hyperbolic, and spherical spaces, with learnable curvature parameters, softmax-normalized learnable geometry weights, and a geometry-aware optimal transport loss that couples the three geometries and preserves class-conditional mass.

3. Theoretical analysis. Theorem 4.1 decomposes the generalization gap into statistical, stability, and geometric distortion terms, and Theorem 4.2 shows that distribution matching in the product space yields a strictly tighter generalization error bound, Δ_product ≤ Δ_Euclid − Lδ, than a single Euclidean space.

4. Empirical validation. Experiments on MNIST, CIFAR-10, and CIFAR-100 showing gains over state-of-the-art distillation baselines, plus ablation and cross-architecture studies and robustness checks across different distribution-matching objectives.

Main Findings

• MNIST at IPC = 1: GeoDM reaches 96.3 ± 0.2 accuracy, versus the best baseline M3D at 94.4 ± 0.2, about a 2% improvement. At IPC = 10 GeoDM reaches 98.6 ± 0.2 and at IPC = 50 99.3 ± 0.1 (a model trained on the whole MNIST set scores 99.6 ± 0.0).

• CIFAR-10: GeoDM attains 74.4 ± 0.3 at IPC = 10 and 78.3 ± 0.2 at IPC = 50, versus 71.8 ± 0.5 and 77.4 ± 0.3 for the NCFM baseline. At IPC = 1 it reaches 51.2 ± 0.2. The paper reports surpassing state of the art by approximately 3% at IPC = 10. Whole-dataset CIFAR-10 accuracy is 84.8 ± 0.1.

• CIFAR-100: GeoDM reaches 38.0 ± 0.4 at IPC = 1 and 49.2 ± 0.3 at IPC = 10, gains of about 3% and 1% over the best previous result. At IPC = 50 it reaches 55.0 ± 0.2. Whole-dataset CIFAR-100 accuracy is 56.2 ± 0.3.

• Gains shrink as IPC grows: The paper states that relative improvement diminishes as images-per-class increases, because having many synthetic samples reduces the need for extra structural information.

• Ablation on CIFAR-10, IPC = 10: Starting from 71.8 ± 0.3, adding the product space alone gives 73.5 ± 0.2; adding curvature and weight adaptation alone gives 72.3 ± 0.1; product space plus curvature/weights gives 73.9 ± 0.2; product space plus OT gives 73.8 ± 0.1; and all three components together give the best result of 74.4 ± 0.2.

• Robustness to the matching objective: Replacing the matching component with a simple DM objective yields up to 2.79% improvement over single-space baselines; using the stronger DSDM formulation still gives roughly 2% gain. The paper argues the benefit comes from product-space modeling itself rather than any one matching instantiation.

• Cross-architecture transfer on CIFAR-10: Synthetic data distilled with ConvNet-3 transfers best under GeoDM. At IPC = 10, GeoDM scores 74.4 ± 0.3 on ConvNet-3, 69.1 ± 0.5 on ResNet-18, and 69.0 ± 0.3 on AlexNet. At IPC = 50, the corresponding scores are 78.3 ± 0.2, 74.5 ± 0.4, and 75.6 ± 0.1. For comparison at IPC = 50, NCFM scores 77.4 ± 0.3 on ConvNet-3, 73.7 ± 0.2 on ResNet-18, and 75.5 ± 0.3 on AlexNet, while M3D drops to 31.6 ± 0.9 on ResNet-18 and 35.6 ± 0.8 on AlexNet.

• Theory result: Theorem 4.2 states there exists δ > 0 such that Δ_product ≤ Δ_Euclid − Lδ. The paper attributes the improvement to hyperbolic embeddings representing hierarchical structure isometrically where Euclidean embeddings incur distortion growing exponentially with radius, and to spherical embeddings preserving geodesic arcs where Euclidean embeddings reduce them to chords (at least a π/2 factor of distortion), while flat data incurs no loss in either case.

Methodology in Plain English

The core idea is to stop pretending data lives in a flat space. Real data often sits on a low-dimensional curved surface — tree-like hierarchies, circles, directions — that a flat Euclidean embedding cannot represent without distortion.

GeoDM embeds every real and synthetic image into three spaces at once: an ordinary Euclidean space, a hyperbolic space (curvature c_H < 0, implemented with a Poincaré ball), and a spherical space (curvature c_S > 0). A Riemannian convolutional network with spherical convolutions and hyperbolic exponential/logarithmic maps produces features for each branch, which are concatenated into a single product-space feature vector.

Two things are made learnable. First, the curvatures of the hyperbolic and spherical parts are optimized along with network weights, kept stable by a curvature loss that penalizes embeddings drifting from the ball radius r_H = 1/√(−c_H) and the sphere radius r_S = 1/√c_S. Second, the three geometries get learnable weights α, β, γ, normalized by softmax so they sum to one, which controls how much each geometry contributes to the product-space distance.

Distribution matching itself follows the NCFM objective in the product space, but the authors emphasize the framework is agnostic to that choice. On top of it, a geometry-aware optimal transport loss measures the squared 2-Wasserstein distance between the embedded real and synthetic distributions. Because the space is curved, gradients for the hyperbolic and spherical components are computed in tangent spaces and mapped back to the manifold with exponential maps, keeping every embedding a valid manifold point. The three terms are combined as L_total = L_DM + λ_OT·L_OT + λ_curv·L_curv.

Experiments use a 3-layer convolutional backbone, T = 10⁴ matching iterations, λ_OT = 2, λ_curv = 1, curvature and geometry weights initialized uniformly at random, evaluation models trained for 1500 epochs, and accuracy as the metric, with three repeats reported as mean and standard deviation.

Why This Matters

Impact on research. The paper argues that geometry is not a cosmetic detail in distillation but a term that appears explicitly in the generalization bound. It opens a route for other condensation and matching objectives to be lifted into curved or product spaces, and it connects the dataset-distillation literature (gradient matching, trajectory matching, distribution matching, coreset selection) to Riemannian representation learning. Results on three standard benchmarks give the claim empirical backing, and the reported robustness across DM variants suggests the benefit is not tied to one loss.

Potential applications (the paper does not evaluate these directly; they follow from the dataset-distillation setting):

• Compressing large image datasets for training on resource-limited hardware such as phones, drones, or embedded sensors.
• Continual or incremental learning, where a compact distilled memory must retain earlier-task information.
• Medical and scientific imaging pipelines, where storage and annotation budgets limit how much data can be kept.
• Rapid architecture search and benchmarking, where cheap surrogate datasets speed up repeated training runs.

Industry relevance. Any team that pays the storage, transfer, or compute cost of training on huge image corpora has a stake in distillation ratios. The paper reports results at ratios ranging from 0.017% (MNIST, IPC = 1) to 10% (CIFAR-100, IPC = 50). The cross-architecture results matter for deployment, since synthetic data is often distilled on one network and deployed on another: GeoDM's distilled CIFAR-10 data retains 69.1 ± 0.5 accuracy on ResNet-18 and 69.0 ± 0.3 on AlexNet at IPC = 10, while several baselines degrade sharply across architectures.

Future Directions

• Scaling beyond CIFAR-scale benchmarks. The experiments cover MNIST, CIFAR-10, and CIFAR-100; whether the product-space advantage holds on larger, higher-resolution datasets is not reported.
• Extending the theory's assumptions. The bounds rest on Assumption 4.1 (data manifold, Lipschitz loss, uniform stability, statistical convergence), so relaxing those conditions and testing the predicted δ gap empirically remain open.
• Choosing and learning the geometry decomposition. The authors fix the dimensions (d_E, d_H, d_S) across datasets rather than learning them, noting that varying dimensionality destabilizes optimization; finding a stable way to learn these is unresolved.
• Applying the framework to other matching objectives and modalities. The paper shows robustness when swapping in DM and DSDM, but the study is limited to image classification, and the truncated ablation table in the provided content leaves the full component-interaction analysis incomplete.

Target Audience

Researchers and graduate students working on dataset distillation, dataset condensation, and efficient training; practitioners of Riemannian or hyperbolic deep learning interested in a concrete vision application; and machine-learning engineers who need compact training sets and want to know how much accuracy current distillation methods give up relative to full-data training. Readers should be comfortable with optimal transport and differential geometry, since the method and the theoretical section use curvature, tangent spaces, and exponential maps throughout.

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