Paper deep dive
MGPO: Manifold-Guided Diffusion Alignment for Task-Aware Dataset Distillation
Yunyi Chen, Chenru Wang, Xinyi Ye, Zexin Zheng, Chi Zhang
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.05252v1
- Canonical: https://arxiv.org/abs/2610.05252v1
Trouble viewing inline? Open PDF directly →
Full Text
101,147 characters extracted from source content.
Expand or collapse full text
MGPO: Manifold-Guided Diffusion Alignment for Task-Aware Dataset Distillation Yunyi Chen †thanks: Equal contribution. Affiliation: AGI Lab, Westlake University Eindhoven University of Technology Chenru Wang11footnotemark: 1 Xinyi Ye Zexin Zheng Affiliation: Guangdong University of Finance and Economics Chi Zhang †thanks: Corresponding author. Abstract Diffusion-based dataset distillation (D) 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 D 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. 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. 1 Introduction Dataset distillation (D) (Wang et al., 2018) has emerged as a critical technique in large-scale machine learning, aiming to synthesize a small, representative subset of data from a massive original dataset. Recently, generative models, particularly diffusion models (Ho et al., 2020; Song et al., 2020), have emerged as effective backbones for this task. By leveraging powerful generative priors, diffusion models can effectively capture the complex underlying distribution of the source data to synthesize diverse and high-fidelity images. Unlike traditional gradient-matching (Zhao et al., 2020) or trajectory-matching (Cazenavette et al., 2022) approaches, which often struggle with optimization stability and limited sample diversity, these generative frameworks naturally overcome such bottlenecks. Consequently, diffusion-based methods have established state-of-the-art (SOTA) performance in dataset distillation (Gu et al., 2024; Su et al., 2024; Wang et al., 2025; Chan-Santiago et al., 2025; Zou et al., 2025; Qi et al., 2025), offering superior results within limited computational budgets. Despite recent advances, diffusion-based dataset distillation under extremely low-budget settings faces two distinct challenges. First, a semantic objective mismatch: likelihood-driven diffusion models prioritize pixel-level density approximation over the discriminative features required by downstream tasks (Biroli et al., 2024; Ho et al., 2020; Karczewski et al., 2025). Second, a geometric coverage loss: relying solely on density or naive classifier guidance often leads to mode collapse, failing to capture the complex geometric manifold structure and structural diversity of each class in the original data distribution. Consequently, generated samples may exhibit high perceptual fidelity yet fail to effectively represent the dataset. These challenges show that diffusion-based distillation must transcend perceptual fidelity by resolving both issues at the same time. An effective distilled dataset requires two orthogonal properties: semantic discriminability, which pushes samples toward highly separable regions to better serve downstream tasks; and geometric preservation, which spreads samples along a sparse geometric skeleton of each class to maintain structural diversity and prevent mode collapse (Guo et al., 2024). Following these principles, we propose Manifold-Guided Policy Optimization (MGPO), a framework that reformulates diffusion-based dataset distillation as a multi-objective reinforcement learning problem. Our core contribution is a Dual-Space Alignment strategy that addresses both challenges in their respective domains. We design a pixel-space discriminative reward to enforce class separability, alongside a latent-space geometric reward guided by a class-wise Minimum Spanning Tree (MST) to encourage geometric coverage. Crucially, MGPO utilizes a decoupled group-wise policy optimization strategy to internalize these heterogeneous, non-differentiable rewards directly into the diffusion weights. This RL formulation eliminates the need for expensive inference-time guidance, ensuring highly efficient sampling while harmonizing semantics and geometry. We further provide an idealized analysis showing that, under smoothness and reach assumptions, a Euclidean MST δ-approximates the class manifold in Hausdorff distance with δ2/τδ^2/τ tightness in the backward direction. This analysis motivates the distance term of the geometric reward; the implemented reward is evaluated empirically. Beyond classification, the flexibility of our RL formulation unlocks a new reward-modular paradigm for dataset distillation. While the MST-based geometric prior is inherently task-agnostic, the outer framework can be adapted to complex downstream tasks, such as object detection and segmentation, by substituting the pixel-space reward model. In such scenarios, the same frozen task model provides the reward during fine-tuning and pseudo-annotates the class-conditioned generated images afterwards, without modifying the core generative architecture. Extensive empirical evaluations across multiple benchmarks demonstrate that MGPO consistently outperforms existing SOTA approaches. For instance, under representative low-budget settings, our method achieves 18.3 mIoU on segmentation tasks, surpassing the strongest baseline by +8.0 mIoU. Our contributions are summarized as follows: • We reformulate diffusion-based dataset distillation as a multi-objective reinforcement learning problem, internalizing task-specific rewards directly into the diffusion weights and eliminating inference-time guidance overhead. • We design a dual-space reward with a pixel-level discriminative term and a latent-level MST geometric term, addressing the semantic mismatch and geometric coverage loss. • We provide an idealized analysis motivating the MST-based reward: under smoothness and reach assumptions, (i) a Euclidean MST approximates the class manifold in Hausdorff distance, and (i) a subsampled MST admits a bound independent of the dataset size. • Extensive experiments show MGPO achieves SOTA performance on standard classification benchmarks and extends to other complex structured tasks like object detection and semantic segmentation. Figure 2: Overview of the proposed MGPO framework. (a) Pixel Reward (rdiscr_disc) is computed from downstream objectives such as classification, detection, and segmentation, encouraging semantic correctness and task consistency. (b) MST Reward (rgeor_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. 2 Related Work 2.1 Diffusion Models for Dataset Distillation Recent work has explored diffusion models as powerful generative backbones for D. Minimax (Gu et al., 2024) fine-tunes pre-trained diffusion models under a minimax objective to balance sample representativeness and diversity. D4M (Su et al., 2024) incorporates text-to-image diffusion with prototype learning to improve intra-class consistency. D3HR (Zhao et al., 2025) addresses distribution mismatch via DDIM inversion (Song et al., 2020) and group-wise sampling, aligning generated and real data distributions. CaO2 (Wang et al., 2025) performs inference-time latent optimization to reduce feature level inconsistencies, while DDVLCP (Zou et al., 2025) leverages vision–language prototypes to enhance semantic coherence. Despite their success, diffusion-based approaches are inherently likelihood-centric, concentrating samples in high-density regions, which preserves perceptual fidelity but neglects boundary and hard examples essential for discriminative training. Most related to our work, ManifoldGD (Roy et al., 2026) explores manifold structure for diffusion-based distillation, but relies on inference-time guidance without formal geometric guarantees. In contrast, MGPO internalizes geometric constraints into the diffusion policy via RL, with an idealized Hausdorff analysis motivating the MST-based reward, while simultaneously resolving the semantic objective mismatch through dual-space alignment. 2.2 Reinforcement Learning in Diffusion Models Reinforcement learning has recently been introduced to align diffusion and flow-based generative models with task-level objectives beyond likelihood-based training. A representative approach is Group Relative Policy Optimization (GRPO) (Shao et al., 2024), which stabilizes policy updates by comparing relative rewards within groups of sampled trajectories. Subsequent works extend GRPO to generative settings. Flow-GRPO (Liu et al., 2025) adapts it to diffusion and flow matching models, while Temp-Flow-GRPO (He et al., 2025) introduces time-dependent credit assignment. Mix-GRPO (Li et al., 2026) combines multiple sampling dynamics, and structured variants such as Branch-GRPO (Li et al., 2025) and Dance-GRPO (Xue et al., 2025) improve scalability through branching and group-wise rollouts. However, existing methods focus primarily on generative quality or preference alignment. In contrast, we leverage reinforcement learning to directly optimize downstream discriminative objectives. More related works see Appendix B.1 3 Preliminary 3.1 Dataset Distillation Dataset distillation aims to synthesize a small and compact dataset =(x~i,y~i)i=1NSS=\( x_i, y_i)\_i=1^N_S that preserves the essential training information of a large-scale real dataset =(xi,yi)i=1NT,T=\(x_i,y_i)\_i=1^N_T, where xix_i denotes an input sample (e.g., an image), yiy_i its corresponding label, and NS≪NTN_S N_T. The goal of dataset distillation is to enable models trained on S to achieve performance comparable to those trained on T, while significantly reducing storage and computational costs. 3.2 DDIM Sampling in Latent Space We adopt a latent diffusion formulation, where z0=E(x0)z_0=E(x_0) denotes the VAE-encoded latent, and diffusion is performed directly on the latent variables ztz_t. The DiT backbone parameterizes the noise prediction ϵθ(zt,t) _θ(z_t,t) in this latent space. We model the reverse diffusion dynamics using the generalized sampling process of Denoising Diffusion Implicit Models (DDIM) (Song et al., 2020). The transition from state ztz_t to zt−1z_t-1 is defined as: zt−1=αt−1(zt−1−αtϵθ(zt,t)αt)+1−αt−1−σt2⋅ϵθ(zt,t)+σtϵt,z_t-1= _t-1 ( z_t- 1- _t _θ(z_t,t) _t )+ 1- _t-1- _t^2· _θ(z_t,t)+ _t _t, (1) where ϵt∼(0,I) _t (0,I) is standard Gaussian noise and αt _t follows the standard variance schedule. This non Markovian formulation unifies deterministic and stochastic sampling through the variance parameter σt=η1−αt−11−αt1−αtαt−1 _t=η 1- _t-11- _t 1- _t _t-1, which is parameterized by a scalar η. The hyperparameter η explicitly controls the degree of stochasticity. When η=0η=0, the random noise term vanishes, reducing the sampling process to the strictly deterministic DDIM formulation. Conversely, when η=1η=1, the process recovers the stochastic Markovian transitions of traditional DDPM (Ho et al., 2020). During the reward driven policy optimization phase, we maintain η=1η=1 to encourage stochastic trajectory exploration. Once the generative policy is fully aligned with our dual space objectives, we set η=0η=0 to perform highly efficient and deterministic dataset synthesis. 4 Method To bridge the gap between likelihood driven diffusion generation and discriminative downstream tasks, MGPO internalizes dual space task objectives into the diffusion backbone via RL. 4.1 Dual-Space Reward Design To align diffusion generation with downstream objectives, we design two complementary rewards in different spaces. The discriminative reward encourages semantic separability in pixel space, while the geometric reward encourages coverage of class-level geometric structure in latent space. Discriminative Reward (rdiscr_disc). To ensure the synthesized samples are informative for training downstream models, we employ a pre-trained proxy classifier fϕf_φ. For a generated sample x~0 x_0 with target label c, the reward is defined as the log probability of the ground-truth class: rdisc(x~0,c)=log(exp(fϕ(x~0)c)∑kexp(fϕ(x~0)k)),r_disc( x_0,c)= ( (f_φ( x_0)_c) _k (f_φ( x_0)_k) ), (2) where fϕ(x~0)cf_φ( x_0)_c denotes the predicted logit for class c. This continuous formulation provides a smooth reward landscape and penalizes samples that lie near ambiguous decision boundaries. Manifold Geometric Reward (rgeor_geo). Likelihood-driven diffusion models tend to concentrate samples in high-density regions, which may overlook geometrically informative regions of the class manifold. To encourage geometric coverage of each class, we introduce a geometric reward grounded in a class-wise Minimum Spanning Tree (MST), which provides a sparse geometric skeleton of the class distribution and is more discriminative than density-based regularization. Concretely, let c=zii=1NcZ_c=\z_i\_i=1^N_c denote the VAE-encoded latent representations of class-c samples. We model the class manifold via a weighted graph Gc=(c,Ec)G_c=(Z_c,E_c) with edge weights w(u,v)=‖u−v‖2w(u,v)=\|u-v\|_2, and define the class-wise MST as c=argmin∑(u,v)∈T⊆Gcw(u,v),T_c= _T G_c _(u,v)∈ Tw(u,v), (3) where T spans all nodes and contains no cycles. To ensure scalability, in practice we construct cT_c from a subsampled latent set on a sparse k-nearest-neighbor graph (see Algorithm 2). Given a generated latent z0z_0 corresponding to x~0 x_0, we project it into the normalized latent space via ugen=gc(z0)u_gen=g_c(z_0), and define the geometric reward as rgeo(x~0,c)=−λdistminv∈c‖ugen−v‖2−λdir(1−|cos(gen,e∗)|),e∗=argmin(i,j)∈ℰc‖umid−vi+vj2‖2. gatheredr_geo( x_0,c)=- _dist _v _c\|u_gen-v\|_2- _dir (1-| (v_gen,d_e^*)| ),\\ e^*= _(i,j) _c \|u_mid- v_i+v_j2 \|_2. gathered (4) where cV_c and ℰcE_c are the node and edge sets of cT_c, the first term is the Euclidean distance from the projected terminal latent to its nearest MST node, and λdir _dir controls the strength of the directional alignment term (analysis in Appendix). The transition direction is gen=gc(z0)−gc(zmid)v_gen=g_c(z_0)-g_c(z_mid), where zmidz_mid is a latent sampled from an intermediate diffusion timestep, and e∗d_e^* is the direction of the MST edge e∗e^* whose midpoint is nearest to the projected intermediate latent umid=gc(zmid)u_mid=g_c(z_mid). The first term encourages generated samples to remain close to the manifold skeleton, while the second encourages alignment with its local direction. By incorporating rgeor_geo into policy optimization, the diffusion model internalizes the geometric prior during fine-tuning, removing the need for inference-time guidance. Theoretical Justification. We provide two idealized results that motivate the MST-based geometric reward and its subsampled construction. Both concern a Euclidean MST built on samples from a smooth class manifold and an edge-distance formulation. They do not cover the VAE/PCA representation, the nearest-node distance used in practice, or the direction term, which we evaluate empirically. Let ℳcM_c be a compact, connected C2C^2 class submanifold with reach τc>0 _c>0, and let the MST nodes cover ℳcM_c with radius δ<τcδ< _c. Theorem 5 (Appendix A.2) shows that dH(ℳc,|c|)≤δd_H(M_c,|T_c|)≤δ, with the sharper backward bound supy∈|c|d(y,ℳc)≤δ2/τc _y∈|T_c|d(y,M_c)≤δ^2/ _c; that is, in this idealized setting the MST is a reach-aware geometric approximation of the class manifold. Theorem 10 (Appendix A.3) shows that an MST built on n i.i.d. samples satisfies dH(|(Vc)|,|(Uc)|)=((log(n/η)/n)1/mc)d_H(|T(V_c)|,|T(U_c)|)=O (( (n/η)/n)^1/m_c ) with probability at least 1−η1-η, where mcm_c is the intrinsic dimension; this bound is independent of the full dataset size NcN_c, so the per-class MST construction cost stays constant as datasets scale. These results together motivate our design of the implemented reward. 4.2 MGPO : Manifold-Guided Policy Optimization We formulate diffusion-based dataset distillation as a sequential decision process (see Appendix A.1), where the diffusion model acts as a stochastic policy that gradually transforms Gaussian noise into a synthetic sample. At each timestep, the policy πθ(zt−1∣zt,c) _θ(z_t-1 z_t,c) predicts the reverse diffusion transition conditioned on the class label c. Following the generalized DDIM formulation introduced in Sec. 3.2, we set η=1η=1 during policy optimization to introduce stochasticity into the sampling process. This injected stochasticity is vital as it enables the policy to explore a broader range within latent space, while task-driven rewards guide the trajectories toward regions that are both discriminative and geometrically structured. Policy Optimization. To optimize the generative policy without training an explicit value network, we adopt a relative policy optimization strategy. For each class condition c, we sample a group of G trajectories z0(i)i=1G\z_0^(i)\_i=1^G where G denotes the rollout group size (see Sec. 5.4), and evaluate them using the dual-space rewards defined in Sec. 4.1. The diffusion policy is updated by maximizing a clipped surrogate objective: (θ)=c∼p(c),ϵ[1G∑i=1Gmin(ρiA^i,clip(ρi,1−ϵ,1+ϵ)A^i)−βKL(πθ∥πθref)],J(θ)=E_c p(c),ε [ 1G _i=1^G ( _i A_i,clip( _i,1-ε,1+ε) A_i )-β\,D_KL ( _θ\,\|\, _ _ref ) ], (5) where ρi=πθ(zt−1(i)∣zt(i),c)πθold(zt−1(i)∣zt(i),c) _i= _θ(z^(i)_t-1 z^(i)_t,c) _ _old(z^(i)_t-1 z^(i)_t,c) denotes the importance sampling ratio and A^i A_i is the aggregated advantage of the i-th trajectory. The MST skeleton and the projection function gc(⋅)g_c(·) are pre-computed following Algorithm 2. MGPO training procedure is summarized in Algorithm 1. Algorithm 1 MGPO Training Procedure 0: Diffusion policy πθ _θ, proxy reward model fϕf_φ, class-wise MST skeleton cT_c and projection function gc(⋅)g_c(·) 1: for each training iteration do 2: Sample class label c∼p(c)c p(c) 3: Generate G trajectories τii=1G\ _i\_i=1^G via SDE and ODE sampling 4: for each trajectory τi=zT(i),…,z0(i) _i=\z_T^(i),…,z_0^(i)\ do 5: Compute geometric reward rgeo(i)r_geo^(i) using latent projection ugen(i)=gc(z0(i))u_gen^(i)=g_c(z_0^(i)) 6: Decode image x~0(i)=D(z0(i)) x_0^(i)=D(z_0^(i)) 7: Compute discriminative reward rdisc(i)r_disc^(i) using (Eq. 2) 8: end for 9: Perform group-wise normalization of rewards (Eq. 6) 10: Aggregate advantages across objectives (Eq. 7) 11: Apply batch-level normalization (Eq. 8) 12: Update diffusion policy using the MGPO objective (Eq. 5) 13: end for Multi-Reward Aggregation. A key challenge in dual-space optimization is that semantic and geometric rewards may lie on different numerical scales. Directly combining them can lead to unstable training. Following GDPO (Liu et al., 2026), we adopt a three-stage advantage normalization strategy. First, for each reward type k∈disc,geok∈\disc,geo\, we perform group-wise normalization across the G rollouts: Ai,k=rk(i)−mean(rk(j)j=1G)std(rk(j)j=1G)+ϵ.A_i,k= r_k^(i)-mean(\r_k^(j)\_j=1^G)std(\r_k^(j)\_j=1^G)+ε. (6) This step removes scale differences while preserving the relative ranking of samples within the group. Second, the normalized advantages are combined through a weighted summation: Aisum=wdiscAi,disc+wgeoAi,geo.A_i^sum=w_discA_i,disc+w_geoA_i,geo. (7) where wdiscw_disc and wgeow_geo weight the discriminative and geometric rewards (Sec. 5.4). Finally, we apply batch-level normalization across all samples in the training batch: A^i=Aisum−mean(Ajsum)std(Ajsum)+ϵ. A_i= A_i^sum-mean(\A_j^sum\)std(\A_j^sum\)+ε. (8) This normalization stabilizes policy gradients and prevents numerical explosion when fine-tuning large diffusion backbones. 4.3 Reward-Modular Extension to Structured Tasks While the geometric reward rgeor_geo in Section 4.1 is task-agnostic, structured tasks such as object detection and segmentation require pixel-space rewards that capture spatial, in addition to semantic, correctness. The modular design of MGPO allows substituting rdiscr_disc with such task-specific rewards while keeping the policy optimization, geometric reward, and overall training pipeline unchanged. We illustrate this with two examples; each task still requires a task-appropriate reward design, but the burden is reduced from modifying the framework to specifying a scalar reward signal. Object Detection. To encourage clear spatial boundaries without composite losses, we derive the reward from a pretrained detector. For an image x~0 x_0 conditioned on class c, we apply Non-Maximum Suppression to filter redundant predictions, and define the reward as the maximum remaining confidence score for class c: rdet(x~0,c)=max(c(x~0)),r_det( x_0,c)= (S_c( x_0) ), (9) where c(x~0)S_c( x_0) denotes the set of post-NMS confidence scores. Image Segmentation. Instead of relying on predefined mask annotations, we extract the spatial probability map from a pre-trained segmentation network. To prevent reward hacking via isolated high-confidence pixels, we average the top-K pixel confidences for class c: rseg(x~0,c)=1K∑k=1KTopK(c(x~0))k,r_seg( x_0,c)= 1K _k=1^KTopK (P_c( x_0) )_k, (10) where c(x~0)P_c( x_0) is the flattened probability map. The hyperparameter K enforces a minimum expected pixel area (analysis in Appendix). In both cases, the reward depends only on the prompted class c. Spatial annotations are obtained by class-conditioned image generation followed by pseudo-annotation with the same frozen task model. For detection, we keep the post-NMS predictions of class c (score threshold 0.05, IoU threshold 0.5) as boxes. For segmentation, we keep the pixels whose argmax prediction is c as the mask; the Top-K operation in Eq. (10) only computes the scalar reward. Images without a valid class-c prediction are kept with empty annotations (detection) or an all-background mask (segmentation). The current formulation therefore does not model multi-class co-occurrence or inter-class relations (see Appendix F). This reward-modular design allows MGPO to extend to new structured tasks through reward specification rather than architectural modification, though designing an effective reward for each task remains a non-trivial step (see Appendix C for further design considerations). 5 Experiments 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 (Ratio) Model Random K-Center Herding DiT Minimax DM IDC-1 D4M DDVLCP ManifoldGD MGPO (Ours) Full 10 (0.8%) ConvNet-6 24.3± 1.1 19.4± 0.9 26.7± 0.5 34.2± 1.1 37.0± 1.0 26.9± 1.2 33.3± 1.1 29.4± 0.9 34.8± 2.4 36.9± 0.6 37.2± 0.8 86.4± 0.2 ResNetAP-10 29.4± 0.8 22.1± 0.1 32.0± 0.3 34.7± 0.5 39.2± 1.3 30.3± 1.2 39.1± 0.5 33.2± 2.1 39.5± 1.5 38.3± 0.4 40.9± 1.1 87.5± 0.5 ResNet-18 27.7± 0.9 21.1± 0.4 30.2± 1.2 34.7± 0.4 37.6± 0.9 33.4± 0.7 37.3± 0.2 32.3± 1.2 39.9± 2.6 39.2± 0.7 40.1± 0.3 89.3± 1.2 20 (1.6%) ConvNet-6 29.1± 0.7 21.5± 0.8 29.5± 0.3 36.1± 0.8 37.6± 0.2 29.9± 1.0 35.5± 0.8 34.0± 2.3 37.9± 1.9 37.7± 0.6 39.6± 1.3 86.4± 0.2 ResNetAP-10 32.7± 0.4 25.1± 0.7 34.9± 0.1 41.1± 0.8 45.8± 0.5 35.2± 0.6 43.4± 0.3 40.1± 1.6 44.5± 2.2 45.6± 0.5 46.1± 0.3 87.5± 0.5 ResNet-18 29.7± 0.5 23.6± 0.3 32.2± 0.6 40.5± 0.5 42.5± 0.6 29.8± 1.7 38.6± 0.2 38.4± 1.1 44.5± 2.0 42.4± 0.4 45.9± 0.2 89.3± 1.2 50 (3.8%) ConvNet-6 41.3± 0.6 36.5± 1.0 40.3± 0.7 46.5± 0.8 53.9± 0.6 44.4± 1.0 43.9± 1.2 47.4± 0.9 54.5± 0.6 51.3± 0.2 55.7± 0.4 86.4± 0.2 ResNetAP-10 47.2± 1.3 40.6± 0.4 49.1± 0.7 49.3± 0.2 56.3± 1.0 47.1± 1.1 48.3± 1.0 51.7± 3.2 57.3± 0.5 57.9± 0.6 58.1± 0.2 87.5± 0.5 ResNet-18 47.9± 1.8 39.6± 1.0 48.3± 1.2 50.1± 0.5 57.1± 0.6 46.2± 0.6 48.3± 0.8 53.7± 2.2 58.9± 1.5 58.2± 0.6 59.7± 0.4 89.3± 1.2 70 (5.4%) ConvNet-6 46.3± 0.6 38.6± 0.7 46.2± 0.6 50.1± 1.2 55.7± 0.9 47.5± 0.8 48.9± 0.7 50.5± 0.4 55.8± 1.7 55.3± 0.6 56.4± 1.3 86.4± 0.2 ResNetAP-10 50.8± 0.6 45.9± 1.5 53.4± 1.4 54.3± 0.9 58.3± 0.2 51.7± 0.8 52.8± 1.8 54.7± 1.6 60.6± 0.3 60.6± 0.1 60.8± 0.4 87.5± 0.5 ResNet-18 52.1± 1.0 44.6± 1.1 49.7± 0.8 51.5± 1.0 58.8± 0.7 51.9± 0.8 51.1± 1.7 56.3± 1.8 60.3± 0.3 61.0± 0.6 61.2± 0.9 89.3± 1.2 100 (7.7%) ConvNet-6 52.2± 0.4 45.1± 0.5 54.4± 1.1 53.4± 0.3 61.1± 0.7 55.0± 1.3 53.2± 0.9 57.9± 1.5 62.7± 1.4 60.6± 0.2 61.3± 1.1 86.4± 0.2 ResNetAP-10 59.4± 1.0 54.8± 0.2 61.7± 0.9 58.3± 0.8 64.5± 0.2 56.4± 0.8 56.1± 0.9 59.5± 1.8 65.7± 0.5 65.2± 0.3 65.8± 0.6 87.5± 0.5 ResNet-18 61.5± 1.3 50.4± 0.4 59.3± 0.7 58.9± 1.3 65.7± 0.4 60.2± 1.0 58.3± 1.2 63.8± 1.3 68.3± 0.4 66.3± 0.6 68.4± 0.2 89.3± 1.2 5.1 Experimental Setup Generative Backbone. For classification experiments, we adopt a pretrained DiT-XL/2 (Peebles & Xie, 2022) as the generative backbone and fine-tune it using DiffFit (Xie et al., 2023), all synthetic images are generated at a resolution of 256×256256× 256. For detection and segmentation, we follow the experimental protocol of UniDD (Qi et al., 2025) and employ Stable Diffusion v1.4 (Rombach et al., 2022) as the generative backbone to synthesize labeled images with structured annotations. Implementation Details. For classification, student networks are trained from scratch for 300 epochs using SGD with cosine learning-rate decay. For downstream tasks, we use Faster R-CNN (Ren et al., 2016) with a ResNet-50 backbone for object detection and LR-ASPP with a MobileNetV3 (Howard et al., 2019) backbone for semantic segmentation, with standard pretrained backbone weights. All experiments are conducted on a single NVIDIA A100 GPU (80GB). To ensure statistical reliability, each synthetic dataset is generated three times and each downstream model is trained with three random seeds; we report mean ± standard deviation. More details see Appendix C. Baselines. For classification, we compare against representative optimization-based and diffusion-based dataset distillation methods, including DM (Zhao & Bilen, 2023), IDC (Kim et al., 2022), SRe2L (Yin et al., 2023), RDED (Sun et al., 2024), Minimax (Gu et al., 2024), D4M (Su et al., 2024), DDVLCP (Zou et al., 2025), CaO2 (Wang et al., 2025), and ManifoldGD (Roy et al., 2026). For segmentation and object detection, we primarily compare with UniDD (Qi et al., 2025), a recent universal dataset distillation framework. Since diffusion-based distillation for structured outputs remains relatively underexplored, we additionally include coreset selection baselines (Random, Uniform, K-Center (Sener & Savarese, 2018), and Herding (Welling, 2009)). Evaluation Protocol. For classification, we evaluate ConvNet-6, ResNetAP-10, and ResNet-18 trained from scratch. For segmentation and object detection, we follow standard training recipes with pretrained backbones as described above. Following prior works, we report the better performance between hard-label and soft-label supervision. 5.2 Results on Classification We first evaluate our approach on classification benchmarks, providing comprehensive comparisons with multiple state-of-the-art dataset distillation methods across four benchmark datasets: ImageWoof (Howard, 2019b), ImageNette (Howard, 2019a), ImageIDC (Kim et al., 2022) and ImageNet-1K (Russakovsky et al., 2015). The results are summarized in Tables 1, 2 and 3. Additional comparison experiment results with CaO2 (Wang et al., 2025) are provided in Appendix D, Table 7. 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 Random DiT Minimax D4M DDVLCP MGPO (Ours) Nette 10 54.2± 1.6 59.1± 0.7 58.6± 0.4 60.9± 1.7 64.8± 3.6 66.9± 2.2 20 63.5± 0.5 64.8± 1.2 70.6± 0.3 66.3± 1.3 71.4± 0.5 73.9± 0.7 50 76.1± 1.1 73.3± 0.9 83.7± 0.4 77.7± 1.1 81.2± 0.8 85.4± 1.1 IDC 10 48.1± 0.8 54.1± 0.4 51.9± 1.4 52.8± 0.5 57.0± 1.4 54.6± 1.7 20 52.5± 0.9 58.9± 0.2 59.1± 3.7 58.5± 0.4 63.3± 1.2 64.8± 0.4 50 68.1± 0.7 64.3± 0.6 69.4± 1.4 69.1± 0.8 71.9± 0.4 77.8± 0.9 Table 3: Comparison of dataset distillation methods on ImageNet-1K under different IPC settings. IPC SRe2L RDED DiT Minimax MGPO (Ours) 10 21.3± 0.6 42.0± 0.1 39.6± 0.4 44.3± 0.5 45.1± 0.6 50 46.8± 0.2 56.5± 0.1 52.9± 0.6 58.6± 0.3 59.1± 0.4 5.3 Result on Detection and Segmentation We further investigate whether the distilled datasets preserve spatial structures that generalize to downstream tasks, particularly semantic segmentation and object detection. Object Detection. Table 4 reports object detection results on Pascal VOC 2012 (Everingham et al., ) and MS-COCO (Lin et al., 2015). Under aggressive compression, our method substantially outperforms coreset selection baselines, which indicates that the distilled images retain object-level localization cues rather than merely category-level appearance statistics. Segmentation. Table 6 summarizes segmentation performance on Pascal VOC (Everingham et al., ). Our method achieves consistently higher mIoU across compression ratios, suggesting that pixel-level structural information is preserved during distillation and remains useful for downstream learning. Table 4: Performance comparison with coreset selection methods. Metrics are reported as mAP and AP50(%) across different selection ratios. Method Pascal VOC MS COCO 0.5% 1% 2% 0.25% 0.5% 1% mAP AP50 mAP AP50 mAP AP50 mAP AP50 mAP AP50 mAP AP50 Random 0.8± 0.2 3.1± 0.4 4.2± 0.5 13.7± 0.6 12.4± 0.4 34.3± 0.5 0.5± 0.1 1.7± 0.3 3.7± 0.2 10.1± 0.3 7.2± 0.8 17.3± 0.9 Uniform 0.9± 0.1 3.4± 0.3 5.7± 0.2 17.7± 0.4 13.8± 0.3 36.2± 0.4 0.8± 0.2 2.4± 0.5 3.4± 0.4 9.5± 0.6 7.4± 0.5 17.6± 0.5 K-Center 0.5± 0.1 2.1± 0.3 3.6± 0.6 12.3± 0.3 10.9± 0.6 29.3± 0.6 0.4± 0.1 1.5± 0.2 3.2± 0.5 9.5± 0.5 6.1± 0.3 15.4± 0.6 Herding 0.6± 0.2 2.4± 0.2 3.5± 0.5 11.9± 0.5 10.4± 0.4 28.7± 0.7 0.5± 0.1 1.8± 0.4 3.5± 0.3 9.7± 0.3 6.7± 0.4 16.3± 0.7 UniDD 8.5± 0.4 22.3± 0.6 16.8± 0.5 38.9± 0.7 23.9± 0.5 48.5± 0.6 4.5± 0.3 10.3± 0.4 7.1± 0.4 16.9± 0.3 10.8± 0.4 22.5± 0.5 MGPO(Ours) 13.4± 0.2 26.3± 0.5 21.1± 0.4 41.8± 0.3 26.9± 0.9 53.1± 1.3 5.8± 0.2 13.9± 0.3 8.3± 0.5 19.9± 0.2 11.7± 0.8 25.4± 0.8 Full Data 51.4± 0.8 / 80.3± 0.4 32.6± 0.7 / 51.4± 0.8 Table 5: Semantic Segmentation results. Methods Semantic Segmentation (mIoU) 3.5% 7% 14% Full Random 6.7± 0.4 9.6± 0.5 11.3± 0.4 Uniform 6.5± 0.6 9.3± 0.4 12.1± 0.5 K-Center 6.1± 0.4 8.5± 0.6 10.4± 0.7 32.9± 0.6 Herding 6.2± 0.7 8.7± 0.4 10.6± 0.3 UniDD 10.3± 0.5 12.3± 0.3 14.9± 0.7 MGPO (Ours) 18.3± 0.5 20.4± 0.4 23.1± 0.6 Table 6: Ablation study of rewards. Method ImageWoof ImageNette IPC = 10 IPC = 50 IPC = 10 IPC = 50 DiT 34.7± 0.5 49.3± 0.2 58.9± 1.2 82.9± 0.1 Pixel Reward 37.0 ± 0.3 50.5± 0.5 64.1± 0.9 84.9± 1.1 Latent Reward 36.8± 0.3 50.0± 0.2 65.9± 0.2 84.9± 1.3 Both (Ours) 40.9± 1.1 57.9± 0.2 66.9± 2.2 85.4± 1.1 5.4 Ablation Study We conduct ablation studies to understand how each reward component contributes to both semantic fidelity and spatial consistency. Specifically, we progressively enable the baseline DiT objective, Pixel Reward, MST Reward, and their combination, with results reported in Table 6. In general, when both rewards are applied, best performance is observed across different IPC settings. We further analyze robustness across training hyperparameters (Figure 4) and finetuning design choices (Figure 4). MGPO remains stable across all settings, yielding the best trade-off. Figure 3: Hyperparameter analysis of MGPO. Figure 4: Ablation studies on key design choices of MGPO finetuning process. 6 Conclusion In this paper, we proposed MGPO , a reward driven RL framework that resolves the fundamental objective mismatch in diffusion-based dataset distillation. By integrating pixel level discriminative and latent level geometric rewards, MGPO effectively aligns the generative process with downstream tasks. Furthermore, our plug-and-play reward design enables the framework to extend beyond standard classification to complex downstream tasks like object detection and semantic segmentation. Extensive experiments show that MGPO achieves state of the art performance across multiple classification benchmarks and successfully reaches state of the art on detection and segmentation tasks. AI Use Statement In this work, we used generative AI tools (large language models) for the following tasks with required disclosure: proof assistance, namely helping to derive and write the proofs of the theoretical results in Appendix A.2 and Appendix A.3 (Lemmas 1–11, Theorems 5 and 10, Proposition 8, and the associated corollaries and remarks), and drafting parts of the manuscript, primarily the related-work discussion (Section 2 and Appendix B.1) and the implementation descriptions in Appendix C. We have not used generative AI tools for research ideation, hypothesis refinement, methodology design, implementation of the method, running experiments, data cleaning, or result interpretation; the remaining required disclosure tasks (e.g., translation and qualitative analysis) are not applicable to this work. Note that the diffusion models used to synthesize distilled datasets are the object of study and are not AI assistance in preparing this paper. Additionally, we used generative AI tools to aid and polish the writing (grammar, wording, and clarity) throughout the paper. We have reviewed all AI-assisted work. Every AI-assisted proof was checked step by step by at least two authors, who verified all assumptions, constants, and cited results against the original sources (e.g., Berenfeld et al. (2022), Niyogi et al. (2008)). All AI-drafted or AI-polished text was revised by the authors, and every citation was checked manually against the referenced paper. All numerical results were produced by the authors’ own code and were not generated or edited by AI tools. We take responsibility for the final content of this work, including text, claims, or artifacts produced with the aid of generative AI. Ethics Statement This work does not involve human subjects, user studies, or the collection of new data. All experiments use publicly available benchmarks (ImageNet-1K and its subsets ImageWoof, ImageNette, and ImageIDC; PASCAL VOC 2012; MS-COCO), used under their respective licenses and terms of use; we do not redistribute the original images. Some of these datasets contain images of people collected from the web, and the synthesized distilled datasets may therefore reflect content and biases present in the source data. In addition, because MGPO optimizes the generator with frozen pretrained reward models (classifier, detector, or segmentation network), any biases in these reward models may be internalized and amplified in the distilled data; we discuss this risk and recommend auditing reward models before deployment in Appendix F. We do not foresee direct harmful applications beyond those common to generative models and dataset compression. The authors declare no conflicts of interest or sponsorship that could have influenced this work. Reproducibility Statement We have made the following efforts to ensure reproducibility. The full method is described in Section 4, including the DDIM sampling formulation (Section 3.2), the reward definitions (Section 4.1), and the complete training procedure (Algorithm 1). The MST construction is given in Algorithm 2, with all preprocessing steps and hyperparameters (e.g., k-N size, PCA dimension, number of sampled latents per class) detailed in Appendix C.1; the latent geometric reward and the pixel-space rewards for classification, detection, and segmentation are specified in Appendices C.2 and C.3 (Algorithms 3 and 4). Experimental settings, including backbones, student architectures, training schedules, evaluation protocols, and the number of runs and random seeds, are given in Section 5.1, and hyperparameter sensitivity analyses are reported in Section 5.4 and Appendix D. For the theoretical results, all assumptions are stated explicitly and complete proofs are provided in Appendices A.1, A.2, and A.3. All datasets are publicly available. The source code will be released upon acceptance. References Berenfeld et al. (2022) Clément Berenfeld, John Harvey, Marc Hoffmann, and Krishnan Shankar. Estimating the reach of a manifold via its convexity defect function. Discrete & Computational Geometry, 67(2):403–438, 2022. Biroli et al. (2024) Giulio Biroli, Tony Bonnaire, Valentin De Bortoli, and Marc Mézard. Dynamical regimes of diffusion models. Nature Communications, 15(1):9957, 2024. Cazenavette et al. (2022) George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A Efros, and Jun-Yan Zhu. Dataset distillation by matching training trajectories. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 4750–4759, 2022. Chan-Santiago et al. (2025) Jeffrey A. Chan-Santiago, Praveen Tirupattur, Gaurav Kumar Nayak, Gaowen Liu, and Mubarak Shah. Mgd3: Mode-guided dataset distillation using diffusion models, 2025. URL https://arxiv.org/abs/2505.18963. (5) M. Everingham, L. Van Gool, C. K. I. Williams, J. Winn, and A. Zisserman. The PASCAL Visual Object Classes Challenge 2012 (VOC2012) Results. http://w.pascal-network.org/challenges/VOC/voc2012/workshop/index.html. Gu et al. (2024) Jianyang Gu, Saeed Vahidian, Vyacheslav Kungurtsev, Haonan Wang, Wei Jiang, Yang You, and Yiran Chen. Efficient dataset distillation via minimax diffusion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 15793–15803, 2024. Guo et al. (2024) Yingqing Guo, Hui Yuan, Yukang Yang, Minshuo Chen, and Mengdi Wang. Gradient guidance for diffusion models: An optimization perspective. arXiv preprint arXiv:2404.14743, 2024. He et al. (2025) Xiaoxuan He, Siming Fu, Yuke Zhao, Wanli Li, Jian Yang, Dacheng Yin, Fengyun Rao, and Bo Zhang. Tempflow-grpo: When timing matters for grpo in flow models, 2025. Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020. URL https://arxiv.org/abs/2006.11239. Howard et al. (2019) Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, Quoc V. Le, and Hartwig Adam. Searching for mobilenetv3, 2019. URL https://arxiv.org/abs/1905.02244. Howard (2019a) Jeremy Howard. Imagenette: A smaller subset of 10 easily classified classes from imagenet, March 2019a. URL https://github.com/fastai/imagenette. Howard (2019b) Jeremy Howard. Imagewoof: a subset of 10 classes from imagenet that aren’t so easy to classify, March 2019b. URL https://github.com/fastai/imagenette#imagewoof. Karczewski et al. (2025) Rafał Karczewski, Markus Heinonen, and Vikas Garg. Diffusion models as cartoonists: The curious case of high density regions, 2025. URL https://arxiv.org/abs/2411.01293. Kim et al. (2022) Jang-Hyun Kim, Jinuk Kim, Seong Joon Oh, Sangdoo Yun, Hwanjun Song, Joonhyun Jeong, Jung-Woo Ha, and Hyun Oh Song. Dataset condensation via efficient synthetic-data parameterization, 2022. URL https://arxiv.org/abs/2205.14959. Li et al. (2026) Junzhe Li, Yutao Cui, Tao Huang, Yinping Ma, Chun Fan, Miles Yang, Zhao Zhong, and Liefeng Bo. Mixgrpo: Unlocking flow-based grpo efficiency with mixed ode-sde, 2026. URL https://arxiv.org/abs/2507.21802. Li et al. (2025) Yuming Li, Yikai Wang, Yuying Zhu, Zhongyu Zhao, Ming Lu, Qi She, and Shanghang Zhang. Branchgrpo: Stable and efficient grpo with structured branching in diffusion models, 2025. URL https://arxiv.org/abs/2509.06040. Lin et al. (2015) Tsung-Yi Lin, Michael Maire, Serge Belongie, Lubomir Bourdev, Ross Girshick, James Hays, Pietro Perona, Deva Ramanan, C. Lawrence Zitnick, and Piotr Dollár. Microsoft coco: Common objects in context, 2015. URL https://arxiv.org/abs/1405.0312. Liu et al. (2025) Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di ZHANG, and Wanli Ouyang. Flow-GRPO: Training flow matching models via online RL. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. URL https://openreview.net/forum?id=oCBKGw5HNf. Liu et al. (2026) Shih-Yang Liu, Xin Dong, Ximing Lu, Shizhe Diao, Peter Belcak, Mingjie Liu, Min-Hung Chen, Hongxu Yin, Yu-Chiang Frank Wang, Kwang-Ting Cheng, Yejin Choi, Jan Kautz, and Pavlo Molchanov. Gdpo: Group reward-decoupled normalization policy optimization for multi-reward rl optimization, 2026. URL https://arxiv.org/abs/2601.05242. Niyogi et al. (2008) Partha Niyogi, Stephen Smale, and Shmuel Weinberger. Finding the homology of submanifolds with high confidence from random samples. Discrete & Computational Geometry, 39(1–3):419–441, 2008. Peebles & Xie (2022) William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022. Penrose (2003) Mathew Penrose. Random Geometric Graphs. Oxford University Press, 2003. Pope et al. (2021) Phil Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. In International Conference on Learning Representations (ICLR), 2021. Qi et al. (2025) Ding Qi, Jian Li, Junyao Gao, Shuguang Dou, Ying Tai, Jianlong Hu, Bo Zhao, Yabiao Wang, Chengjie Wang, and Cairong Zhao. Towards universal dataset distillation via task-driven diffusion. In 2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 10557–10566, 2025. doi: 10.1109/CVPR52734.2025.00987. Ren et al. (2016) Shaoqing Ren, Kaiming He, Ross Girshick, and Jian Sun. Faster r-cnn: Towards real-time object detection with region proposal networks, 2016. URL https://arxiv.org/abs/1506.01497. Rombach et al. (2022) Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 10684–10695, June 2022. Roy et al. (2026) Ayush Roy, Wei-Yang Alex Lee, Rudrasis Chakraborty, and Vishnu Suresh Lokhande. Manifoldgd: Training-free hierarchical manifold guidance for diffusion-based dataset distillation. arXiv preprint arXiv:2602.23295, 2026. Russakovsky et al. (2015) Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211–252, 2015. doi: 10.1007/s11263-015-0816-y. Sener & Savarese (2018) Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: A core-set approach, 2018. URL https://arxiv.org/abs/1708.00489. Shao et al. (2024) Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024. URL https://arxiv.org/abs/2402.03300. Song et al. (2020) Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv:2010.02502, October 2020. URL https://arxiv.org/abs/2010.02502. Su et al. (2024) Duo Su, Junjie Hou, Weizhi Gao, Yingjie Tian, and Bowen Tang. Dˆ4: Dataset distillation via disentangled diffusion model. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 5809–5818, 2024. Sun et al. (2024) Peng Sun, Bei Shi, Daiwei Yu, and Tao Lin. On the diversity and realism of distilled dataset: An efficient dataset distillation paradigm. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2024. Wang et al. (2025) Haoxuan Wang, Zhenghao Zhao, Junyi Wu, Yuzhang Shang, Gaowen Liu, and Yan Yan. Cao2: Rectifying inconsistencies in diffusion-based dataset distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), p. 4722–4731, October 2025. Wang et al. (2018) Tongzhou Wang, Jun-Yan Zhu, Antonio Torralba, and Alexei A Efros. Dataset distillation. arXiv preprint arXiv:1811.10959, 2018. Welling (2009) Max Welling. Herding dynamical weights to learn. In Proceedings of the 26th annual international conference on machine learning, p. 1121–1128, 2009. Xie et al. (2023) Enze Xie, Lewei Yao, Han Shi, Zhili Liu, Daquan Zhou, Zhaoqiang Liu, Jiawei Li, and Zhenguo Li. Difffit: Unlocking transferability of large diffusion models via simple parameter-efficient fine-tuning, 2023. URL https://arxiv.org/abs/2304.06648. Xue et al. (2025) Zeyue Xue, Jie Wu, Yu Gao, Fangyuan Kong, Lingting Zhu, Mengzhao Chen, Zhiheng Liu, Wei Liu, Qiushan Guo, Weilin Huang, et al. Dancegrpo: Unleashing grpo on visual generation. arXiv preprint arXiv:2505.07818, 2025. Yin et al. (2023) Zeyuan Yin, Eric Xing, and Zhiqiang Shen. Squeeze, recover and relabel: Dataset condensation at imagenet scale from a new perspective. Advances in Neural Information Processing Systems, 36:73582–73603, 2023. Zhao & Bilen (2023) Bo Zhao and Hakan Bilen. Dataset condensation with distribution matching. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, p. 6514–6523, 2023. Zhao et al. (2020) Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen. Dataset condensation with gradient matching. arXiv preprint arXiv:2006.05929, 2020. Zhao et al. (2025) Lin Zhao, Yushu Wu, Xinru Jiang, Jianyang Gu, Yanzhi Wang, Xiaolin Xu, Pu Zhao, and Xue Lin. Taming diffusion for dataset distillation with high representativeness. arXiv preprint arXiv:2505.18399, 2025. Zou et al. (2025) Yawen Zou, Guang Li, Duo Su, Zi Wang, Jun Yu, and Chao Zhang. Dataset distillation via vision-language category prototype. arXiv preprint arXiv:2506.23580, 2025. Appendix Contents of Appendix Appendix A Theoretical Analysis A.1 Reward-Aligned Diffusion as Policy Optimization Let pθ(x0)p_θ(x_0) denote the data distribution induced by a the current diffusion model, and let pref(x0)p_ref(x_0) denote the distribution of the frozen pretrained model Standard diffusion training maximizes likelihood and therefore prioritizes perceptual fidelity. However, in dataset distillation, our objective is not merely realism, but the generation of samples that are informative for downstream tasks. To incorporate task objectives, we consider a reward-reweighted target distribution pβ(x0)∝pref(x0)exp(βR(x0)),p_β(x_0)\; \;p_ref(x_0) \! (β R(x_0) ), (11) where R(x0)R(x_0) denotes a task-specific reward and β>0β>0 controls alignment strength. This formulation shifts probability mass toward task-relevant regions of the data manifold. We view diffusion sampling process as a trajectory τ=xT,xT−1,…,x0τ=\x_T,x_T-1,…,x_0\ induced by a parameterized policy pθ(τ)p_θ(τ). Under this perspective, aligning diffusion generation with task objectives amounts to optimizing the expected terminal reward over trajectories, regularized towards the pretrained model: β(θ)=τ∼pθ[R(x0)]−1βKL(pθ(x0)∥pref(x0)).J_β(θ)=E_τ p_θ [R(x_0) ]\;-\; 1β\,KL (p_θ(x_0)\,\|\,p_ref(x_0) ). (12) Over all distributions of x0x_0, the maximizer of Eq. (12) is exactly the reward-reweighted density in Eq. (11), since β=1βlogZβ−1βKL(pθ∥pβ)J_β= 1β Z_β- 1βKL(p_θ\,\|\,p_β) with Zβ=pref[exp(βR)]Z_β=E_p_ref[ (β R)]. Policy optimization over the diffusion parameters therefore approximates reward-aligned sampling within the model class, without modifying the inference procedure. This policy view naturally enables the incorporation of heterogeneous reward signals. In our setting, the reward function will consist of (i) a discriminative term defined in pixel space and (i) a geometric term defined in latent space, allowing alignment of both semantic and structural properties. A.2 Theoretical Justification of MST as a Manifold Skeleton Setup. Let ℳ⊂ℝdM ^d be a compact, connected, C2C^2-embedded submanifold with reach τ≔reach(ℳ)>0.τ (M)>0. Let U=u1,…,uN⊂ℳU=\u_1,…,u_N\ be a finite point set on the manifold, and define its covering radius δ≔supx∈ℳmin1≤i≤N‖x−ui‖2.δ _x \, _1≤ i≤ N\|x-u_i\|_2. (13) Let MST=(U,EMST)T_MST=(U,E_MST) denote the Euclidean minimum spanning tree on U, and write |MST|≔⋃(ui,uj)∈EMST[ui,uj]⊂ℝd|T_MST| _(u_i,u_j)∈ E_MST[u_i,u_j]\;⊂\;R^d for its geometric realization, where [ui,uj][u_i,u_j] denotes the closed line segment in ℝdR^d between uiu_i and uju_j. Scope. The analysis in this subsection and in Appendix A.3 is idealized: it concerns a Euclidean MST built on points sampled from a smooth class manifold, with the distance to the tree realization |MST||T_MST| (i.e., to the nearest edge segment). It motivates the geometric reward but does not by itself cover the implemented reward, which uses the nearest-node distance in the VAE/PCA representation, a kkNN-restricted MST, and a direction term. proposition 8 below states what carries over to the implementation. The goal of this subsection is to establish that, under the regime δ<τδ<τ, the MST realization |MST||T_MST| approximates the manifold ℳM in the Hausdorff metric at rate δ in the forward direction and rate δ2/τδ^2/τ in the backward direction. The argument proceeds via four lemmas: (i) a one-sided Hausdorff bound from the manifold to any spanning tree (lemma 1); (i) the existence of a short-edged spanning tree on U (lemma 2); (i) a bottleneck property of the MST (lemma 3); (iv) a chord-to-manifold distance bound based on convexity defect (lemma 4). These ingredients are then combined in theorem 5. Lemma 1 (One-sided Hausdorff bound). For any spanning tree T on the vertex set U, supx∈ℳd(x,||)≤δ. _x \,d(x,|T|)\;≤\;δ. Proof. Fix x∈ℳx . By the definition of δ in eq. 13, there exists ui∈Uu_i∈ U such that ‖x−ui‖2≤δ.\|x-u_i\|_2≤δ. Since uiu_i is a vertex of T, we have ui∈||u_i∈|T|. Therefore d(x,||)≤‖x−ui‖2≤δ.d(x,|T|)\;≤\;\|x-u_i\|_2\;≤\;δ. Taking the supremum over x∈ℳx proves the claim. ∎ Lemma 2 (Existence of a short-edged spanning tree). There exists a spanning tree on U whose every edge has Euclidean length at most 2δ2δ. Proof. For each i∈1,…,Ni∈\1,…,N\, define the closed manifold cap Ai≔B¯(ui,δ)∩ℳ.A_i B(u_i,δ) . By the definition of δ, every x∈ℳx lies within distance δ of some uiu_i, hence ℳ=⋃i=1NAi.M\;=\; _i=1^NA_i. Claim: the intersection graph of Aii=1N\A_i\_i=1^N is connected. Suppose, for contradiction, that the intersection graph is disconnected, so there is a partition 1,…,N=I⊔J,I,J≠∅,\1,…,N\=I J, I,J≠ , satisfying Ai∩Aj=∅for all i∈I,j∈J.A_i∩ A_j= all i∈ I,\ j∈ J. Define X≔⋃i∈IAi,Y≔⋃j∈JAj.X _i∈ IA_i, Y _j∈ JA_j. Each AiA_i is closed in ℳM (as the intersection of a closed Euclidean ball with ℳM), so X and Y are closed in ℳM. They are nonempty, disjoint, and satisfy ℳ=X∪YM=X∪ Y. In particular, X=ℳ∖YX=M Y is also open in ℳM, so X is a nonempty proper clopen subset of ℳM. This contradicts the connectedness of ℳM, establishing the claim. Now construct a graph G on vertex set U by joining uiu_i and uju_j whenever Ai∩Aj≠∅A_i∩ A_j≠ . By the claim, G is connected. For any such edge (ui,uj)(u_i,u_j), choose x∈Ai∩Ajx∈ A_i∩ A_j; the triangle inequality yields ‖ui−uj‖2≤‖ui−x‖2+‖x−uj‖2≤ 2δ.\|u_i-u_j\|_2\;≤\;\|u_i-x\|_2+\|x-u_j\|_2\;≤\;2δ. Thus every edge of G has length at most 2δ2δ, and any spanning tree of G inherits this property. ∎ Lemma 3 (Bottleneck bound on MST edge length). Let ℓMST≔max(ui,uj)∈EMST‖ui−uj‖2. _MST _(u_i,u_j)∈ E_MST\|u_i-u_j\|_2. Then ℓMST≤2δ _MST≤ 2δ. Proof. By Lemma 2, there exists a spanning tree on U whose maximum edge length is at most 2δ2δ. Every Euclidean minimum spanning tree is also a minimum bottleneck spanning tree, i.e., it minimizes the maximum edge length among all spanning trees. Hence ℓMST≤ 2δ.∎ _MST\;≤\;2δ. Lemma 4 (Chord distance to manifold via convexity defect). Let ℳ⊂ℝdM ^d have reach τ>0τ>0. Suppose a,b∈ℳa,b satisfy ‖a−b‖2=ℓ<2τ.\|a-b\|_2= <2τ. Then for every y∈[a,b]y∈[a,b], d(y,ℳ)≤τ−τ2−ℓ24≤ℓ24τ.d(y,M)\;≤\;τ- τ^2- ^24\;≤\; ^24τ. (14) Proof. Let t≔ℓ/2<τt /2<τ. The smallest closed ball in ℝdR^d containing a,b\a,b\ has radius t, namely the ball centered at the midpoint of [a,b][a,b]. Since [a,b]=conva,b,[a,b]=conv\a,b\, every y∈[a,b]y∈[a,b] belongs to the t-local convex hull of ℳM in the sense of (Berenfeld et al., 2022, Definition 4.1). Therefore d(y,ℳ)≤hℳ(t),d(y,M)\;≤\;h_M(t), where hℳh_M denotes the convexity defect function of ℳM. For a set of reach τ, the quarter-circle bound of (Berenfeld et al., 2022, Definition 4.1 and the bound stated immediately thereafter) states hℳ(t)≤τ−τ2−t2,0≤t<τ.h_M(t)\;≤\;τ- τ^2-t^2, 0≤ t<τ. Substituting t=ℓ/2t= /2 yields the first inequality of eq. 14. For the second inequality, observe that for any 0≤s≤τ20≤ s≤τ^2, τ−τ2−s=sτ+τ2−s≤sτ,τ- τ^2-s\;=\; sτ+ τ^2-s\;≤\; sτ, where the equality follows from rationalizing the numerator and the inequality from τ2−s≥0 τ^2-s≥ 0. Setting s=ℓ2/4s= ^2/4 gives τ−τ2−ℓ24≤ℓ24τ,τ- τ^2- ^24\;≤\; ^24τ, which completes the proof. ∎ We are now ready to state and prove the main result of this subsection. Theorem 5 (MST–manifold Hausdorff bound). Suppose δ<τδ<τ. Then dH(ℳ,|MST|)≤δ.d_H\! (M,\ |T_MST| )\;≤\;δ. (15) More precisely, the two one-sided bounds supx∈ℳd(x,|MST|) _x d\! (x,\ |T_MST| ) ≤δ, \;≤\;δ, (16) supy∈|MST|d(y,ℳ) _y∈|T_MST|d (y,\ M ) ≤τ−τ2−δ2≤δ2τ \;≤\;τ- τ^2-δ^2\;≤\; δ^2τ (17) hold simultaneously. Proof. The forward bound eq. 16 is exactly lemma 1. For the backward bound eq. 17, by lemma 3, ℓMST≤ 2δ< 2τ, _MST\;≤\;2δ\;<\;2τ, so lemma 4 applies to every edge [ui,uj]∈EMST[u_i,u_j]∈ E_MST. Each point y on such an edge therefore satisfies d(y,ℳ)≤τ−τ2−‖ui−uj‖224.d(y,M)\;≤\;τ- τ^2- \|u_i-u_j\|_2^24. Since the function ℓ↦τ−τ2−ℓ2/4 τ- τ^2- ^2/4 is monotone non-decreasing on [0,2τ)[0,2τ), and ‖ui−uj‖2≤ℓMST≤2δ\|u_i-u_j\|_2≤ _MST≤ 2δ, we have d(y,ℳ)≤τ−τ2−ℓMST24≤τ−τ2−δ2≤δ2τ,d(y,M)\;≤\;τ- τ^2- _MST^24\;≤\;τ- τ^2-δ^2\;≤\; δ^2τ, where the last inequality follows from lemma 4 applied with s=δ2s=δ^2. Taking the supremum over y∈|MST|y∈|T_MST| yields eq. 17. To deduce the combined Hausdorff bound eq. 15, note that dH(ℳ,|MST|)=maxsupx∈ℳd(x,|MST|),supy∈|MST|d(y,ℳ).d_H\! (M,\ |T_MST| )\;=\; \! \ _x d\! (x,|T_MST| ),\ _y∈|T_MST|d(y,M) \. By eqs. 16 and 17, both terms inside the maximum are at most maxδ,τ−τ2−δ2. \! \δ,\ τ- τ^2-δ^2 \. It remains to show τ−τ2−δ2≤δτ- τ^2-δ^2≤δ whenever δ<τδ<τ. Equivalently, τ−δ≤τ2−δ2τ-δ≤ τ^2-δ^2, which after squaring becomes (τ−δ)2≤τ2−δ2,(τ-δ)^2\;≤\;τ^2-δ^2, i.e., −2τδ+2δ2≤0-2τδ+2δ^2≤ 0, i.e., δ≤τδ≤τ, which holds by assumption. Hence dH(ℳ,|MST|)≤δ.∎d_H\! (M,\ |T_MST| )\;≤\;δ. Remark 6. The bounds in theorem 5 are asymmetric: the forward direction (manifold to MST) decays linearly in δ, while the backward direction (MST to manifold) decays quadratically as δ2/τδ^2/τ. The asymmetry arises because the MST realization includes line segments that may bow off the manifold, but the bowing distance is controlled by the manifold’s reach via lemma 4. This quadratic backward rate is what enables lemma 3 to feed into a tight bound: a generic spanning tree on U would only give a linear (δ)O(δ) bound in both directions. Remark 7 (Geometric approximation, not topology preservation). Theorem 5 is a statement about Hausdorff closeness only. The MST is a one-dimensional acyclic graph, so |MST||T_MST| has trivial first homology and cannot, in general, preserve cycles, higher-dimensional homology, or the intrinsic dimension of ℳM. For example, for points sampled on a circle the MST is a path, which breaks the cycle. We therefore use the MST as a sparse geometric skeleton that encourages coverage of the class support, not as a topology-preserving representation. What carries over to the implemented reward. The proofs of lemmas 1, 2 and 3 use only that ℳM is compact and connected; smoothness and reach enter only through lemma 4, i.e., the backward bound eq. 17. This allows us to state a bound for the nearest-node distance used in practice, directly in the deployed representation. Proposition 8 (Nearest-node distance in the deployed representation). Let K⊂ℝDK ^D be compact and connected, let U=u1,…,uN⊂KU=\u_1,…,u_N\⊂ K have covering radius δK≔supx∈Kmini‖x−ui‖2 _K _x∈ K _i\|x-u_i\|_2, and let T be the Euclidean MST on U. Define the node distance Dnode(y)≔mini‖y−ui‖2D_node(y) _i\|y-u_i\|_2 and the edge distance Dedge(y)≔d(y,||)D_edge(y) d(y,|T|). Then: (i) supx∈KDnode(x)≤δK _x∈ KD_node(x)≤ _K; (i) for every y∈ℝDy ^D, Dedge(y)≤Dnode(y)≤Dedge(y)+δKD_edge(y)\;≤\;D_node(y)\;≤\;D_edge(y)+ _K; (i) if V⊆UV U, then for every y∈ℝDy ^D, 0≤DnodeV(y)−DnodeU(y)≤maxu∈Uminv∈V‖u−v‖2≤δK(V)0≤ D^V_node(y)-D^U_node(y)≤ _u∈ U _v∈ V\|u-v\|_2≤ _K(V), where δK(V) _K(V) is the covering radius of V over K. In particular, with K=gc(c)K=g_c(X_c) the image of a compact, connected class support cX_c under the continuous encoder-and-whitening map gcg_c (Appendix C.1), (i)–(i) hold in the VAE/PCA space where the reward is computed, without requiring gcg_c to be isometric or invertible. Proof. (i) is the definition of δK _K. For (i), the left inequality holds because every node lies on |||T|. For the right inequality, let p∈[ui,uj]p∈[u_i,u_j] be a point of |||T| nearest to y. The nearer endpoint of this edge is within ‖ui−uj‖2/2≤ℓMST/2\|u_i-u_j\|_2/2≤ _MST/2 of p, and the arguments of lemmas 2 and 3, which use only compactness and connectedness of K, give ℓMST≤2δK _MST≤ 2 _K. Hence Dnode(y)≤‖y−p‖2+δK=Dedge(y)+δKD_node(y)≤\|y-p\|_2+ _K=D_edge(y)+ _K. For (i), the left inequality holds since V⊆UV U. For the right one, let u∗∈Uu ∈ U be nearest to y and v∗∈Vv ∈ V nearest to u∗u ; then DnodeV(y)≤‖y−u∗‖2+‖u∗−v∗‖2≤DnodeU(y)+maxu∈Uminv∈V‖u−v‖2D^V_node(y)≤\|y-u \|_2+\|u -v \|_2≤ D^U_node(y)+ _u∈ U _v∈ V\|u-v\|_2, and the last quantity is at most δK(V) _K(V) because U⊂KU⊂ K. Finally, the continuous image of a compact connected set is compact and connected. ∎ Remark 9 (What remains empirical). Proposition 8 does not assume smoothness or positive reach of K, so the backward bound eq. 17 is not claimed in the deployed space. It also does not cover the direction term, which is evaluated empirically, or the kkNN restriction used to build the MST: if the symmetrized kkNN graph contains a Euclidean MST of U, then its MST is a Euclidean MST of U, and otherwise the two trees can differ. Finally, the relation between distances in gc(c)g_c(X_c) and downstream utility is not established theoretically and is assessed through the experiments. A.3 Subsample-MST Approximation and Scalability Setup. We adopt the notation of section A.2. Let ℳc⊂ℝdM_c ^d denote the class-c data manifold, assumed to be a compact, connected, C2C^2-embedded submanifold with reach τc:=reach(ℳc)>0 _c:=reach(M_c)>0 and intrinsic dimension mcm_c. Let μc _c be a probability measure on ℳcM_c with density (with respect to the volume form) bounded below by ρmin>0 _ >0. For any finite point set W⊂ℳcW _c, define its covering radius δ(W):=supx∈ℳcminw∈W‖x−w‖2,δ(W):= _x _c\, _w∈ W\|x-w\|_2, and let (W)=(W,EMST(W))T(W)=(W,E_MST(W)) be the Euclidean minimum spanning tree on W, with geometric realization |(W)|:=⋃(u,v)∈EMST(W)[u,v]⊂ℝd.|T(W)|:= _(u,v)∈ E_MST(W)[u,v]\;⊂\;R^d. The Hausdorff distance between two compact sets A,B⊂ℝdA,B ^d is dH(A,B):=maxsupa∈Ainfb∈B∥a−b∥2,supb∈Binfa∈A∥a−b∥2.d_H(A,B):= \! \ _a∈ A _b∈ B\|a-b\|_2,\ \ _b∈ B _a∈ A\|a-b\|_2 \. Let Uc=u1,…,uNc⊂ℳcU_c=\u_1,…,u_N_c\ _c denote the full real sample set for class c, drawn i.i.d. from μc _c, and let Vc⊆UcV_c U_c with |Vc|=n|V_c|=n be a uniform random subset. By the exchangeability of i.i.d. samples, VcV_c is itself distributed as an i.i.d. sample of size n from μc _c, a characterization we adopt in the analysis below. The aim of this subsection is to show that, with high probability, the subsample MST |(Vc)||T(V_c)| approximates the full MST |(Uc)||T(U_c)| in Hausdorff distance at a rate determined only by the intrinsic dimension mcm_c, justifying the practical use of small random subsamples in the geometric reward computation (Algorithm 1 of the main paper). Theorem 10 (Subsample-MST approximation). Let ℳcM_c, μc _c, UcU_c, and VcV_c be as above. Then there exist constants Cℳc>0C_M_c>0 (depending only on ℳcM_c and ρmin _ ) and n0=n0(ℳc,τc)n_0=n_0(M_c, _c) such that, for every n≥n0n≥ n_0 and every η∈(0,1)η∈(0,1), with probability at least 1−η1-η, δ(Vc)≤Cℳc(log(n/η)n)1/mc<τc,δ(V_c)\;≤\;C_M_c\! ( (n/η)n )^\!1/m_c\;<\; _c, (18) and consequently dH(|(Vc)|,|(Uc)|)≤ 2Cℳc(log(n/η)n)1/mc.d_H\! (|T(V_c)|,\ |T(U_c)| )\;≤\;2\,C_M_c\! ( (n/η)n )^\!1/m_c. (19) Moreover, the geometric reward rgeor_geo defined in Eq. (4) of the main paper satisfies |rgeoVc(x~0,c)−rgeoUc(x~0,c)|≤ 2Cℳc(log(n/η)n)1/mc+λdir, |r_geo^V_c( x_0,c)-r_geo^U_c( x_0,c) |\;≤\;2\,C_M_c\! ( (n/η)n )^\!1/m_c+ _dir, (20) for every generated latent x~0 x_0, where rgeoWr_geo^W denotes the geometric reward computed using the MST on point set W. The first term captures the contribution of the distance component of rgeor_geo and vanishes as n→∞n→∞; the second term is a worst-case bound on the direction component (see remark 12). The proof proceeds in three steps: (i) a probabilistic bound on the covering radius of the random subsample (lemma 11); (i) a deterministic Hausdorff bound between the two MSTs via the triangle inequality, leveraging theorem 5; (i) a Lipschitz argument transferring the geometric bound to the distance component of the geometric reward, combined with a worst-case bound on the direction component. Step 1: Covering radius of a random subsample. We first invoke a classical result on the convergence of empirical covering radius on a manifold. Lemma 11 (Covering radius rate, (Niyogi et al., 2008; Penrose, 2003)). Under the assumptions of theorem 10, there exists a constant Cℳc>0C_M_c>0 depending only on ℳcM_c and ρmin _ such that, for every η∈(0,1)η∈(0,1), with probability at least 1−η1-η, δ(Vc)≤Cℳc(log(n/η)n)1/mc.δ(V_c)\;≤\;C_M_c\! ( (n/η)n )^\!1/m_c. (21) Proof. Cover ℳcM_c by (ρ)N(ρ) geodesic balls of radius ρ, where (ρ)≤c1ρ−mcvol(ℳc)N(ρ)≤ c_1\,ρ^-m_c\,vol(M_c) for some constant c1=c1(ℳc)c_1=c_1(M_c), by standard volume comparison on bounded-reach submanifolds. For ρ<τc/2ρ< _c/2, each such ball has μc _c-mass at least ρminc2ρmc _ \,c_2\,ρ^m_c by the lower density assumption combined with the small-ball volume bound on ℳcM_c. Setting ρ:=(log(n/η)c3n)1/mcρ:=\! ( (n/η)c_3\,n )^\!1/m_c for an appropriate constant c3c_3, and applying a union bound over the cover, the probability that some ball contains no point of VcV_c is at most (ρ)(1−ρminc2ρmc)n≤η.N(ρ)\, (1- _ \,c_2\,ρ^m_c )^n\;≤\;η. On this event, every x∈ℳcx _c lies within ρ of some vi∈Vcv_i∈ V_c, yielding eq. 21 with Cℳc:=c3−1/mcC_M_c:=c_3^-1/m_c. ∎ Choosing n0n_0 large enough that the right-hand side of eq. 21 is strictly less than τc _c, we obtain eq. 18 for all n≥n0n≥ n_0. Step 2: Hausdorff distance between the two MST realizations. Apply theorem 5 to the point set W=UcW=U_c and to W=VcW=V_c. Since Vc⊆UcV_c U_c, the covering radius is monotone non-increasing under enlargement of the point set, i.e., δ(Uc)≤δ(Vc)δ(U_c)≤δ(V_c). Hence, on the event of lemma 11, δ(Uc)≤δ(Vc)<τc,δ(U_c)\;≤\;δ(V_c)\;<\; _c, so theorem 5 applies to both sets and gives dH(ℳc,|(Uc)|)≤δ(Uc),dH(ℳc,|(Vc)|)≤δ(Vc).d_H\! (M_c,\ |T(U_c)| )≤δ(U_c), d_H\! (M_c,\ |T(V_c)| )≤δ(V_c). The triangle inequality for Hausdorff distance yields dH(|(Vc)|,|(Uc)|) d_H\! (|T(V_c)|,\ |T(U_c)| ) ≤dH(|(Vc)|,ℳc)+dH(ℳc,|(Uc)|) \;≤\;d_H\! (|T(V_c)|,\ M_c )+d_H\! (M_c,\ |T(U_c)| ) ≤δ(Vc)+δ(Uc) \;≤\;δ(V_c)+δ(U_c) ≤ 2δ(Vc). \;≤\;2\,δ(V_c). Combined with eq. 18, this establishes eq. 19. Step 3: Stability of the geometric reward. Fix a generated latent x~0 x_0 with class label c, and let ugen=gc(z0)u_gen=g_c(z_0) denote its projected representation as in Sec. 4.1 of the main paper. The geometric reward (Eq. (4) of the main paper) admits the decomposition rgeoW(x~0,c)=−DW−λdirAW,r_geo^W( x_0,c)=-\,D_W\;-\; _dir\,A_W, where DW:=mine∈EMST(W)dist(ugen,e),AW:=1−|cos(gen,eW∗)|,D_W:= _e∈ E_MST(W)dist(u_gen,e), A_W:=1-| (v_gen,v^*_e_W)|, and eW∗∈EMST(W)e_W^*∈ E_MST(W) is the MST edge nearest to ugenu_gen in W. Distance term. Since |(Vc)||T(V_c)| and |(Uc)||T(U_c)| are compact subsets of ℝdR^d, and the map w↦dist(ugen,w)w (u_gen,w) is 11-Lipschitz, |DVc−DUc|=|infw∈|(Vc)|‖ugen−w‖−infw∈|(Uc)|‖ugen−w‖|≤dH(|(Vc)|,|(Uc)|). |D_V_c-D_U_c |\;=\; | _w∈|T(V_c)|\!\|u_gen-w\|- _w∈|T(U_c)|\!\|u_gen-w\| |\;≤\;d_H\! (|T(V_c)|,\ |T(U_c)| ). (22) Direction term. Since |cos(⋅,⋅)|∈[0,1]| (·,·)|∈[0,1], we have the worst-case bound |AVc−AUc|≤ 1. |A_V_c-A_U_c |\;≤\;1. (23) This bound cannot in general be improved through Hausdorff-distance arguments alone: for class manifolds of intrinsic dimension mc≥2m_c≥ 2, two MSTs whose geometric realizations are close as point sets may nonetheless contain edges whose directions span different subspaces of the local tangent plane, so the nearest-edge direction eW∗v^*_e_W is not a Hausdorff-continuous functional of |(W)||T(W)| (see remark 12). Combining the two, |rgeoVc−rgeoUc|≤|DVc−DUc|+λdir|AVc−AUc|≤dH(|(Vc)|,|(Uc)|)+λdir. |r_geo^V_c-r_geo^U_c |\;≤\;|D_V_c-D_U_c|+ _dir\,|A_V_c-A_U_c|\;≤\;d_H\! (|T(V_c)|,\ |T(U_c)| )+ _dir. Substituting eq. 19 yields eq. 20, which completes the proof of theorem 10. ∎ Implemented distance term. For the nearest-node distance used in practice, proposition 8(i) gives, deterministically and in the deployed representation, 0≤DnodeVc(ugen)−DnodeUc(ugen)≤δ(Vc),0\;≤\;D^V_c_node(u_gen)-D^U_c_node(u_gen)\;≤\;δ(V_c), so the distance component of the implemented reward changes by at most λdistδ(Vc) _dist\,δ(V_c) under subsampling. Where the assumptions of lemma 11 hold, δ(Vc)≤Cℳc(log(n/η)/n)1/mcδ(V_c)≤ C_M_c( (n/η)/n)^1/m_c with probability at least 1−η1-η. This bound does not depend on the MST edges, so it is unaffected by the kkNN restriction; the direction term remains uncontrolled, as discussed next. Remark 12 (On the direction component). The bound eq. 20 retains a constant term λdir _dir that does not vanish with δ(Vc)→0δ(V_c)→ 0. This reflects an intrinsic limitation: while the distance component DWD_W depends only on |(W)||T(W)| as a point set and is therefore controlled by Hausdorff distance, the direction component AWA_W depends on the orientation of the nearest MST edge, which is not preserved under Hausdorff convergence on manifolds of intrinsic dimension ≥2≥ 2. In our experiments we set λdir=0.6 _dir=0.6 (chosen via ablation; cf. Fig. 3(b) of the main paper), so the absolute size of the direction contribution remains small relative to the dynamic range of the reward, preserving training stability. A finer analysis of the direction component, e.g. in expectation over ugenu_gen, is left to future work. The stability of the geometric reward guarantees that the excess-risk bound of the main paper carries over with δ replaced by δ(Vc)δ(V_c), up to the additive direction-component constant. Corollary 13 (Excess risk under subsampled MST). Under the assumptions of theorem 10 together with those of the main-paper excess-risk bound, the classifier h^S h_S trained on the MGPO-distilled set generated using the subsample-MST satisfies, with probability at least 1−η−δ01-η- _0, Rμ(h^S)−minh∗∈ℋRμ(h∗) R_μ( h_S)- _h^* R_μ(h^*)\; ≤ 2Lℓ(ε+δ(Vc)2τc+λdir)+4ℜNS(ℓ∘ℋ) ≤\;2\,L_ \! ( + δ(V_c)^2 _c+ _dir )+4\, R_N_S( ) +2LℓCℳNT−1/m+(Blog(1/δ0)NS). +2\,L_ \,C_M\,N_T^-1/m+O\! (B (1/ _0)N_S ). where δ(Vc)≤Cℳc(log(n/η)/n)1/mcδ(V_c)≤ C_M_c ( (n/η)/n )^1/m_c. Proof. Substitute the reward stability bound from eq. 20 into the excess-risk decomposition: the δ-dependent term is replaced by δ(Vc)δ(V_c), and the direction-component constant λdir _dir contributes additively through the Lipschitz constant LℓL_ . The result follows from a union bound over the high-probability events of theorem 10 and the main-paper excess-risk bound. ∎ Corollary 14 (Sample complexity for fixed precision, independent of NcN_c). Fix any target covering radius δtarget∈(0,τc) _target∈(0, _c) and any η∈(0,1)η∈(0,1). Then with probability at least 1−η1-η, n=Θ(log(1/η)+mclog(1/δtarget)δtargetmc)n\;=\; \! ( (1/η)+m_c (1/ _target) _target^m_c ) random samples suffice to ensure δ(Vc)≤δtargetδ(V_c)≤ _target. In particular, the required n depends only on the manifold parameters (τc,mc,ρmin)( _c,m_c, _ ) and the target precision δtarget _target, and is independent of the full dataset size NcN_c. Proof. Set δ(Vc)=δtargetδ(V_c)= _target in eq. 21 and solve for n: δtarget=Cℳc(log(n/η)n)1/mc⟹nlog(n/η)=Θ(δtarget−mc). _target=C_M_c ( (n/η)n )^1/m_c\; \; n (n/η)= ( _target^-m_c ). Inverting gives the stated scaling. None of the constants involved depend on NcN_c. ∎ Remark 15 (Practical implications for ImageNet-scale data). The empirical intrinsic dimension of natural image manifolds, estimated by maximum-likelihood methods, lies in the range mc∈[26,43]m_c∈[26,43] for ImageNet-scale data (Pope et al., 2021), with class-conditional manifolds typically exhibiting smaller intrinsic dimension due to reduced intra-class variation. By corollary 14, the subsample size n required to reach a fixed geometric precision δtarget _target is independent of NcN_c, so a single choice of n generalizes across datasets of different scale. In our implementation, taking n=200n=200 random samples per class gives a ratio of the two covering-radius rates of (Nc/200)1/mc(N_c/200)^1/m_c (ignoring logarithmic factors), which evaluates to ≈1.07≈ 1.07–1.131.13 on ImageNet-1K (Nc≈1300N_c≈ 1300, mc∈[26,43]m_c∈[26,43]). This is a heuristic comparison of upper-bound rates rather than a bound on the realized covering radii, and the intrinsic-dimension estimates refer to image space rather than the VAE/PCA space. It is consistent with, but does not by itself establish, the empirical scalability of MGPO to ImageNet-1K (Table 3). Appendix B Related Work B.1 Optimization-Based Dataset Distillation Dataset distillation was originally introduced as a meta-learning problem that directly optimizes a small set of synthetic examples such that models trained on them achieve comparable performance to training on the full dataset (Wang et al., 2018). A dominant line of follow-up work formulates distillation as bi-level optimization, where the outer loop updates synthetic data and the inner loop trains a network on these synthetic samples(Zhao et al., 2020). Subsequent studies reduce accumulative error in optimization process by matching training trajectory. Later, SRe2L (Yin et al., 2023) series of work further unleashed the ability of decoupling the distillation process by squeeze and recover approach. Besides, an alternative line of works matches higher-level statistics such as distribution or feature in representation space to avoid expensive unrolling, improving efficiency while sometimes sacrificing fidelity of training dynamics. Nevertheless, these optimization-based methods often suffer from heavy compute or memory overhead due to unrolling, sensitivity to architecture choices, and difficulties in scaling to higher resolutions or large numbers of classes. B.2 Diffusion Models for Dataset Distillation Recent advances have explored diffusion models as powerful generative backbones for dataset distillation, enabling the synthesis of compact yet highly representative datasets. As a pioneering work, Minimax Diffusion (Gu et al., 2024) leverages pre-trained diffusion models and further fine-tunes them under a minimax criteria to jointly enhance sample representativeness and diversity. Building upon this idea, D4M (Su et al., 2024) incorporates pre-trained text-to-image diffusion models with prototype learning, clustering latent features to learn category prototypes and thereby improving intra-class consistency and visual fidelity. To mitigate distribution mismatch and stochasticity in latent sampling, D3HR (Zhao et al., 2025) proposes a domain mapping framework that applies DDIM inversion (Song et al., 2020) to project VAE latents into a Gaussian space with improved normality, combined with a group-based sampling strategy to better align the generated data distribution with the real one. More recently, CaO2 (Wang et al., 2025) rectifies feature and class level inconsistencies by inference time latent optimization. DDVLCP (Zou et al., 2025) extends diffusion-based distillation to the vision–language setting by constructing paired image–text prototypes using large language models, which effectively enhances semantic coherence and reduces artifacts such as missing or illogical objects. Despite their success, diffusion-based approaches are inherently likelihood-centric, concentrating samples in high-density regions, which preserves perceptual fidelity but neglects boundary and hard examples essential for discriminative training. B.3 Reinforcement Learning in Diffusion models Reinforcement learning has recently been adopted to align diffusion and flow based generative models with task level objectives that are difficult to encode through likelihood based training alone. A representative line of work builds upon Group Relative Policy Optimization (GRPO) (Shao et al., 2024), which updates generative policies by comparing relative rewards within groups of sampled trajectories, enabling stable optimization without explicit value function estimation. Following the introduction of GRPO, several variants have been proposed to adapt it to different generative settings. Flow-GRPO (Liu et al., 2025) extends group relative optimization to diffusion and flow matching models by introducing controlled stochasticity into the sampling process. Temporal variants such as Temp-Flow-GRPO (He et al., 2025) further incorporate time dependent credit assignment to emphasize critical denoising stages. Mix-GRPO (Li et al., 2026) combines multiple sampling dynamics within a unified optimization framework to balance exploration and efficiency. More recently, structured rollout strategies such as Branch-GRPO (Li et al., 2025) and Dance-GRPO (Xue et al., 2025) have been proposed to improve scalability and stability through branching trajectories and group wise reward aggregation. Despite these advances, existing GRPO-based methods primarily focus on improving generative quality or preference alignment. In contrast, our work leverages reinforcement learning to explicitly incorporate downstream discriminative objectives into the diffusion generation process, addressing the objective mismatch between likelihood-driven diffusion models and dataset distillation for classification. Appendix C Implementation Details Algorithm 2 Per-class MST Construction in Latent Space 0: Real dataset T, VAE encoder E(⋅)E(·), max samples per class NcN_c, PCA dimension d, kNN size k 0: Class-wise MST skeleton cc\T_c\_c and PCA statistics (μc,Wc)( _c,W_c) 1: for each class c do 2: Sample up to NcN_c real images xi\x_i\ from class c 3: Encode latents zi=E(xi)z_i=E(x_i) and flatten to vectors 4: Stack all latent vectors into matrix Z 5: Compute centered features Zc=Z−μcZ_c=Z- _c 6: Perform PCA projection to obtain normalized embeddings ui=Wc(zi−μc)u_i=W_c(z_i- _c) 7: Construct a k-nearest-neighbor graph over ui\u_i\ 8: Compute the minimum spanning tree cT_c from the sparse adjacency graph 9: end for C.1 MST Construction Details The class-wise MST skeleton is pre-computed once before MGPO fine-tuning. For each class c, we sample up to NcN_c real images from the training set. All images are resized to 256×256256× 256, center-cropped, converted to tensors, and normalized to the range [−1,1][-1,1]. Each image is then encoded by a frozen Stable Diffusion VAE encoder. We use the mean of the encoder posterior as the latent representation and divide it by the standard latent scaling factor 0.182150.18215. The resulting latent tensor is flattened into a vector. Stacking all vectors from class c produces a latent matrix Z∈ℝN×D,Z ^N× D, where N is the number of collected samples and D is the flattened latent dimension. Class-wise PCA normalization. To obtain a compact and numerically stable embedding space, we perform PCA on the centered latent matrix. Let μc=1N∑i=1Nzi _c= 1N _i=1^Nz_i be the class mean and define the centered matrix Zc=Z−μc.Z_c=Z- _c. We compute a low-rank PCA decomposition of ZcZ_c and retain the top d principal components. Let V∈ℝD×dV ^D× d denote the projection matrix and S∈ℝdS ^d the singular values. To avoid numerical instability, singular values are lower-bounded by a small constant ε . Each latent vector z is then projected into a normalized embedding space by u=gc(z)=(z−μc)Vdiag(S)−1.u=g_c(z)=(z- _c)Vdiag(S)^-1. This transformation simultaneously performs dimensionality reduction and whitening, making distances more comparable across directions. In our implementation, the PCA dimension is set to d=128d=128. Sparse kkNN graph construction. After obtaining the normalized embeddings uii=1N\u_i\_i=1^N, we construct a sparse k-nearest-neighbor graph in the embedding space. For each node uiu_i, we identify its k nearest neighbors according to Euclidean distance. An adjacency matrix A is then constructed where Aij=‖ui−uj‖2,if uj is among the k nearest neighbors of ui,+∞,otherwise.A_ij= cases\|u_i-u_j\|_2,&if u_j is among the k nearest neighbors of u_i,\\ +∞,&otherwise. cases The graph is symmetrized by assigning the same edge weight to both AijA_ij and AjiA_ji once either direction is observed. This sparse construction significantly reduces graph density and improves scalability. In practice, we use k=10k=10. Minimum spanning tree extraction. Given the sparse adjacency matrix, we compute the minimum spanning tree using a standard graph solver. The resulting tree is represented as an edge set ℰc=(i,j),E_c=\(i,j)\, which connects all nodes while minimizing the total edge weight. Ideally, a spanning tree over N nodes contains N−1N-1 edges. The extracted tree therefore provides a sparse skeleton of the class manifold that preserves global connectivity while avoiding redundant edges. Saved statistics. For each class, we store the PCA mean μc _c, the projection matrix V, the singular values S, the normalized node embeddings, and the MST edge list. These quantities define both the class-wise tree cT_c and the projection function gc(⋅)g_c(·) used during MGPO training. When a generated latent z0z_0 is obtained at the end of the diffusion trajectory, it is projected using the same transformation gc(⋅)g_c(·) and evaluated according to its geometric relation to the nearest edge of the corresponding MST skeleton. C.2 Latent Geometric Reward Details We provide additional implementation details for the latent geometric reward rgeor_geo used in MGPO . The reward is evaluated on a two-stage diffusion trajectory that combines a stochastic SDE transition with a deterministic ODE refinement. This design enables both stable policy optimization and geometry-aware trajectory supervision. Latent projection. To measure geometric consistency with the class manifold, both ztmz_t_m and z0z_0 are projected into the class-specific PCA space used for MST construction. Let gc(⋅)g_c(·) denote the projection function defined in Section C.1. For a flattened latent vector vec(z)vec(z), the projection is u=gc(vec(z)).u=g_c(vec(z)). Thus we obtain utm=gc(vec(ztm)),u0=gc(vec(z0)).u_t_m=g_c(vec(z_t_m)), u_0=g_c(vec(z_0)). Let c=(c,ℰc)T_c=(V_c,E_c) denote the class-wise MST skeleton constructed from real data. Distance-to-manifold term. The first component of the latent reward encourages the terminal latent to stay close to the class manifold represented by the MST nodes. Let vi\v_i\ denote the nodes in cV_c. The distance penalty is defined as ℒdist=mini‖u0−vi‖2.L_dist= _i\|u_0-v_i\|_2. This term measures how far the generated latent deviates from the learned manifold skeleton. Direction alignment term. To additionally constrain the diffusion trajectory, we measure whether the local latent transition direction aligns with the direction of nearby MST edges. The sample transition direction in the projected space is dsample=u0−utm‖u0−utm‖2.d_sample= u_0-u_t_m\|u_0-u_t_m\|_2. For each MST edge (i,j)∈ℰc(i,j) _c, we define mij=vi+vj2,dij=vj−vi‖vj−vi‖2.m_ij= v_i+v_j2, d_ij= v_j-v_i\|v_j-v_i\|_2. We first find the nearest edge midpoint to the intermediate latent (i∗,j∗)=argmin(i,j)∈ℰc‖utm−mij‖2.(i ,j )= _(i,j) _c\|u_t_m-m_ij\|_2. The direction penalty is then defined as ℒdir=1−⟨dsample,di∗j∗⟩.L_dir=1- d_sample,d_i j . This term becomes small when the local diffusion update follows the tangent direction of the class manifold. Final latent reward. The geometric reward is defined as the negative weighted sum of the two penalties: rgeo=−λdistℒdist−λdirℒdir.r_geo=- _distL_dist- _dirL_dir. Therefore, higher reward values correspond to smaller manifold deviation and better alignment with the manifold geometry. C.3 Pixel Classifier Reward Details We provide additional implementation details for the pixel-space discriminative reward rdiscr_disc used in MGPO. The pixel reward evaluates the semantic consistency of generated samples in pixel space and is computed from the decoded image x~0=(z0) x_0=D(z_0) corresponding to the terminal latent produced by the diffusion trajectory. Classification. For standard image classification, the discriminative reward is derived from a pretrained proxy classifier fϕf_φ. Given a generated image x~0 x_0 and target class c, the reward is defined as the log-probability of the ground-truth class: rdisc(x~0,c)=log(exp(fϕ(x~0)c)∑kexp(fϕ(x~0)k)),r_disc( x_0,c)= ( (f_φ( x_0)_c) _k (f_φ( x_0)_k) ), (24) which provides a smooth reward landscape and penalizes samples near ambiguous decision boundaries. Object Detection. For object detection, the discriminative reward is derived from a pretrained Faster R-CNN (Ren et al., 2016) with a ResNet-50 FPN backbone. Given a generated image x~0 x_0 conditioned on class c, the reward is the maximum post-NMS confidence score for class c, as defined in eq. 9 of the main paper. The full computation is described in Algorithm 3. Algorithm 3 Detection Reward Computation 0: Generated image x~0 x_0, target class c, pretrained detector fϕf_φ (Faster R-CNN) 0: Scalar reward rdetr_det 1: proposals←fϕ.region_proposals(x~0)proposals← f_φ.region\_proposals( x_0) 2: detections←NMS(proposals,τiou=0.5,τscore=0.05)detections (proposals,\ _iou=0.5,\ _score=0.05) 3: c←score(d)∣d∈detections,class(d)=cS_c←\score(d) d ,\ class(d)=c\ 4: if c=∅S_c= then 5: rdet←−1.0r_det←-1.0 ⊳ no valid detection for class c 6: else 7: rdet←max(c)r_det← (S_c) ⊳ eq. 9 of main paper 8: end if 9: return rdetr_det Semantic Segmentation. For semantic segmentation, the discriminative reward is derived from a pretrained LR-ASPP (Howard et al., 2019) with a MobileNetV3-Large backbone. Given a generated image x~0 x_0 conditioned on class c, the reward is the mean confidence over the top-K activated pixels of the spatial probability map for class c, as defined in eq. 10 of the main paper. The TopK aggregation prevents reward hacking by enforcing a minimum expected object area, discouraging the policy from exploiting isolated high-confidence pixels. The full computation is described in Algorithm 4. Algorithm 4 Segmentation Reward Computation 0: Generated image x~0 x_0, target class c, pretrained segmentation network fϕf_φ (LR-ASPP), TopK fraction α 0: Scalar reward rsegr_seg 1: c←fϕ(x~0)[c]∈ℝH×WP_c← f_φ( x_0)[c] ^H× W ⊳ softmax probability map for class c 2: ←flatten(c)∈ℝH⋅Wp (P_c) ^H· W 3: K←⌊α⋅H⋅W⌋K← α· H· W ⊳ number of top pixels 4: if max()<ϵ (p)<ε then 5: rseg←−1.0r_seg←-1.0 ⊳ no activation for class c 6: else 7: rseg←1K∑k=1KTopK()kr_seg← 1K _k=1^KTopK(p)_k ⊳ eq. 10 of main paper 8: end if 9: return rsegr_seg Appendix D Additional Experiments D.1 Experiment Results on Different Student Model 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. Backbone IPC SRe2L (Yin et al., 2023) Minimax (Gu et al., 2024) RDED (Sun et al., 2024) CaO2 (Wang et al., 2025) MGPO (Ours) ResNet-18 10 20.2± 0.2 40.1± 1.0 38.5± 2.1 45.6± 1.4 47.3± 0.8 50 23.3± 0.3 67.0± 1.8 68.5± 0.7 68.9± 1.1 69.7± 1.6 ResNet-50 10 17.3± 1.7 37.3± 1.1 29.9± 2.2 40.1± 0.1 41.3± 1.1 50 24.8± 0.7 64.3± 0.9 67.8± 0.3 68.2± 1.1 69.0± 0.4 ResNet-101 10 17.7± 0.9 34.2± 1.7 31.3± 1.3 36.5± 1.4 37.8± 1.6 50 21.2± 0.2 62.7± 1.6 59.1± 0.7 63.1± 1.3 63.5± 1.5 As shown in Table 7, we report the performance of different distillation methods under the soft-label based training protocol. Experiments are conducted on ImageWoof across multiple backbone architectures, including ResNet-18, ResNet-50, and ResNet-101, under various IPC settings. Across all settings, MGPO consistently achieves the best performance. Notably, the performance gains remain stable across different backbone capacities, indicating that the distilled datasets produced by MGPO generalize well to diverse student architectures. In particular, under low IPC setting (IPC = 10), MGPO boosts ResNet-18 performance from 45.6% (CaO2) to 47.3%, and ResNet-50 accuracy from 40.1% to 41.3%. Under higher IPC (IPC = 50), MGPO continues to outperform strong baselines such as Minimax (Gu et al., 2024), RDED (Sun et al., 2024), and CaO2 (Wang et al., 2025). This suggests that our reward-guided distillation framework is able to capture class-discriminative information that transfers effectively across models of varying depth and capacity. D.2 Hyperparameter Analysis We further study the sensitivity of our method to several key hyperparameters. All experiments are conducted on ImageWoof with 10 IPC distilled datasets. Specifically, we analyze the parameters involved in the construction of the MST that models the latent geometric structure. By varying these parameters while keeping other settings fixed, we evaluate the robustness of our framework and examine how they influence the quality of the generated synthetic data and the downstream performance (Figure 5). 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. Overall, MGPO exhibits stable performance across a broad range of hyperparameter settings. The results suggest that the method is not overly sensitive to individual parameters, and moderate values generally provide the best balance between optimization stability and reward effectiveness. Appendix E Visualization E.1 Results on Classification Figure 6, 7, and 8 demonstrate the visualization results on ImageWoof, ImageNette and ImageIDC, respectively. 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). E.2 Results on Detection Figure 9 demonstrates the visualization results on detection results. Figure 9: Samples generated by the proposed MGPO (Object Detection). E.3 Results on Segmentation Figure 10 demonstrates the visualization results on segmentation results. Figure 10: Samples generated by the proposed MGPO (Segmentation). Appendix F Limitations and Broader Impacts Limitations. While MGPO establishes an efficient and adaptable paradigm for task aware dataset distillation, it still possesses certain limitations. First, from a theoretical perspective, as discussed in remark 12, the direction component of the geometric reward relies on local tangent approximations which may exhibit fluctuations on high intrinsic dimension manifolds. Second, although MGPO eliminates the computational bottleneck of inference time guidance, the reinforcement learning policy optimization phase introduces additional offline training overhead compared to purely training free selection methods. Finally, scaling the class wise MST construction to datasets with millions of categories might require more advanced hierarchical clustering algorithms to maintain strict scalability. Broader Societal Impacts. Our research presents significant positive societal impacts by heavily compressing the data required to train deep learning models. This inherently reduces the carbon footprint and energy consumption associated with large scale model training, democratizing artificial intelligence by allowing researchers with limited computational resources to train state of the art models. However, as a potential negative impact, if the pre-trained reward models (such as the proxy classifier or detector) contain inherent biases regarding demographic attributes, the policy optimization process might internalize and amplify these biases within the distilled dataset. Therefore, practitioners must ensure that the reward models used for distillation are thoroughly audited for fairness and equity prior to deployment in sensitive applications.