Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow
Machine Learning
Summary
The authors focus on improving one-step generative models, which create data in a single step from random noise, to better match desired outcomes guided by rewards. They use a mathematical concept called Wasserstein Gradient Flow to smoothly adjust the model's output distributions without needing gradient information from the rewards, allowing flexibility with different types of rewards. Their method also helps prevent common problems like the model cheating or producing limited variety. Tests on image datasets show their approach aligns generated data more closely with given rewards than other methods.
Authors
Hoseong Hwang, Woorim Han, Joungin Chun, Jinseong Park, Jaewoong Choi
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.