PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.05252v1cs.CVcs.AI

TL;DR

Diffusion-based dataset distillation can prioritize density over downstream discriminative needs and fail to cover class-manifold diversity. MGPO optimizes pixel-space and latent-space rewards through reinforcement learning; it reports state-of-the-art results across tasks, including a low-budget segmentation gain of +8.0 mIoU over the strongest baseline.

  • Problem

    Likelihood-driven diffusion prioritizes density approximation over task-discriminative features, while density-based or naive classifier guidance can miss class-manifold structure and diversity.

    Source: §1 · p. 2

  • Method

    MGPO reformulates dataset distillation as multi-objective reinforcement learning, combining a pixel-space discriminative reward with a latent-space geometric reward guided by a class-wise MST.

    Source: §1 · p. 2

  • Results

    MGPO reports state-of-the-art performance across classification, detection, and segmentation, including 18.3 mIoU on low-budget segmentation, +8.0 mIoU over the strongest baseline.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The reward-modular design extends MGPO beyond classification to object detection and semantic segmentation.

    Source: §6 · p. 9

    Policy optimization adds offline training overhead compared with purely training-free selection methods.

    Source: F LIMITATIONS AND BROADER IMPACTS · p. 27

Abstract

from arXiv · show

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.

1 INTRODUCTION

Diffusion-based dataset distillation can yield perceptually faithful samples that miss downstream discriminative needs and class-manifold diversity. MGPO addresses both with rewards optimized through reinforcement learning.

1 INTRODUCTION

  • Likelihood-driven diffusion prioritizes density approximation over task-discriminative features, while density-based or naive classifier guidance can miss manifold structure and class diversity.Generated samples may remain perceptually faithful yet fail to represent the dataset effectively.Source: §1 · p. 2
  • MGPO reformulates distillation as multi-objective reinforcement learning, combining a pixel-space discriminative reward with a latent-space, class-wise MST geometric reward.The rewards are internalized into diffusion weights, eliminating inference-time guidance overhead.Source: §1 · p. 2
  • Under smoothness and reach assumptions, the analysis gives a Hausdorff approximation bound for a Euclidean MST and a subsampling bound independent of dataset size.The Hausdorff analysis motivates the geometric reward's distance term.Source: §1 · p. 2
  • For detection and segmentation, MGPO can substitute a task-specific pixel-space reward while retaining its geometric prior and core generative architecture.The same frozen task model supplies the fine-tuning reward and pseudo-annotations for generated images.Source: §1 · p. 2
  • MGPO reports state-of-the-art performance across benchmarks; under representative low-budget segmentation settings, it achieves 18.3 mIoU, surpassing the strongest baseline by +8.0 mIoU.The framework is also reported to consistently outperform existing state-of-the-art approaches.Source: §1 · p. 2
  • The MST reward aligns generated latents with a manifold representation built by encoding real features, projecting with PCA, and using a kNN graph followed by an MST.The figure describes this alignment as promoting representativeness and structural consistency.Source: Figure 2 · p. 3

2 RELATED WORK

Prior work applies diffusion to dataset distillation and reinforcement learning to generative alignment, but the cited approaches leave distinct gaps that MGPO targets.

2.1 DIFFUSION MODELS FOR DATASET DISTILLATION

  • Prior diffusion distillation methods address representativeness, intra-class consistency, distribution mismatch, or semantic coherence, but remain likelihood-centric; ManifoldGD uses inference-time guidance without formal geometric guarantees.MGPO instead internalizes geometric constraints through reinforcement learning and motivates its MST reward with Hausdorff analysis.Source: §2.1 · p. 3

2.2 REINFORCEMENT LEARNING IN DIFFUSION MODELS

  • Existing reinforcement-learning methods for diffusion focus mainly on generative quality or preference alignment, whereas MGPO uses reinforcement learning to optimize downstream discriminative objectives directly.The related work names GRPO and subsequent diffusion or flow-based variants as examples.Source: §2.2 · p. 3

3 PRELIMINARY

Dataset distillation seeks a compact synthetic set that preserves training information. The paper uses latent-space DDIM sampling, with stochasticity controlled by a schedule parameter.

3.1 DATASET DISTILLATION

  • Dataset distillation synthesizes a small set whose training performance can be comparable to training on the large real dataset, while reducing storage and computational costs.The notation distinguishes synthetic samples and labels from the corresponding real dataset.Source: §3.1 · p. 4

3.2 DDIM SAMPLING IN LATENT SPACE

  • The method performs diffusion on VAE-encoded latents using DDIM; η controls stochasticity, with η = 1 for reward-driven exploration and η = 0 for efficient deterministic synthesis after alignment.The latter setting removes the random noise term and yields deterministic DDIM sampling.Source: §3.2 · p. 4

4 METHOD

MGPO reformulates diffusion-based dataset distillation as reinforcement learning that internalizes task objectives in the diffusion model. It combines pixel-space discriminative and latent-space geometric rewards, then adapts the reward design for structured tasks.

4 METHOD

  • MGPO internalizes dual-space task objectives in the diffusion backbone through reinforcement learning.The framework is framed as a way to bridge likelihood-driven generation and discriminative downstream tasks.Source: §4 · p. 4, §4.1 · p. 4

4.1 DUAL-SPACE REWARD DESIGN

  • A proxy classifier rewards target-class log probability, while a latent class-wise MST reward encourages manifold proximity and local-direction alignment.The MST is built from class-wise VAE latents, with a subsampled construction on a sparse k-nearest-neighbor graph for scalability.Source: §4.1 · pp. 4–5
  • Under smoothness and reach assumptions, the Euclidean MST approximates the class manifold in Hausdorff distance, and the subsampled construction has a dataset-size-independent bound.The analysis motivates the reward; it does not cover the practical VAE/PCA representation, nearest-node distance, or direction term.Source: §4.1 · p. 5

4.2 MGPO : MANIFOLD-GUIDED POLICY OPTIMIZATION

  • Group-wise reward normalization addresses scale differences before weighted advantage combination, while stochastic sampling supports latent-space exploration during policy optimization.Batch-level normalization is also applied to stabilize gradients and prevent numerical explosion when fine-tuning large diffusion backbones.Source: §4.2 · pp. 5–6

4.3 REWARD-MODULAR EXTENSION TO STRUCTURED TASKS

  • Detection uses maximum post-NMS class confidence, while segmentation averages top-K class-pixel confidences as rewards for generated images.The same frozen task model supplies rewards and pseudo-annotations; the current formulation does not model multi-class co-occurrence or inter-class relations.Source: §4.3 · pp. 6–7

5 EXPERIMENTS

The experiments compare MGPO with dataset-distillation and coreset baselines across classification and structured tasks. Detection, segmentation, and reward ablations report advantages for MGPO or its combined rewards.

5.1 EXPERIMENTAL SETUP

  • Classification uses DiT-XL/2, while detection and segmentation use Stable Diffusion v1.4 with Faster R-CNN and LR-ASPP downstream models.Baselines include optimization- and diffusion-based methods for classification and UniDD and coreset selection for structured tasks; evaluation reports the better hard- or soft-label result. Each synthetic dataset is generated three times and each downstream model trained with three seeds.Source: §5.1 · pp. 7–8

5.2 RESULTS ON CLASSIFICATION

  • Classification comparisons cover ImageWoof, ImageNette, ImageIDC, and ImageNet-1K, with results reported in Tables 1–3.The tables compare dataset-distillation methods across IPC settings, with Table 1 also varying backbone settings.Source: Table 1 · p. 7, §5.2 · p. 8, Table 2 · p. 8

5.3 RESULT ON DETECTION AND SEGMENTATION

  • Under aggressive compression, MGPO substantially outperforms coreset selection in object detection and achieves consistently higher mIoU across segmentation compression ratios.The reported comparisons cover Pascal VOC 2012 and MS-COCO for detection, and Pascal VOC for segmentation.Source: §5.3 · p. 8, Table 4 · p. 9, Table 5 · p. 9

5.4 ABLATION STUDY

  • Applying both Pixel Reward and MST Reward yields the best performance across the tested IPC settings.The ablation progressively compares the baseline DiT objective, each reward, and their combination.Source: Table 6 · p. 9, §5.4 · p. 9
  • MGPO remains stable across the analyzed training hyperparameters and finetuning design choices, with the authors reporting the best trade-off across settings.Figure 3 covers hyperparameter analysis; the prose also describes Figure 4 as analyzing finetuning choices.Source: §5.4 · p. 9, Figure 3 · p. 9

6 CONCLUSION

The paper concludes that MGPO aligns diffusion-based dataset distillation with downstream tasks using pixel-space and latent-space rewards, with reported state-of-the-art results across classification and structured tasks. Its appendix frames this alignment as reward-guided policy optimization of diffusion sampling.

6 CONCLUSION

  • MGPO combines pixel-level discriminative and latent-level geometric rewards, and the paper reports state-of-the-art performance across classification, detection, and segmentation.The conclusion describes the reward design as plug-and-play for structured downstream tasks.Source: §6 · p. 9
  • Frozen pretrained reward models may have biases that are internalized and amplified in distilled data; the authors recommend auditing them before deployment.Source: ETHICS STATEMENT · p. 10
  • The paper provides method, MST, reward, experiment, and proof details for reproducibility, and states that source code will be released upon acceptance.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 REWARD-ALIGNED DIFFUSION AS POLICY OPTIMIZATION

  • The appendix models diffusion sampling as policy optimization over trajectories, regularized toward the pretrained model to approximate reward-aligned sampling without changing inference.The reward-reweighted target distribution shifts probability mass toward task-relevant regions of the data manifold.Source: §A.1 · pp. 14–15
  • The policy view accommodates heterogeneous rewards: a pixel-space discriminative term and a latent-space geometric term align semantic and structural properties.Source: §A.1 · p. 15

A.2 THEORETICAL JUSTIFICATION OF MST AS A MANIFOLD SKELETON

This section analyzes the MST as a geometric, not topology-preserving, manifold skeleton under idealized assumptions and presents subsampling guarantees for that analysis. It clarifies that the subsampling comparison concerns covering-radius upper-bound rates, not realized radii, and does not establish empirical scalability to ImageNet-1K.

A.2 THEORETICAL JUSTIFICATION OF MST AS A MANIFOLD SKELETON

  • When δ < τ, the idealized Euclidean MST approximates the manifold in Hausdorff distance at rate δ forward and δ^2/τ backward.The analysis concerns distance to the tree’s edge segments, not the implemented nearest-node, kNN-restricted reward with its direction term.Source: §A.2 · pp. 15, 17, 18
  • The Hausdorff bound establishes geometric closeness, not topology preservation: an MST can break cycles and cannot generally preserve higher-dimensional homology or manifold dimension.The paper therefore uses the MST as a sparse coverage skeleton rather than a topology-preserving representation.Source: §A.2 · p. 18
  • For a compact, connected class-support image in deployed representation, nearest-node distance bounds hold without requiring the encoder-and-whitening map to be isometric or invertible.These guarantees do not establish that representation distances predict downstream utility.Source: §A.2 · p. 18
  • The deployed-space guarantees do not cover the direction term or necessarily the kNN-restricted MST, and the link between representation distances and downstream utility remains empirical.The restricted graph yields the Euclidean MST only when it contains one; otherwise the trees can differ.Source: §A.2 · p. 18

A.3 SUBSAMPLE-MST APPROXIMATION AND SCALABILITY

  • With high probability, a random subsample MST approximates the full-sample MST at a rate governed by intrinsic dimension, and fixed-precision sample size is independent of full dataset size Nc.The required sample size depends on manifold parameters and target precision, not Nc.Source: §A.3 · pp. 19, 21
  • Under subsampling, the implemented distance-reward change is bounded by λdist δ(Vc), but Hausdorff convergence does not control the nearest-edge direction term on manifolds of intrinsic dimension at least two.The direction contribution retains a constant worst-case bound; experiments set λdir = 0.6, while finer analysis is left to future work.Source: §A.3 · pp. 20–21
  • ≈1.07–1.13: the estimated ratio of covering-radius rates for 200 samples per class on ImageNet-1K (Nc≈1300, mc∈[26, 43]).This compares upper-bound rates while ignoring logarithmic factors; the intrinsic-dimension estimates are for image space, not VAE/PCA space. The heuristic is consistent with, but does not establish, MGPO’s empirical scalability to ImageNet-1K.Source: §A.3 · p. 22

B.1 OPTIMIZATION-BASED DATASET DISTILLATION

The related-work section traces dataset distillation from optimization-based methods to diffusion backbones, then distinguishes MGPO’s use of reinforcement learning for downstream discriminative objectives.

B.1 OPTIMIZATION-BASED DATASET DISTILLATION

  • Optimization-based distillation uses synthetic-data optimization, trajectory matching, or decoupled squeeze-and-recover approaches, but can incur compute, memory, architecture-sensitivity, and scaling costs.Some approaches instead match distributional or representation statistics to reduce expensive unrolling, sometimes at the cost of training-dynamics fidelity.Source: §B.1 · p. 22

B.2 DIFFUSION MODELS FOR DATASET DISTILLATION

  • Diffusion distillation methods use minimax objectives, prototypes, DDIM inversion, latent optimization, or image-text prototypes, yet remain likelihood-centric and can neglect discriminative boundary and hard examples.The passage describes Minimax, D4M, D3HR, CaO2, and DDVLCP as examples of these approaches.Source: §B.2 · p. 22

B.3 REINFORCEMENT LEARNING IN DIFFUSION MODELS

  • GRPO compares rewards within sampled-trajectory groups, while later variants adapt it to diffusion, time-dependent credit assignment, or mixed sampling dynamics.The passage also cites branching and group-wise reward aggregation as efforts to improve scalability and stability.Source: §B.3 · pp. 22–23
  • MGPO differs from prior GRPO-based work by using reinforcement learning to incorporate downstream discriminative objectives into diffusion generation.Its per-class MST construction samples real images, encodes and PCA-projects their latents, then builds a kNN graph and its MST.Source: §B.3 · p. 23

C IMPLEMENTATION DETAILS

The implementation builds class-wise latent MSTs, scores generated trajectories against that geometry, and uses task-specific pixel-space rewards for classification, detection, and segmentation.

C.1 MST CONSTRUCTION DETAILS

  • Before MGPO fine-tuning, each class’s sampled real images are encoded with a frozen Stable Diffusion VAE, projected into normalized PCA space, and used to build a sparse k-nearest-neighbor graph and MST skeleton.The PCA dimension is 128.Source: §C.1 · pp. 23–24

C.2 LATENT GEOMETRIC REWARD DETAILS

  • The latent geometric reward combines terminal-latent distance to MST nodes with alignment between the diffusion transition and a nearby MST-edge direction.The reward evaluates a two-stage stochastic-SDE and deterministic-ODE trajectory; higher values indicate smaller manifold deviation and better geometric alignment.Source: §C.2 · pp. 24–25

C.3 PIXEL CLASSIFIER REWARD DETAILS

  • Pixel-space rewards use a proxy classifier’s target-class log-probability, a detector’s maximum class confidence, or segmentation confidence averaged over top-activated pixels.The segmentation Top-K aggregation discourages exploiting isolated high-confidence pixels by enforcing a minimum expected object area.Source: §C.3 · p. 25

D.1 EXPERIMENT RESULTS ON DIFFERENT STUDENT MODEL

MGPO outperforms comparison methods across tested ImageWoof settings and student backbones, while its MST hyperparameters show stable performance across a broad range.

D.1 EXPERIMENT RESULTS ON DIFFERENT STUDENT MODEL

  • 47.3% versus 45.6% for ResNet-18 and 41.3% versus 40.1% accuracy for ResNet-50 at IPC = 10, with MGPO reported best across tested settings.The comparison is with CaO2 under the soft-label training protocol; reported gains remain stable across ResNet-18, ResNet-50, and ResNet-101.Source: §D.1 · pp. 25–26, Table 7 · p. 27

D.2 HYPERPARAMETER ANALYSIS

  • MGPO exhibits stable performance across a broad range of MST hyperparameter settings, with moderate values generally balancing optimization stability and reward effectiveness.The analysis varies kNN neighbor count, maximum sampled latent features per class, and PCA projection dimension on ImageWoof at 10 IPC.Source: §D.2 · pp. 26–27, Figure 5 · p. 27

E.1 RESULTS ON CLASSIFICATION

The appendix reports generated-sample visualizations for classification, object detection, and segmentation.

E.1 RESULTS ON CLASSIFICATION

  • Figures 6–8 visualize classification results on ImageWoof, ImageNette, and ImageIDC.Source: §E.1 · p. 27

E.2 RESULTS ON DETECTION

  • Figure 9 visualizes generated samples for object detection.Source: §E.2 · p. 27

E.3 RESULTS ON SEGMENTATION

  • Figure 10 visualizes generated samples for segmentation.Source: §E.3 · p. 27

F LIMITATIONS AND BROADER IMPACTS

MGPO's stated constraints include geometric-reward approximation, offline optimization overhead, and MST scalability; its broader impacts include resource savings and potential reward-model bias amplification.

Limitations

Broader Societal Impacts

Loading 2610.05252v1…