TL;DR — RL fine-tuning of diffusion models is well-studied, but flow matching models (Stable Diffusion 3, Flux) remain underexplored due to unique challenges: no likelihood, policy collapse risk, high compute. ORW-CFM-W2 solves all three — online reward-weighting for alignment, Wasserstein-2 regularization for diversity, and a tractable bound that makes it efficient.

Why Flow Matching Is Hard to Fine-Tune

No Likelihood

Unlike diffusion models, flow matching has no tractable likelihood — standard KL regularization cannot be computed directly.

Policy Collapse

Without diversity constraints, reward maximization causes the model to collapse onto a few high-reward modes, losing generative quality.

No Human Data

Existing fine-tuning methods require expensive human-curated datasets or preference annotations. ORW-CFM-W2 needs neither.

Three Key Contributions

1

Online Reward-Weighted Mechanism

Integrates RL into flow matching via reward weighting — guides the model to prioritize high-reward regions in data space, without requiring reward gradients or filtered datasets.

2

Wasserstein-2 Regularization

Derives a tractable upper bound for W2 distance in flow matching models — prevents policy collapse and maintains generation diversity throughout continuous optimization.

3

Unified RL Perspective

Establishes a connection between flow matching fine-tuning and traditional KL-regularized RL, enabling controllable reward-diversity trade-offs and deeper understanding of the learning behavior.

Experiments

SD3
Successfully fine-tunes
Stable Diffusion 3
0
Human-collected data
required
Spatial understanding
& compositional tasks
SOTA
Orders of magnitude
less data than baselines

Tasks: target image generation, image compression, text-image alignment. Consistently achieves optimal policy convergence while allowing controllable diversity-reward trade-offs.

Results on SD3

On Stable Diffusion 3, ORW-CFM-W2 improves spatial and positional understanding on complex compositional prompts that the base model fails on, compares favorably against multiple baselines for target image generation, supports multi-reward optimization across diverse prompts and reward signals, and produces images with richer semantic detail and better semantic correspondence.

Unified RL Framework

We establish a formal connection between flow matching optimization and KL-regularized reinforcement learning, providing theoretical grounding for the reward-diversity trade-offs realized by ORW-CFM-W2.

Controlled Studies: the α Trade-off

Controlled experiments on MNIST and CIFAR isolate the reward–diversity trade-off set by the W2 regularization strength α. Raising α preserves more diversity at some cost in reward, and the trade-off is smooth and controllable rather than a cliff — so α behaves as a usable knob. The regularization also prevents the mode collapse that unconstrained reward-weighting falls into, and the same control carries over to text-conditional CIFAR.

Alpha Control: Reward–Diversity Knob

The regularization strength α provides a smooth knob between pure reward maximization (α=0) and full diversity preservation (α=1): α=0 gives pure reward and mode collapse, α=0.3 is a balanced recommended setting, and α=0.8 preserves high diversity. The same knob applied to SD3 spans from the base model through balanced fine-tuning to full reward optimization.

Publication Journey

Oct 2024
Submitted to ICLR 2025
Submitted to The Thirteenth International Conference on Learning Representations.
Feb 2025
Accepted at ICLR 2025 (Poster)
Accepted as a poster. OpenReview.
Apr 2025
Presented at ICLR 2025 · Singapore
ICLR 2025.

BibTeX

@inproceedings{fan2025online,
  title={Online Reward-Weighted Fine-Tuning of Flow Matching
         with Wasserstein Regularization},
  author={Jiajun Fan and Shuaike Shen and Chaoran Cheng
          and Yuxin Chen and Chumeng Liang and Ge Liu},
  booktitle={The Thirteenth International Conference on
             Learning Representations},
  year={2025},
  url={https://openreview.net/forum?id=2IoFFexvuw}
}