💻 computer science

Geometry-Aware Dataset Condensation for Diffusion Model Training

This paper proposes Geometry-Aware Dataset Condensation (GADC), a method that reformulates real subset selection as a geometry-aware distribution alignment problem using one-sided partial optimal transport and semantic regularization to construct compact datasets that preserve the geometric structure and distributional fidelity required for effective diffusion model training.

Original authors: Xiao Cui, Yulei Qin, Mo Zhu, Wengang Zhou, Hongsheng Li, Houqiang Li

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiao Cui, Yulei Qin, Mo Zhu, Wengang Zhou, Hongsheng Li, Houqiang Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you want to teach a master chef (a Diffusion Model) how to cook a perfect meal. Traditionally, you'd give them a massive library of millions of recipes and ingredients (the Full Dataset). While this works, it takes forever to read through everything, costs a fortune in storage, and requires a huge kitchen to process it all.

Dataset Condensation is the idea of creating a tiny, perfect "cheat sheet" of just a few dozen recipes that teaches the chef everything they need to know, without the bloat.

However, previous attempts to make these cheat sheets had two big problems:

Synthetic Cheats: Some tried to invent new recipes from scratch. The result? The recipes looked weird and tasted fake, confusing the chef.
Bad Selection: Others tried to just pick the "best" existing recipes. But they picked them based on a single, simple score (like "how hard is this recipe?"). This missed the big picture, leaving out important variations and creating a biased menu.

This paper proposes a new way to build that cheat sheet, called Geometry-Aware Dataset Condensation. Here is how it works, using simple analogies:

1. The Problem: The "One-Dimensional" Mistake

Imagine the full dataset is a giant, complex city with neighborhoods, parks, and skyscrapers (the Data Distribution).

Old Method (D2C): This method tried to pick the best houses by ranking them on a single straight line based on "difficulty." It's like trying to pick the best houses in a 3D city by only looking at their height. You might pick a tall, empty skyscraper and miss a cozy, essential cottage. You lose the shape of the city.
The Goal: We need to pick a small group of houses that perfectly represents the entire shape of the city, keeping the parks, the streets, and the neighborhoods intact.

2. The Solution: "One-Sided Partial Transport"

The authors use a mathematical tool called Optimal Transport, which is like a logistics company trying to move cargo from a warehouse (the Full Dataset) to a new, smaller warehouse (the Condensed Subset).

The Old Way (Balanced Transport): The old rules said, "You must move every single grain of sand from the big warehouse to the small one, perfectly matching the weight."
- The Flaw: Since the small warehouse is tiny, this forces the logistics company to drag heavy, useless sand from the edges of the city (low-density areas) just to fill the weight quota. This distorts the map.
The New Way (One-Sided Partial Transport): The authors say, "We only need to move the important cargo. We don't need to move the sand from the empty, low-density outskirts."
- The Benefit: This allows the small warehouse to focus entirely on the "core" of the city—the busy streets and popular neighborhoods. It ensures the small subset captures the true geometry (the shape and structure) of the original data without being dragged down by noise.

3. The Safety Net: "Statistical Regularization"

Just moving the cargo isn't enough; we need to make sure the new warehouse still feels like the original city. The authors add two "safety nets":

Mean-Variance Check: They ensure the average "height" and "spread" of the buildings in the small warehouse match the big city. If the big city has a mix of tall and short buildings, the small one must have that same mix.
Confidence Check: They make sure the selected houses are clearly recognizable. If a house looks like a blurry mess that could be a barn or a garage, they reject it. This ensures the "chef" isn't confused by ambiguous examples.

4. The Strategy: "Greedy Construction + Swap Refinement"

How do you actually pick these specific houses? You can't check every possible combination (there are too many!). So, they use a two-step strategy:

Greedy Building: Start with an empty lot and add one house at a time, always picking the one that improves the map the most right now. It's like building a puzzle piece by piece.
The Swap: Once the puzzle is built, they look for mistakes. "Hey, this house in the corner isn't working; let's swap it with that house outside." They keep swapping until the map is as perfect as possible.

The Results

When they tested this method on ImageNet (a massive database of 1.4 million images) to train AI image generators:

Better Quality: The AI generated images that looked much sharper and more diverse (lower "FID" scores) compared to previous methods.
Efficiency: They could train the AI using only 0.8% of the original data (10,000 images instead of 1.4 million) and still get better results than using random chunks of the full data.
Speed: The process of selecting these 10,000 images was much faster than previous methods.

In Summary:
This paper teaches us that to train a powerful AI on a small dataset, you shouldn't just pick the "hardest" or "easiest" examples. Instead, you should mathematically select a tiny group of images that perfectly preserves the shape, structure, and diversity of the original massive dataset, ignoring the empty, noisy edges. It's like curating a museum exhibit that captures the soul of an entire art history collection in just one room.

Technical Summary: Geometry-Aware Dataset Condensation for Diffusion Model Training

1. Problem Statement

Dataset condensation aims to construct compact datasets from real data via synthesis or selection to reduce the storage and computational costs associated with training deep learning models. While existing condensation methods have shown success in discriminative tasks (e.g., classification), they are ill-suited for training diffusion models.

Synthesis Limitations: Methods that generate synthetic data via continuous pixel optimization often produce low-fidelity samples. Diffusion models are highly sensitive to noise and structural distortions, making these synthetic samples unsuitable for authentic modeling.
Selection Limitations: Existing real subset selection methods typically rely on fixed criteria or heuristic rankings (e.g., D2C ranks images by diffusion difficulty along a one-dimensional axis). These approaches fail to preserve the complex distributional geometry required by diffusion likelihood objectives. They often lack a principled optimization objective aligned with diffusion training, leading to subsets that are poorly aligned with the full data distribution and fail to capture the necessary manifold structure.

The core challenge is to select a compact, real subset that preserves the geometric structure and distributional fidelity of the full dataset, specifically optimized for the likelihood-based training objectives of diffusion models.

2. Methodology

The authors propose Geometry-Aware Dataset Condensation (GADC), which reformulates real subset selection as a geometry-aware distribution alignment problem. The framework consists of three main components:

A. One-Sided Partial Optimal Transport (POT)
To address the capacity mismatch between a small subset and the full dataset, the authors employ one-sided partial optimal transport.

Concept: Unlike balanced OT which enforces full mass alignment, one-sided POT allows the source mass (selected subset) to be fully transported while relaxing the target mass (full dataset) under a capacity constraint.
Mechanism: This allows unmatched mass in low-density peripheral regions, concentrating the transport on high-density, geometrically stable, dominant regions of the data manifold.
Implementation: The problem is solved efficiently using a "dummy-source" reformulation and entropic regularization, optimized via Sinkhorn iterations. This enables flexible mass allocation where the subset focuses on representative regions while maintaining necessary coverage.

B. Statistical and Semantic Regularization
To complement the geometric alignment and ensure distributional fidelity, two lightweight regularization terms are introduced:

Mean-Variance Regularization ( $L_{sta}$ ): Aligns the first- and second-order statistics (mean and variance) of feature representations between the selected subset and the full dataset to preserve global distributional fidelity.
Confidence Regularization ( $L_{conf}$ ): Ensures semantic consistency by penalizing the selection of samples with low predicted class probabilities. This prevents the inclusion of unreliable geometric anchors that could compromise the alignment.

C. Two-Stage Discrete Optimization Strategy
Since the selection of a fixed-size subset is a combinatorial problem, the authors propose an efficient two-stage solver:

Stage I (Greedy Geometry-Guided Selection): An incremental construction process that selects samples to minimize the composite objective (POT loss + regularizations), establishing broad manifold coverage.
Stage II (Swap-Based Refinement): A refinement phase that performs pairwise swaps between selected and unselected samples to correct early myopic choices, further optimizing the global geometric alignment and consistency.

3. Key Contributions

Problem Reformulation: The paper reformulates dataset condensation for diffusion models as a distribution alignment problem, moving away from heuristic ranking or scalar-based selection. It introduces a one-sided POT objective with statistical regularization to select representative real subsets.
Optimization Framework: It proposes a two-stage discrete optimization framework (greedy construction followed by swap-based refinement) that efficiently solves the alignment problem, avoiding the limitations of fixed-criterion or ranking-based sampling.
Empirical Performance: Extensive experiments demonstrate that the method achieves superior fidelity and distributional coverage compared to prior approaches across various diffusion variants (DiT, SiT), subset sizes, image resolutions, and training rounds.

4. Experimental Results

The method was evaluated on ImageNet-1K using DiT-L/2 and SiT-L/2 architectures with data budgets ranging from 0.8% (10K images) to 8% (100K images).

Performance Metrics: The proposed method consistently achieved the lowest Fréchet Inception Distance (FID) and highest Inception Score (IS), Precision, and Recall compared to baselines including Random Sampling, K-Center, Herding, CCS, DQ, and D2C.
- Example: On ImageNet 256×256 with a 10K data budget (0.8%), the method achieved an FID of 3.43, outperforming D2C (4.20) and Random (35.86).
- Example: At 512×512 resolution with a 10K budget, the method achieved an FID of 6.17, significantly better than D2C (14.8).
Robustness: The improvements persisted across different diffusion variants (DiT and SiT), different evaluation protocols (10K vs. 50K generated samples), and extended training iterations (up to 300K).
Efficiency: The selection process is computationally efficient. On a single GPU, the method took 5.5 hours to select a 10K subset from ImageNet, compared to 41.9 hours for D2C and 30.4 hours for DQ.

Ablation Studies:

Removing the POT loss ( $L_{OT}$ ) degraded FID and reduced precision/recall, confirming the necessity of geometric alignment.
Replacing one-sided POT with balanced OT degraded performance, validating the benefit of allowing unmatched mass in peripheral regions.
Removing statistical ( $L_{sta}$ ) or confidence ( $L_{conf}$ ) regularization led to worse FID or IS, respectively, highlighting the importance of global statistics and semantic clarity.

5. Significance and Claims

The paper claims that by treating dataset condensation as a geometry-aware distribution alignment problem, it successfully bridges the gap between data selection and the specific requirements of diffusion model training.

Data Efficiency: The method enables high-quality generative modeling on condensed datasets, significantly reducing the storage and computational costs associated with training diffusion models.
Fidelity: Unlike synthetic generation methods, the approach retains the high-fidelity structure of real samples while ensuring the subset aligns with the full data distribution's geometry.
Practicality: The proposed two-stage optimization strategy is scalable and computationally efficient, making it a practical solution for resource-constrained settings.

The authors conclude that their method provides a reliable data substrate for diffusion-style generative training, offering a complementary, data-centric path to improving efficiency alongside model-side optimizations.