这篇论文提出了WGF方法，能微调单步生成模型且无需奖励梯度，解决了非可微奖励问题。
研究人员提出了一种通过Wasserstein梯度流(WGF)对单步生成模型进行奖励引导微调的新方法。该方法在CIFAR-10和ImageNet 256×256数据集上进行了测试，使用JPEG压缩率、类别概率、黑白和CLIP对齐等多种奖励函数。实验表明，该方法比基线模型实现了更好的奖励对齐效果，同时避免了奖励黑客和模式崩溃问题。
Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow
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.