MGPO: Manifold-Guided Diffusion Alignment for Task-Aware Dataset Distillation

Yunyi Chen; Chenru Wang; Xinyi Ye; Zexin Zheng; Chi Zhang

cs.CVOct 4, 2026

MGPO: Manifold-Guided Diffusion Alignment for Task-Aware Dataset Distillation

Authors: Yunyi Chen, Chenru Wang, Xinyi Ye, Zexin Zheng, Chi Zhang

Organizations: AGI Lab, Westlake University · Eindhoven University of Technology · Guangdong University of Finance and Economics

Abstract

Diffusion-based dataset distillation (DD) suffers from a fundamental objective mismatch: likelihood-driven diffusion models prioritize density approximation over the discriminative decision boundaries required for downstream tasks. Beyond semantic mismatch, relying solely on density also leads to geometric coverage loss, where generated samples collapse into a few high-density modes and fail to cover the manifold's structural diversity. We propose Manifold-Guided Policy Optimization (MGPO), which reformulates DD as a multi-objective reinforcement learning problem and achieves Dual-Space Alignment via a pixel-space discriminative reward and a latent-space geometric reward guided by a class-wise Minimum Spanning Tree (MST). The discriminative reward enforces class separability, while the MST-based geometric reward encourages generated latents to cover a sparse geometric skeleton of each class, jointly addressing both failure modes. We further provide an idealized analysis that motivates the MST-based reward, including a Hausdorff approximation bound and a subsampling bound independent of the dataset size. The reward-modular design extends to structured tasks such as object detection and segmentation by substituting the frozen task reward model. Extensive experiments show MGPO consistently outperforms existing methods, including a +8.0% mIoU gain on segmentation under low-budget settings.

Figures & tables

Figure 1: t-SNE of latent distributions on ImageWoof. Blue density: real data; orange dots: Minimax Diffusion; red dots: MGPO (ours); red stars with green edges: class-wise MST skeleton. MGPO covers a broader area along the MST skeleton, indicating improved manifold coverage.

Figure 2: Overview of the proposed MGPO framework. (a) Pixel Reward ( $r_{\text{disc}}$ ) is computed from downstream objectives such as classification, detection, and segmentation, encouraging semantic correctness and task consistency. (b) MST Reward ( $r_{\text{geo}}$ ) is constructed in latent space. Real dataset features are encoded by encoder and projected via PCA to obtain a structured latent representation. A kNN graph is built to approximate local neighborhood structure, followed by a Minimum Spanning Tree (MST) to capture the global manifold topology. The generated latent samples are aligned with this manifold structure, promoting both representativeness and structural consistency.

IPC (Ratio)	Model	Random	K-Center	Herding	DiT	Minimax	DM	IDC-1	D 4 M	DDVLCP	ManifoldGD	MGPO (Ours)	Full
10 (0.8%)	ConvNet-6	24.3 $\pm$ 1.1	19.4 $\pm$ 0.9	26.7 $\pm$ 0.5	34.2 $\pm$ 1.1	37.0 $\pm$ 1.0	26.9 $\pm$ 1.2	33.3 $\pm$ 1.1	29.4 $\pm$ 0.9	34.8 $\pm$ 2.4	36.9 $\pm$ 0.6	37.2 $\pm$ 0.8	86.4 $\pm$ 0.2
	ResNetAP-10	29.4 $\pm$ 0.8	22.1 $\pm$ 0.1	32.0 $\pm$ 0.3	34.7 $\pm$ 0.5	39.2 $\pm$ 1.3	30.3 $\pm$ 1.2	39.1 $\pm$ 0.5	33.2 $\pm$ 2.1	39.5 $\pm$ 1.5	38.3 $\pm$ 0.4	40.9 $\pm$ 1.1	87.5 $\pm$ 0.5
	ResNet-18	27.7 $\pm$ 0.9	21.1 $\pm$ 0.4	30.2 $\pm$ 1.2	34.7 $\pm$ 0.4	37.6 $\pm$ 0.9	33.4 $\pm$ 0.7	37.3 $\pm$ 0.2	32.3 $\pm$ 1.2	39.9 $\pm$ 2.6	39.2 $\pm$ 0.7	40.1 $\pm$ 0.3	89.3 $\pm$ 1.2
20 (1.6%)	ConvNet-6	29.1 $\pm$ 0.7	21.5 $\pm$ 0.8	29.5 $\pm$ 0.3	36.1 $\pm$ 0.8	37.6 $\pm$ 0.2	29.9 $\pm$ 1.0	35.5 $\pm$ 0.8	34.0 $\pm$ 2.3	37.9 $\pm$ 1.9	37.7 $\pm$ 0.6	39.6 $\pm$ 1.3	86.4 $\pm$ 0.2
	ResNetAP-10	32.7 $\pm$ 0.4	25.1 $\pm$ 0.7	34.9 $\pm$ 0.1	41.1 $\pm$ 0.8	45.8 $\pm$ 0.5	35.2 $\pm$ 0.6	43.4 $\pm$ 0.3	40.1 $\pm$ 1.6	44.5 $\pm$ 2.2	45.6 $\pm$ 0.5	46.1 $\pm$ 0.3	87.5 $\pm$ 0.5
	ResNet-18	29.7 $\pm$ 0.5	23.6 $\pm$ 0.3	32.2 $\pm$ 0.6	40.5 $\pm$ 0.5	42.5 $\pm$ 0.6	29.8 $\pm$ 1.7	38.6 $\pm$ 0.2	38.4 $\pm$ 1.1	44.5 $\pm$ 2.0	42.4 $\pm$ 0.4	45.9 $\pm$ 0.2	89.3 $\pm$ 1.2

Table 1: Comparison of SOTA dataset distillation methods on ImageWoof under various IPC and backbone settings. The best results are marked as bold, and the second are underlined.

	IPC	Random	DiT	Minimax	D 4 M	DDVLCP	MGPO (Ours)
Nette	10	54.2 $\pm$ 1.6	59.1 $\pm$ 0.7	58.6 $\pm$ 0.4	60.9 $\pm$ 1.7	64.8 $\pm$ 3.6	66.9 $\pm$ 2.2
	20	63.5 $\pm$ 0.5	64.8 $\pm$ 1.2	70.6 $\pm$ 0.3	66.3 $\pm$ 1.3	71.4 $\pm$ 0.5	73.9 $\pm$ 0.7
	50	76.1 $\pm$ 1.1	73.3 $\pm$ 0.9	83.7 $\pm$ 0.4	77.7 $\pm$ 1.1	81.2 $\pm$ 0.8	85.4 $\pm$ 1.1
IDC	10	48.1 $\pm$ 0.8	54.1 $\pm$ 0.4	51.9 $\pm$ 1.4	52.8 $\pm$ 0.5	57.0 $\pm$ 1.4	54.6 $\pm$ 1.7
	20	52.5 $\pm$ 0.9	58.9 $\pm$ 0.2	59.1 $\pm$ 3.7	58.5 $\pm$ 0.4	63.3 $\pm$ 1.2	64.8 $\pm$ 0.4
	50	68.1 $\pm$ 0.7	64.3 $\pm$ 0.6	69.4 $\pm$ 1.4	69.1 $\pm$ 0.8	71.9 $\pm$ 0.4	77.8 $\pm$ 0.9

Table 2: Comparison on ImageNette and ImageIDC under different IPC settings. Results are obtained on ResNet-18. The best results are marked as bold and the second are underlined.

IPC	SRe 2 L	RDED	DiT	Minimax	MGPO (Ours)
10	21.3 $\pm$ 0.6	42.0 $\pm$ 0.1	39.6 $\pm$ 0.4	44.3 $\pm$ 0.5	45.1 $\pm$ 0.6
50	46.8 $\pm$ 0.2	56.5 $\pm$ 0.1	52.9 $\pm$ 0.6	58.6 $\pm$ 0.3	59.1 $\pm$ 0.4

Table 3: Comparison of dataset distillation methods on ImageNet-1K under different IPC settings.

Method	Pascal VOC						MS COCO
	0.5%		1%		2%		0.25%		0.5%		1%
	mAP	AP 50	mAP	AP 50	mAP	AP 50	mAP	AP 50	mAP	AP 50	mAP	AP 50
Random	0.8 $\pm$ 0.2	3.1 $\pm$ 0.4	4.2 $\pm$ 0.5	13.7 $\pm$ 0.6	12.4 $\pm$ 0.4	34.3 $\pm$ 0.5	0.5 $\pm$ 0.1	1.7 $\pm$ 0.3	3.7 $\pm$ 0.2	10.1 $\pm$ 0.3	7.2 $\pm$ 0.8	17.3 $\pm$ 0.9
Uniform	0.9 $\pm$ 0.1	3.4 $\pm$ 0.3	5.7 $\pm$ 0.2	17.7 $\pm$ 0.4	13.8 $\pm$ 0.3	36.2 $\pm$ 0.4	0.8 $\pm$ 0.2	2.4 $\pm$ 0.5	3.4 $\pm$ 0.4	9.5 $\pm$ 0.6	7.4 $\pm$ 0.5	17.6 $\pm$ 0.5
K-Center	0.5 $\pm$ 0.1	2.1 $\pm$ 0.3	3.6 $\pm$ 0.6	12.3 $\pm$ 0.3	10.9 $\pm$ 0.6	29.3 $\pm$ 0.6	0.4 $\pm$ 0.1	1.5 $\pm$ 0.2	3.2 $\pm$ 0.5	9.5 $\pm$ 0.5	6.1 $\pm$ 0.3	15.4 $\pm$ 0.6
Herding	0.6 $\pm$ 0.2	2.4 $\pm$ 0.2	3.5 $\pm$ 0.5	11.9 $\pm$ 0.5	10.4 $\pm$ 0.4	28.7 $\pm$ 0.7	0.5 $\pm$ 0.1	1.8 $\pm$ 0.4	3.5 $\pm$ 0.3	9.7 $\pm$ 0.3	6.7 $\pm$ 0.4	16.3 $\pm$ 0.7

Table 4: Performance comparison with coreset selection methods. Metrics are reported as mAP and AP 50 (%) across different selection ratios.

Table 7

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Backbone	IPC	SRe 2 L ( Yin et al., 2023 )	Minimax ( Gu et al., 2024 )	RDED ( Sun et al., 2024 )	CaO 2 ( Wang et al., 2025 )	MGPO (Ours)
ResNet-18	10	20.2 $\pm$ 0.2	40.1 $\pm$ 1.0	38.5 $\pm$ 2.1	45.6 $\pm$ 1.4	47.3 $\pm$ 0.8
ResNet-18	50	23.3 $\pm$ 0.3	67.0 $\pm$ 1.8	68.5 $\pm$ 0.7	68.9 $\pm$ 1.1	69.7 $\pm$ 1.6
ResNet-50	10	17.3 $\pm$ 1.7	37.3 $\pm$ 1.1	29.9 $\pm$ 2.2	40.1 $\pm$ 0.1	41.3 $\pm$ 1.1
ResNet-50	50	24.8 $\pm$ 0.7	64.3 $\pm$ 0.9	67.8 $\pm$ 0.3	68.2 $\pm$ 1.1	69.0 $\pm$ 0.4
ResNet-101	10	17.7 $\pm$ 0.9	34.2 $\pm$ 1.7	31.3 $\pm$ 1.3	36.5 $\pm$ 1.4	37.8 $\pm$ 1.6
ResNet-101	50	21.2 $\pm$ 0.2	62.7 $\pm$ 1.6	59.1 $\pm$ 0.7	63.1 $\pm$ 1.3	63.5 $\pm$ 1.5

Appendix

Table 7: Performance comparison with soft-label-based baselines on ImageWoof under different IPC and backbone settings. The best results in each setting are highlighted in bold.

Figure 5: Ablation studies on the hyperparameters used in MST construction, including the number of neighbors k in the kNN graph, the maximum number of sampled latent features per class, and the PCA projection dimension.

Figure 6: Comparison between samples generated by Minimax Diffusion (left) and generated by the proposed MGPO method (right) on ImageWoof (Classification).

Figure 7: Comparison between samples generated by Minimax Diffusion (left) and generated by the proposed MGPO method (right) on ImageNette (Classification).

Figure 8: Comparison between samples generated by Minimax Diffusion (left) and generated by the proposed MGPO method (right) on ImageIDC (Classification).

Figure 9: Samples generated by the proposed MGPO (Object Detection).

Figure 10: Samples generated by the proposed MGPO (Segmentation).

Explore similar work

DMGD: Train-Free Dataset Distillation with Semantic-Distribution Matching in Diffusion Models
May 5, 2026Qichao Wang, Yunhong Lu, Hengyuan Cao +2Diffusion-Based Dataset Distillation Dataset Distillation
Zhejiang University · Shanghai Institute for Advanced Study-Zhejiang University · Shanghai Institute for Mathematics and Interdisciplinary Sciences
D3S2: Diffusion-Guided Dataset Distillation for Semantic Segmentation
May 24, 2026Wenjie Zheng, Haoji Hu, Jiali Lu +2Efficient Semantic Segmentation Semantic Segmentation
Zhejiang University, China
Dataset Distillation Based on Saliency-Driven Prototype Alignment
Jul 28, 2026Yawen Zou, Wenqi Cai, Guang Li +3Diffusion-Based Dataset Distillation Dataset Distillation
University of Toyama · Hokkaido University · University of Fukui

May 5, 2026, cs.CVJ/K move · Enter open · S save

DMGD: Train-Free Dataset Distillation with Semantic-Distribution Matching in Diffusion Models

Qichao Wang, Yunhong Lu, Hengyuan Cao, Junyi Zhang+1

Zhejiang University · Shanghai Institute for Advanced Study-Zhejiang University · Shanghai Institute for Mathematics and Interdisciplinary Sciences

Diffusion-Based Dataset Distillation Dataset Distillation Diffusion Models Imagenet

Open paper AskPDF arXiv

Dataset distillation enables efficient training by distilling the information of large-scale datasets into significantly smaller synthetic datasets. Diffusion based paradigms have emerged in recent years, offering novel perspectives for dataset distillation. However, they typically necessitate additional fine-tuning stages, and effective guidance mechanisms remain underexplored. To address these limitations, we rethink diffusion based dataset distillation and propose a Dual Matching Guided Diffusion (DMGD) framework, centered on efficient training-free guidance. We first establish Semantic Matching via conditional likelihood optimization, eliminating the need for auxiliary classifiers. Furthermore, we propose a dynamic guidance mechanism that enhances the diversity of synthetic data while maintaining semantic alignment. Simultaneously, we introduce an optimal transport (OT) based Distribution Matching approach to further align with the target distribution structure. To ensure efficiency, we develop two enhanced strategies for diffusion based framework: Distribution Approximate Matching and Greedy Progressive Matching. These strategies enable effective distribution matching guidance with minimal computational overhead. Experimental results on ImageNet-Woof, ImageNet-Nette, and ImageNet-1K demonstrate that our training-free approach achieves significant improvements, outperforming state-of-the-art (SOTA) methods requiring additional fine-tuning by average accuracy gains of 2.1%, 5.4%, and 2.4%, respectively.

Abstract

Explore similar work

DMGD: Train-Free Dataset Distillation with Semantic-Distribution Matching in Diffusion Models

D3S2: Diffusion-Guided Dataset Distillation for Semantic Segmentation

Dataset Distillation Based on Saliency-Driven Prototype Alignment