Research map / Reinforcement learning / RL from feedback
Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ORW-CFM-W2, an online reward-weighted conditional flow matching method with a Wasserstein-2 bound regularizer to fine-tune flow models on arbitrary rewards while avoiding policy collapse.
Problem
Fine-tuning continuous flow-based generative models to arbitrary user-defined rewards is hard: likelihoods are costly in CNFs, KL is intractable, and unregularized online reward weighting collapses to a delta distribution.
Approach
Reward-weighted conditional flow matching trained on self-generated samples (online), plus a tractable W2 distance upper bound (Monte Carlo over vector-field differences to a frozen reference model) that prevents policy collapse; avoids likelihoods, filtered datasets, and differentiable rewards.
Key results
- On SD3 text-image alignment, ORW-CFM-W2 reaches CLIP score 35.46 with diversity 4.21, versus RAFT+W2 30.88 and ReFT+W2 32.03 (diversity 2.81 and 3.63).
- Without W2 regularization, ORW-CFM reaches CLIP score 33.63 but diversity drops to 3.47 from the SD3 baseline's 4.78, and the paper reports policy collapse to near-identical images.
- On MNIST, reward reaches near-optimal (1) quickly; increasing entropy coefficient tau makes policy greedier and less diverse, while W2 coefficient alpha controls diversity.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use online reward weighting instead of the RAFT or ReFT online baseline | SD3 fine-tuning · SD3 text-image alignment, “a cat in the sky” | CLIP Score: 29.30 -> 33.63; 29.32 -> 33.63 | Helped |
| add W2 regularization (ablation) | ORW-CFM · SD3 text-image alignment, “a cat in the sky” | CLIP Score: 33.63 -> 35.46 | Helped |
| add W2 regularization (ablation) | ORW-CFM · SD3 text-image alignment, “a cat in the sky” | Diversity Score: 3.47 -> 4.21 | Helped |
| add W2 regularization (ablation) | RAFT · SD3 text-image alignment, “a cat in the sky” | Diversity Score: 2.05 -> 2.81 | Helped |
| add W2 regularization (ablation) | RAFT · SD3 text-image alignment, “a cat in the sky” | CLIP Score: 29.30 -> 30.88 | Helped |
| add W2 regularization (ablation) | ReFT · SD3 text-image alignment, “a cat in the sky” | Diversity Score: 3.26 -> 3.63 | Helped |
| add W2 regularization (ablation) | ReFT · SD3 text-image alignment, “a cat in the sky” | CLIP Score: 29.32 -> 32.03 | Helped |
| remove W2 regularization (ablation) | ORW-CFM-W2 · SD3 text-image alignment, “a cat in the sky” | Diversity Score: 4.21 -> 3.47 | Hurt |
| increase the W2 coefficient α (ablation) | ORW-CFM-W2 · CIFAR-10 image compression | reward | Hurt |
| increase the W2 coefficient α (ablation) | ORW-CFM-W2 · MNIST target image generation | diversity | Helped |
| increase the entropy coefficient τ (ablation) | ORW-CFM without W2 · MNIST target image generation | diversity | Hurt |
| set τ=0 (ablation) | ORW-CFM · MNIST target image generation | reward | No clear effect |
Surprising findings
- Unregularized online reward-weighted methods (including RAFT and ReFT) collapse to nearly identical images despite high reward, which the paper frames as a theoretical delta-distribution limit.
Limitations
- Theoretical results assume perfect learning of the reward-weighted distribution at each epoch and Lipschitz vector fields.
- W2 bound uses Monte Carlo estimates and a scaling constant gamma; practical distance estimation details are limited.
- Evaluation on a single SD3 prompt set and a CLIP-based diversity metric; no human evaluation.
Experimental setup
Datasets: MNIST CIFAR-10 DrawBench prompts
Benchmarks: MNIST even-digit generation CIFAR-10 image compression DrawBench spatial prompts
Metrics: CLIP score diversity score (mean pairwise CLIP embedding distance) reward W2 distance
Baselines: RAFT ReFT SD3 RAFT+W2 ReFT+W2 DPO DDPO
Builds on: conditional flow matching reward weighted regression DDPO DPOK ReFT
Closest papers by meaning
- WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps (NeurIPS 2026 · similarity 0.91)
- Flow-GRPO: Training Flow Matching Models via Online RL (NeurIPS 2025 · similarity 0.90)
- Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control (ICLR 2025 · similarity 0.90)
- Efficient Adjoint Matching for Fine-tuning Diffusion Models (NeurIPS 2026 · similarity 0.89)
- Value Gradient Guidance for Flow Matching Alignment (NeurIPS 2025 · similarity 0.89)
- Adaptive Divergence Regularized Policy Optimization for Fine-tuning Generative Models (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.