Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization

Authors: Jiajun Fan, Shuaike Shen, Chaoran Cheng, Yuxin Chen, Chumeng Liang, Ge Liu
Primary area: reinforcement learning
Keywords: Flow Matching, Reinforcement Learning, Wasserstein Regularization, Exploration-Exploitation Trade-off, Fine-Tuning, Generative Model

Abstract:
Recent advancements in reinforcement learning (RL) have achieved great success in fine-tuning diffusion-based generative models. However, fine-tuning continuous flow-based generative models to align with arbitrary user-defined reward functions remains challenging, particularly due to issues such as policy collapse from overoptimization and the prohibitively high computational cost of likelihoods in continuous-time flows. In this paper, we propose an easy-to-use and theoretically sound RL fine-tuning method, which we term Online Reward-Weighted Conditional Flow Matching with Wasserstein-2 Regularization (ORW-CFM-W2). Our method integrates RL into the flow matching framework to fine-tune generative models with arbitrary reward functions, without relying on gradients of rewards or filtered datasets. By introducing an online reward-weighting mechanism, our approach guides the model to prioritize high-reward regions in the data manifold. To prevent policy collapse and maintain diversity, we incorporate Wasserstein-2 (W2) distance regularization into our method and derive a tractable upper bound for it in flow matching, effectively balancing exploration and exploitation of policy optimization. We provide theoretical analyses to demonstrate the convergence properties and induced data distributions of our method, establishing connections with traditional RL algorithms featuring Kullback-Leibler (KL) regularization and offering a more comprehensive understanding of the underlying mechanisms and learning behavior of our approach. Extensive experiments on tasks including target image generation, image compression, and text-image alignment demonstrate the effectiveness of our method, where our method achieves optimal policy convergence while allowing controllable trade-offs between reward maximization and diversity preservation.

Decision: accept (poster)
Reviewer scores: 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730681913930

thread_type:
official_review

<thread_object>
creator: reviewer_s526

<content>
summary:
The paper introduces a new online reinforcement learning (RL) fine-tuning method for continuous flow-based generative models, named Online Reward-Weighted Conditional Flow Matching with Wasserstein-2 Regularization (ORW-CFM-W2). It addresses the challenges of policy collapse and high computational costs associated with traditional fine-tuning methods. The authors propose integrating RL within the flow matching framework, utilizing an online reward-weighting mechanism to focus on high-reward regions and a Wasserstein-2 distance regularization to balance exploration and exploitation. The paper provides theoretical analyses and empirical results across various tasks, demonstrating the effectiveness of the proposed method in achieving optimal policy convergence with controlled trade-offs between reward maximization and the generation capacity.

strengths:
1.The paper effectively identifies and addresses some key issues in fine-tuning continuous flow-based generative models, such as policy collapse and computational inefficiency.
2. The introduction of the online reward-weighting mechanism and Wasserstein-2 distance regularization is well-suited for flow matching models that balances exploration and exploitation, mitigating the policy collapse problem.
3. The theoretical analyses are rigorous and provide a solid foundation for the proposed method. The empirical results across various tasks are convincing and demonstrate the method's effectiveness.

weaknesses:
Potential Overemphasis on Theoretical Analysis:
While the theoretical underpinnings are robust, the paper might overly focus on the theoretical aspects at the expense of practical considerations. Balancing the presentation (e.g.section 4.6 to the appendix) to include more  case studies could make the findings more relatable to a broader audience.
Lack of Comparative Analysis with Other Regularization Techniques:
The paper introduces W_2 distance regularization but does not compare its effectiveness with other potential regularization methods. Including such comparisons could strengthen the paper's contribution by positioning it within the broader landscape of regularization strategies.
Narrow Empirical Validation:
The empirical validation is commendable, but the paper could benefit from testing the method across a wider range of datasets (e.g. CeleA face dataset) and tasks to further establish the generalizability and robustness of the approach.

questions:
1. What kind of reward functions can be fine-tuned without collapsing by the proposed W_2 regularization method?  
2. Is this method capable of performing fine-grained fine-tuning tasks, such as controlling specific semantic parts of images?
3. Why not use W_1 distance for regularizing?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712377075

thread_type:
official_review

<thread_object>
creator: reviewer_tfnw

<content>
summary:
The paper presents a method to perform reward finetuning of flow-based models. The idea starts with the reward-weighted version of the standard flow matching loss (i.e., doing simple importance sampling) and, to remove the dependency on pretraining datasets and to perform online training, changes the sampling distribution from the original data distribution to the finetuned sampling policy. Such a strategy proves very prone to overfitting, as the finetuned distribution collapses into a single mode if it is trained for too many epochs. Therefore, the authors further proposes to regularize the sampling policy to be not too far away from the pretrained one (using a Wasserstein distance). The paper discusses some theoretical results like the asymptotic behaviors of the proposed methods and empirically show that the proposed method can be applied to finetuning of flow matching models pretrained on MNIST and CIFAR-10.

strengths:
The paper theoretically analyzes probably one of the most intuitive methods of reward reweighting, and by introducing a regularization loss on the finetuned distribution, shows that this naive method can be extended to the online setting. To support the claims, the paper does sufficient amount of experiments on different small-scale image datasets and different reward functions. Especially, the paper shows that their online variant is better than the offline one.

Compared to baselines like DDPO that requires one specify the number of sampling steps for finetuning, the proposed method finetunes a model in the way very similar to flow matching -- to sample images from a "data" distribution and some random t in [0,1], and to compute the matching loss.

weaknesses:
The paper does not compare the proposed method against any other methods, for instance DDPO (running PPO for diffusion finetuning). While one may argue that DDPO is not designed for continuous flow models, one eventually samples from CNFs with some discretization and can therefore construct an MDP for DDPO finetuning, and not to mention some more recent methods. On flow matching, there is a very recent paper [1] that does reward finetuning for flow matching (though this paper should be considered as a concurrent one). There also exist some more recent papers for reward finetuning to compare with, and I feel that showing at least one of them will be great.

The proposed method seems to be a bit sensitive (in theory) to hyperparameter tuning due to its online nature. It is a bit unsatisfactory that the resulted distribution (Eqn 12 in the paper) is dependent on the number of epochs. While in practice it is not a super big concern, an objective that guarantees convergence to a specific distribution (e.g. P_pretrained(x) * exp(lambda * r(x)) / Z) is generally considered better.

Many of the baselines are tested on large-scale models like StableDiffusion, and many of them can converge in a reasonably fast speed on simple reward functions like Aesthetic Score used in DDPO. The paper fails to show results in these more realistic settings (though it probably requires some compute, but one might be able to find a smaller model to do experiments).

[1] Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control
. Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, Ricky T. Q. Chen. https://arxiv.org/abs/2409.08861

questions:
Besides the points raised in the weakness section:

1. It is probably better to also show quantitative metrics like diversity scores (e.g., feature pairwise distances) and FID scores.
2. In Eqn 10, it is probably more aesthetic to write \theta_\text{ft} and \theta_\text{ref} (for the subscripts), instead of \theta_{ft} and \theta_{ref}.
3. W2 distance is great, but I wonder if it makes a big difference if one instead uses KL divergence (both theoretically and empirically).
</content>
</thread_object>

</thread>
<thread>
creation date:
1731124966973

thread_type:
official_review

<thread_object>
creator: reviewer_dxdz

<content>
summary:
This work introduces a way to finetune conditional flow matching models to maximize some user-defined reward function. Specifically, the paper combines two techniques: (1) reward-weighted conditional flow matching; and (2) a constraint that bounds the pretrained model and the finetuned model. The work gives some theoretical analyses to justify the proposed method is grounded and some experiments also show its effectiveness.

strengths:
- The problem of finetuning conditional flow matching models is of general interest to the community. How to preserve the generation diversity and avoid model collapse is a challenging problem.

- Combining the reward-weighted matching loss and the Wasserstein distance regulatization seems to be empirically effective. The experimental results look good.

- There are quite a few theoretical justifications for the proposed method. Although I didn't check carefully, I find them to be quite reasonable claims.

weaknesses:
- The contribution of the paper seems ad-hoc to me. There is quite little connection between the reward-weighted matching loss and the Wasserstein regularization. I find both techniques independent of each other, so I find the motivation of the work quite weak. Could the author elaborate more on why these two techniques should be used together (other than empirically well-performing)?

- Given that the reward-weighted matching loss and the Wasserstein regularization are unrelated contributions, I will be interested to see how much gain each individual component contribute to the performance gain? Could the authors conduct some ablation study?

- I find it less convincing for the performance gain, since there is no compelling baselines for comparison. For example, the paper claims that the Wasserstein regularization performs well. How about other discrepancy measure? How is the Wasserstein distance a good choice here? I think more discussion on the motivatiojn will help the reader gain more insights. 

- Whlle I am no expert in this domain, I am wondering whether there are other stronger baselines to compare to. The problem this paper studies doesn't seem to be new, so I think there will be some other finetuning methods for comparison, say [Imagereward: Learning and evaluating human preferences for text-to-image generation, NeurIPS 2023].

- The experiments are relatively small-scaled. I don't know how the proposed method scales with the size of the model/dataset. Could the authors conduct some experiments to study the scaling performance of this finetuning technique?

questions:
See the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732085815219

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Overall Response to Major Concerns and Feedbacks

comment:
We sincerely thank all reviewers for their thoughtful and constructive feedback. After careful consideration of the reviewers' comments, we have comprehensively addressed three main shared concerns in our revision:

1. **Extended Empirical Evaluation on Large-Scale Models** (Reviewers tfNW, dXDz, S526): Multiple reviewers requested demonstrations on larger models and stronger baseline comparisons. We address this through `extensive experiments on online fine-tuning of Stable Diffusion 3 (SD3)` in Section 5.3 and Appendix A. We first note that many existing RL fine-tuning methods cannot be directly applied to flow matching (FM) models due to intractable ELBO and computationally expensive likelihood/KL calculations in continuous-time ODE-based FM models. Therefore, we compare against applicable baselines RAFT and ReFT that don't require likelihood calculations. Our comprehensive evaluation demonstrates: (i) `superior spatial relationship alignment` while maintaining image quality (Figure 5), (ii) `effective prevention of policy collapse` through W2 regularization with detailed ablation studies (Figure 6), (iii) strong `adaptability across multiple reward architectures` (HPS-V2, Pick Score, Alpha CLIP, Figure 9), and (iv) successful handling of complex compositional prompts (Figure 10). Most importantly, our framework achieves these results while preserving the simplicity of the original CFM loss and continuous-time ODE properties, making it highly scalable and easily applicable to any flow matching architecture from TorchCFM to SD3. All of this is achieved through training on self-generated data without requiring  manually collected datasets (i.e., `theoretically-guaranteed collapse-free online fine-tuning of FM while training on self-generated data`).
2. **Theoretical Motivation and Necessity of W2 Regularization**: Regarding the theoretical connection between reward-weighted matching and W2 regularization (Reviewer dXDz) and the choice of W2 distance (Reviewer S526, dXDz, tfNW), we emphasize that combining reward-weighted CFM with W2 regularization is theoretically motivated and necessary. Our analysis proves that online reward-weighted fine-tuning `without regularization inevitably collapses to a delta distribution (Lemma 1)`, necessitating effective regularization to maintain diversity. While KL divergence is common in previous works, it requires expensive ODE simulation with Hutchinson traces estimator for flow matching models, making `KL computationally intractable`. Unlike diffusion models which can leverage variational bounds, flow matching lacks tractable ELBO formulations to connect with its vector field loss. To address this fundamental challenge, for the first time, `we derive  a tractable upper bound for W2 distance (Theorem 3)` in flow matching fine-tuning  that enables practical regularization via vector field loss directly, avoiding expensive likelihood calculations while effectively `preventing policy collapse`.
3. **Improved Paper Structure** (Reviewer S526): Following suggestions about presentation balance, we have restructured the paper to improve accessibility while maintaining technical rigor. We moved theoretical analyses of RL perspectives to the appendix while `expanding Section 5.3 and Appendix A with comprehensive case studies on online fine-tuning of SD3`. The enhanced experimental section now clearly demonstrates our method's advantages through diverse examples: superior spatial relationship control (Figure 5), detailed ablations validating W2 regularization's effectiveness (Figure 6), strong adaptability across reward architectures (Figure 9), and complex semantic control capabilities (Figure 10). We also improved `notation aesthetics` following Reviewer tfNW's suggestions.

These improvements have strengthened both the theoretical foundations and practical impact of our work. We are grateful for the reviewers' detailed feedback that helped us make the paper more accessible to a broader audience while preserving its theoretical rigor. We hope our revisions have thoroughly addressed all concerns, and we welcome any further discussion.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732476898453

thread_type:
official_comment

<thread_object>
creator: area_chair_qcem

<content>
comment:
Dear Reviewers,

This is a gentle reminder that the authors have submitted their rebuttal, and the discussion period will conclude on November 26th AoE. To ensure a constructive and meaningful discussion, we kindly ask that you review the rebuttal as soon as possible (if you have not done so already) and verify if your questions and comments have been adequately addressed.

We greatly appreciate your time, effort, and thoughtful contributions to this process.

Best regards, 
AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a novel method for fine-tuning flow matching models trained with a Wasserstein objective using reinforcement learning that addresses two key challenges: policy collapse and computational tractability. It does so by first showing two complementary observations: (a) online reward-weighted training of flow matching models inevitably leads to policy collapse, (b) a tractable upper bound exists for Wasserstein distance in flow matching that enables practical regularization without expensive likelihood calculations. By combining these insights, the authors finally yield a method that achieves optimal policy convergence while maintaining generation diversity. The paper's key strengths lie in its comprehensive theoretical analysis with clear proofs and guarantees and strong empirical results demonstrating better text-image alignment and diversity preservation compared to baselines like RAFT and ReFT. The reviewers highlighted some weaknesses with regards to the choice of Wasserstein distance being orthogonal to the finetuning procedure, as well as limited comparisons with related approaches.