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

Thank you for your detailed and constructive review. We are grateful for your recognition of our work's theoretical rigor and its contributions to addressing key challenges in flow matching fine-tuning, particularly regarding policy collapse and computational efficiency. We’re happy to address your concerns with the following:

Q: Is this method capable of performing fine-grained fine-tuning tasks, such as controlling specific semantic parts of images? Include more case studies could make the findings more relatable to a broader audience.

A: Thank you for the great suggestion. We comprehensively demonstrate our ORW-CFM-W2 method's strong capabilities in fine-grained semantic control through extensive experiments with Stable Diffusion 3 (SD3) [6].

As shown in Section 5.3, Figure 5 shows precise control over complex spatial relationships between objects like "a banana on the left of an apple" and "a cat on the left of a dog", achieving better positional accuracy and image quality compared to baselines like RAFT [1] and ReFT [2]. The ablation studies in Figure 6 using the challenging "a cat in the sky" prompt demonstrate that while methods without W2 regularization exhibit clear signs of policy collapse (generating nearly identical images), our approach with W2 regularization successfully maintains generation diversity while achieving high semantic accuracy.

In our more comprehensive SD3 experiments (Appendix A), our method's fine-grained control capabilities are further validated through diverse reward architectures. Figure 9 demonstrates consistent performance across multiple text-image alignment reward models (HPS-V2, Pick Score, and Alpha CLIP) for the challenging "train on top of a surfboard" prompt, showing robust spatial understanding while maintaining natural variations in train appearances and ocean conditions. Most notably, Figure 10 exhibits sophisticated semantic control with complex compositional prompts like "A stack of 3 cubes. A red cube is on the top, sitting on a red cube. The red cube is in the middle, sitting on a green cube. The green cube is on the bottom." with specific color and position requirements, where our method successfully controls multiple attributes simultaneously while maintaining image quality and stylistic diversity. These comprehensive results validate our theoretical framework's ability to enable precise semantic control while preventing the policy collapse issues that typically challenge online fine-tuning methods [1] [2].

Q: What kind of reward functions can be fine-tuned without collapsing by the proposed W2 regularization method?

A: Our method is designed to fine-tune flow matching models with arbitrary reward functions without requiring gradients of rewards [5] or filtered datasets [3], as comprehensively demonstrated through our experiments across different scales. In the experiments section, we showcase this versatility through controlled convergence using varying $\tau$ and $\alpha$ for classifier-guided digit generation (Figures 2-3), exhibiting precise control over digit categories while preserving diversity. Figure 4 validates our method with compression-based rewards for image optimization, where we demonstrate clear reward-distance trade-offs while maintaining generation quality. Figure 5 shows successful optimization with CLIP-based similarity rewards for text-image alignment, demonstrating superior performance compared to previous methods like RAFT [1] and ReFT [2]. Figure 6 further demonstrates through ablation studies that methods incorporating W2 regularization (including RAFT+W2, ReFT+W2, and ours) successfully prevent policy collapse while maintaining performance.

In our SD3 experiments (See Appendix A), we further demonstrate this reward-agnostic property across diverse reward architectures. Figure 9 validates our method's adaptability across different reward architectures including HPS-V2, Pick Score, and Alpha CLIP. Figure 10 exhibits success with complex compositional tasks using CLIP rewards, where we effectively optimize multiple attributes (colors, positions, objects) while maintaining semantic coherence. The W2 regularization prevents collapse across all these reward functions through the tractable bound derived in Theorem 3, as particularly evident in our results on challenging tasks like spatial understanding (Figure 7), preventing policy collapse (Figure 8), multi-scale rewards (Figure 9) and multi-attribute control (Figure 10), where the method maintains both high rewards and generation diversity, when KL-divergence and W1-divergence are prohibited due to constraints mentioned above.

Q: Why choose W2 instead of other distance for regularizing (e.g., KL, W1)?

A: Our choice of W2 distance is both theoretically motivated and practically necessary to address policy collapse in online fine-tuning of flow matching models. Finding a computationally tractable divergence regularization in ODE-based flow matching models is non-trivial. While KL divergence is widely used for fine-tuning diffusion models and LLMs [3, 4], calculating KL divergence for flow matching models is infeasible, as it requires solving intricate transport dynamics and tracking probability flows across continuous time, which is computationally costly (i.e., exact likelihood computation requires expensive ODE simulation with Hutchinson traces estimator, detailed in Appendix B.2.4). Unlike diffusion models that get around this with variational bounds, there is no established relationship between vector field loss and ELBO in continuous-time ODE-based flow matching. Similarly, W1 distance faces intractability issues as the true marginal vector field is intractable (i.e., Theorem 2 of [7] does not hold for W1).

To address these challenges, we are the first to derive a computationally tractable upper bound for W2 distance in flow matching (Theorem 3). This bound only requires calculating the difference between vector fields, avoiding expensive likelihood calculations needed for KL divergence, thus enabling effective constraint of the discrepancy between fine-tuned and reference models while maintaining computational efficiency. Our theoretical analysis shows that this bound effectively prevents policy collapse while preserving the generation capacity of the pre-trained model.

Our experimental results strongly validate this choice, particularly in our comprehensive experiments with SD3 in Appendix A. Figure 8 provides detailed ablation studies demonstrating that our W2-regularized approach successfully prevents policy collapse while maintaining semantic accuracy for challenging prompts like "a cat in the sky", where baseline methods without regularization generate nearly identical images. The effectiveness of W2 regularization is further evidenced in Figure 7 for spatial relationship control and Figure 10 for complex compositional prompts, where our method achieves precise semantic control while maintaining natural variations in generation. These results demonstrate that our theoretically-derived W2 bound provides a practical and effective solution for flow matching fine-tuning that enables both high performance and stable learning.

Q: 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.

A: Following your kind suggestion, we have restructured the paper to improve accessibility while maintaining technical rigor. We have moved the "RL Perspective of Online Fine-Tuning Method" (previously Section 4.6) to the appendix, allowing us to dedicate more space in the main text to empirical evaluations that demonstrate the practical implications of our work. Specifically, we have significantly expanded Section 5.3 and Appendix A with comprehensive case studies on online fine-tuning of Stable Diffusion 3 (SD3), where Figures 5, 7 demonstrate superior performance in spatial relationship control compared to baselines like RAFT [1] and ReFT [2], Figures 6, 8 provide detailed ablation studies of policy collapse prevention, Figure 9 validates our method's adaptability across different reward architectures (HPS-V2, Pick Score, and Alpha CLIP), and Figure 10 showcases complex semantic control capabilities.

Our expanded experimental section demonstrates that our method not only has strong theoretical foundations (as shown by Theorem 3, 5) but also provides practical solutions to real challenges in fine-tuning large generative models. The addition of diverse case studies, from spatial understanding to multi-attribute control, makes our findings more accessible to readers while complementing our theoretical contributions. These comprehensive results, particularly with SD3, help readers better understand both the theoretical innovations and practical benefits of our approach compared to existing methods [1, 2].

References

[1] Dong, Hanze, et al. "Raft: Reward ranked finetuning for generative foundation model alignment." arXiv preprint arXiv:2304.06767 (2023).

[2] Huguet, Guillaume, et al. "Sequence-Augmented SE (3)-Flow Matching For Conditional Protein Backbone Generation." arXiv preprint arXiv:2405.20313 (2024).

[3] Rafailov, Rafael, et al. "Direct preference optimization: Your language model is secretly a reward model." Advances in Neural Information Processing Systems 36 (2024).

[4] Fan, Ying, et al. "DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models." Advances in Neural Information Processing Systems 36 (2024).

[5] Domingo-Enrich, Carles, et al. "Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control." arXiv preprint arXiv:2409.08861 (2024).

[6] Esser, Patrick, et al. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, March 2024." URL http://arxiv. org/abs/2403.03206.

[7] Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).

We deeply appreciate your thorough evaluation of our paper and the recognition of our theoretical contributions and empirical effectiveness. Thanks for noting that our method preserves the simplicity of the original flow-matching framework, making it easy to use while ensuring theoretical soundness. We are pleased to address each of your concern in detail:

Q: Why not compare with methods like DDPO and other recent baselines?

A: Thank you for the suggestion and we now provide additional results on text-to-image alignment tasks of Stable Diffusion3 (SD3) with comparison to more recent baselines. Meanwhile, we'd like to clarify several important perspectives regarding choice of baselines.

Technical Limitations of Existing RL Fine-tuning Methods for Flow Matching:

Most existing RL fine-tuning methods cannot be directly applied to fine-tune continuous-time ODE-based flow matching models due to intractable calculation of ELBO, KL divergence and expensive calculation of likelihood with Hutchinson trace estimator. Methods like DDPO [4] and DPOK [7] rely heavily on computationally tractable MDP trajectory likelihood calculations and ELBO approximations, which are computationally intractable or ill-defined for ODE-based flow matching models (detailed in Appendix B.2.4). Besides, the continuous-time nature of flow matching invalidates discrete timestep-based policy updates used in DDPO [4]. While discretizing ODE is possible, derivation of corresponding MDP trajectory likelihood (i.e., $\log p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{c}\right)$[4]) in continuous-normalizing-flow is non-trivial and itself could lead to a new research work that deviates significantly from what's proposed in DDPO [4]. Additionally, KL divergence [7] is also intractable in flow matching, necessitating our novel and tractable W2 regularization approach.

Regarding the concurrent work of Adjoint Matching [5], while theoretically interesting, it has certain practical limitations - it requires differentiable rewards, lacks theoretical convergence analysis, hasn't been open-sourced, and hasn't demonstrated effectiveness on SOTA FM architectures like SD3 [6]. It also employs a complex optimization procedure with stochastic optimal control. In contrast, our method works with arbitrary rewards, and preserves the simplicity of original CFM training loss, while providing both theoretical guarantees and strong empirical results on SD3.

Comprehensive Evaluation with Applicable Baselines:

Given these technical constraints, we focused our comparisons on reward-based methods that can be adapted to flow matching - specifically RAFT [1] and ReFT [2], which don't rely on likelihood calculations. Though extending these methods [1] [2] to flow matching models was far beyond the empirical scope/contributions of their original work, we implemented them for SD3 fine-tuning to provide fair and meaningful comparisons. To the best of our knowledge, we are the first to demonstrate successful online fine-tuning of SD3 models. Our comprehensive experiments show our method's advantages across multiple dimensions: superior handling of spatial relationships (Figure 5), strong adaptability across different reward models including HPS-V2, Pick Score, and Alpha CLIP (Figure 9), and successful management of complex semantic relationships (Figure 10). We also provide ablation results in Figure 6 to demonstrate unique contribution of ORW and W2 regularization.

Q: Why use W2 distance instead of KL divergence?

A: Our choice of W2 distance is both theoretically motivated and practically necessary for fine-tuning flow matching models. As proven in Lemma 1, online reward-weighted fine-tuning without regularization will inevitably collapse to a delta distribution centered at maximum reward, causing policy collapse. This necessitates effective regularization to maintain diversity.

However, finding a suitable and computationally tractable divergence regularization in ODE-based CFM is non-trivial. While KL divergence is widely used in previous works for fine-tuning diffusion models and LLMs [3, 7], calculating KL divergence for flow matching models requires solving intricate transport dynamics and tracking probability flows across continuous time, which is computationally intractable (detailed in Appendix B.2.4). Unlike diffusion models which can leverage variational bounds [4], continuous flow matching lacks tractable ELBO formulations, making KL-based approaches computationally infeasible.

We address this challenge by deriving a computationally tractable upper bound for W2 distance in flow matching (Theorem 3). To the best of our knowledge, we are the first to derive such a tractable W2 bound for online fine-tuning of flow matching models. Our experimental results strongly validate this choice - Figure 4 demonstrates how W2 regularization enables controlled divergence from the reference model while maintaining diversity. This is particularly evident in our extensive experiments with SD3 [6], where our method achieves better spatial relationship control (Figure 5) compared to RAFT [1] and ReFT [2]. The ablation studies in Figure 6 show that all methods without W2 regularization exhibit clear signs of policy collapse, generating nearly identical images. Additional experiments in Appendix A further demonstrate our method's effectiveness across different reward architectures (Figure 9) and complex compositional prompts (Figure 10), consistently maintaining both high-quality generation and semantic accuracy while preventing policy collapse.

Q: Regarding the dependency on number of epochs in the resulted distribution (Equation 12).

A: Equation 12 establishes a general mathematical framework for analyzing the convergent behavior of our proposed ORW-CFM-W2 method for different cases, and its limiting case depends on the configuration of key parameters ($\tau, \alpha$). To the best of our knowledge, we are the first to demonstrate the theoretical analysis of convergent behavior of online reward-weighted fine-tuning methods for CFM, applicable to both cases with W2 (Theorem 5) and without W2 regularization (Theorem 2 and Case 2 of Theorem 5) . Through both theoretical analysis (Theorem 5) and comprehensive experiments, we have demonstrated two important limiting cases, where the learned policy will converge to a specific distribution independent of the number of epochs:

1. When $\tau=0$ or $\alpha \to \infty$ (Case 1 in Theorem 5), the model maintains behavior similar to the reference model, as validated by Figure 2 ($\tau=0$);
2. Conversely, when $\alpha=0$ or $\tau \to \infty$ (Case 2 in Theorem 5), the distribution collapses to a delta distribution maximizing rewards, as shown in Figure 3 ($\alpha=0$) and Figure 5, 6 (without W2).

We further provide another interesting case in appendix (Case 6 in App. C.7) and prove the existence of the best trade-off parameter for balancing reward optimization term and divergence. Although not all limiting cases have closed-form solution, the practical significance of our theoretical framework is validated through experiments where the tradeoffs can be visually confirmed (Figures 2-10).

Q: Results on larger-scale models like StableDiffusion.

A: Thank you for your suggestions. We have now added extensive experiments on Stable Diffusion 3 (SD3) [6] in Section 5.3 and Appendix A to demonstrate our method's effectiveness on large-scale FM models. Our experiments show that ORW-CFM-W2 successfully fine-tunes SD3 on challenging spatial relationship prompts (Figure 5), achieving better positional relationship control and semantic understanding compared to RAFT [1], ReFT [2], and baseline SD3 [6] while maintaining image quality. The ablation studies in Figure 6 validate our theoretical predictions, demonstrating that W2 regularization effectively prevents policy collapse in large models - all methods without W2 regularization exhibit clear signs of policy collapse, generating nearly identical images. Even without W2 regularization, our base ORW-CFM method outperforms both RAFT and ReFT.

Additionally, our comprehensive experiments in Appendix A showcase the method's versatility and robustness. We demonstrate strong adaptability across different reward architectures (Figure 9) including HPS-V2, Pick Score, and Alpha Clip. The method also successfully handles complex compositional prompts (Figure 10) with multiple interleaved requirements spanning spatial relationships, color specifications, and object attributes. These results demonstrate that our theoretically-grounded approach not only scales effectively to large models but also enables stable fine-tuning while maintaining both generation quality and semantic accuracy.

Q: Notation suggestion for equation 10: "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}."

A: Thank you for this suggestion. We agree that using $\theta_{\text{ft}}$ and $\theta_{\text{ref}}$ would improve clarity and aesthetics. We have updated this in our paper.

References

[1] Dong, Hanze, et al. "Raft: Reward ranked finetuning for generative foundation model alignment." arXiv preprint arXiv:2304.06767 (2023).

[2] Huguet, Guillaume, et al. "Sequence-Augmented SE (3)-Flow Matching For Conditional Protein Backbone Generation." arXiv preprint arXiv:2405.20313 (2024).

[3] Rafailov, Rafael, et al. "Direct preference optimization: Your language model is secretly a reward model." Advances in Neural Information Processing Systems 36 (2024).

[4] Black, Kevin, et al. "Training diffusion models with reinforcement learning." arXiv preprint arXiv:2305.13301 (2023).

[5] Domingo-Enrich, Carles, et al. "Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control." arXiv preprint arXiv:2409.08861 (2024).

[6] Esser, Patrick, et al. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, March 2024." URL http://arxiv. org/abs/2403.03206.

[7] Fan, Ying, et al. "DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models." Advances in Neural Information Processing Systems 36 (2024).

Thank you very much for your kind response and your positive feedback. We’re happy to address your lingering concerns with the following:

Q: Quantitative results on large-scale experiments.

A: Thank you for your suggestion. We now include Table 1 (also put it below for ease of reading) to quantitatively evaluate our method's performance and effectiveness. Our evaluation focuses on two key metrics: CLIP scores to measure reward optimization (higher indicates better text-image alignment) and diversity scores computed as the mean pairwise distance between CLIP embeddings of generated samples (higher indicates more diverse outputs). We use Euclidean distance to calculate the distance between any two embeddings when calculating mean pairwise distance. As shown in Table 1, even without W2 regularization, our ORW-CFM method significantly outperforms baseline methods RAFT [1] and ReFT [2] in both CLIP scores and diversity metrics. The addition of W2 regularization ('+W2' variants) further helps preserve generation diversity across all methods, with our full approach (ORW-CFM-W2) achieving the best alignment score while maintain diversity comparable to the base SD3 model.

Regarding FID scores, we note that FID is usually used to measure fitness to a given training dataset distribution, and it is only applicable when evaluating models whose training set is available. However, SD3 does not disclose its training data and we are online fine-tuning method without ground-truth datasets, and therefore FID may not be a proper metric to use here (note that FID was not evaluated in previous online fine-tuning works like DPOK [7] /DDPO [4] or concurrent work like Adjoint Matching [5] either).

Q: About dependency on number of epochs and hyperparameter tuning.

A: Thank you for this insightful observation. Besides the theoretically analyzed limiting cases, equation (12) in Theorem 5 provides intuitive guidance for practical parameter selection, and we provide extensive experimental demonstrations of the effective controlling behavior for non-extreme value cases. Our analysis clearly shows how $\tau$ and $\alpha$ intuitively control the convergent behavior: $\tau$ determines the algorithm's preference for reward maximization (as shown in Figure 2's impact on policy collapse and reward optimization), while $\alpha$ maintains diversity and prevents collapse (demonstrated in Figures 3-4's reward-diversity trade-off). This understanding allows practitioners to adjust parameters purposefully rather than through random guesses.

Though deriving convergent distributions for all general cases is challenging, our empirical results demonstrate that our method achieves stable convergence and controllable fine-tuning across moderate hyper-parameter values. As shown in Figures 2-4, our method trains reliably to convergence without requiring early stopping, and practitioners can customize/adjust the convergent policy by adjusting $\tau$ and $\alpha$ according to their needs - increasing $\tau$ for stronger reward maximization or $\alpha$ for greater diversity. The convergence behavior is validated by extensive experiments across different tasks, from target image generation to text-image alignment, showing consistent and predictable behavior that follows our theoretical intuition.

Q: The "ad hoc" issue. "Use of W2 distance is ad-hoc, and indeed the authors show that W2 distance works well even if we use ReFL…That's said, the fact that the W2 regularization is a bit ad-hoc does not greatly compromise the contribution of this paper. But I would like to suggest that the authors lower their tone, instead of claiming that W2 is something very tied to the proposed online reward matching method."

A: Thank you for your insightful feedback. As we pointed out in our response to Reviewer dXDz, we have not claimed that W2 should be tied to ORW-CFM. Instead, we have written the paper with the goal of introducing two separate contributions: ORW-CFM as an effective online reward optimizer for CFM, and W2 as an effective and general regularizer to prevent policy collapse in online fine-tuning of flow matching models.

This separation is more evident from our independent introduction of ORW-CFM (see Abstract lines 20–21; Introduction lines 107–111; and theoretical results in Section 4.3 and App. C.3, which show the learning/convergent behavior for ORW-CFM alone) and W2 (see Abstract lines 22–23; Introduction lines 112–117; and theoretical results in Section 4.4 and App. C.4, which show the tractable W2 bound for flow matching alone). If there are parts of our writing that didn't make this separation of contributions clear, we're happy to make adjustments to further clarify it.

Our ablation studies showed that ORW-CFM by itself outperforms the baselines (Figure 6, Table 1), and that W2 is effective in preventing policy collapse in online fine-tuning. Moreover, the combination of the two creates a synergy that leads to the optimal performance of ORW-CFM-W2. Therefore, we believe that our demonstration of the strength of both components, especially the general effectiveness of W2 regularization, and their synergistic advantage, should enhance the standing of our paper (i.e., strengthen the significance of our methods), rather than "downplaying" it.

Q: On the baseline issue.

A: Thank you for suggesting the possibility of implementing MDP-based RL method as future directions for camera-ready. As discussed in our first response, we believe the conversion of deterministic neural ODE to stochastic MDP while preserving the original probability path is a non-trivial topic, and categorizing it as a "naive baseline" might be a bit over-reaching, as the scope of which could lead to a new research study in our opinion. We're also not sure if such modification will still lead to a valid continuous normalizing flow and whether we could still consider it an FM fine-tuning method. We'll try our best to include discussions regarding the feasibility of such a method in our final version.

References

[1] Dong, Hanze, et al. "Raft: Reward ranked finetuning for generative foundation model alignment." arXiv preprint arXiv:2304.06767 (2023).

[2] Huguet, Guillaume, et al. "Sequence-Augmented SE (3)-Flow Matching For Conditional Protein Backbone Generation." arXiv preprint arXiv:2405.20313 (2024).

[3] Rafailov, Rafael, et al. "Direct preference optimization: Your language model is secretly a reward model." Advances in Neural Information Processing Systems 36 (2024).

[4] Black, Kevin, et al. "Training diffusion models with reinforcement learning." arXiv preprint arXiv:2305.13301 (2023).

[5] Domingo-Enrich, Carles, et al. "Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control." arXiv preprint arXiv:2409.08861 (2024).

[6] Esser, Patrick, et al. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, March 2024." URL http://arxiv. org/abs/2403.03206.

[7] Fan, Ying, et al. "DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models." Advances in Neural Information Processing Systems 36 (2024).

Q: How about other discrepancy measures? Why is Wasserstein distance a good choice?

A: Our choice of W2 distance is theoretically motivated and necessary. While KL divergence is widely used in other methods [3] [7], KL is computationally intractable for continuous-time ODE-based flow matching models (Appendix B.2.4) and there are no tractable ELBO alternatives either. To overcome this, for the first time, we derive a tractable upper bound for W2 distance (Theorem 3) in online fine-tuning of flow matching that enables practical regularization via vector field loss directly.

Our experiments strongly validate this choice - as shown in Figure 4, varying the W2 regularization coefficient $\alpha$ provides explicit control over the trade-off between reward maximization and diversity. The reward-distance curves demonstrate that as $\alpha$ increases, our method explores optimal solutions within a constrained neighborhood of the pre-trained model, preserving diversity while still optimizing the reward. Furthermore, our ablation studies in Figure 6 show that methods without W2 regularization (RAFT [1], ReFT [2]) exhibit clear policy collapse, while our approach with W2 regularization successfully maintains generation diversity without sacrificing semantic alignment. The results across various reward models in Figure 9 further demonstrate how our tractable W2 bound enables stable fine-tuning regardless of the underlying reward mechanism.

References

[1] Dong, Hanze, et al. "Raft: Reward ranked finetuning for generative foundation model alignment." arXiv preprint arXiv:2304.06767 (2023).

[2] Huguet, Guillaume, et al. "Sequence-Augmented SE (3)-Flow Matching For Conditional Protein Backbone Generation." arXiv preprint arXiv:2405.20313 (2024).

[3] Rafailov, Rafael, et al. "Direct preference optimization: Your language model is secretly a reward model." Advances in Neural Information Processing Systems 36 (2024).

[4] Black, Kevin, et al. "Training diffusion models with reinforcement learning." arXiv preprint arXiv:2305.13301 (2023).

[5] Domingo-Enrich, Carles, et al. "Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control." arXiv preprint arXiv:2409.08861 (2024).

[6] Esser, Patrick, et al. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, March 2024." URL http://arxiv. org/abs/2403.03206.

[7] Fan, Ying, et al. "DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models." Advances in Neural Information Processing Systems 36 (2024).

[8] Xu, Jiazheng, et al. "Imagereward: Learning and evaluating human preferences for text-to-image generation." Advances in Neural Information Processing Systems 36 (2024).

We sincerely thank you for your insightful review and recognition of our empirical and theoretical contribution. We are happy to address each of your concern with the following:

Q: Why should reward-weighted matching loss and Wasserstein regularization be used together (other than empirically well-performing)? What's the theoretical connection? The contribution of the paper seems ad-hoc.

A: Thank you for your recognition of our empirical result. We'd like to clarify that our contributions include introducing and theoretically grounding both online reward-weighted CFM as an effective RL loss and W2-regularizer as an effective divergence/diversity control mechanism for flow matching (none of them is ad-hoc, but theoretically motivated and necessary). To the best of our knowledge, both components are novel and non-trivial (detailed below) in flow-matching online fine-tuning tasks. We do not claim that they should be always used together, but their combination is theoretically motivated, and we derived a clean and interpretable convergent distribution for ORW-CFM-W2 (Theorem 5) which demonstrates the tradeoffs between optimization and diversity controlled by key coefficients $\tau$ and $\alpha$ (i.e., controllable and interpretable convergent/learning behavior as Figures 2-4). We has provided extensive experiment results to illustrate this tradeoff (Figures 2-4), and now provide further ablation studies to show the effectiveness of both components in achieving superior online fine-tuning performance on the text-to-image alignment tasks of SD3 (Figure 5-6).

Motivation and Significance of Online-Reward-Weighting (ORW-CFM). Obtaining theoretically guaranteed online RL fine-tuning methods for continuous-time ODE-based flow matching models is non-trivial due to the intractable ELBO and costly exact likelihood calculation in FM. Our online reward-weighted CFM loss is not a simple application of RL algorithms that leverage likelihood/ELBO, since the equivalence of CFM loss and ELBO cannot be established easily as in diffusion, and thus theoretical guarantees are non-trivial and not granted for free. We do provide novel theoretical analysis of ORW-CFM's convergent behavior (Theorems 2, 5), optimal convergence (Lemma 1), and its equivalence to the optimal policy as KL-regularized RL (Appendix C.9), all without relying on explicit likelihood/ELBO. We prove that although $\tau$ can control the convergent speed, ORW alone may eventually collapse to delta distribution at optimal reward (Lemma 1), motivating the use of W2-regularizer to further prevent policy collapse.

Motivation and Necessity for W2 Regularization. Finding a tractable divergence regularization in flow matching is non-trivial since the intractable ELBO and costly exact likelihood computation make KL-divergence impractical (See Appendix B.2.4). To address this fundamental challenge, we derive a tractable upper bound for W2 distance (Theorem 3) that enables tractable and efficient divergence regularization of flow matching via vector field loss directly. The W2 regularization is not an independent add-on but a theoretically motivated solution to the policy collapse problem inherent in our online reward-weighted approach for ODE-based flow matching.

Motivation and Necessity for Combining ORW-CFM with W2. The W2 regularization is a theoretically motivated solution to the policy collapse phenomenon in our ORW-CFM, which is evident by Theorem 5-our key theoretical result on the convergent distribution after combining ORW and W2. We show that their combo nicely results in a Boltzmann distribution whose energy is controlled by both coefficients from ORW and W2, balancing reward optimization and divergence regularization in an explicit and interpretable manner (Theorem 5). We provided extensive experimental results on how controlling $\tau$ and $\alpha$ leads to different tradeoffs (figure 2-4), and now include more benchmarks and ablations on Stable Diffusion3 (Figures 5-10) to show the effectiveness of ORW and W2-regularizer independently, as well as combinatorially.

In short, our online reward weighting enables us to achieve theoretically interpretable and controllable learning/convergent behavior (Theorem 5, Figures 2-6), optimal convergence (Lemma 1, Figure 2-6), ELBO/likelihood-free online RL fine-tuning (Theorems 5), and equivalent optimal policy with KL-regularized RL (App. C.9) - while preserving the simplicity of original CFM loss and continuous-time ODE property, making it easily integratible with any existing FM works like TorchCFM and SD3 (Sec. 5.3 and App. A). Additionally, our tractable W2 regularization (Theorem 3) effectively handles policy collapse (Lemma 1), allowing convergence to optimal policies that preserve diversity without collapsing (Theorem 5, Figure 6).

Q: How about comparison with more stronger baselines and how does the method scale to larger models/datasets? Whether there are other stronger baselines to compare to?

A: We now provide extensive text-to-image alignment experiments on state-of-the-art large-scale flow matching model Stable Diffusion 3 [6] (Section 5.3, Appendix A) to demonstrate both our superiority over baselines and excellent scalability. We emphasize that many existing RL fine-tuning methods (e.g., DDPO [4]) cannot be directly applied to flow matching models due to intractable calculation of ELBO, KL divergence and costly likelihood computation in ODE-based continuous-time flow models. Concurrent work like Adjoint Matching [5] requires differentiable rewards limiting its applicability, and has not provided open-source implementation. Nevertheless, we implement and extend two popular reward-based methods that don't rely on likelihood calculations, RAFT [1] and ReFT [2], as our baselines. RAFT has been used for diffusion [1] and ReFT has been adopted in aligning flow-matching models for protein structures [2]. We are, to our knowledge, the first to establish successful online fine-tuning results on SD3 with these baselines and our ORW-CFM-W2 (noting that ReFL of Image Reward [8] is an offline RL method that rely on offline datasets and has not yet been applied to flow matching online fine-tuning).

Our comprehensive experiments demonstrate superior performance across multiple tasks - Figure 5 shows our superior performance on spatial relationship control compared to baselines (RAFT[1] and ReFT [2]), Figure 6 provides ablations showcasing the effectiveness of ORW/W2 independently and the global superiority of combining them, Figure 9 validates strong adaptability across multiple reward models (HPS-V2, Pick Score, Alpha CLIP), and Figure 10 showcases successful handling of complex semantic relationships. Our method uniquely provides a theoretically sound approach for fine-tuning flow matching models with arbitrary reward functions while preventing policy collapse through W2 regularization. These experiments demonstrate that our approach scales effectively to state-of-the-art models while maintaining computational efficiency, handling complex tasks from spatial control to compositional generation while preserving diversity and quality.

Q: How much gain does each component contribute? Where are the ablation studies?

A: We now do provide comprehensive ablation studies through a series of controlled experiments on both small-scale and large-scale models that demonstrate the necessity and contribution of each component.

For online reward-weighting (ORW), Figure 2 quantitatively demonstrates how the entropy coefficient $\tau$ affects convergence behavior - as $\tau$ increases, the policy optimizes reward more aggressively but at the cost of diversity. We also show in Figure 6 that even without W2 regularization, our ORW-CFM method achieves the best semantic alignment compared to other baselines like RAFT [1] and ReFT [2], validating the effectiveness of our online reward-weighting mechanism. However, consistent with our theoretical prediction in Lemma 1, online fine-tuning methods without regularization exhibit clear policy collapse. This necessitates tractable regularization to maintain diversity.

For W2 regularization, Figure 3 demonstrates how the coefficient $\alpha$ controls the diversity-reward trade-off - from $\alpha=0$ to $\alpha=0.8$, the diversity of generated samples increases without significantly compromising performance. Figure 4 provides quantitative trade-off curves showing how $\alpha$ enables explicit control over the balance between reward maximization and divergence from the pre-trained model. Our ablation studies in Figure 6 further validate this - methods incorporating W2 regularization (RAFT+W2, ReFT+W2, and our ORW-CFM-W2) successfully prevent collapse while maintaining high-quality generation, validating the effectiveness of our W2 regularization , wherein our combined approach achieving the best balance between semantic alignment and diversity preservation. The results empirically validate both the individual contributions of ORW and W2 regularization as well as their synergistic effects in enabling efficient online fine-tuning of flow matching models.

Thanks for the suggestion of Reviewer tfNW, we have now added comprehensive quantitative evaluation in Table 1 (we also put it below for ease of reading) that systematically compares our method's performance on SD3. Our evaluation focuses on two key metrics: CLIP scores to measure reward optimization (higher indicates better text-image alignment) and diversity scores computed as the mean pairwise distance between CLIP embeddings of generated samples (higher indicates more diverse outputs). Using these metrics, we demonstrate significant advantages over baselines. Our ORW-CFM-W2 approach achieves the highest CLIP score while maintaining strong diversity comparable to the base SD3 model. Notably, even without W2 regularization, our ORW-CFM method outperforms both RAFT and ReFT on these metrics. The addition of W2 regularization helps preserve generation diversity across all methods, with our combined approach striking the best balance between alignment quality and output diversity.

We believe our quantitative results effectively address the comments regarding larger scale experiments and further validate our method's effectiveness at larger scales, and we thank the reviewers for their constructive suggestions.

We extend our deepest gratitude to all reviewers for their invaluable time and effort throughout the review and discussion process. We are encouraged that our responses, revisions, and additional empirical validation effectively addressed all concerns, leading to consistently positive assessments of our work. The recognition of both our theoretical foundations and practical contributions has been especially encouraging, with Reviewer dXDz noting our "reasonable theoretical justifications" and "good experimental results," Reviewer tfNW emphasizing that our work is "definitely valuable to the community" and recognizing our "sufficient small-scale experiments," and Reviewer S526 affirming that "this paper is definitely valuable" while highlighting our successful handling of "policy collapse and high computational costs."

Following the suggestions of all reviewers (Reviewers dXDz, tfNW, S526), we expanded our empirical validation by extending our method to online fine-tuning of large-scale flow matching models like SD3. We further appreciate the constructive and fruitful discussions with Reviewer tfNW, leading to enhanced clarity through additional qualitative and quantitative experiments with recent baselines. Once again, we sincerely thank the reviewers for their dedicated engagement, which has significantly improved the quality and clarity of our work.