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

arXiv:2608.29647 · cs.LG · Submitted 2026-08-30 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow".

Jane: The paper was written by Hoseong Hwang, Woorim Han, Joungin Chun, Jinseong Park and Jaewoong Choi from Sungkyunkwan University, Seoul, Republic of Korea and Korea Institute for Advanced Study, Seoul, Republic of Korea.

Tom: Stay tuned as we take you through the paper and discuss its implications.

The Core Idea: Tom: So, let's talk about what the paper actually does. They're using this concept called Wasserstein Gradient Flow or WGF to model how the generative model learns to move from its initial state toward a target that is defined by a reward function.

Jane: It’s essentially guiding the distribution of probability mass in a controlled way, Tom. Think of it like taking an image that's randomly generated and smoothly pulling it towards the "dog" class, without just forcing it into one single specific sample.

Lu: The authors define this target distribution as being weighted by an exponential function of a reward signal, which is represented by nu r. This isn't just a simple average; it’s a targeted evolution toward that specific reward-weighted goal.

Meng: And they are using the JKO scheme to make this continuous flow practical, right? That's crucial for implementation. It discretizes the process so that, instead of solving one massive equation, we take small, manageable steps in time.

Lalam: This is a huge step because it suggests that AI can learn not just from what *is* there, but from what we *want* to be there—a shift in the definition of learning itself.

The Methodology and Innovation: Tom: The authors say their method is a big improvement over existing approaches, and they claim it handles things that usually cause problems in AI training. They call this being able to handle non-differentiable rewards.

Jane: That's a game changer, Tom. Many fine-tuning methods require gradients through the reward function, but if you are using something like JPEG file size as the measure of quality—which is non-differentiable—those standard methods fail completely.

Lu: The paper’s innovation is that it bypass those gradients by using a specific mathematical framework called Unbalanced Optimal Transport or UOT. This allows the model to map from the current state to the next without needing continuous feedback on the reward function itself.

Meng: From an engineering standpoint, this also means we' don't need to worry about gradient vanishing when dealing with those complex rewards; we just apply a practical semi-dual formulation instead, which is much more robust for large-scale deployment.

Lalam: It’ provides a way for the AI to learn that aligns with human concepts—like "this image should be highly compressed" or "this image should look like a dog"—without having to mathematically define the exact gradient of that concept.

The Results and Performance: Tom: Now, let's talk about what the results actually show. They tested this method on various tasks, including incompressibility and class probability on both CIFAR-ten and ImageNet two hundred fifty-six times two hundred fifty-six.

Jane: And the evidence shows that our method achieves a much better reward alignment than the baselines, which is pretty impressive given how complex those reward signals are. It’s not just achieving the highest score; it's doing it while maintaining fidelity.

Lu: The authors specifically highlight that this approach mitig reward hacking and mode collapse, which is a major weakness in many other reinforcement learning methods where the AI finds one high-reward spot and ignores everything else.

Meng: The performance on ImageNet is particularly striking; they are achieving high aesthetic scores while maximizing the target rewards, which means we’ aren't just getting flashy results but *good* results.

Lalam: It shows that this method allows the AI to be both ambitious in its pursuit of a goal and also disciplined enough to stay true to the quality of its foundational training.

Conclusion and Final Thoughts: Tom: So, as we wrap up this discussion, we're looking at a major breakthrough with "Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow." It offers a stable way to train fast generative models that align precisely with complex human desires.

Jane: It’s clear that the core idea is creating a smooth, controlled evolution of the probability distribution, rather than relying on iterative steps.

Lu: I'm really excited about the potential for how many different creative tasks this could apply to—the possibilities are endless once we have this stable guidance mechanism.

Meng: I think the practical takeaway is that we can now build deployment pipelines that don't sacrifice speed for quality, which is a massive win for industry.

Lalam: For culture, it means AI isn' no longer just a tool; it becomes a highly focused partner in shaping what we want to see and achieve.

Tom: Well, I think that’s a perfect way to wrap up the discussion on this fantastic paper today! Thank you all for joining us.

Jane: It was truly insightful, Tom. We'll be back next week with another exciting paper.

Hoseong Hwang, Woorim Han, Joungin Chun, Jinseong Park, Jaewoong Choi

Sungkyunkwan University, Seoul, Republic of Korea · Korea Institute for Advanced Study, Seoul, Republic of Korea

cs.LG

Submitted: 2026-08-30

Updated: 2026-08-30

Comments: 14 pages, 9 figures

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 81/100

The gist: This paper introduces a novel framework for guiding one-step generative models using reward signals, specifically employing Wasserstein Gradient Flow (WGF) principles.

Key concepts

Wasserstein Gradient Flow (WGF)
WGF is a concept used to model how a generative model learns by smoothly moving its probability distribution from an initial state toward a target. This controlled evolution guides the AI's output based on what the defined reward function determines as desirable.
Reward-Guided Fine-Tuning
This process trains an AI model not just on existing data, but specifically to align its output with human desires or goals. It uses a reward signal to guide the probability distribution toward a targeted, desired outcome rather than relying on simple averages.
Unbalanced Optimal Transport (UOT)
UOT is a mathematical framework that allows the model to learn without needing continuous feedback (gradients) on the reward function. This capability bypasses standard training limitations when dealing with complex quality measures that are not mathematically smooth.

Terminology

Summary

This paper introduces a novel framework for guiding one-step generative models using reward signals, specifically employing Wasserstein Gradient Flow (WGF) principles. This method is crucial because it allows for the fine-tuning of generative models toward complex, non-differentiable objectives—such as achieving specific visual qualities or adhering to external rewards—by minimizing the distance between the generated samples and a desired target distribution defined by these rewards.

Core Methodologies for Reward Guidance

The framework evaluates several strategies for incorporating reward signals into the fine-tuning process. These include:

  • Winner Strategy: This approach Selects the single highest-reward sample x+ and optimizes the direct reconstruction objective: LSFT-Winner = E T theta(z +) - x+ squared. This focuses on optimizing the model to directly reconstruct the best possible sample.

  • Top-k Mean Strategy: This method constructs a reward-weighted blended target from the top- k highest-reward samples x i i=1 k. The objective minimizes the distance between the mean predicted output and this blended target: LSFT-TopK = E T theta(z) - 1 over k sum i=1 k w i x i squared.

  • DPO Baseline: The authors compare against the DPO baseline, which utilizes a specialized loss function: LDPO = -E[sigma(beta DPO [(theta(i-) - old(i-)) - (theta(i+) - old(i+))])].

Training and Evaluation Protocols

The experiments are conducted on two datasets: CIFAR-10 (32 times32) and the high-resolution ImageNet (256 times256). For ImageNet, the backbone architecture used is SiT-XL/2. To manage scale discrepancies across different reward functions, the raw reward r is standardized using a formulation defined as w = beta (r-mu) / (sigma+ epsilon), where mu and sigma are the exponential moving average (EMA) of the mean and standard deviation, respectively. The training process for the proposed method involves a total of 16k iterations with a batch size of 32, utilizing Adam optimization with specific learning rate schedules.

Comparison to Fixed vs. Updated Distribution

A key ablation study compares two modes of updating the model: Updated and Fixed. The results demonstrate that while the Fixed approach might initially move toward the target distribution faster, it ultimately degrades into blurry samples with degraded structure. In contrast, the Updated method successfully produces clear dogs while preserving sharpness, indicating superior long-term sample quality preservation.

Performance Across Multiple Reward Tasks

The framework is tested across various reward objectives, including:

  • Incompressibility and Compressibility (representing non-differentiable reward optimization).

  • Class 5 (Dog) (representing differentiable reward optimization).

  • B&W and CLIP-red tasks, which test the model's ability to adhere to specific stylistic or semantic constraints.

The quantitative results show that the proposed method achieves consistent reward improvement across all tasks throughout training. Conversely, the DPO baseline is noted to converge quickly in the early stage but subsequently saturates or even declines thereafter, highlighting the stability and robustness of the proposed reward-guided approach.

Improvements for AI systems

Based on a thorough review of this paper's methodologies concerning reward-guided generative model fine-tuning, I have identified several critical improvements that must be integrated into any state-of-the-art generative AI system. These improvements primarily focus on stabilizing training, enhancing sample quality retention, and improving the efficiency of complex reward optimization tasks.

Here are the specific improvements and the capabilities they enable:

Improvement: Implement a standardized, hybrid reward weighting strategy that combines elements of Direct Preference Optimization (DPO) with robust sample aggregation techniques like the Top-k Mean Strategy.

  • Mechanism: Instead of relying solely on binary preference comparisons (as in standard DPO), the system must calculate a weighted target by minimizing the distance between the predicted mean output (T theta(z)) and a blended, weighted target derived from the k highest-reward samples (1 over k sum i=1 k w i x i).

  • Capability: This allows the model to effectively optimize for complex, multi-faceted rewards (e.g., combining incompressibility and compressibility metrics) without being destabilized by the variance of a single optimal sample. It provides superior convergence stability compared to Winner Strategies and better generalization than pure DPO when multiple reward signals are active.

The resulting AI system will be a Highly Robust, Multi-Objective Generative Framework capable of:

  1. Superior Fine-Tuning: Achieving target rewards consistently across diverse, independently defined metrics (e.g., simultaneously optimizing for both incompressibility and a specific class label like 'Dog').

  2. High Fidelity at Scale: Generating photorealistic and structurally coherent samples at high resolutions (256x256) without the degradation or blurring typical of standard optimization methods, while maintaining structural similarity to professional pre-trained models.

  3. Adaptive Optimization: Automatically selecting the optimal reward aggregation strategy (Top-k Mean, Winner, or DPO) based on whether the reward signal is highly variable, single-sample dependent, or purely comparative.

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.

Sources

Related papers