PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.29647v1cs.LG

TL;DR

One-step generators make inference fast, but existing reward-fine-tuning methods depend on multi-step trajectories unavailable to them. This paper uses Wasserstein Gradient Flow with JKO-based transport updates and importance sampling, requiring no reward gradients. Across synthetic and image benchmarks, it reports better reward alignment than baselines while maintaining smooth distributional updates.

  • Problem

    Existing fine-tuning methods rely on intermediate states from multi-step sampling trajectories, leaving reward-guided fine-tuning for one-step generators insufficiently addressed.

  • Method

    The method evolves a pre-trained one-step generator toward a reward-weighted distribution using WGF, JKO discretization, cumulative-map reparameterization, and importance sampling.

  • Results

    Experiments on 2D synthetic data, CIFAR-10, and ImageNet 256×256 show better reward alignment than baselines across diverse reward tasks.

  • Takeaways & Limitations

    The framework supports reward-guided one-step fine-tuning with smooth controlled updates and both differentiable and non-differentiable rewards without reward gradients.

  • Takeaways & Limitations

    The method assumes access to a pre-trained one-step generator, and the paper notes higher FID on two ImageNet tasks when reward-weighted targets are far from the data distribution.

Abstract

from arXiv · show

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.

1 Introduction

One-step generators reduce inference cost through direct noise-to-data mappings, but existing reward-fine-tuning methods rely on multi-step trajectories unavailable in one-step models. The paper addresses this gap with WGF-based fine-tuning and reports better reward alignment across synthetic and image benchmarks.

  • Motivation: One-step models reduce generation to a direct mapping from noise to data, motivating reward fine-tuning with fewer forward passes.Recent fine-tuning work targets one- and few-step frameworks to reduce computational cost.
  • Research gap: Existing diffusion fine-tuning methods cannot directly apply because one-step generators lack intermediate denoising states and generally use linear generation paths.Diffusion-based methods assume nonlinear iterative trajectories that one-step models do not produce.
  • Approach: The proposed method uses Wasserstein Gradient Flow to optimize one-step model distributions through smooth, controlled dynamics in probability space.Unlike diffusion fine-tuning, it does not rely on iterative inference trajectories.
  • Approach: The method discretizes WGF with the JKO scheme and uses semi-dual optimization plus importance sampling to derive a practical training objective.The resulting algorithm supports both differentiable and non-differentiable rewards.

2 Preliminaries

The preliminaries frame generation as transport from Gaussian noise to data and introduce WGF as distributional steepest descent. JKO discretization provides a tractable approximation to this continuous evolution.

  • Probability flows: Probability-flow models transform Gaussian noise into data through a continuous ODE governed by a learned velocity field.Sampling requires solving the flow from t = 0 to t = 1.
  • One-step models: Numerical integration requires tens to hundreds of network evaluations, motivating models that learn the flow map directly for one-step generation.Consistency Trajectory Models and MeanFlow parameterize mappings between times rather than repeatedly solving the ODE.
  • One-step models: The paper assumes a pre-trained one-step map T0 whose pushforward of a Gaussian prior approximately matches the data distribution.Formally, (T0)#µ ≈ ν with µ Gaussian.
  • WGF: Wasserstein Gradient Flow evolves probability densities along the steepest descent direction of a functional under the 2-Wasserstein metric.For an f-divergence target functional, the flow moves an initial distribution toward the target distribution.
  • JKO scheme: The JKO scheme approximates continuous WGF by iteratively minimizing a time-discretized objective, with step size h converging to the continuous trajectory as h → 0.Each update balances proximity to the current distribution with descent toward the target.

3 Method

The method formulates reward-guided fine-tuning as WGF toward a reward-weighted target, then makes each flow step trainable with JKO, semi-dual optimal transport, reparameterization, and importance sampling. Wasserstein regularization yields controlled updates without reward gradients or iterative inference at deployment.

  • 3.1 Problem Formulation: Fine-tuning evolves the pre-trained generator distribution toward a reward-weighted target by minimizing an f-divergence through WGF.The target may use differentiable or non-differentiable rewards, without requiring direct target samples.
  • 3.2 Fine-tuning via JKO Scheme: The JKO scheme replaces continuous WGF with iterative distributional updates whose Wasserstein term keeps successive distributions close while divergence drives reward improvement.As h → 0, the discrete sequence converges to the continuous WGF trajectory.
  • 3.2 Fine-tuning via JKO Scheme: Each JKO step has a semi-dual unbalanced optimal transport formulation that directly parameterizes a transport map and potential function.This avoids explicitly modeling the evolving probability density.
  • 3.2 Fine-tuning via JKO Scheme: Reparameterization trains the cumulative map Tk instead of simulating every intermediate JKO distribution, requiring samples only from the fixed initial distribution.Tk represents the fine-tuned generator after successive refinements from the uncontrolled distribution toward the reward-weighted target.
  • 3.2 Fine-tuning via JKO Scheme: Importance sampling estimates reward-target expectations using pre-trained-distribution samples weighted by exp(β · r(y))/Mβ.The normalization constant is estimated with a moving average over mini-batches.
  • 3.2 Fine-tuning via JKO Scheme: Updating the reference distribution during training reduces mismatch between the fixed pre-trained distribution and the current generator, improving optimization stability.The algorithm alternates generator transport-map and discriminator potential-function updates.
  • 3.3 Advantages of WGF for Fine-tuning: Because rewards enter only as scalar importance weights, the method requires no reward gradients and handles both differentiable and non-differentiable rewards.This supports black-box rewards such as compression-based objectives.
  • 3.3 Advantages of WGF for Fine-tuning: The learned transport map realizes each WGF descent step directly, avoiding velocity-field learning and iterative ODE solving at inference time.The map moves samples from the current distribution toward the reward-weighted target.

4 Related Works

Related work covers one-step generative models and reward fine-tuning for diffusion or flow models. The paper evaluates reward tasks qualitatively on CIFAR-10, including incompressibility and class targeting.

  • One-Step Generative Models: Consistency Models learn mappings from points along an ODE trajectory to its endpoint, enabling one-step training, while CTM generalizes this framework.These models belong to the broader effort to reduce iterative generation.
  • CIFAR-10 evaluation: CIFAR-10 evaluation reports final reward and FID for incompressibility, compressibility, and Class 5 tasks, with FID computed from 50k samples.The table separates non-differentiable and differentiable reward settings.
  • CIFAR-10 evaluation: CIFAR-10 qualitative samples compare incompressibility and Class 5 ("Dog") steering.The incompressibility task is shown on the left and Class 5 on the right.
  • Reward Fine-Tuning: Diffusion and flow reward-fine-tuning methods align models to rewards while attempting to preserve the pre-trained distribution.DDPO may overfit to a few high-reward samples, whereas DPOK adds KL regularization to preserve diversity.

5 Experiments

The experiments evaluate reward-guided fine-tuning across 2D synthetic data, CIFAR-10, and ImageNet 256×256, covering non-differentiable and differentiable rewards. Across these settings, the WGF method provides more stable reward–fidelity trade-offs and better reward alignment than the evaluated baselines.

  • Experimental Setup: Experiments span 2D synthetic data, CIFAR-10, and ImageNet 256×256 with JPEG compression rewards, class probability, CLIP-red, and Black-and-White rewards.FID is used on CIFAR-10 and ImageNet, with ImageNet additionally evaluated using Aesthetic score.
  • 2D Synthetic Data: The WGF method progressively captures all eight target modes while avoiding the structural collapse and mode averaging observed with SFT, RWR, and DPO.The 2D experiment tracks training from 5k to 25k iterations toward an 8-Gaussian reward-weighted target.
  • CIFAR-10: The method improves non-differentiable CIFAR-10 rewards while preserving semantic structure, whereas SFT and RWR fail on incompressibility and DPO produces noise and blur.For the differentiable Class 5 task, the method reaches reward −0.87 versus −7.71 for the pretrained model, while baseline improvements involve fidelity costs or fail to improve reward.
  • CIFAR-10: At a comparable reward level of 1.40, our method attains an FID of 9.32 against 89.79 and 138,89 for the compared 20-step DDPM methods.The comparison uses different pretrained models, with FID values of 3.90 and 29.90 before fine-tuning.
  • ImageNet 256×256: On ImageNet, our method achieves the highest reward on three of four tasks and the best aesthetic score on all four tasks.It maintains aesthetic score near the pretrained value on Black-and-White and CLIP-red while achieving higher rewards than DPO.

6 Conclusion

The paper presents WGF-based reward-guided fine-tuning for one-step generators, using JKO discretization to produce stable updates without reward gradients. It reports stronger reward alignment than baselines while acknowledging an untested prompt-conditional text-to-image setting.

  • The method is the first WGF-based reward-guided fine-tuning framework for one-step generative models.
  • It discretizes distributional evolution with the JKO scheme to move a pretrained distribution toward a reward-weighted target.
  • The resulting objective requires no reward gradients and therefore applies to non-differentiable rewards.
  • Experiments show superior reward alignment over baseline approaches on synthetic and real datasets, while the method mitigates reward hacking.
  • The paper does not evaluate prompt-conditional text-to-image generation, leaving that application for future work.
  • The algorithm initializes current and reference generators from the pretrained generator, repeatedly updates the current model, and then refreshes the reference generator.

B.1 Rewards

The experiments cover differentiable and non-differentiable rewards, generate candidate groups from perturbed noise, and compare several reward-driven fine-tuning baselines.

  • Rewards: The evaluation uses JPEG (in)compressibility, CIFAR-10 dog-class probability, ImageNet CLIP-red alignment, and Black-and-White rewards.
  • Rewards: JPEG incompressibility rewards file size directly, encouraging high-frequency detail, while compressibility uses the negative JPEG file size.
  • Rewards: The JPEG objectives are strictly non-differentiable because JPEG compression is a discrete, black-box operation.
  • Rewards: The class-probability reward uses a pretrained ResNet-56 classifier to target CIFAR-10’s dog class, class 5.
  • Rewards: CLIP-red measures cosine similarity to the text prompt “red,” whereas Black-and-White penalizes RGB-channel variance to encourage grayscale images.
  • Candidate Generation: Candidate generation samples Gaussian noise, forms perturbed intermediate latents, and produces candidate pairs with associated rewards.
  • Baselines: SFT uses highest-reward candidates, DPO contrasts highest- and lowest-reward candidates, and RWR weights all candidates by exponentiated rewards.

B.3 Training Details

Training uses fixed datasets, seeds, and hardware across CIFAR-10 and ImageNet experiments, with longer high-resolution training and reward standardization for ImageNet.

  • Training Details: All experiments use CIFAR-10 32×32 and ImageNet 256×256, train on one NVIDIA A100 GPU, and evaluate with fixed seeds.
  • CIFAR-10 Experiments: CIFAR-10 models train for 5,000 iterations, with a 2,000-iteration discriminator warmup for the Class 5 reward experiment.
  • ImageNet 256×256 Experiments: ImageNet experiments use the SiT-XL/2 backbone and standardize rewards using exponential moving averages of reward mean and standard deviation.
  • ImageNet 256×256 Experiments: The proposed ImageNet models train for 16k iterations with batch size 32, Adam optimization, cosine learning-rate scheduling, and generator-weight EMA.
  • Qualitative Comparison: At 800 iterations, both CIFAR-10 proposals remain near the pretrained distribution; at 5,000, Updated preserves sharp dog images while Fixed becomes washed out.
  • Baselines: DPO uses group size 8 for ImageNet, trains for 16k iterations, and applies gradient clipping and reward-scale-dependent skip thresholds.

C Additional Results

On ImageNet 256×256, the proposed method improves reward consistently across training, whereas DPO improves early but later saturates or declines.

  • The proposed method achieves consistent reward improvement across all ImageNet 256×256 tasks throughout training.
  • DPO converges quickly early in training but later saturates or even declines.

C.2 Ablation Study on Fixed vs. Updated

The Updated pretrained distribution preserves sample quality during training, whereas Fixed initially advances faster but later produces blurry samples.

  • Fixed initially moves toward the target distribution more quickly but degrades into blurry samples as training proceeds.
  • Updated retains the sample quality of the pretrained model throughout training.

C.3 Additional Qualitative Results

Additional qualitative comparisons show generated samples across multiple reward tasks on CIFAR-10 and ImageNet 256×256.

  • The qualitative comparisons cover samples generated by the proposed method and baselines on CIFAR-10 and ImageNet 256×256.
  • Figure 8 compares pre-trained samples with CIFAR-10 samples optimized for incompressibility, compressibility, and Class 5 (Dog) rewards.
  • Figure 9 compares pre-trained ImageNet 256×256 samples with samples fine-tuned for each reward target under matched class labels and seeds.
Loading 2608.29647v1…