Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow

• URL: http://arxiv.org/abs/2608.29647v1
• Date: Sun, 30 Aug 2026 08:18:30 GMT
• Title: Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow
• Abstract summary: We propose a reward-guided fine-tuning of a one-step generative model via Wasserstein Gradient Flow.<n>We derive a practical training method that requires no reward gradients, thereby handling both non-differentiable and differentiable rewards.<n>Our method provides smooth and stable reward-guided distributional updates while mitigating reward hacking and mode collapse.
• Score: 10.90140528360878
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: To mitigate the time complexity of generative models, one-step generative models have recently emerged through direct mapping from noise to data in a single forward pass. However, the reward-guided fine-tuning method of one-step generative models remains largely unexplored. To address this, we consider one-step generators from an optimal transport view, investigating Wasserstein Gradient Flow (WGF) for modeling smooth and controlled distributional evolution in probability space. We then propose a novel reward-guided fine-tuning of a one-step generative model via WGF. We derive a practical training method that requires no reward gradients, thereby handling both non-differentiable and differentiable rewards. Moreover, our method provides smooth and stable reward-guided distributional updates while mitigating reward hacking and mode collapse. Experiments on 2D synthetic data, CIFAR-10, and ImageNet 256$\times$256 with diverse rewards, including JPEG (in)compressibility, class probability, Black-and-White and CLIP alignment, show that our method achieves better reward alignment compared to baselines.

Related papers

• Perceptual Flow Matching for Few-Step Generative Modeling [64.26537634881244]
  Perceptual Flow Matching is a framework for few-step generation in flow-matching models.<n>We show that PFM produces high-quality few-step generators when supervised in an appropriate representation space.
  arXiv  Detail & Related papers  (2026-07-03T17:55:01Z)
• Aligning Few-Step Generative Models by Amortizing Sample-based Variational Inference [60.82875447631739]
  We introduce FAV, Few-step Generative Models Alignment via Sample-based Variational Inference.<n>FAV cast alignment as sampling from a reward-tilted distribution anchored to a reference distribution.<n>We evaluate FAV on two domains: robotics manipulation and image generator alignment.
  arXiv  Detail & Related papers  (2026-05-26T05:02:49Z)
• One-Step Generative Modeling via Wasserstein Gradient Flows [56.266175641614986]
  W-Flow is a framework for training a generator that transforms samples from a simple reference distribution into samples from a target data distribution in a single step.<n>We set a new state of the art for one-step ImageNet 256$times$256 generation, achieving 1.29 FID, with improved mode coverage and domain transfer.
  arXiv  Detail & Related papers  (2026-05-12T08:29:44Z)
• Scaling Categorical Flow Maps [29.25832277355597]
  Continuous diffusion and flow matching models could represent a powerful alternative to autoregressive approaches for language modelling.<n>We show how to generate diverse, high-quality text in as few as $4$ inference steps while maintaining near-data-level token entropy.<n>We uncover some of the challenges that arise from training these models at scale, and we provide prescriptive insights on loss weighting and time scheduling.
  arXiv  Detail & Related papers  (2026-05-08T14:50:57Z)
• Unlearning for One-Step Generative Models via Unbalanced Optimal Transport [11.695528552671371]
  UOT-Unlearn is a novel plug-and-play class unlearning framework for one-step generative models.<n>Our method formulates unlearning as a principled trade-off between a forget cost, which suppresses the target class, and an $f$-divergence penalty.
  arXiv  Detail & Related papers  (2026-03-17T13:19:04Z)
• FAIL: Flow Matching Adversarial Imitation Learning for Image Generation [52.643484089126844]
  Post-training of flow matching models-aligning the output distribution with a high-quality target-is mathematically equivalent to Imitation learning.<n>We propose Flow Matching Adrial Learning (FAIL), which minimizes policy-expert divergence through adversarial training without explicit rewards or pairwise comparisons.
  arXiv  Detail & Related papers  (2026-02-12T16:36:33Z)
• Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding [84.3224556294803]
  Diffusion models excel at capturing the natural design spaces of images, molecules, DNA, RNA, and protein sequences. We aim to optimize downstream reward functions while preserving the naturalness of these design spaces. Our algorithm integrates soft value functions, which looks ahead to how intermediate noisy states lead to high rewards in the future.
  arXiv  Detail & Related papers  (2024-08-15T16:47:59Z)
• Score-Based Generative Modeling through Stochastic Differential Equations [114.39209003111723]
  We present a differential equation that transforms a complex data distribution to a known prior distribution by injecting noise. A corresponding reverse-time SDE transforms the prior distribution back into the data distribution by slowly removing the noise. By leveraging advances in score-based generative modeling, we can accurately estimate these scores with neural networks. We demonstrate high fidelity generation of 1024 x 1024 images for the first time from a score-based generative model.
  arXiv  Detail & Related papers  (2020-11-26T19:39:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.