Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization
Authors: Jiajun Fan
Shuaike Shen
Chaoran Cheng
Yuxin Chen
Chumeng Liang
Ge Liu
Venue: ICLR 2025 Poster (2025)
Keywords: Flow Matching
Reinforcement Learning
Wasserstein Regularization
Exploration-Exploitation Trade-off
Fine-Tuning
Generative Model
Primary area: reinforcement learning
TL;DR: We propose a novel and theoretically sound method for fine-tuning flow matching generative models using a reward-weighted mechanism and Wasserstein-2 regularization to optimize user-defined rewards while preventing overoptimization.

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.

Published as a conference paper at ICLR 2025
ONLINE REWARD-WEIGHTER MATCHING WITH WASSERSTEIN REGULARIZATION
OF FLOW
Jiajun Fan' Shuaike Shen?i, Chaoran Cheng', Yuxin Chen Chumeng Liang"1 Ge Liu'
1 University of Illinois Urbana-Champaign, 2 *Zhejiang University.
3 Tsinghua University, 4 University of Southern California.
{ fjiajunf,, chaoran7, geliuj @ @illinois.ee.. edu
2,3,4 (shenshuaike256, itsme. yxchen, caradryan2022 (atprraaann22222mm1112 com
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 func- tions remains challenging, particularly due to issues such as policy collapse from overoptimization and the prohibitively high computational cost of likelihoods continuous-time flows. In this paper, we propose easy-to-use and theoretically sound RL fine-tuning method, which we term Online Reward- Weighted Condi- Our
tional Flow Matching with W Regularization method integrates RL into the flow matching framework to fine-tune generative models with arbitrary reward functions, without relying gradients of rewards or filtered datasets. By introducing an online reward-weightin 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-Leiber (KL) regularization and offering 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 conver- gence while allowing controllable trade-offs between reward maximization and diversity preservation.
mechanism, our
INTRODUCTION
Generative models have achieved remarkable success in producing high-fidelity data across various domains, including text and images (Ouyang et al.) 2022a Esser et al.) 2024). Among these, continuous flow-based models have gained attention for their more concise and flexible design of
the ODE-based denoising process to model complex data distributions (Tong et al. 2222aa] Lipman ct al. 2023). However, fine-tuning such flow-based models to align with arbitrary user-defined reward objectives remains challenging. While recent advancements in reinforcement learning (RL) have demonstrated considerable success in fine-tuning diffusion based generative models (Rafailoy al. 2023). their application to continuous flow-based models remains underexplored. The primary challenges lie in lacking computationally efficient divergence measurement and tractable methods to avoid overoptimization (Black et al. 2024) and policy collapse in fine-tuning continuous time flow matching models. Particularly, traditional policy gradient methods (Black et al.) 2024) struggle in this continuous domain due to the high computation cost of computing the exact ODE likelihood for continuous normalizing flows (CNFs), which requires solving intricate transport dynamics and tracking probability flows with model divergence calculation across time (See App. B.2.4.
"Corresponding authors.
Work done research intems University of Illinois Urhama-Champaign.

---

Published as a conference paper at ICLR 2025
Furthermore, existing diffusion-basse based RLHF methods, such DPO (Rafalov et al. 2023) and DDPO (Black et al. 2024), face sigmificant limitations when applied to CNFs. DPO relies on pairwise reward comparisons and filtered datasets, making it unsuitable for optimizing arbitrary reward functions. DDPO adopts a PPO-style policy optimization that depends on frequent likelihood calculations, which computationally intractable for CNFs. Moreover, methods for diffusion often rely tractable ELBO likelihood approximation or MDP formulation, which does not apply to CNFs. To tackle these challenges, we propose a novel RL fine-tuning framework that bypasses the need for explicit likelihood estimation by focusing on wwarrd-wiihhte objectives tailored specifically for continuous flow-based models, which we call Reward- Weighted Conditional Flow Matching (RW-CFM). Our approach significantly reduces computational overhead while maintaining optimal performance.
In the meanwhile, many existing reight-weighted RI. fine-tuning methods require train on offline datasets manually collected by hu- mans (Huguet et al.) 2024) or derived from pre-tree tre-trained models (Lee et al.] 2023). While this approach ensures stable learning, restricts the model's ability to explore optimally in an online setting, inducing what is known as the online-offline gap (Tang et al. (024) (See Fig. [ad Figs. 2). This limitation reduces the model's flexibility and can lead to sub-optimal convergence. In contrast, online RL enables the generative model to continuously update its training distribution, allowing for greater adaptability and diverse exploration within the reward space. Despite the success in traditional RL domains, the convergence behavior of online RL in fine-tuning generative mod- els, particularly continuous flow-based models, has not yet been widely studied. To leverage the advantages of online RL for generative models, we propose online seward-weighter Conditional Flow Matching (ORW-CFM), In this paper,
The TaadEFM
Sampling (ty
Dountry Conrel
Chille kovari Warhed Fax- Tuaking
Gib
1ib
ClIP
Choite
Ovengrosies
Figure 1: A General Architecture of Our Method.
fine-tuning
method called Online theoretically demonstrate that in online reward-weighted settings, the continuous updates of the training data distribution of the fine-tuned generative model based on reward feedback can cause the learned policy collapse into a delta distribution (i.e.. becoming extremely exploitative; see Lemma Th. Similar phenomena have been experimentally observed in mode collapse of GAN (Goodfellow jtt al. 2014)) and over-optimization in diffusion models (Black et al. 2024). Nevertheless, this issue has not yet been fully addressed theoretically demonstrated in previous works (Black et al., 2024).
To address these challenges and avoid policy collapse, we introduce a divergence regularization approach to bound the distance between the fine-tuned model and the pre-trained reference model into our online RL method. This regularization ensures that the policy updates remain within a controllable and stable range, preventing collapse while still allowing the model to explore and improve. One of the major challenges in fine-tuning CNFs lies in finding a computationally tractable distance measure between different flow matching models, especially considering the KI. divergence used in classic RL methods is inefficient in CNFs. To overcome this, we instead propose to use the Wasserstein-2 (W2) distance (Arjovsky et al.) 2017) measure the distance between the fine-tuned model and the reference model. In this paper, we derive a tractable upper bound of W2 distance in flow matching seward-distanee trade-off
(See Theorem 3) and incorporate into our online RL method, inducing in fine- fine-tuning flow matching models. We demonstrate both theoretically and empirically that W2 distance regularization allows for controllable and stable fine-tuning process, enabling our agent to balance the reward objective and generative diversity.
In summary, for the first time, we propose an online RL. method to fine-tune continuous flow-based generative models towards arbitrary reward objectives with sufficient theoretical guarantees and detailed analysis of the convergence behavior in both online and offline settings. Our approach can leverage classifier-uuided reward models, compression rate, and text-image similarity score with CLIP models (Radford et al. 2021b), eliminating the need for cine-con likelihood calculations that involve the integration of the model divergence across the entire continuous trajectory. Our novel
2

---

Published as a conference paper at ICLR 2025
framework significantly reduces computational costs and enhances the feasibility of using RL to fine-tune continuous flow-based generative models. The main contributions of our paper are:
We propose the Online Reward- Weighted Conditional Flow Matching (ORW-CFM) method, novel RL approach for fine-tuning continuous flow-based generative models. Our method enables optimization with user-defined reward functions without the need for likelihood calculation (Black et al.. 2024), filtered dataset (Rafailov et al., (202)) or differentiable reward (Clark et al.) (024). expanding the applicability of flow-based models various downstream generative tasks.
A Novel RL-Based Fine- Tenn-TTning Framework for CNFs:
Tractable Wasserstein Regularization: To prevent policy collapse and maintain the di- versity of the fine-tuned models, we introduce Wasserstein-2 (W2) distance regularization into online method, and derive tractable upper bound of it in flow matching. This regularization effectively balances exploration and exploitation by bounding the distance between the fine-tuned model and the pre-trate pre-trained reference model, inducing a theoretically and empirically controllable convergent behavior of fine-tuning flow matching models.
Theoretical Analysis and RL Connections: We provide comprehensive theoretical anal- yses of our method, including convergence properties and the induced data distributions. We establish connections between our approach and traditional RL algorithms with KL regularization, offering new insights into fine-tuning flow matching generative models.
Through extensive experiments on tasks such as target image generation, image compression, and text-image alignment, we empirically demonstrate the effectiveness of our method. Our approach achieves optimal policy convergence, allowing controllable trade-offss between reward maximization and generative diversity.
Empirical Validation:
2
RELATED WORKS
Conditional Flow Matching Conditional Flow Matching (CFM) (Lipman et al. 2023 Tong et al., 2024a) and related methods have been pivotal transforming simple data distributions into complex, task-specific ones. [ipman et al. (2023) and Tong al. (2024a) focus training flow-based models by conditioning the generative flow on known distributions, allowing for more precise control over the generative process. However, while these models achieve great success in generation tasks, how to fine-tune FM models to fit arbitrary user-defined objectives has not yet been widely studied.
Fine-tuning From Human Feedback Recent works typically leverage DPO (Rafallov al. 2023) or policy gradient RI. methods (Black et al. 2024 Fan et al. 2023) to achieve RLHF (Ouyang al. 2022a).. However, since the likelihood of the generative policy cannot be easily calculated in continuous-time flow models, none of these approaches can be easily adapted to fine-tune flow matching models. Additionally, DPO (Rafailov et al. 2023) relies heavily on a filtered dataset, which does not work with general reward model, while DPOK (Fan al. 2023) relies on KL divergence, which is computationally inefficient and lacks effective evidence lower bound (ELBO) in FM.
Fine-Tuning Generative Models with RL RL has been increasingly adopted fine-tune generative models, allowing them to adapt to specific downstream tasks by optimizing for user-defined reward objectives. Black al. (2024); Fan et al. (2023); Lee et al. (2023) used RL to fine-tune diffusion models to align generated data with task-specift rewards. Lee et al. (2023) proposed an offline rewght-weighted method, while Black et al. (2024) proposed an online version without divergence regularization, which either failed to converge into optimal policy or collapsed into an extremely greedy policy without diversity (Black et al. 2024). Other reward-weine -wighttee methods like ReFT, almost limited to the offline settings with given datasets and lack theoretical analysis of the convergent behavior. Importantly, how to adopt online RI. methods to fine-tune FM models has not yet been widely studied in both theory and practice.
ReST (Huguet et al. 22244) Gülçchree et al. 2023)
3

---

Published as a conference paper at ICLR 2025
3 PRELIMINARIES
3.1 FLOW MATCHING MODELS
Flow Matching (FM) is promising technique for training generative models by transforming simple initial distribution p (zo) like Gaussian noises, into a complex target distribution p (z1). which represents real-world data (Esser et al. 2024 This transformation works continuously over time, governed by a time-dependent vector field ue (2)) and the flow ordinary differential equation (ODE): de = uy (x). The goal of FM is to learn a vector field that transports samples from the base zt
distribution the target distribution by minimizing the difference between the learned vector field and the true dynamics of the probability flow. This process can be formalized by the Flow Matching loss: CFM(0) (0) 10=22(0.1),,=ppp.. (2)]l] (o x) (=)) However, the original FM objective is generally intractable because the time-varying distribution pr (x) and the true vector field u,(xr) are often intractable. To address this, Conditional Flow Matching (CFM) has been proposed (Lipman al. 2023). which simplifies the task by conditioning the flow on target samples to derive a tractable objective. The CFM loss is defined as follows:
(CFM(()
(t, z) - us (x x1]ll
wherein uy (x 211) becomes tractable by defining explicit conditional probability paths from xo to 3i, such as OT-paths (Tong et al. 2024a) or linear interpolation paths (Lipman et al. (023). this
paper, we adopt OT-paths for all experiments based on TorchCFM (Tong et al. 2024b) and Diffusers (von Platen et al. 2022).
3.2 REWARD WEIGHITE REGRESSION
Reward Weighted Regression (RWR) (Peters & Schaal 2007 Black al. 2024) is approach used in RL for policy optimization. In this framework, actions that yield higher rewards are assigned greater importance, and the policy is updated by re-weighting actions based on the rewards they produce. The key idea behind RWR is to adjust the probability of selecting actions proportionally to their associated rewards, encouraging the agent focus on high-reward actions.
Specifically, the policy x(a | s) is updated by re-weighting the actions via Tnew (a s) оx Tota s) exp(T $ r(s, a)), where r(s, a) is the reward for taking action state s, and is temperature parameter that controls the trade-off between exploration and exploitation. Higher rewards lead to higher probabilities for the corresponding actions, allowing the agent improve its policy over time by favoring actions that maximize cumulative rewards: x arg max, Ex 1E "r((s,, at)].
METHODS
4.1 PROBLEM STATEMENT
In this paper, we aim to fine-tune continuous flow-based models using a reward-weighted RL method (Peters& & Schaal, 2007 Peng al. (019) optimize user-defined reward function r (x1). Specifically, our objective is to find the optimal generative policy model, denoted as T²,, which maximizes the expected reward throughout data generation:
arg max ((T) = arg max (ay [r (x1)]
*
where the policy и refers to the parameterized flow matching data generative model p" (æ1). deter- mining how data points (e.g., images) are generated, which is learned solely from scalar rewards, without ground-truth labels or filtered dataset. In fine-tuning tasks, we may also consider constrained optimization, where the distance between the pre-trained model and fine-tuned model will be con- strained (Ouyang et al. 2022a). Considering the computational inefficiency of KL in continuous flow models, we instead propose Wasserstein-2 distance to achieve that.
4.2 FINE-TUNING FLOW MATCHING MODELS WITH OFFLINE REWARD- WEIRD-WIIHTING
To find the optimal data generation policy in equation 2 in this section, we extend the standard CFM framework by introducing rewird-weighting mechanism focus the model's learning on more
4

---

Published as a conference paper at ICLR 2025
important or highly rewarded regions of the data space, which thus improve its performance towards user-defined reward objectives r(z1).
Let X [ Rd be the data space, and let g(x1) be the probability density function (pdf) of the random variable 21 E X. We consider continuous-time flow over t E [0. 1|, where t is sampled from the uniform distribution O( 1). The conditional pdf of time t, given 21, is denoted by px (1). The true conditional vector field is u,(x xx and the model's 's learned vector field is 2, (z 0). where 0 represents the model parameters. In CFM (Tong al. 2024a), the goal is to learn (z 0) such that it closely approximates (x (1)) with equal importance via equation | Based on that, we introduce weighting function w(x1) 7(x1) : X + [0, oo) into CFM loss, inducing reward-weighted CFM method (RW-CFM) as equation 3 Then, we can have the Theorem (Proof in App. C.2
Theorem 1. Let w X + 00, xc) be measurable weighting function such that O < 2 Sx w (x1) q (x1) dx, oc and t((11) r(æ1)- Define the reward weighted CFM loss function:
CRW-CFM(0) = an (( [ (21) [1 us (2c 21)11²]
(3)
where w (x1) is weighting function, (I1) is the original data distribution (e.g., data distribution induced by well-learned pre-trained models), pt. (x x1) is the conditional distribution at time 1, (er;;; 0) is the model's 's vector field, and us (ar xn) is the true vector field conditioned on Then, training the model using this RW-CFM loss leads the model to learn a new data distribution:
nnow (x1) w (11)
(4)
Z
where Z = [w (x1) q (x1) dx is the normalization constant.
obvious that the model trained with the RW-CFM loss will focus on high-reward regions compared to the original data distribution (. However, since the training data is sampled from a static dataset (21). which is equivalent to an offline RL setting (Black et al. 2024), the model may not find the optimal data distribution p" (i.e., optimal policy x²) that maximizes the expected reward (i.e., online and offline gap, see Fig- (2. To address this, we propose using the fine-tuned flow matching models as the data sampling policy (i.e., g(z1) = p² (z))),, enabling an online generation of training data and thereby inducing an online RW-CFM method.
Based on Theorem
4.3 FINE-TUNING FLOW MATCHING MODELS WITH ONLINE REWARD- WIGHTING
As discussed, mitigate the online-offline gap and achieve better performance, we can instead adopt an online RL algorithm, where the fine-tuned FM models are used as data generation/sampling policy to continuously generate new training data and explore high-value regions, inducing what we call the Online Reward-Weighted Conditional Flow Matching (ORW-CFM) method:
CORW-CFM (0) = (0.])=/1ppp (11) (xsl)))) w (z1) [1: 0) - u (a ==))))]]
(5)
In the online setting, the model can adaptively sample from the current learned data distribution, and fine-tune itself based on the newly generated samples, inducing an online policy iteration that
promotes more efficient exploration of the high-reward areas in the data space. In general, the data distribution learned by the ORW-CFM method follows Theorem 2 (Proof in App. C.3):
Theorem 2 (Online Reward Weighted CFM), Let g (x1) be the initial data distribution, and w (xx) ) be a non-negative weighting functton integrable with respect to q (x1), and propertional to user-definer reward r ((x) Assume that at each epoch n, the model vo perfectly learns the distribution implied by the CRWWCFFM loss in equation 5
Then, the learned data distribution after N epochs is given by:
W (x1)) (x1)
T (xx)
(6)
ZN
where ZN = Sx (21) 9 (x)) dx, is the normalization constant ensuring IN (x1) ) is a valid probability distribution.
5

---

Published as a conference paper at ICLR 2025
Based on Theorem 2 easily derive the limiting behavior of ORW- CFM- method (Proof in App. C.3) as the training converges (i.e., the learned generative policy collapses into an extremely explottativ//rredyy policy without diversity the case where 0 in Fig. 3):
Lemma (Limiting Case). We now consider the limiting cases where N 00:
w (21) & r(z1) for all X and attains its maximum rewards 21, then N + oc, the learned data distribution qø (z1) converges to a Dirac delta function centered at x
lim T (x1) = 8 6 - x;)) N
From Lemma T in the online setting, training progresses, the ORW-CFM method may converge the optimal policy that yields maximum reward without any diversity, potentially leading to a policy collapse problem (See Figs. 3 6 and I1 where a 0):
Corollary (Policy to a lack of diversity in generated samples, similar to mode collapse in GANs (Goodfellow al. 2014), overoptimization in fone--utuing diffusion models (Black et al. 2024 or policy collapse in RL
Overemphasizing high-reward regions can lead
Haarmoja et 2018
Sutton & Barto 1998).
Based on Lemma and Theorem assuming an exponential weighting function w(x1) =
(z1)** ) exp(N 7 * r(z1)) where can control the conver- gence/collapse speed learned policy, and r + co also induces a policy collapse problem (See introduce Wasserstein-2 regularization into our method, which can bound the distance between the current learned model and the pre-trained model. This helps maintain diversity and prevents the model from collapsing into extremely greedy policy, inducing an exploration-exploitatiee trade-off in fine-tuning flow matching models.
exp(T * r(x1)) case, we can have (21)
Fig. 2). To mitigate the policy collapse problem,
4.4 WASSERSTEIN-2 REGULARIZATIO IN FLOW MATCHING
To address the policy collapse problem in the ORW-CFM method, we introduce Wasserstein-2 (W2) regularization, which helps maintain diversity by bounding the distance between the current learned model and a pre-trained reference model. This regularization induces an trade-off, allowing the model to avoid focusing solely on high-reward regions at the expense of diversity. general, we can write the Wasserstein-2 Distance as follows (Arjovsky et al. 2017): Definition 4.1 2 Distance), Given two probability measures and on R., the squared Wasserstein- 2 distance between μ and v is defined as:
inf f W (H, rEIl(u.)) JRExR"
| v1l
(8)
where M( v) denotes the set of all couplings of p and v; that is, all joint distributions 7y on IR" x IR with marginals μ and v.
However, it is intractable to directly introduce W2 distance into our loss function. Therefore, we instead derive and introduce its upper bound as our distance regularizer (Proof in App. C.4.
Theorem (W2 Bound for Flow Matching). We consider two flow matching models parameterized by 02 and Oz inducing time-evolving data distributions pt (x) and f (xx), respectively. The models define vector fields eth (t, x) and o (t, x) that transport an initial distribution po to final distri- butions pe: (x:) and p (1:) at time t 1. Assume that @ x) is Lipschitz continuous in x with Lipsehitz constant Then, the squared Wasserstein-2 distance between the distributions pht and p induced by the flow matching models at time I. = 1 satisfies:
W3 PT pT 2L J." E Egpp? [1100 -00,,))) (s) ds
Since the integral can be approximated by Monte Carlo sampling, this bound allows us constrain the discrepancy between the learned and reference models, effectively preventing the model from collapsing into greedy policy. By controlling the Wasserstein-2 distance, the model maintains balance between exploring new areas of the data space and exploiting known high-reward regions in constrained neighhorhood from the pre-trained model, thus preserving diversity.

---

Published as a conference paper at ICLR 2025
4.5 FINE-TUNING FLOW MATCHING MODELS WITH W2 REGULARIZATION
The offline version of the rewght-weighted method can be considered as the first step of the online version (Black et al. 2024). For simplicity, we only derive the ORW-CFM method with W2 distance regularization (i.e., ORW-CFM-W2 method) in this section, and left the offline version (i.e., RW-CFM with W2 distance regularization, namely RW-CFM-W2 method) in App. C.5 for ease of reading.
To further improve the stability of the ORW-CFM method and prevent policy collapse. we incorporate W2 regularization into our method, which has not been widely studied (Black al. 2024). This regularization introduces balance between exploration and exploitation by penalizing divergence from the reference model, thus maintaining diversity the learned model. The W2 distance bound can be directly incorporated into the ORW-CFM loss, inducing an ORW-CFM-W2 method:
LORW-CFM-W2
(I1) | (x, in (z =1)
(10)
+ | ~ 1)
wherein is trade-off coefficient, x (x, (2; 04) and you is the reference model. Then, the learned data distribution after epochs follows the Theorem (Proof in App. C.
Theorem 4. Under the online RL setup with DRW-CFM- loss, the data distribution after epochs evolves according to the following rules:
2e (x1) x [w ( (a11 exp (-3Da-1 (x1))]
(11)
||0000 (1, - (1, |//. 3 = ya, >> 0 is a scaling
where, Dn-1 (21) ((lla))
constant, gm-1 denotes the fine- fine-tuned model parameters after epoch 1
The W2 regularization term prevents the model from solely collapsing into high-reward regions and ensures diversity in the learned policy (See cases where a > 0 in Figs and also inducing the exploration-explootttioo trade-off.
In practice, we often use an exponential function or Boltzmann distribution to formulate the weighting function wr(211) = exp(r r(x)) )) where is entropy coefficient that controls the entropy of the induced policy (Haarnoja et al. 2018). Thus, we can derive the induced data distribution for the exponential case under the ORW-CFM-W2 loss (Proof in App. C.7].
Theorem Given w (21) = exp (T r (xr)), the induced duta distribution after N epochs of training under ORW-CFM CFM-W2 Loss as equation 70 is:
N
de (x1) x exp
(Nr(xi) B DD (x1) (x11)
(12)
n=1
[| [10000( (t,x)
where, D= (x1) = (x,a-)((a(os)
-a, 7 > 0 is a scaling And (t,
constant, 8t gm-1 denotes the model parameters after epoch n
Based on this, we can derive the learning behavior of ORW-CFM-W2 under two limiting cases:
Case 1 (Dominant Regularization Term: + 20 or 0). Dn (x1) grows rapidly with N or is large, since a, the regularization term dominates, namely EN D (21) > T (21).
The induced distribution is heavily penalized for deviations from the reference model, potentially leading to a distribution similar to the initial distribution q (x1) (See
Fig. 2
Case 2 (Dominant Reward Term: a = 0 or T + oo). If D (x1) grows slowly with N or a is small, since x a, the reward term dominates the regularization term, namely 7 r (x1) >
S D (11). The induced distribution q (æ1) concentrates on the xn maximizing r (#1).
similar the case without considering 1 (x1) ) as Lemma [ 0 in Fig.
Though the Boltzmann distribution is commonly used for the weighting function to handle the normalization constant Z, the derivation is similar the exponential form. For completeness, provide the derivation of the Boltzmann form weighting function in App. C.8 for ease of reading.
interpret ORW-CFM- method from the perspective of policy iteration in RL (Sutton & Barto 1998). which can connect our method to KL-constrainedd policy optimization (Ouyang al. 2002a Peng al. 2019) See App. C.9 for more detailed discussion.
Interestingly,

---

Published as a conference paper at ICLR 2025
5 EXPERIMENTS
We empirically evaluate our ORW-CFM-WW W2 method in equation 10 by fine-tuning small-scale FM models with U-net architecture from TorchCFM (Tong et al. 2024a] and large-scale FM models like Stable Diffusion (SD3) (Esser et al. 2024) from diffusers (von Platen et al. (0022) for reproducibility In all experiments,
use the exponential weighting function described in Theorem [ We also explore the impact of each component of our method and conducted detailed analyses of convergent behavior various hyperparameter settings. More experimental details can be found in App. D and with our Pseudocode in App. In general, we focus on the following questions:
Optimal Policy Convergence: Can our online RL methods converge to an optimal genera- tive data distribution that maximizes various user-defined reward objectives Theorem and Lemma while the offline baseline method converges to a sub-optimal?
derived in
Reward-Diversity Trade-off: Does W2 regularization effectively prevent policy collapse by maintaining a balance between reward maximization and diversity as Theorem 5
Stable and Controllable Policy Optimization: Can our method achieve a stable and con- trollable fine-tuning process of FM models in different settings, and control the convergent behavior of learned policy by adjusting entropy coefficient т and W2 coefficient a??
1.0
6
5
8
3
0
1 2
4 6 g
0 6 D C
6 0 2
0 4 o² 6
6
6
0
2
8
9
6 2 6
6 6 6 b 6 6 6 6 6
100 Y=0
3
6
O 6
è
lo
6
lo
6.7
5 7
2
6 4
/
A
6
6
6.6
optina
1
2
4
+
6
1 9
1
3
7
3
2
0 2 2
6 lo
2000 4000 soco aoboo 105co
8
0
9
2
D
6
6 6
6
l
Training Stop
8
8
0
O
0
2
6 8
6
4
(a) Reward Curve
(c) 0.01
(d)
10
Figure 2: Learning and generated images in target image generation task with different while a = 0. As increases, the convergent policy becomes increasingly greedy, and the diversity decreases. remains the similar distribution pre-trained model (See Theorem
1.04
6 à 6 99
6
6
6
6
6
6
6
0.9
6
6
6
6
6
6
6
8
6
6
6
6
6
8
6 6
6
L 6 6
6
6
6
6
2
6
6
4
6
0
6
l
6
4
6.7-
6
6
9
6
à
4
0
9
7
7
6
b
0
6
6.6
9
9
4
6
6
D
9999
6
6
6
4
6
6 6
6
3
6
6
6
0
6
2000 4000 8000 5000 10000
Trairing Stop
6 6
b 6
6
b
U
D
6
6
6
5
0
D
(a) Reward Curve
(b) 0
(c)
0.3
(d)
0.8
Figure 3: Learning curve and generated images in target image generation task with different a while 0.8, the diversity of the convergent generative distribution increases
10. From a without sacrificing too much performance. a 0 collapse to a Delta Distribution as Lemma
5.1 TARGET IMAGE GENERATION
At first, we evaluate our ORW-CFM method, with and without W2 regularization, by fine-tuning pre-trained model of MNIST (LeCCnn et al. 1998) to generate only even numbers using reward signals. The reward function is defined as (x1) Peven (z1) Podd (æ1), calculated by pre-trained binary classifier. An optimal expected return of indicates the model exclusively generates even numbers.
Figs. 2 and demonstrate that our method quickly reaches near-optimal performance across different settings of the entropy coefficient 7 and regularization coefficient a. From Fig. 2 increasing 7 leads to more greedy, wwwardmmxxiiizing policy at the cost of diversity, consistent with the theoretical limits in Theorem 5 and Lemma] [] while - = remains a similar performance as the pre- Besides, from Fig. W2 regularization controlled by a remains diverse by controlling how far the
pre-trained trained model.
8

---

Published as a conference paper at ICLR 2025
fine-tuned model diverges from the pre-trained model, enabling balance between exploration and exploitation. In general, ORW-CFM, especially with W2 regularization, achieves a stable fine-tuning process via maximizing the rewards in constrained neighborhood of the pre-trained model, inducing a flexible balance between reward and diversity, avoiding extremely exploitative policies.
5.2
IMAGE COMPRESSION
(a) a =
(b) a = 0.5
(c) a:
(d) pre-trained t-tirrined model
Reward Distance Trade-Off
Reward Distance Trade-Off
Convergent W2 Distance
Convergent Rowards
0.7
0.10
0.8-
pre-trained model
0.08
pre-trained model
0.9
0.06
-1.0
0.04
-1.1
0.02
-1.3
-1.2
Finc Tuning
0.00
(e) Reward Curve
(l) W2 Curve
Figure 4: Neward-Diversite Trade-off and W2 Distance Control a in Image Compression Task
(Black et al. 2024) of CIFAR-10 (Krizhessky & Hinton 2009) with 1. W2 distance is estimated
by its upper bound (see App. C.4. Each point in (c) and (f) corresponds to one group of experiments with a varying from 0 1. As increases, the final reward decays while distance between the fine-tuned model and the reference model become closer, inducing controllable fine-tuning process.
The effect of W2 regularization and exploration-exploottion trade-off is particularly evident in the image compression task using the model pre-trained in CIFAR-10, as shown in Figs. 4 In Fig. see that the introduction of W2 regularization prevents the policy from collapsing into extremely greedy solution. As increases. the model explores optimal solutions within a constrained neighborhood of the pre-trained model, preserving diversity while still optimizing the reward.
The reward-distance trade-off curve in Fig. further highlights this balance. Varying a provides an explicit trade-off between maximizing the reward and minimizing the divergence from the pre-trained model. As decreases, the policy becomes more exploitative, potentially collapsing into Delta distribution, as described in Lemma Conversely, increasing a encourages broader exploration of the data space, ensuring a more diverse generative policy. Therefore, proper W2 regularization effectively balances exploration and exploitation, avoiding policy collapse.
5.3 TEXT-IMAGE ALIGNMENT
Aligning Large-Scale Flow Matching Models To further validate our method's effectiveness, we
evaluate our BRW-CFM-VV W2 method on fine-tuning Stable Diffusion (SD3) (Esser al.) (024)..
Besides, we adopt LoRA (Hu et (022) for efficient fine-tuning. As shown in Figure we compare
our method against RAFT (Dong al. | (2023) and ReFT (Huguet al. 2024) on spatial relationship prompts used in DPOK (Fan et al. 2023). Our approach demonstrates better positional relationship
control while maintaining image quality and semantic coherence. On challenging prompts like "a banana on the left of an apple" and "a cat on the left of dog". our method consistently generates images with correct object placement and natural appearances, outperforming previous approaches in terms of both alignment accuracy and generation quality. Results these complex compositional prompts further demonstrate our approach's 's capability in handling multiple semantic constraints while maintaining generation diversity. See App. A more details and additional results.

---

Published as a conference paper at ICLR 2025
Proepes
huxaza apple
Alant racktt
lapepp ce up ot's
SDS
REFT
Ouis CORW-CMM
Figure 5: General Comparison of Different Fine-tuning Methods on SD3 via CLIP Rewards.
Tree Alguneer, A Car to The Sky
RAFT
ReFT
BeFT
WW2
Figure 6: Ablation Studies Text-Image Alignment Tasks of SD3, All use CLIP rewards.
Ablation Results The ablation studies in Figure 6 demonstrate several key findings that align with our theoretical predictions. Our ORW-CFM method, even without W2 regularization (a = 0), achieves better semantic alignment than other online methods like RAFT and ReFT, though all methods without W2 regularization exhibit clear signs of policy collapse, generating nearly identical images despite achieving high rewards- empirically validating our theoretical prediction in Lemma
I The introduction of W2 regularization (RAFT+W2, ReFT+W2, and our ORW-CFM-W method, 1) successfully prevents this collapse while maintaining performance, validating our tractable W2 distance bound in Theorem 3 Our method achieves the best balance between reward optimization and diversity preservation, with generated images showing consistent semantic understanding of "cat in sky" while maintaining natural variations in compositions. See App. A.2 for more details.
6 CONCLUSION
In this paper, we present the first theoretically-grounded online RL framework for fine-tuning flow matching models Online Reward-Weighee Conditional Flow Matching with Wasserstein- 2 Regularization (ORW-CFM-W2) without requiring calculations of ELBO, exact likelihood or KL divergence, which are intractable ODE-based flow matching models. Our work provides the first theoretical analysis of convergent behavior in flow matching online fine-tuning, proving that unregularized methods inevitably converge to a Delta distribution (Lemma []. inducing policy collapse that maximizes rewards at the cost of diversity. This theoretical discovery motivates our second major contribution the derivation a tractable Wasserstein-2 regularization that can be directly computed from vector fields, providing an efficient way prevent collapse while maintaining diversity. Extensive experiments across target image generation, image compression, and text-image alignment demonstrate that our method achieves optimal convergence while maintaining stable learning, with ablation studies validating both the online reward-weigs weighttng mechanism for optimal convergence and W2 regularization for diversity preservation. See App. for more discussions.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Martín Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial net- works. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 214 223. PMLR, 2017. URL http://proceedii:
mlr.. www.s/rssseev...m httppprs//////rrroooy..llaa html
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, OpenReview.net, 2024. URL https:://ppeneellee net Forum mineeumm?iiYCWWW (CGrJDG
Kevin Clark, Paul Vicol, Kevin Swersky, and David J. Fleet. Directly fine-tuning diffuusio models on differentiable rewards. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL httttttts///opeeeeeiii net
Carles Domingo-Enrinn Ennich, Michal Drozdzal, Brian Karrer, and Ricky T. Q. Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. CoRR, abs/a4b9 08861, 2024. doi: 1044550/ARXIV..409.058
URL https://doi. org/
10. 18350/arXiv.24099
Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. RAFT: reward ranked finetuning for gener- ative foundation model alignment. Trans. Mach. Learn. Res., 2023, 2023. URL https
//openreview net.:/wwweeeeeeeee 2ott:///oormmm22mdmm55557771
Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O. Stanley, and Jeff Clune. First return, then explore. Nat., 590(7847):580-55888 2021. doi: (0008//41586-020-0115555
|
7/doi oogoi.org/// 10. 10100048111555
020-03157-1
Patrick Esser, Sumith Kulal, Andreas Blattmann. Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion En- glish, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthe- sis. In Forty-first International Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL htlpttts:: wwtpp//pwwwwww..p///ppppp net/ forum? FPrUhsQJ5B
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Picter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. DPOK: reinforcement learning for fine-tuning text-to-image diffuusion models. CoRR, abs/2305. 63881, 2023. doi: 10.48550/ARXIV. 2305.16381. URL htttps:/:///to org/ 48550/ arxiv 2305 16381
lan J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In Zoubin Ghahra- mani, Max Welling, Corinna Cortes, Neil D. Lawrence, and Kilian Q. Weinberger (eds.), Advances in Neural Information Processing Systems 27: Annual Conference on Neural In- formation Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada, pp. 2672 2680, 2014. URL SolLtttp///proceeinnng nwwrips www.rips.ccpppper//....... 4/hhhhhh Cca3eeb122f4 61 1016611fffo0
Abstrrccttbbatttttcoccc html
Çaglar Gülçehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, Wolfgang Macherey, Arnaud Doucet, Orhan Firat, and Nando de Freitas. Reinforced self-training (rest) for language modeling. CoRR, abs/2308 08998, 2023. doi: 10.48550/ARXIV.33888. 0889988 htttss://oii.... 10. 18550//aXX30 208. 08998
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforeement learning with stochastic actor. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüsssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings
11

---

Published as a conference paper at ICLR 2025
of Machine Learning Research, pp.
PMLR, 2018. URL http::/ptcceedd.ggg
mlr. www.spressse...... Teteess/v0//hae hwwp///r///aaaaaaaaaaaaa ja 8l. html
Geoffrey Hinton. Training products of experts by minimizing contrastive divergence. Neural https://///d
Comput., 14(8):1711-1800, 2002. doi: 10. 10.116208996660766110000000
org
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In The Tenth Interna- tional Conference on Learning Representations, ICLR 2022, Virtual Event. April 25-29, 2022. OpenReview.ne,, 2022. httLhtttss///peeeeeevee
Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Laufer, Pablo Lemos, Riashat Islam, Cheng-Hao Liu, Jarrid Rector-Brooks, Tara Akhound-Sadegh, Michael M. Bronstein, Alexander Tong, and Avishek Joey Bose. Sequennce-augmeetted se(3)-flow matching for conditional protein backbone generation, CoRR, abt/22405.222113333 2024. doi:
10 10.48500ARKI2244022033
URL
20313
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications -imistics-Siimuatiion and Computation, 18(3):059-10076 1989.
Steven Kapturowski, Georg Ostrovski, John Quan, Rémi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. JRLLLttps://operev..ee net/ fortm/formmsm/mo 1ttp://orruuun.rrllyyy..."""
James Kirkpatrick, Razvan Pascanu, Neil C. Rabinowitz, Joel Veness, Guillaume Desjardins, An- drei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabskk-Barwinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forgetting in neural networks. CoRR, abs/1612.00776, 2016. URL http: http://axxxv.oogg/aaa/. 00796
Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick-a-pic: An open dataset of user preferences for text-to-image generation. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 / 16, 2023, 2023, UtL http: http::/papppps.eeeee nips wwtppcerpppppppeeepppppppeee...e fecaape_ffilee files html [3accb010bb1100653105555------
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical Report 0, University of Toronto, Toronto, Ontario, 2009. URI. htt.::tttt///wwt cs
foronto.edu//krrt.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proc. IEEE, 86(1 11:22)8--3224 1998. doi: 10.1 10.1109/5.726791 URL
https://doi wwtp:o...rgg/.... 1109/5.../.... 726791
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, et al. A tutorial on energy-based learning. Predicting structured data, 10), 2006.
Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to-imaae models using human feedback. CoRR, 121222200.12192 2023. doi: 10.4850/ARXXVV232222199 URL https
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL
https:://openreview net www.:oorrmmoooooo
12

---

Published as a conference paper at ICLR 2025
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F. Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9. 2022, 2022a. URL http: / nips. cc wwwppppppppp.pcccccccccccccc..ee http:pppprrrieesppppppeeeeee htwpse/aappp//peeee /002 2022//222/2222 2022/hash/ /hash/ html
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F. Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December cc/ 9, 2022, httpp//aaatppppppee......sos... wwtppppppppppppppppppp.p..eeeeee hwwppppppppp..oceeecc.ecceepeeee Seapperrfillsee files/ http:eeeeeieeeeeeeeeeeeeee
hash/
html
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage regression: Simple and scalable off policy reinforcement learning. CoRR, abs/1910.00177, 2019. URL
http:://axiii org/ abs 19101 00177
Jan Peters and Stefan Schaal. Reinforcement learning by reward-weighted regression for operational space control. In Zoubin Ghahramani (cd.), Machine Learning, Proceedings of the Twenty-Fourth
International Conference (ICMI. 2007), Corvallis, Oregon, USA, June 20-24, 2007, volume 227 of ACM International Conference Proceeding Series, pp. 745-750 ACM, 2007. doi: 144// 1273496.1273590. URL luttps://dot http://....011111111111111111399 1273590
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 8748 P763. 2021a. URL radford21a h
mlr. www...reeesvvv3.3
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML, 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 8748-8763, PMLR, 2021b. radFord21a.hee html
mmw.eiiinnn.mm.......mm.mm.m mlr. wttp:///ws........s/////./...
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URLhttp://papers wipppprrs..iipp...ppppppppp http://s.ccppppppeeeeeee wwt.:ffilleeeeeeeeee / 8 hww.0//eeedocee744ee5eeaae
Aatpreeettteeetatatattt------ Cew.acttt-ooffeeeneeee html
Abst
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells III, and Alejandro F Frangi (eds.), Medical Image Computing and Computer-Assisted Intervention MICCAI 2015 18th International Conference Munich, Germany, October 9, 2015, Proceedings, Part III. volume 9351 of Lecture Notes in Computer Science, pp. 234 241. Springer, 2015. doi: 10.1007/ 978-3-319-24574 28. URL https:///ol.ooggo..... doi org /
28
13

---

Published as a conference paper at ICLR 2025
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to- image diffusion models with deep language understanding. hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 2022, 2022. URL http: http:://papers..eteeeeeee nips. wwtppppr.ccpppppppppppppcppp..e fi wwwps//aaeeeeeeeeeeeaeeeaaaeaeee hhahh/ html
In Sanmi Koyejo, S. Mo-
John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICMI. 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 100-11897.JMMR... 2015. mllr. wtt.:/ppwwsw////eeemee.m http:///////hhummaamaaam .html
URL hatt:::/proceedinggee
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/1707 7004776 2017. URL http::/arxiv, www.iiiiioorgaaaa /bs/
1707. 06347
Ilia Shumailov, Zakhar Shumaylov, Yiren Zhao, Nicolas Papernot, Ross J. Anderson, and Yarin Gal. AI models collapse when trained on recursively generated data. Nat., (31(2)))) 759, 2024. doi:
-88-44456--2--7566666 Y. URL httts:://ol.orqq// s41586 024-0066--67
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F. Christiano. Learning to summarize with human feed- back. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Anmual Con- ference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. /RRL::ttpps:/prrceedddng neurips. hccppapper/222//aaaaaaaa (2860088166222999933888888888888
html
Zeyi Sun, Ye Fang, Tong Wu, Pan Zhang. Yuhang Zang, Shu Kong, Yuanjun Xiong, Dahua Lin, and Jiaqi Wang. Alpha-clip: A CLIP model focusing on wherever you want. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2024, Seattle, WA, USA, June 16-22, 2024, pp. 13019-13029, IEEE, 2024. doi: 10. 10.199/0VPR2273322040033333 /ww.//ooii.oo bttpes:/Bboo..gr...0110/////22/
URL
2024. 2123.01237
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: An introduction. IEEE Trans. Neural Nenworks, 15):1054- 1998, doi: 10.1109/TNN.1998.72222 URL httpp:://dbi
/tgg/10.110
1 1998 92]
Yunhao Tang, Zhaohan Daniel Guo, Zeyu Zheng, Daniele Calandriello, Yuan Cao, Eugene Tarassov, Rémi Munos, Bernardo Ávila Pires, Michal Valko, Yong Cheng, and Will Dabney, Understanding the performance gap between online and offline alignment algorithms. CoRR, abs/2405.08848, 2024. doi: 10.4850/ARXIV.2405.0888 24C5. 08448
URL ht doi, org/10 48500/8255554/2
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Trans. Mach. Learn. Res., 2024, 2024a. URL https //openreview. net (openrevieenneer//ormm--------- fewee//ormmidddceeeeee
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-frie free schrödinger bridges via score and flow matching. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li (eds.), International Conference Artificial Intelligence and Statistics, 2-4 May 2024, Palau de Congressos, Valencia, Spain, volume 238 of Proceedings of Machine Learning Research, pp. 1279-1287. PMLR, 2024b. URL https: / http://e/oceeee//o ngs ml www.mmr...rs....mm prw.sssss//ee..... perees/vve v ton./////..oo/og...... html
14

---

Published as a conference paper at ICLR 2025
Patrick Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj. Dhruv Nair, Sayak Paul, William Berman, Yiyi Xu, Steven Liu, and Thomas Wolf. Diffusers: state-of-the-at diffusion models. petps://ggthhbe wwtp://ww....gggggaee..ce.
diffusers 2022.
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text- to-image synthesis. CoRR, abs/2306.093411 09341, 2023. doi: 1048550/ARXXV.33000094441 URL
https::/dde .rgg10...... www88500aaiiii 2306 09341
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko. Moritz Hardt, and Sergey Levine (eds.), Advances Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023, NewrlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. U nips wwtpscppppppppeeeeeeeeeeeeeee.ee 036400orooa5551111abbb52222222 ottp:://eeeaoaaattoaoaaaaaa--aaa Conference.htmm
/iless/paapee//222////eeeee Abstract
Daniel M. Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul E. Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. CoRR, abs/1909.08593, 2019, URL http://arxiv org/ abs/ 909. 08593
15

---

Published as a conference paper at ICLR 2025
A TEXT-IMAGE ALIGNMENT EXPERIMENTS OF STABLE DIFFUSION 3
In this section, to further empirically demonstrate our performance on fine-tuning large-scale genera- tive models like the stable diffusion series (Esser et al. 2024), we conduct extensive experiments on fine-tuning Stable Diffusion 3 (SD3) (Esser et al. (024), which adopts flow matching architecture (Lipman al. 2023 Esser al., 2024p. For reproducibility, we built our code on the open-sourced diffuser codebase (von Platen et al. 2022), and adopted LoRA (Hu et al.) (022) and fp16 precision to reduce GPU memory requirements. Unless otherwise specified, we use a 1, T 0.1 for all experiments in this section.
To verify the effectiveness of our method, we primarily use the raw CLIP score (i.e., logit) (Radford et al.] 20) as our rewards for text-image alignment for SD3 fine-tuning. We also demonstrate our method's adaptability by extending to various other text-image alignment reward models, including
HPS-V2 (Wu al. 2023). Pick Score (Kirstain al. 2023). and Alpha Clip (Sun et al., 2024).
Notably, our method learns entirely from selfrgene data (Dong et al. 2023; Shumailov et al.
2024) without requiring manually collected (Ouyang et al.) (002aa) or filtered datasets (Rafailov
ct al. 2023). while effectively preventing policy collapse through our Eteoretically-derive W2
regularization (See App. C.4)
A.1 TEXT-IMAGE ALIGNMENT WITH SPATIAL UNDERSTANDING
A banana on the left of an apple
A on the left of a dog.
A tennis racket undernetth traffic
A laptop top of teddy bear
Prompts
SD3
RAFT
RcFT
Ours (ORW-CFM
-W2)
Figure 7: General comparison of positional relationship understanding between different fine-tuning methods on SD3 Esser et al. 2024). We compare our proposed ORW-CFM-W2 method against RAFT (Dong et al. 2023), ReFT (online variants, training on self-generated data) (Huguet al.] 2024), and the original SD3 model [Esser et al., 2024) across various positional relationship prompts from DrawBench (Saharia et al. 2022) used in DPOK (Fan al. 2023). Each column shows generations for specific prompt testing spatial understanding (e.g.. "left ot "on top of", "unddrneath"), The results demonstrate that our method achieves better positional relationship control while maintaining image quality and semantic coherence. This suggests that our Wasserstein-2 regularized online reward-weiihted approach effectively fine-tunes flow-based models to better understand and generate complex spatial relationships between objects. Model generations shown are sampled with the same random seeds across methods for fair comparison.
In the first experiments, we aim to improve the model's 's understanding of spatial relationships between objects, challenging task that requires both semantic understanding and precise positional control.
To demonstrate the effectiveness our method, we use the raw CLIP score (i.c., logits) (Radford) al. 2021a Xu et al. 2023 Dong et al. 2023) directly our text-image alignment reward model,
16

---

Published as a conference paper at ICLR 2025
eliminating the need for probability conversion complex reward shaping (Kapturowski al. 2019). This straightforward approach further demonstrates our method's 's ability to optimize directly from raw similarity scores, making it more practical for real-world applications (Dong et al. et al. 2021a).
2023 Radford
Figure| (i.e., Same Figure 5 but with more detailed discussion for ease of reading.) compares our method against several baselines: RAFT (Dong et al.) 2023), ReFT (online variants, training on self-generated data) (Huguet al. 2024). and the original SD3 model. The results demonstrate several key advantages of our approach:
1. Spatial Semantic Coherence/Alignmenn Our method shows superior understanding of spatial relationships ("left of", "on top of", "underneath")) while maintaining high image quality. This validates our theoretical analysis that ORW-CFM-W2 can effectively optimize arbitrary reward functions without compromising generation quality.
2. Controlled Optimization: As shown in Figure 3 and Figure 7 online RL methods trained on self-generated data (Dong et al. 2023 Huguet et al. 2024 Black et al., 2024 Shumailov et al. 2024) normally face great challenges of policy collapse, where models converge to limited, homogeneous outputs that maximize rewards but lack diversity a phenomenon theoretically characterized in Lemma I To address this challenge, we introduce W2 regularization, which effectively prevents over-optimization and policy collapse (Lemma demonstrates that this regularization enables controlled trade-off between reward maximization and output diversity (See Figure A). The experimental results validate this prediction, showing consistent spatial relationships while maintaining natural variations in object appearances, styles, and compositions.
D. Our theoretical analysis in Theorem [
3. Easy-to-Us Use and Stable Fine-tuning Method: Our method achieves these improvements without requiring filtered datasets (Rafailov et al., 2023), likelihood calculations (Black ett al.) 2024) differentiable rewards Ceomingo-EEriiee et al. 2024). demonstrating the practical advantages of our reward-we -weihhee flow matching framework. This aligns with our theoretical framework that bypasses the computational challenges of likelihood estimation in continuous flow models.
In general, these results further validate our theoretical contributions in a practical, large-scale setting. The successful fine-tuning of SD3 demonstrates that our method can effectively balance between reward maximization and diversity preservation, as predicted by off analysis (Theorem (5). Furthermore, the stable learning behavior supports our theoretical analysis of the convergence properties under W2 regularization. These experiments also highlight the practical significance of our two key innovations: the online reward-weighting mechanism and the W2 regularization. The online approach allows the model to continuously explore and improve spatial understanding, while the W2 regularization prevents the collapse into simplistic or degenerate solutions a common challenge in fine-tuning large generative models (Dong et al. 2023 Black
trade-
et al. 2024 Shumailov et al. 2024
A.2 ABLATION STUDIES AND POLICY COLLAPSE ANALYSIS
In the second part of SD3 fine-tuning experiments, to empirically validate our theoretical analysis of policy collapse in online RI. fine-tuning as Lemma I and demonstrate the effectiveness of our W2 regularization via ablation, we further fine-tune SD3 models align out-of-distribution (OOD) prompt "a cat in the sky". This setup provides an ideal test case for examining both the policy collapse phenomenon and our method's ability maintain diversity while optimizing rewards in different online methods. The ablation results have been concluded in Figure 8
Effectiveness of Online Reward Weighting (ORW-CFM) The ablation studies in Figure demonstrate that our ORW-CFM method, even without W2 regularization (a = 0), achieves better semantic alignment than other online methods like RAFT (Dong et al. 2023) and ReFT (Huguet et al.) 2024). This superior performance validates our theoretical framework for online reward-weit weightee fine-tuning of flow matching models (See Lemma Optimal Convergent Behavior in Lemma II. The method successfully optimizes for the desired semantic content, achieving higher-quality generations that better match the prompt "cat in sky" compared to baseline approaches.
17

---

Published as a conference paper at ICLR 2025
Text- Text-Image Alignment: Cat in Sky
SD3
RAFT +W2
RAFT
ReFT +W2
ReFT
Ours (no W2)
Ours
Figure 8: Comparison of different fine-tuning methods fine-tuming SD3 for the prompt "a cat in the sky" (reward also the CLIP score Radford et al. (2021a); Xu et al. (2023): Dong et al. (2023). The results demonstrate that policy collapse frequently occurs during online RL fine-tuning (Black et al., Dong et al.) 2023 [hhumallov et al. 2024). The baseline SD3 Esser et al. (224 shows limited spatial/semantic understanding for unusual prompts (i.e., hard
2024
our normal life). While RAFT (Dong et al. 2023) and ReFT (Huguet et al. 2024) without W2 regularization achieve high rewards (e.g., successful cat placement), they exhibit clear signs of policy collapse, generating nearly identical images with limited diversity. Our method without W2 regularization (Ours no W2) similarly suffers from policy collapse (also see Figure 2]. aligning with our theoretical prediction of the convergent behavior in Lemma T In contrast, methods incorporating W2 regularization (RAFT+W2 ReFT+W2, and Ours) maintain generation diversity while successfully positioning cats in the sky. Our approach particularly good at balancing reward optimization with diversity preservation, demonstrating the effectiveness of our theoretically-derivee W2 regularization bound (Theorem in preventing policy collapse while maintaining high-qualij, generations. The results empirically validate our theoretical analysis of the reward-di wwwarddiiversity trade-off in Theorem
Policy Collapse Phenomenon However, in Figure g all methods without W2 regularization (RAFT, ReFT, and Ours without W2) exhibit clear signs of policy collapse (Shumailov et al. (2024) Black cet al. 2024), generating nearly identical images despite achieving high rewards. This empirically validates theoretical prediction in Lemma that unregularized online reward- weighted methods naturally converge to a delta distribution at the maximum reward point, sacrificing diversity for reward maximization.
Effectiveness of W2 Regularization Based on Figure and Table I the introduction of W2 regularization (RAFT+W2 ReFT+W2, and ORW-CFM-W method, the effectiveness of our theoretically-derieed tractable W2 distance bound from Theorem 3 This regularization successfully prevents policy collapse while maintaining performance across all methods, enabling controlled divergence from the reference model. The consistent improvement in output diversity across different baseline methods validates our theoretical framework for controlling model behavior through W2 regularization in flow matching architectures.
1) demonstrates
Optimal Balance Through ORW-CFM-W2 In Figure 8 our ORW-CFM-W2 method achieves the best balance between reward optimization and diversity preservation generated images show consistent high-level semantic understanding of "cat in sky" while maintaining natural variations in cat appearances, sky conditions, and overall compositions. While both RAFT+W2 and ReFT+ W2 show diversity improvements from W2 regularization, their outputs still exhibit less variation in style
18

---

Published as a conference paper at ICLR 2025
and composition compared to ORW-CFM-W. W2. This improved performance aligns with our theoretical analysis in Theorem | regarding the reward- -warr-diversity trade-off in the exponential weighting case. The results also validate our optimization framework's ability to achieve stable convergence without collapsing into extreme exploitation.
Theoretical and Practical Validation These experimental results in Figure 8 provide strong empirical validation of both key components of our approach. First, they demonstrate the ORW- CFM method's s ability achieve optimal policy convergence, though with potential collapse unregularized (as predicted by our theoretical analysis of online reward-weighted learning). Second, they show how W2 regularization effectively prevents this collapse while preserving output quality and diversity- a capability derived from tractable W2 distance bound for flow matching. The comprehensive improvement across different baseline methods (Dong et al., 2023 Huguet et al.)
2024) illustrates how our theoretical framework enables stable fine-tuning that maximizes rewards while maintaining the rich generative capabilities of flow-based models (Lipman et al. 2023
Esser
al. 2024).
Table 1: Performance and Diversity comparison of different fine-tuning methods on text-image alignment using SD3. 'CLIP Score' measures alignment with the text prompt using the CLIP model's similarity score (higher is better). 'Diversity Score' quantifies the diversity of generated images by measuring 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. Our method achieves the highest CLIP score while maintaining strong diversity comparable to the base SD3 model. The addition of W2 regularization ('+W2' variants) helps preserve generation diversity across all methods, with our full approach (Ours) achieving the best balance between alignment (35,46) and diversity (4.21) compared to baseline methods RAFT and ReFT. All scores are calculated by averaging across 64 samples.
Method
CLIP Score t Diversity Score t
Effectiveness of ORW-CFM without W2
SD3 (Baseline)
28.69
4.78
ORW-CFM (Ours w/o W2)
33.63
3.47
RAFT
29.30
2.05
ReFT
29.32
3.26
Effectiveness of W2 regularization
ORW-CFM-W2 (Ours)
35.46
4.21
RAFT + W2
30.88
2.81
ReFT + W2
32.03
3.63
A.3 ADAPTABILITY ACROSS DIFFERENT REWARD MODELS
In addition to the three reward functions in Section 5 to demonstrate the broad applicability and reward-agnosite nature/property of our methods, we further evaluate our performance using three different text-image alignment reward models: HPS-V2 (Wu et al. 2023). Pick Score (Kirstain et al. 2023), and Alpha Clip (Sun al., 2024). Using the challenging Spatial prompt "a train on top of surfboard," we examine how our method adapts to different reward signals/models while maintaining generation quality and semantics coherence.
In general, the results shown in Figure demonstrate several key strengths of our approach:
Reward Adaptability: Our method shows consistent performance improvement across all three reward models, successfully positioning trains on surfboards while maintaining realistic scene com- position. This empirically validates our theoretical framework's 's ability to optimize arbitrary reward functions without requiring reward-specific modifications (Dong et al.) 2023) filtered datasets (Rafallov et al., 2023). The high-quality generations across different reward models demonstrate that our W2 regularization (Theorem 3) provides effective stabilization regardless of the underlying reward mechanism (It is worth mentioning that appropriately increasing the a value can be used to adjust the distance from the reference model, which is also demonstrated in Figure
19

---

Published as a conference paper at ICLR 2025
Text-Image Alignment: A train on top of a surfboard
HPS-V2
Pick Score
Alpha Clip
Figure 9: Comparison of fine-tune taeett methods for the prompt "a train top of surfboard", demon- strating our method's adaptability across different reward models (HPS- V2 Wu et al. (2023). Pick Score Kirstain et al.] (202)), and Alpha Clip scores (Sun al. (2024)). The results validate theoretical analysis that ORW-CFM-W2 can effectively optimize diverse reward functions while maintaining generation quality through W2 regularization. 1) Top row: Using HPS-V2 reward shows our method successfully aligns locomotive placement while preserving realistic wave interactions. 2) Middle row: Pick Score reward further demonstrates robust spatial understanding across varying train models and surfboard configurations. 3) Bottom row: Alpha Clip reward showcases consis- tent performance even with different text-image alignment rewards, validating the reward-agnostic pature/ of our approach. This comprehensive evaluation across multiple reward models empirically confirms our thcoretical prediction that the W2 regularization (Theorem 3 enables stable fine-tuning regardless of the underlying reward mechanism. The consistent high-quality generations across all three reward models demonstrate the practical advantages of our reward- -wightee flow matching framework, which achieves strong performance without requiring filtered datasets (Rafailov et al. 2023), likelihood calculations (Black et al.) 2024)). reward-specte optimizations (Dong et al., 2023 Huguet al., 2024), or differentiable rewards (Domingo-Ennichh et al. 2024). These results particularly highlight how our method's reward-diversity trade-off (Theorem 5) generalizes effectively
across different reward models while preventing policy collapse (Lemmal
Spatial Understanding: Across all reward models, our method demonstrates robust spatial under- standing:
1. HPS-V2: Achieves precise locomotive placement while capturing realistic wave interac- tions.
2. Pick Score: Maintains consistent spatial relationships across varying train models and surfboard configurations.
3. Alpha Clip: Shows stable performance even with different similarity metric framework.
Diversity Preservation:
Consistent with our theoretical analysis in Theorem []] the generated images maintain natural variations train appearances, surfboard designs, and ocean conditions while adhering the spatial constraints. This demonstrates how our W2 regularization effectively prevents policy collapse (Lemma D regardless of the reward models used.
Stable Fine-tuning:
The consistent quality/alignment improvement across reward models validates our theoretical contribution regarding the convergence behavior of ORW-CFM-WW W2 (See App. C.6). As predicted by our analysis of the exponential weighting case (Theorem 5) the method achieves
20

---

Published as a conference paper at ICLR 2025
stable learning without collapsing into extreme exploitation, even when faced with different reward scales and architectures. It also worth mentioning that the reward model is very important for the
success of the fine-tuming task (Dong et al. 2023 Black al. 2024 Rafailov et al. 2023 Xu et al. 2023), so how to select the best reward model still needs further exploration.
All in all, the successful adaptation across multiple reward models showcases how online reward- weighting mechanism and W22 regularization work together to enable robust, stable fine-tuning without requiring reward-specific optimizations. This generality makes our method particularly valuable for real-world applications where different reward models might be needed for different tasks or domains (Black al. 2024 Shumailov et 2024).
Furthermore, these results address a key challenge in flow matching fine-tuning: the ability to optimize diverse reward objectives without compromising generation quality or diversity while training on self-generated data. Our method's consistent performance improvement across reward models demonstrates that the reward-distance trade-off formulation (See Theorem (5) provides a easy-to-use and theoretically-gronnee diversity, regardless of how the reward is computed.
approach to balancing task-specific optimization with output
A.4 COMPOSITIONAL AND COMPLEX SEMANTIC UNDERSTANDING
A stack of 3 plates. A blue plate is on the top, sitting on blue plate, The blue plate is in the middle, sitting on green plate. The green plate on the bottom.
stack of cubes. A red cube is on top, sitting a cube. The red cube is sitting green cube. The green cube on the bottom.
An emoji of baby panda wearing red hat, green gloves, red shirt, and green punts
Prompts
middle,
Samples
ORW-CFM-W2
Aligned Samples
(Ours)
Figure 10: Additional evaluation of ORW-CFM-W2's semantic understanding and compositional generation capabilities on complex structured prompts from DrawBench (Saharia et af.) 2022) used DPOK (Fan et 2023). Each column demonstrates our method's ability to handle facetef requirements including spatial relationships, color specifications (red, green), and detailed object attributes. The figure compares baseline SD3 (Esser al. 2024) samples with our DRW-CFM W2 results across three challenging compositional prompts: stacked colored cubes, stacked colored plates, and a customized panda emoji. Our method shows consistently better semantic alignment while maintaining generation diversity, validating our theoretical framework's 's ability to prevent mode collapse (Shumailov et al., 2024) through W2 regularization.
In the fourth experiments of fine-tuning SD3 models (Esser et al.) 2024), to further validate our method's capability in handling complex semantic relationships and demonstrate the practical benefits of our theoretical framework, we evaluate ORW-CFM- W2 on a set of challenging compositional prompts that test multiple aspects of semantic understanding. These experiments not only demonstrate our method's 's effectiveness on large-scale flow matching models like SD3, but also empirically validate our theoretical predictions about controlled optimization and diversity preservation. We use CL.IP
21

---

Published as a conference paper at ICLR 2025
score text-image alignment reward model (Radford et al. 2021a Xu et al. 2023 2023).
Dong et al.
Compositional Attribute Control: As shown in Figure 10 method successfully handles prompts with multiple interleaved requirements spanning spatíal relationships, color specifications, and object attributes. In the cube-stacking example, the model demonstrates precise control over both spatial ("on top", "in the middle", "on the bottom") and color attributes (two red cubes on green cube), while maintaining physically plausible stacking arrangements. This validates that our ORW-CFM-W2 can effectively achieve fine-grained fine-tuning without falling into degenerate solutions.
Spatial Color Relationships: The stacked plates example particularly highlights our method's 's ability to jointly optimize multiple constraints positional relationships (top, middle, bottom) and color attributes (blue plates top/middle, green plate on bottom). The consistent high-quality results validate our theoretical prediction that our methods enables stable fine-tuning even when dealing with complex, spatially color constraints.
Multi-attribute Character Specification: The panda emoji example demonstrates fine-grained con- trol over multiple clothing attributes (red hat, green gloves, red shirt, green pants) while maintaining the core semantic concept and cartoon-style consistency. This success handling detailed attribute specifications without semantic drift empirically supports our analysis of the reward-dive diversistyy trade-off (Theorem showing how our method balances attribute-level precision with overall coherence.
In general, across all test cases, our method achieves notably better semantic alignment compared to the baseline SD3 samples while maintaining image quality and stylistic diversity. This improvement particularly evident in cases requiring precise attribute combinations spatial relationships, where baseline samples often exhibit attribute mixing or spatial inconsistencies. The results demonstrate that approach effectively addresses the policy collapse problem (Lemma T through W2 regularization, enabling precise semantic control without sacrificing generation diversity. These results offer strong empirical validation of our thcoretical framework. The successful handling of complex compositional prompts demonstrates that our online reward-weighting mechanism and W2 regularization work together enable robust, stable fine-tuning large-scale flow matching models. This generality makes our method particularly valuable for real-world applications requiring precise semantic control and attribute specification while maintaining natural variations the generated outputs. All in all, the consistent performance across diverse prompt types demonstrates our ability to optimize multiple semantic objectives without compromising generation quality or diversity. Our method's success in handling complex compositional prompts (Saharia et al.) [2022] Fan et al.) 2023) while maintaining coherent generations provides strong empirical support for our theoretical contributions while demonstrating the practical utility of our approach in real- world applications.
theoretically-grounee
22

---

Published as a conference paper at ICLR 2025
B BACKGROUND
B.1 REINFORCEMENT LEARNING
B.1.1 RL FORMULATION
In RL (Sutton & Barto, 1998), an agent interacts with an environment by taking actions A in states S, aiming maximize cumulative rewards over time. The agent's behavior is govemned by a policy Ta s), which represents the probability of taking action a given state s. The main goal of RL is optimize the policy maximize the expected return:
( = E
Σ (s at) ===
(13)
where y E [0, 1] is the discount factor that prioritizes early rewards, I' is the time horizon (which could be infinite), (st, az) is the reward at time step
The core problem in RL is how to learn the optimal policy r"(a s) that maximizes the expected return (TT). special case of RL, where there is no states but only actions for each agent, is known as the multi-armed bandit problem (Sutton & Barto 1998).
B.1.2 REWARD- WEIGITED REGRESSION
In Reward- Weighted Regression (RWR), the policy update done by re- re-weighting the probability of actions based on the rewards they produce. The updated policy is given by:
Tnew (a 8) x Wold (a s) exx(px * r(s,a))
(14)
where is temperature parameter controlling the trade-off between exploration and exploitation.
B.1.3 ADVANTAGE- WEIGHTED REGRESSION
Advantage- Weighted Regression (AWR) refines RWR by weighting actions based on their advantage, which measures how much better an action is compared the expected value of other actions:
New (a 8) x Told s) exp(rA(s,a))
(15)
where A(s,a) r(s, a) (s) is the advantage function.
In practice, we can also call the data-collection policy, the behavior policy μ (Sutton Barto 1998.
Schulman al. 2017 Peng al., 2019). and the optimal target policy x².
Then, the update rules of AWR becomes:
Tnew (a 8) x u(a s) exp(TA(ssa))
(16)
And the optimal policy is therefore given by (Peng et al. 2019).
(a where 7(s) is the partition function that normalizes the conditional action distribution.
((s) u(a s) exp (T A" (st, a:))
(17)
B.2
FLOW MATCHING
B.2.1
FL.ww MATCHING FORMULATIOO
Flow Matching (Lipman al. 2023 is methodology for training continuous normalizing flows, which are generative models that transform simple base distribution into complex target distribution through continuous, invertible mapping. The core idea of FM is define continuous flow that
23

---

Published as a conference paper at ICLR 2025
gradually transports samples from initial distribution Po(z) (usually simple distribution like standard Gaussian) to a target distribution (z) (the data distribution we aim to model).
This transformation is governed by ordinary differential equation:
dr ux(x))
(18)
where E Idd is a data point in d-dimensional space, / E [0, 1] is the time parameter, and ue (2) is time-dependent vector field that dictates the dynamies of the flow at each point in time and space.
The solution to this ODE defines flow (x) that maps point a at time & to another point at time f + öt. The transformation over the entire time interval [0, 1] effectively transports the entire distribution (z pl (z).
The evolution probability density px(x) over time is described by the continuity equation:
(x) (pr (z))+(z))
Öt
=
(19)
where V. denotes the divergence operator. The goal in FM is to learn the vector field uz (x(2)) such that the flow induced by (z(z) transforms po into P1 (z) as progresses from 0 to
B.2.2 CONDITIONAL FLOW MATCHING
Conditional Flow Matching (Lipman al. 2023) an extension of FM that introduces conditioning to simplify the learning process. Instead of modeling the entire distribution transformation directly, CFM focuses on modeling the conditional dynamics of the flow given a target sample æ1 from data distribution.
In CFM, we consider pairs (10,,21), where xo ~ (x) and ø1 ~ (ar). The conditional flow transports to 1 over time & [0, using vector field (z æ1) that depends æ
dr us (x 21)
(20)
dt
The objective CFM is to learn a parametric model ug (1, that approximates u, (x 21). The loss function used to train the model is:
(CF = Er, [100 - - (a =1)l"²]
(21)
where the expectation is taken over ~ U(0, 1), ~ (x(), and x ~ Pe (x x1), the conditional distribution at time given T1.
By conditioning on æ CFM reduces the complexity of the learning problem, as the model only needs learn how transform 2g Iy for given pairs, rather than modeling the entire joint distribution.
B.2.3
OT-CFM
Our baseline code is built upon torchefm (Tong et al. 2024a), where we use OT-CFM loss to learn flow matching models (i.e., OT-paths). Optimal Transport theory provides framework for finding most efficient way to transport one probability distribution another (Tong et al. 2024a). OT-CFM integrates optimal transport into the conditional flow matching framework to define vector fields based optimal transport maps.
In OT-CFM (Tong et al., 2024a), the vector field ue (a x1) is derived from the optimal transport plan that minimizes the cost of transporting (z) Pı (x)) The cost is typically defined in terms of the expected squared Euclidean distance between transported points.
By using optimal transport, OT-CFM ensures that the flow from æ1 follows the most efficient path, which can improve the convergence and quality of the generative model. The OT-CFM loss function remains similar to the CFM loss but incorporates the optimal transport-based vector field:
24

---

Published as a conference paper at ICLR 2025
(OT-CEM(() =E , [lloo(t, oT (æ [x))l"
(22)
where ur (x (11) represents the optimal transport vector field.
B.2.4 LIKELIHOOD CALCULATION
As special case of CNF, flow matching models share the benefit of exact likelihood calculation. Specifically, the learned vector field induces a probability path induced by the pushforward:
pr = (P
(23)
where Po is the initial noise distribution and (x) == 2, is the flow governed by the learned vector
field and the flow ODE as dæz such shange-of-probabine
dr (22). The continuity equation in equation alloww us calculate
term with the estimation of the model divergence as (Lipman et al. 2023):
Po (21) log (o(o)) J. (iv(v))(x))d
(24)
The trajectory is generated reverse through time from t = 1 (arbitrary data) to t = O via the learned vector field. Direct calculation of model divergence div(u:) prohibitively expensive. Instead, Hutchinson's 's trace estimator is used to efficiently approximate the model divergence as (Hutchinson) 1989):
div(e:)(z) Ee- (e)())]
(25)
In contrast to diffusion models that rely on variational bounds that can be calculated in closed form as the summation of KI. divergence (Ho et al., 2020), the likelihood for flow-based models requires expensive simulation and the calculation of vector- model divergence. Therefore, the likelihood-basee al. (2024); Schulman et al. (2017); Rafailov et al. (2023); Fan et al. (2023) is non-trivial and computationally expensive for flow models.
product to estimate the RL formulation for diffusion models Black
25

---

Published as a conference paper at ICLR 2025
C PROOFS
C.I INDUCED DATA DISTRIBUTION
In practice, generative model, we are normally very interested in the induced data distribution (new((1)) ) after learning under design loss function L. In order to facilitate the subsequent derivation, this section, we first derive general Theorem of Induced Data Distribution.
Theorem (Induced Data Distribution). Given the loss function:
Liorat (0) = Emira((as)
L (21)
(26)
where i. (x1) any per-sample loss function, the induced data distribution over an is:
Pnow (x1) x (x)) exp (-7I -x1)(x))
(27)
with > O being a positive constant.
Proof of Theorem 6 from RI. Perspective. In fact, most previous RLHF Ziegler et al. 2019 Stien non et al. 2020 Ouyang al. 2022b) and RL works (Peng al. 2019) have found that assume reward model r(x) that captures human preferences, with the goal modifying the base generative model ((x)) such that it generates the following tilted distribution:
Pnew (x1) p² (x11 g(x1) exp(T * r(r))
(28)
just set the reward model as the negative of per-sample loss i. (x1), namely:
r(æ1) - (x1)
(29)
and set 7 which is constant, and we can obtain:
Phew (x1) = p" (x1) x (x1) (->ā(xi)
(30)
Then, the proof concludes.
Proof of Theorem 6 from Energy Function Perspective. In standard EBMs, the probability density over 31 is defined using an energy function E (x1):
PERM (II) = exp (-3E (zi))
(31)
where 3 > is the inverse temperature, and Z is the partition function ensuring normalization:
Z f exp( (x1)) dri
(32)
This formulation assumes a uniform base measure over Z1-
We consider the per-sample loss L (x1) with the energy function E (x1), and set B = = Therefore:
Pnew (x1)
7 exp (-7i.
(33)
When the data distribution q (x1) is not uniform. when we have prior knowledge about 21 like ((x)), is appropriate to include q (æ,) in the model (LeCun et al. 2006). This leads to modified probability density:
26

---

Published as a conference paper at ICLR 2025
Pnew (21)= = z (x1) exp (-mī(xz)
(34)
Here, q (21) acts a base measure or prior distribution, and L (x1) corresponds the energy function E (x1). The inclusion of (I1) is necessary to accurately represent the underlying distribution of the data when it is not uniform.
Remark. The total loss function is:
Liottl (0) -0)=/a (
(35)
Minimizing Croat (0) with respect to corresponds to adjusting the model parameters to reduce the expected loss over the data distribution (z1).
We can interpret this minimization as seeking a distribution Prew (x)) that places higher probability on samples with lower L (x1). This aligns with the principle of maximum likelihood estimation in EBMs, where we aim maximize the likelihood of observed data under the model.
Then, the proof concludes.
Proof of Theorem 6 from Bayesian Inference Perspective. In Bayesian inference, the posterior dis- tribution Pnew (x1 D) over latent variable 2ı given data D is proportional to the product of the prior q (x1) and the likelihood p (D) Φ1):
Pnew (0)
D) = z (x1) a1)
(36)
where the normalization constant:
Z = ( x1) dr
(37)
In statistical modeling and machine learning, the loss function Í. (x1) quantifies the discrepancy between the model's predictions and the observed data. It is common to relate the loss function to the likelihood function through the negative log- ligliiclihood:
L (xx) - log (D) x1)
(38)
7
Rearranging this equation gives:
p (D x1) exp (->i (rii(z)) (z:))
(39)
The posterior distribution becomes:
Phew (x1) 7 (21) (x1)exx (->L(z))
(40)
where:
z z= ==/0(1)ess (xIL(x)) da
(41)
which ensures that Pnew (x1) integrates to one.
Then, the proof concludes.
27

---

Published as a conference paper at ICLR 2025
Corollary 2 (Limiting Behavior of Well Learned Model). As L (x1) O. the exponential term approaches 1:
lim exp (-)ī())) =1
(42)
-(m))-0
Thus, Pnew (x1) oc q (21), meaning the induced distribution is the same as the data sampling distribution q(21).
Corollary 3 (Weighted Prior Distribution). Considering a weighed prior sampling distribution Pess(z1)= ((())q((1)))) given the loss function:
Liotal (0) (p1
L (21))
(43)
where L (x1) any per-sample loss function, the induced data distribution over æ1 is:
Pnew (I1) x Pess (x1) exp (-rL(x1))
(44)
x w (x1) exp (->I (211)
with > O being a positive constant.
Proof of Corollary data distribution:
Substituting
(x11) Peps (x1) into the Theorem 6 we can have the induced
Pnew (x1) Peff (x) exp (-7[(())
(45)
ww(x1 (21) exp (-rā (z1)
then, the proof concludes.
Corollary 4 (Constant Loss). If (x1) c.for all Iı, then:
Dnew (x1) x q (I1) exp(- q (211 constant.
(46)
Thus, Pnew (x1) is proportional to (x1), and the induced distribution remains unchanged.
28

---

Published as a conference paper at ICLR 2025
C.2 REWARD WEIGHTED CONDITIONAL FLOW MATCHING
We start from a conditional flow matching setting, where model parameterized by 0 defines a vector field vo (1, that transports samples from base distribution target distribution over time E [0, 1]. The standard CFM loss function is as follows:
(0F (0) (=)- (el [Iva(t, u (<: x1)1l
(47)
which aims to align (o(t, x:) with the true conditional vector field (x(x)))) by minimizing the expected squared error over the data distribution ((11).
In practical applications, it often desirable to focus the model's learning capacity on specific regions of the data space. This can be achieved by incorporating weighting function w (x1) ) into the loss function, effectively reweighting the data distribution. However, the theoretical implications of this modification on the learned distribution have not been fully studied.
In this part, we rigorously analyze the impact of using a weighted regression loss in CFM training. We prove that incorporating weighting function into the CFM loss leads the model to learn new distribution over 21. given by pdew (21) ""[a; (o(m1) where q (x1) the original data distribution and Z is normalization constant. Our analysis provides theoretical support for the use of weighting functions to control the focus of generative models, with implications for target data generation, RLHF Ouyang et al. 2022a) and RL fine-tuning (Black et al., 2024 Rafailvv et al. 2023).
2
Let X S Rdd denote the data space, and let q (x1) be the probability density function (pdf) of the random variable 2, E X. We consider a continuous-time flow over E [0, 1], where / is sampled from the uniform distribution (0 1]. The conditional pdf of x at time given æı is denoted by Pt (x x.).
The true conditional vector field is ue (x | æ1). and the model's vector field is (2; 0), where represents the model parameters. The goal is to learn (z(2; 0) such that it approximates un x1)
We introduce a weighting function w X [0, oo), which is measurable and integrable with respect to (Z1).
Then, we can have the following Theorem (For ease of reading, we rewrite the Theorem I to be proved in this section as follows):
Theorem. Let w X' - (o) be a measurable weighting function such that 0 < =
Sx w (I1) 9 (I1) dx1 co. Define the reward weighted CFM loss function:
LRW-CFM Ee-2200.1), (ampp)))))))))) (x1) (x; 0) - us (x x 21)1
(48)
where w (x1) a weighting function, q (x1) is the original data distribution, Ps (x x1) is the conditional distribution at time t, v2 f) is the model's vector field, and ue (x: x2) is the true vector field conditioned on I1. Then, training the model using this weighted loss function leads the model to learn a new data distribution:
(11)
now (21)=
(49)
Z
where Z f (x1) (x1) dx is the normalization constant.
Proof of Theorem 7 We first aim to show that incorporating the weighting function w (I1) into the loss function effectively changes the distribution over x1 from q (x1) to new (x1).
Lemma 2. The weighting function modifies the data distribution over æ1 from q(x1) to:
w (x1) (x1)
now (x))
(50)
7
Proof of Lemma
The weighted loss function can be explicitly written as:
LRw-Crx(0)- S/ J (2z) / [lx((z;0) us (x 21)ll (x) xx) drdxydt
(51)
29

---

Published as a conference paper at ICLR 2025
Then, define the reweighted distribution pne" (a1) as:
p guw (x1) (21) (211 where = / w (21) (211) dri
(52)
Z
Since (x1) 0 and Z < 00, (xi) is valid pdf over X.
Substituting prw (x1) into the loss function, we have:
CRW-CFM(0)= = Z p now (x1) [loc(z;0) - ue (2 x1)ll Pe (or x1) drdridt
(53)
This can be expressed as expectation:
LRW-CFM(0) 21. (2)) [1oe((;;)) 0) ue (x mx))la]
(54)
The gradient of (RW is:
VOLRW-CFM(0) = 20, --ppe (x)) (s(a))
Vo ()) - (x =2))
(55)
Since constant with respect to does not affect the optimization process. Therefore, minimiz- ing (FW) is equivalent to minimizing the expected loss under the distribution new (x1)
set Pefs(a1)= w( w(x1)g(o))
Remark. According Corollary
we have:
Pnew (x1) x Pess (21) exp -xb((1))
(56)
x w (II) (x1) exp (-7I (21))
When the model fits the data well (iee, attain convergence), and i. (xI) 0, then according to Corollary 2 we have:
Pnew (x1) w (21) a (an) exp (-21((11)
(57)
w ( (z1)
Remark. During optimization, the model minimizes the expected squared difference between the model's vector field and the true vector field, but crucially, this expectation is now taken with respect to the reweighted distribution nrew (ee(x1), (x1). As a result, the model prioritizes accurately modeling the flow dynamics in regions of higher weight, effectively learning generate data according to new ( ).
The proof concludes.
Based on Lemma it's obvious that Theorem [
Based on equation 55 the reward weighted conditional flow matching loss in equation 48 actually
learn a weighted distribution new (x1):
new (x1) 9
(58)
Z
Remark (Normalization Constant Z Z a finite normalization constant ensuring new (xx) is valid probability distribution (i.e.. integrates to 1). In practice, during training, the normalization constant can be ignored because does not affect the optimization process since is constant with respect to 0.
Lemma 3 (Finite Normalization Constant Z). The normalization constant Z is finite if w (x1) is integrable with respect to q (x1).
30

---

Published as a conference paper at ICLR 2025
Proof of Lemma 3
Since w (z1) is weighting function and (z1) is probability density function,
the integral:
2 z= w (z1) (211 dæ1
(59)
is finite if w (I1) does not introduce singularities and is integrable over X'
Remark (Data Traversal). In fact, to optimize the reward weighted CFM loss in equation 48 requires the agent nearly traverse almost possible samples æ1, which may also similar the exploration problem in RL, namely how to explore the data space X. Therefore. maintaining the ability to contimuously explore the data space important for RL method (Ecoffer et al. 2027 Haarnoja et al. 2018).
31

---

Published as a conference paper at ICLR 2025
C.3 ONLINE REWARD WEIGHTED CONDITIONAL FLOW MATCHING
In App. C.2 we have discussed offline version of Reward Weighted Conditional Flow Matching, which is learned by optimizing the following loss:
CRW-CrM 10 (a (asp(l))))) w (ari) [1o, 0) de (ar an |11]
(60)
wherein q (x1) ) is the data distribution approximated by the training dataset or sampled from pre-
trained model, Pr (x: x11) is the conditional pdf of x at time t given 21 (Tong et al. 2024a).
We consider N epoch online fine-tuning scenario, where our vector field model va, is initialized by pre-trained model o that can capture the origin data distribution ((21), namely the data distribution pr (x1) induced by pre-trained model Uew can meet p (x1) R2 g(z1). At each epoch of training, the agent will sample data from reweighted data distribution (x*(x1 induced by U0,, and use the newly sampled data train the model under the Reward Weighted CFM Loss, then the online reward weighted CFM loss can be writen as:
CORW-CFM = Ex~ee(o, (o.-lu(j. axppp to (211 (1; 0) (x =1)1a]
(61)
wherein nhh (ar1) is the re- weighted data distribution given by fine-tuned vector fields von- Then, can have the following (For ease reading, we rewrite the Theorem [ass follows):
Theorem (Online Reward Weighted CFM). Let q (I1) be the initial data distribution, and w (x1) be non-negative weighting function integrable with respect to q (a1). Assume that at each epoch the model ve perfectly learns the distribution implied by the Online Reward- Weighted CFM loss. Then, the data distribution after N epochs is given by:
(311 (311)* N
(31)
q (x1)
(62)
ZN
where ZN = -xx - (x1) a (xz) dx, is the normalization constant ensuring d (xx) is valid probability distribution.
Proof of Theorem
We can use induction to present the proof.
Example 1 (Base Case N = 1). We first consider the basic case, where N 1, namely single epoch of learning, which equals to offline RW-CFM. Then, The initial data distribution is g8 (x1) q (x ).
According to Theorem 7 after applying the online reward weighted CFM loss, we can obtain the new distribution is:
dd (x1) = w (z (1.)
(63)
Z1
where Z₁ = fx w (x1) (x1) da
Assume that after 1 epochs, the distribution is:
n- 9 1 (21) W (21) (211)=11 q (x11
(64)
Zn-1
with Zn-1 = fx (z1)"=1 (x)) dxi.
According to Theorem 1 at epoch n, the weighted loss leads effective distribution:
w (I1) d (x1)
z (x1)
(65)
Zm
where Z, fx w (x1) 1/1 (xx dri Substituting d0 2-1 (x1)
32

---

Published as a conference paper at ICLR 2025
w (21)
w(xx )"- Zn-1
(((())))) 9(01)
w (211) (211
lø (x1)
(66)
Ze
Zn Zn
Define Zn Zn-12m, then:
w (x1)) (x1)
96 (x1)
(67)
Ze
By induction, Zn IIk= Zk. and write:
w (x1)) (x1)
2 (x1) =
(68)
Il Zx
Then, the proof concludes.
Corollary 5 (Normalization Constant of ORW-CFM), The normalization constant after N epochs is: ZN / w (21) q (x1) dr
(69)
Proof of Corollary 5 By induction on the normalization constants Zn and the fact that o (z1) must integrate to , we have:
Zk / (xx)) (21) dz
Z
(70)
Remark (RL Perspective of ORW-CFM). From a RL. perspective, this iterative reweighting process can be interpreted as analog policy iteration:
Policy Representation: The data distribution d (x1) represents the policy at epoch n.
Policy Improvement: The application of the weighting function w (x1) x r(æ1) corresponds evaluating and improving the policy based the rewards.
Value Estimation: The weighting function be seen the expected return value function guiding the policy update.
Exploration Exploitation Trade-off: Similar to the exploration-exploicai trade-off, there is a balance between focusing on high-reward regions and maintaining diversity.
This process also aligns with the concept of importance sampling in RL, where the probability of selecting certain actions (data points) is adjusted based their estimated value, and enhance the high-rewardee probability area and reduce the cow-rewarded probability area.
Corollary 6 (Convergence Behavior). According to 2 as N - oo, the distribution q (x1) may become increasingly concentrated on the subset of X' where (I1) is maximized.
Corollary (Mode Collapse visk/Overopiimizzation) Overemphasis on high-reward regions can lead to lack of diversity in generated samples, similar to mode collapse in GANs (Goodfellow et al. 2014) or overoptimization in fine-tuning diffusion models (Black et al. 2024).
In this paper, we introduce two methods to handle the Overoptimization and ease the mode collapse risk online RW-CFM algorithms. The first is to introduce a W2 distance bound between fine- tuned model On and reference model O which works similar to the KL bound in previous RL (Schulman et al. 2017) and RLHF methods (Black et al. 2024 Rafailov et al., 2023). The second to incorporate entropy regularization (i.e., w(x1) exp(T x1)) smoothing the weighting function (i.e., (xx)) = Softmax(r * #z 22)) can mitigate excessive concentration or extremely greedy policy.
33

---

Published as a conference paper at ICLR 2025
Lemma (Limiting Case). We now consider the limiting cases where N -
If w (æ1) > 0 for all æ1 E X and attains its maximum at oi, then as - cc, the distribution N dg (I1) converges Dirac delta function centered at 2
lim 16 N (21) 8 (x1 xi)
(71)
Proof of Lemma
/ Given that (21) > 0 for all xι X and attains its maximum xj. we define:
(a,)
<(x1) w (x])
(72)
which means, cx1)) 1 since w (at) w (zt) and S cx1)) for 21 * 2j, since w (x1) < w (æt).
Then, we can rewrite 4/ (a1) using < as:
v (x1) [w (zt)] N (21) v g (x1)
(73)
Zn
Then for * 2i, we have: ( 0 as N 00 since c(x1) 1. Thus % (x1)
0, V. * 27
And for Iı = zf. we can have <(xt)* = 1, and % (x1)
2N
And we can have the normalization constant follows:
ZN / w (x1)" q (x1) [w (xi)) / (x11) q (x1)) dr
(74)
Similarly, we can obtain N ~ m1 (x7)) a (x7) as N + 00.
Then, we can have the limit behavior: For 21 * xf:
[w (xj))* (21) (x1)
q (21)) g
40 N
(I1)
0.
(75)
Iw (x+)) a (z1)
(zf)
For 21 = xi
w [2t)] (zt)
/6 (x1)
(76)
(w (z+) (zt)
Then, we can have:
lim ded (21) 8 (zı x;))
(77)
According to Lemma iteratively utilize the ORW-CFM loss to fine-tune the flow matching model without bounding the distance between reference model and fine-tuned model may lead us to greedy policy over 21, which may induce the overoptimization problem (Black al. 2024). In this paper introduce W2 Distance into the ORW-CF training to bound the distance between reference model and fine-tuned model.
34

---

Published as a conference paper at ICLR 2025
C.4 WASSERSTHIN-2 2 DISTANCE FOR FLOW MATCHING
We consider two flow matching models parameterized by 0₁ and 02, inducing time-evolving data distributions pt (x) and p? 8a (I), respectively. The models define vector fields e (1,x:) and v (t, x)
that transport an initial distribution (z to final distributions pi (ar) and pi (æ:) at time = 1.
We aim to bound the Wasserstein-2 distance between p| (a:) and p| (a)) in terms of the difference between e and y²
First, we can write down the Wasserstein-2 Distance as follows (For ease of reading, we rewrite the Theorem to be proved in this section as follows):
Definition C.1 (Wasserstein-2 Distance). Given two probability meusures iu and v on R", the squared Wasserstein- 2 distance between μ and v/ is defined
| 1 J
W (2 v)
inf >EIl((p.))
(78)
JR
where [1 v) denotes the set of all couplings of y and v; that is, all joint distributions IR' x R with marginals and v.
Theorem (W2 Bound for Flow Matching). Assume that (l, Lipschitz continuous in æ with Lipschitz constant L. Then, the squared Wasserstein-2 distance between the distributions ph and pi induced by the flow matching models at time t = 1 satisfies:
W2 (p''.pf' S 25. S/ B Benpe (s, D (s, (1]) ds
(79)
Before providing the full proof of Theorem Lemma:
3 we first want to prove the Differential Flow Discrepancy
Lemma 5 (Differential Flow Discrepancy Lemma).
Let of; (æo) be the flow map induced by v 02 (ro) satisfies the differential inequality: starting from To. The difference A (zo) oi 9 (zo) oi
| (xo)]l - 0000 (1 (2o) - 00 (t, 08 (zo)) II (2o)
(80)
dt
Proof of Lemma
Let A. (zo) 0/2 (xo) 0. (xo). Then:
d A. (xo) = u (t.00t (zo)) 00 (1, (zo)
(81)
dt
Adding and subtracting 0z (4 (00)
d Ae (ax) = n/ (t (zo) - 22 (t (mo)) +o02 (t. (zo)) - (t.0² (2o)) (82)
dt
6.(()
8y(1)
Assumption 1 (Lipschitz Continuity). Assume that wh is Lipschitz continuous in with Lipschitz constant L,
110% - (1.2il S Lolkæ vll
(83)
By the Lipschitz continuity of Oz
[()(())] s L |I (fo)ll
(84)
Thus:
A (2o)
$ (1,5(1)) + I. IA, (2o)
(85)
di
35

---

Published as a conference paper at ICLR 2025
Then, the proof of Lemma Sonclludes
Proof of Theorem
Using Lemma 5 define f(t) ||, (zo) |I. Then:
f(t) S [1,. L
(86)
Multiplying both sides by e Ltt
d
di (e- IM f(1) S -- (+)l
(87)
Integrating from 0 to
f()- f(0) - / (
c Le -
ds
(88)
Assumption Assume that the initial noise distribution 2a ~ p(xo) are shared two flow matching models, thus f(0) = Ao (zo)ll =
Since f(0) 0:
f(t) s etu S "e-==110(00) (1o(((ls)
(89)
Att=1:
1(1) S b (- ebe 11002 (s, (zo)) - 0² (s. (fo)) ds
(90)
Taking the expectation over up po and applying Jensen's inequality:
Exo [f(1)]] S E22L S22 (( [[l00 (s,
00 (s
(91)
Lemma Let o: (ao) be the flow map induced by v starting from Zo. The flow matching
models define vector fields n (t,r) and e (t,x) that transport an initial distribution po(x) to final distributions p (x) and 82 (x) at time t = 1. Let A. (xo) 0.: (2o) of² (æo) and f(t) || (I Then, the squared Wasserstein-2 distance between the distributions p1
0n and p induced
02
by the flow matching models at time t = satisfies:
W2 pi .p²)) S Ezo [f(1)
(92)
According to Lemma we have:
-0= [ Banpe [s,)) (s)))))))"((s) (s, ds
w22 pt. .Pt)) s Ezg [f(1)]] s 2r.
(93)
Then, the proof Theorem] concludes.
Proof of Lemma
Recall the definition of the Wasserstein-2 2 distance:
/ x || dz
W
91 yell (pi x2))
inf
(94)
We construct specific coupling 7 between pf and pfa by mapping the same initial sample zo ~ Po through both flow maps:
01 881 (xo).
36

---

Published as a conference paper at ICLR 2025
02 21 07 (zo).
This coupling 7 is defined via the joint distribution of (z9) . By the definition of the Wasserstein-2 distance, for any coupling 7
W pft.pp"' - 1 x= | ull dx,,))
(95)
Using our constructed coupling:
f JaMXa | y112d v) -Ezo ||09 (xo) 03e (200)]²)
Ex [f(1)²]
(96)
Thus, have:
W3 pp=) (/²) s EEo [f(1)
(97)
Then, the Proof concludes.
Corollary 8 (Uniform Vector Field Deviation Bound). If there exists constant such that for all E [0,1] 1]:
E anp [1120(0..2) (s, (s, S²
(98)
then the Wasserstein-2 distance between the induced distributions is bounded by:
W2 pp PpT)) S eE
(99)
Proof of Corollary 8 Given the assumption in Corollary g from Theorem 3
w33 (p) s 222 S( ed 22
(100)
Thus:
W2
8: pi .p²))
Ec
(101)
L
Corollary 9 (W2 Bound for Flow Matching Fine-tuning). Considering have two flow matching
models, the fine-tunee model daa and the referencee/re-treun -triied model ont According to Theorem
3 we bound the Wasserstein-2 (W2) distance between the distributions po.. (21) and po (I1)
iduucdd by these two flow matching models:
(zo))l1] S 2bx S (xemp((s)))) [110t((s)0) Pna (8, [))] ds
Exzo [|A
(102)
Proof of Corollary g Set 0 = 0n and 02 = frrs according to the Theorem
3 the proof concludes.
37

---

Published as a conference paper at ICLR 2025
Corollary 10 (Monte Carlo Sampling for W2 Bound in Flow Matching). However, the integral of time in Theorem is normally intractable and hard to obtain. In practice, we can use Monte Carlo sampling approximate the integral:
J.
E (app((0))) (1,)0 e (x, ds
(t, EU, (0.1))) (0 ( (t, 0 (t.
(t. 2 1)
According to Corollary
10
obtain the W2 bound RW-CFM loss follows:
LRW-CFM CPF--CFFW-2 W2 = how.://....arrrgttaeeeeeeeeaaa (x1) [[ue 0) ue (xr xr x=))l
* (4, 10 (4, (1,=))11
wherein On = 0 is our fine-tuned model, and Oer is reference/pre-trainee model, a is the trade-off coefficient, indicates how close want the final convergence point to be to the reference model.
Also, according to Corollary T0 we can obtain the W2 bound ORW-CFM loss as follows:
W2
(w (21) [1;; 0) (x 21]l ?
+ a [e, (t, 2:) Uo_,, (t, x)) ||
wherein the trade-off coefficient, indicates how close we want the final convergence point be to the reference model (See the Divergence Controlled by in 4
38

---

Published as a conference paper at ICLR 2025
C.5 RW-CFM WITH W2 DISTANCE BOUND
In this section, we provide detailed analysis of how the inclusion of a Wasserstein regularization term in the Reward-Wel -erihheedd Conditional Flow Matching (RW-CFM) loss affects the induced data distribution Pnew (x1). Specifically, we prove that the new distribution becomes proportional to (21) q (z1) exp (-aD (z1)), where D (z1) divergence measure between the fine-tuned model and the reference model.
Theorem 7 (Learned Distribution Induced by RW-CFM-W2 Loss). Let w X + [0, oa) be measurable weighting function, q (x1) be the original data distribution, and Oet be fixed reference model. Consider the loss function:
LRw-CFM-w2 (0) = E, 1) ((1)) (1/21) (a1) lo (z; 0) - (x x1)1l +o Io (2;; 0) de (r;; Pny )1
(103)
where us (x 21) is the true vector field conditioned on Iι, ve(2; 0) is the model's ss vector field, and 0 is the regularization coefficient. Then, minimizing LRw-crrwz (0) over 0 induces data distribution:
Pnew (I1) x w (x11) (21) exp (-3D (x1))
(104)
where D (x1)= (e
(tt) n x)) (1 -a, O is scaling constant.
Proof Theorem 7 We can re-write the RW-CFM-W2 Loss follows:
LRW.CFM.W2 (0) Er w (21) [1;, 0) un (x x1)l + a 1I (1, amt (1, { (105)
where a controls the regularization strength. Furthermore, we can write the loss function as an expectation over 21
Ciotal (0) / (x1) [w (x1) Lerm (x1) + aeLivg (x1) dæ
(106)
where the first term gives Letn (x1) (ppp)())) [1lua(2; 0) us (x m1)1 and the second
term gives Lreg (I1) = (x)
(t,) poor (t, x1]]
Assumption 3. Assuming that the expectation Leog (x1) has an upper bound e
Leeg (I1) SE, VT1
(107)
Since we" re interested how Leog (21) varies with æi, we can let D (x1) = Lrog (21).
Then, our total loss is:
Liaatl (0) - (x1) [w (x1) Lem (x1) + aD (x1) dz
(108)
We consider the optimization over 0, which indirectly affects the distribution Prew (x1) that the model learns.
Then, consider the model as implicitly assigning higher importance to values of Φ that minimize the total loss per 21
7. (m1) (अ1) Teetn (x1) aD (x1)
(109)
39

---

Published as a conference paper at ICLR 2025
Remark (Probabilistic Perspective). In optinization, samples 21 with lower L (21) have higher impact on training. To make this explicit, consider a probabilistic model where the likelihood of xı is proportional exp (->L (x1), where 7 is positive constant Hiinton 2002).
Therefore, the induced distribution over 21 becomes (detailed proof can see App. C.I
Phew (x1) (x11 exp (-7 [w (x1) Letm (x1) aD (x1)))
(110)
We can also induce the similar results from the perspective of energy function (LeCun et al. 2006
Lemma (Energy Function Perspective). The per-sample loss contribution L (x1) can be interpreted as an energy function E (z1), and the induced distribution follows a Boltzmann distribution LeCun
al. 2006);
Pnew (x1) x (21) exp (-E
(111)
(31) exp (-7 (1r (x1) Tegfon (x1) aD (21)))
Proof of Lemma 7 By interpreting the per-sample loss an energy E (x1). ) we have:
(xx) (x1) w (x1) Letm (x1) + aD (x1)
(112)
The induced distribution is then:
Prew (x1) x (I1) exp (-E (x1))
(113)
q (x1) exp (-> (w (x1) Teem (x1) aD (x1)))
This aligns with the Boltzmann distribution in statistical mechanics, where higher energy states (higher loss) are less probable.
Besides, based on Theorem we can have:
Puew (x1) (I1) exp (-y+ L (01))
(114)
(x1) exp (-> (w (x1) Letm (x1) aD (z1)))
Assumption 4. Assuming Lejin (x1) is approximately constant or negligible compared to (21).
namely:
When the model fits the data well, Lefin (x1) small. (i.e., w (æ1) dominates)
The term w (21) Legi (x1) becomes proportional to w (x1)-
Remark, Generally speaking, if we fine-tune from a well-learned pre-trained model, the model can usually fat the data well, Lifin (x1) is rather small. In practice, (x1) can generally dominates.
Thus, we can write:
Pnew (x1) x q (x1) (xz) exp (-3D (x1))
(115)
where 3 = -yck absorbs the constants. Then the proof concludes.
40

---

Published as a conference paper at ICLR 2025
Remark (Importance Sampling Perspective). We can also derive similar results from the perspective of importance sampling. We start from the total loss function:
Lional (0) f w (xi) (Z1)
Lefin (x1) +
D (z1)
da
w (x1)
(116)
Tuefa (211 w (x1) (211)
We denote the reward weighted distribution as pwn
(w(()) w (x1) (21)
(117)
Then the loss function can be re-write as follows:
Liotal (8) = -Ezmpa (ar) [ngin (x1) + a
((1))
(118)
w (x1)
Then the induced data distribution Pnew (x) becomes:
Pnew (Ø1) Pw (x1) (->) Lejm (x1) + (x1)
D (21)
(119)
Under Assumption 4 we can obtain:
Pnew (xI) x Pro (II) exp
7 a
(x1)
w (x1)
3
x Pro (x1) exp
w (I1) (
(120)
x w(x)) (I)) exp
(x1)
In practice, we can assume that Va1, 3 > (x11), then B dominates, and we can obtain:
Pnew (x1) w(x1) (21) exp (-3D (z1))
(121)
Also, can set Oi w(x1), then:
exp (-w w (x1) D (x1) exp (-BoD (x1))
B
(122)
wherein Bo = Cgg *
Corollary 11. As a + 0, - 0, the induced distribution Pnew (x1) approaches w (r1) (x1) /2. recovering the standard RW-CFM result.
Proof of Corollary 77
When a + 0, + 0, the exponential term exp (-BD (x))) + 1, so:
Prew (x1) (x1) (21)
(123)
Z
where Z the normalization constant ensuring that Pnew (x1) integrates to
Corollary 12. As + oc, + oo, the induced distribution Peew (x1) concentrates on minimizing D (I1), effectively aligning with the reference model's 's distribution (See cases with increased a, such as a 10, Fig. 4
Proof of Corollary 12 When + 00, B + cc, the exponential term exp (-3D (x1) )) approaches
zero unless (x)) 0. Therefore, Prew (x)) concentrates on the set of x where v, (x; f matches we (a x1). Namely, the fine-tuned model remains unchanged from the reference model.
n
41

---

Published as a conference paper at ICLR 2025
Remark. The theorem demonstrates that incorporating the Wasserstein regularization term into the RW-CFM loss results in an induced data distribution that not only depends on the reward weighting (21) but also exponentially penalizes data points based reference model (See Trade-off curve via traversing a Fig. 4 This provides a mechanism to control the trade-off between fitting the reward weighted data and adhering to the behavior of pre-trained or reference model (i.e., reward-diversiee trade-off).
their divergence D (21) from the
The divergence D (x1) acts as measure of how much the true dynamics we (xr x1) differ from the reference model's 's dynamics us (2; 0mf) ) along trajectories starting from It. By adjusting o, we can control the extent which the model prioritizes matching the reference model versus fitting the wwightee data.
42

---

Published as a conference paper at ICLR 2025
C.6 ORW-CFM WITH W2 DISTANCE BOUND
Similar to Theorem 7 we can also derive the induced data distribution learned under the online reward weighted CFM loss and W2 Distance Bound.
According to Theorem
3
We can directly introduce the W2 Distance Bound into ORW-CFM loss
equation 5
LORW- =
21) (ur (21) [1 On) We (s ==)
11o, (f,22) (t, (124)
wherein is trade-off coefficient.
For case reading, we rewrite the Theorem to be proved in this section as follows:
Theorem. Under the online reinforcement learning setup with ORW-CFM-W loss, the data distri- bution after n epochs evolves according to:
Le (x1) x [w (II) 1i' (xx1 exp (=)()
exp
(125)
1100 (1,,)) oes 100 (1,20] ||]]]] = yoa, 0- denotes the model
((ppa) parameters after epoch n 1.
where, (2-1((x1)
Proof of Theorem 4 Similar to the proof of we also proceed by induction.
Case (Base Case (n 1)). At 1, d (#1) is obtained by training on data sampled from q (œ1).
From previous analysis, we have:
db (z1) w (x1) (2i) exp (-3D" (I1)
(126)
where D (I1) measures the discrepancy between the initial model (before training) and the reference model.
Assume that at epoch n 1, data distribution is qp 1-1 (x1). During epoch n, the model trained
using data sampled from 93 1-1 (x1). By applying ORW-CFM-W2 Loss, compared to RW-CFM-W2 Loss, the loss function remains the same, but the data distribution changes to % n-1 (z1).
The total loss associated with epoch is:
L (z1) w (x1) Lii (x1) a1" (x1)
(127)
where L (21) is the data fitting loss of CFM epoch n, and D (x1) is the discrepancy from the reference model using parameters g--1!
Similarly, assume that the probability of being selected for training at epoch is propor- tional to exp (->i" (z1)). Therefore:
o (x1) % 2-1 (011(1))
(x1) exp (-rin (m1))
(128)
Substituting in (x1):
1 (x1) -0- (m1) exp (-> [w (a1) Di (x1) a -xx))
(129)
Assuming Lam (x1) is approximately constant or negligible compared to a (xI) (i.e., well learned), simplify:
d (x1) (x1) - exp) (-3D=-1 (a1))
(130)
43

---

Published as a conference paper at ICLR 2025
Then, the induction step completes. This concludes the proof.
Besides, we can also derive the proof via the Importance Sampling Perspective as equation 16 process is very similar to the offline cases. We just need to substitute L" (x1) into equation 16
The
take
out the weights w(xx) to form wxx)) wxxx)a"--<x (21) and reuse Theorem 6
Corollary 13 (Recursive Formulation). The data distribution after N epochs training under LORW-CFM-W2 is given by:
u (a1) & (we (x1)) exp -3 3 Da-1 (x1)
(131)
Proof of Corollary 13 By recursively applying Theorem
we have:
N (x1) w (1n N- -1 (x1) exp (-3D N-1 (x1))
w (x1) (w (x1) q exp (-3DN-2 (x1))) exp (-8DN-1(z) (z1))
(132)
N
& (21)) 9 (x1) exp -B Σ D=-1 (x1)
Corollary 14 (Convergence Behavior). D (211) decreases with each epoch and (x1) is constant (i.e., w (I1) = w ). then ui (x1) converges towards distribution that places higher probability on data points where the model aligns closely with the reference model (See cases Fig. 2
As D" (x1) decreases, the exponential term exp (-8 EN D (m1)
Proof of Corollary 14
increases for those #1 where the model aligns better with the reference model. With (z1) constant, the weighting does not favor any specific 2i, so the distribution q (æ1) is increasingly influenced by the cumulative discrepancy term, leading to higher probabilities for data points where the model matches the reference model.
44

---

Published as a conference paper at ICLR 2025
C.7 EXPONENTIAL WEIGHT FUNCTION
As we discussed, normally we can have w(x1) r(zi) and w(x1) F(r(x))) (). In practice, normally have several form of weighting function to enable w(r) > (Peng et al. 2019 Peters & Schaall 2007 Black et al.) (024), which maps from reward function to weight. In this section, based on the previous derivation, we focus on a special cases fine-tuning flow matching model with ORW-CFM-WV W2 loss, where w(x1) exp(T r(z1)), wherein the temperature coefficient.
Then simply have the following Theorem (For of reading, rewrite the Theorem to be proved in this section as follows):
Theorem. Given w (x) ) = exp (Tr (21)), the induced data distribution after N epochs of training under ORW-CFM-W2 Loss is:
g (x1) x exp (Nr(c1)) ED D (mi) g
(133)
Proof of Theorem 5 According to Corollary 13 we can re-write the final induced data distribution
learned under CRW-CFM-W2 loss function
/ (11) (ne (21)) (xp -B
Σ D (21)
(134)
Let r (111 be a reward function, and t > 0 temperature coefficient. The weighting function becomes:
w ( exp (r
(135)
Substituting into equation 134
N (11) x q exp (NN) 8 qe
Σ Dー- (z1)
(136)
Then, the proof concludes.
Corollary 15 (Limiting Behavior). If 3 E (z1) grows proportionally with N, iee. BEN D= (x1) N8 (x1), the induced data distribution becomes:
det (xx) x exp (N (rr (æ1) - 5 (z1))) (æ1)
(137)
This suggests that the effective reward is adjusted by 8 (x1) due the W2 distance bound.
Proof of Corollary [5 Given the proportional growth, we can write:
n- (x1) = (m))
(138)
Substituting into the expression for o (x1):
N (x1) eNp(Nr (r
(x)) (x))) (xx) (139)
45

---

Published as a conference paper at ICLR 2025
Corollary 16 (The Effect of W2 Discrepancy). The W2 distance bound D (x1) serves regularization tern that penalizes the discrepancy between the current model -0" and a reference model elt This term affects the induced data distribution by effectively adjusting the reward:
i ( 5(21)
The induced distribution then becomes:
q (x1) & (T
(140)
This adjustment means that the W2 distance bound reduces the attractiveness certain samples based on how much the model deviates from the reference model when conditioned on those samples reward-distance (See reward- distance reward-off in Figs.
According to different hyper-parameter settings, we discuss the limiting behavior of the method of the following cases:
Case (Dominant Reward Term). Df (x1) grows slowly with N or 3 is small, the reward
D= (x1). The induced distribution 4o (xr) concentrates on the 2 maximizing r. (II), similar to the case without considering D= (x1) (See 10 in Fig. 2
dominates the regularization term, namely T*T (x1) > 3 S D=
Case 5 (Dominant Regularization Term). If D (r1) grows rapidly with N or 3 is large, the
regalarization term dominates, namely 3 , D=-1 > + (x1). The induced distribution heavily penalized for deviations from the reference model, potentially leading to a distribution similar the initial distribution (21) (See = 10 in Fig. 3
Case (Balanced Terms). When both terms are balanced, namely T * (11) B B D (æ1). the induced distribution reflects trade-off between maximizing rewards and staying close to the reference model. The model achieves a compromise, focusing on high-reward samples that do not deviate significantly from the reference.
Lemma 8 (Growth Rate of D (x1)). Assuming that D (I1) S Duax for all and æι, the total W2 distance bound over N epochs satisfies:
8
D= S NBDDma
(141)
n=1
Proof of Lemma
Since D (x1)
Dmax summing over n
3
D=-1 (xI) S BNN max
(142)
Based on the above derivation, we can discuss the regularization effect D (xi) and how to tune the trade-off coefficient 3. The W2 distance bound term acts regularizer that:
Penalizes Model Drift: Prevents the model from deviating excessively from the reference model o.
Encourages Smooth Updates: Ensures that updates to the model parameters are gradual, promoting stability during training.
Balances Exploration and Exploitation: Mitigates the risk of overoptimization high-reward samples by considering the model's 's consistency.
Besides, the parameter controls the strength of the regularization:
Large B Strong regularization, the model remains close the reference, potentially underexploiting high-reward regions and inducing more diversitylexploration (See 7 in Fig.
46

---

Published as a conference paper at ICLR 2025
Small 3: Weak regularization, the model prioritizes reward maximization, possibly the cost of stability and induce the over-optimization problem (Black et al. (024) (See T = in Fig. IT
Corollary 17 (Optimal 3 for Exploration-Explootaiio Exploitation Trade-off). An optimal 3 exists that balances the influence of the reward term and the W2 distance bound, satisfying:
3r TF
(143)
where F is the average reward and D is the average W2 distance bound.
Proof. Setting the magnitudes of the reward and regularization terms equal, we have:
NTF = N TF
(144)
D
where is the average reward and D the average W2 distance bound.
Remark. The regularization terms crucial for ensuring stable learning dynamics, preventing catastrophic forgetting (Kirkpatrick et al. (2016). Besides, Regularization terms also offer a trust
region constraint like TRPO (Schulman al. 2015) or policy update in proximal area like PPO Schulman et al. 2017).
47

---

Published as a conference paper at ICLR 2025
C.8 BOLTZMANN WEIGHT FUNCTION
In this section, based on the previous derivation, we focus on a special cases of fine-tuning flow matching model with ORW-CFM-W2 loss, where w(x1) x Softmax(( r(z1)), wherein the temperature coefficient.
According Corollary 1 we can also re-write the final induced data distribution learned under ORW-CFM-W2 loss function as:
g (<1) & lno (x1))* q (xp))epp (-8 (-8 E Da- (a1)
(145)
Definition C.2 (Boltzmann Weighting Function). The Boltzmann weighting function over the set of all possible is defined as:
w (21) softmax (rr (x1)) exp (Tr (x1))
(146)
Z
where R the temperature parameter, Z = Ez, exp (rr (x/)) is the partition function ensuring that En, w (I1) = 1.
Then, we can have the following:
Theorem 8. With (x1) softmax (Tr (21)), the induced data distribution after N epochs is:
exp (Tr (xi)))
N
N
1 (x1) x
q (21) exp
exp
- Σ Dz (x11)
(147)
2
nml
Proof of Theorem 8 Starting from the recursive formulation:
g (x1) x (x1) ge (x1) exp (-3DN-1(am)
(148)
By induction, assume the expression holds for N 1:
exp (rr (ar)))
N-1 2 (x)) -B Σ D (m11
% N-1 (x1) -x1)
(149)
n=1
At epoch N
got (a1) x (21) 90 N-1 (1BDD)-
exp (Tr (x)()
exp (Tr (x1))
N-1
-3 > Dr- Dm-2 (m1)
N
a (z)) exp
2
7
(150)
exp (Tr (x1))
N -8
q (xx1 exp
Da D (x1)
Z
Thus, the expression holds for N. Then, the proof concludes.
Corollary 18 (Simplifying Induced Distribution). The induced data distribution can be expressed as:
N6 (z1) exp (Tr
Σ D=-1 (x1) g (x1)
(151)
n=1
up normalization constant independent of g1
48

---

Published as a conference paper at ICLR 2025
Proof of Corollary tionality constant:
18
Since Z is constant with respect to 21,
absorb Z into the propor-
T6 (x1) x exp (NTr (x1))
(x))
(- Σ D (11))
7.N
(152)
N x exp NTr (x1) - D Da-1
(21) - a (x1).
Despite the presence of the partition function Z, the induced data distribution q3 (x1) under softmax weighting ultimately depends on the unnormalized exponentiated rewards, similar to the exponential weighting case. The key difference lies in the normalization:
Exponential Weighting: Does not normalize w (æ1), leading to potential numerical instabil- ity. Softmax Weighting: Normalizes (x1) each epoch, ensuring Ex, (z1) 1.
However, when considering the evolution over multiple epochs, the normalization constants ZZ N accumulate but do not affect the proportionality of q (21) with respect ø1. And, in practice.
softmax weighting normally requires large batch size for efficient learning.
Corollary 19 (Exploration Exploitation Trade-Off). Proper tuning of T and 3 is essential to balance exploration (sampling diverse Fy and exploitation (focusing on high-reward x ):
High t Increases emphasis on high-reward samples.
High B: Increases the penalty for deviations from the reference model.
In practice, the Boltzmann Weight Function help slabilize the learning process and keep numerical stability. However, in practice, when the batch size is smaller, the Boltzmann Weight Function methods normally fail to learn optimal results in similar training steps as Exponential weighting function methods.
Remark.
49

---

Published as a conference paper at ICLR 2025
C.9 RL PERSPECTIVES OF ORW-CFM- W2
In this section, we start to offer a different perspective of the learning behavior of models under ORW-CFM-W2 loss from the perspective of RL.
We first do the following parallels: 1) Data Points I1 as Actions: In our setup, each data point 31 corresponds
action in RL. 2) Discrepancy Term as Regularizer: The discrepancy term D (I1) acts as regularizer, similar to entropy regularization or KL divergence penalties in RL.
Then, define the objective that RL agents want to optimize: Definition C.3 (RL Objective of ORW-CFM).
I(0) = (es)) |r (x1)]
x DKI. (40 (211) 1106 11---1 (21))
(153)
where 9% (x1) the data distribution at epoch n. 93 o (x1) is the data distribution at epoch
DKL. acts as a regularizer to keep the new distribution close to the previous one. X controls the strength of the regularization.
We obtain the following Theorem (For ease of reading, rewrite the Theorem C.9 to be proved in this section as follows):
Theorem. RL settings, We aim to find qb (x1) that maximizes J(0). This involves solving:
1
L (x1) = arg max Eai [r (xx1) X DKI. (q (x11 11 1/001 (x1))
(154)
9
And the optimal data distribution (policy) % (#1) can be obtained
g (x1) x d --1 (x1) exp (Ar (x1))
(155)
Proof of Theorem C.9 The optimization problem can be solved using variational methods. The KI. divergence between (z1) and 9g -1 (x1) is:
DKL (a (x1) 11 1-1 (x1)) = =f (z1) log
q (51)
dx
(156)
4 n-1
(x1)
Introducing a Lagrangian enforce the normalization constraint (x1) dri = 1
C= f a (z1)
(x11) 1 log
1
9 (
dx d (f1())
(157)
d0 - (m))
Then, take the functional derivative of C with respect q (x1) and it to zero:
5L = r (21)
1 - -1
(og
q (x1)
+ 1) =
(158)
8g (x1)
(21)
Simplifying:
q (11)
(1
r (x1)
log
+1) =
(159)
98
n
(xz)
A
Rewriting the equation:
q (x1)
log (x1)
= X
(160)
where C C= 1/2 + 17 is constant independent of 21
50

---

Published as a conference paper at ICLR 2025
Exponentiating both sides:
9 (x1)
= exp (Ar (211 XC)
(161)
9g -- (21)
Thus:
1 (21) 0/ 1/= (211) exp (Ar (
(162)
Since -AC normalization constant, we have:
/ (x1) x %0 (11) (Ar (x1))
(163)
Then, the proof concludes.
Similarly, if we introduce the Regularization Term Da (æ1). we can have the following Theorem (For ease of reading, we rewrite the Theorem C.9 to proved this section as follows):
Theorem. Considering an additional discrepancy term D (21) that acts as a regularizer We can include this term in the objective function:
= =
[r (x1) - 3Dn (x1) - x DKL (q (x1) 11% 1 (21))
(164)
Then, the Rl. agents aim to find qg (x1) that maximizes J()). This involves solving:
g (I1) = arg max
Max.xaxxaa Being [r (x1) 3-0- - - DKL (g (x1) 1165 (x1))
(x1)
(165)
And the optimal data distribution (policy) oB (x1) can be obtained as:
1 (x1)
dg n-1 (x1) exp (Ar (x1) x3D==1(x1))
(166)
Proof of Theorem
C.9
Following similar steps, the functional derivative becomes:
8L
1
Iog g (x1) + +1) = 0
= r (x1) 3D" (x1) - A % n (z1) 8q (z1)
(167)
Simplifying:
log q (x1)
(r((1))
(168)
46 (x1)
Exponentiating:
d (x1) % nn1 (z) exp (Ar (21) >3D -180 (a)) e AC
-AC
(169)
Therefore, the updated data distribution is:
z (x1) x ob 001 ex: (ar - -x3==11xx)
(170)
Then, proof concludes.
Corollary 20 (ORW-CFM-W2 As Reinforcement Learning). Based on Theorem C.9 we can rewrite the induced optimal policy in equation 164 as follows:
1g (x1)
qb 0-1 (xr exp (Ar (x1) A3D (x1))
(171)
51

---

Published as a conference paper at ICLR 2025
By setting 3' = * w((1) exp(Ar(z1)). we can obtain:
d (x1) w(x1) 1 (x1) (-B' 1"-1 (21))
(172)
which is the same form of induced data distribution over 21 in Theorem 4
Therefore, ORW-CFM with W2 Distance Regularization is actually solving a Reinforcement Learning Problem, where the objective is trade-off between reward and discrepancy.
The update rule for % (21) mirrors the policy update in RL methods like AWR, where the new policy proportional to the old policy multiplied by the exponentiated advantage.
The discrepancy D (x1) acts as a penalty term, discouraging the model from deviating too far from the reference model, similar to how KL divergence or entropy regularization is used in RL
(Peng et al., 2019 Schulman et al. (017)) RL.F (Black et al.] 2024 Rafailov et al. 2023) to maintain policy stability and keep the fine-tuned model close to pre-trained model.
Remark (Connection to AWR). AWR (Pngg et 2019) also updates the policy by weighting actions according to the exponentiated advantage. Similarly, our update rule similarly weights data points T1 according exp (X r (I1)).
Remark (Connection PPO). PPO (Schulman el al. 2017 a clipped objective to prevent large policy updates. In our setup, the regularization term D (x1) prevents the data distribution from changing too drastically, Besides, D (x1) can also used limit the distance between ref model and fine-tuned model.
Remark (Exploration and Exploitation Trade-Off). The term exp (A x (21)) encourages the model to prioritize data points with higher rewards. While the regularization term exp (-3' D (x1)) ensures the model does not deviate excessively from the reference, promoting stability and diversity (e.g., avoid the policy collapses (Haarnoja et al. (2018)
Corollary (Learning Flow Matching as Reinforcement Learning). In general, by adopting reinforcement learning perspeclive, we have derived the evolution of the data distribution (x1) under self-training with the RW-CFM-W2 loss. The derivation shows that:
d (x1) x di 1--1
(173)
This update rule mirrors policy updates in RL methods like Advantage Weighted Regression, where the new policy is updated based on the exponentiated advantage while incorporating a regularization term ensure stability.
52

---

Published as a conference paper at ICLR 2025
D EXPERIMENTAL DETAILS
D.1 IMPLEMENTATION DETAILS
To ensure reproducibility, we built our code upon the open-source flow matching code base called TorchCFM (Tong et al., 2024a b) to demonstrate that our method is general enough to be applied to fine-tuning existing flow matching based methods, which provides efficient and theoretical sound tools for training continuous normalizing flow models. TorchCFM offers simulation-free training objective for Conditional Flow Matching (CFM), enabling conditional generative modeling and accelerating both training and inference processes.
We utilized the cxisting unconditional flow matching training baselines provided by TorchCFM on datasets such as MNIST and CIFAR- 10. Our implementations of the RW-CFM, ORW-CFM, and ORW-CFM-W2 W2 methods are based on the TorchCFM codebase, and we adopt the same network archi- tecture U-Net (Ronneberger al. 2015) for the image generation tasks. The pre-trained models used in our experiments are trained with the standard CFM loss as defined in Equation equation on CIFAR-10 and MNIST datasets.
For the image compression task, we directly adopt the reward function from DDPO (Black et al. 2024). In the even-number generation task, trained binary classifier distinguish between and odd digits, which serves as the reward signal. For the text-image alignment task, we employ the CLIP model from OpenAl (Radford et al.. 20 calculate the similarity between text and images; the CLIP model is pre-trained and provides a probability score representing the alignment between the generated image and the given text prompt.
The reward functions for each task can be found in App. F And the hyper-parameterd can be found in App.
D.2 RESOURCES USED
All experiments were conducted system with worker equipped with an 8-core CPU and, an NVIDIA A10 GPU, and memory of 32 GB. The MNIST task cost about 4-8 hours for training, while image compression and text-image alignment task costs about 2-3 days for training.
53

---

Published as a conference paper at ICLR 2025
E HYPER-PARAMETERS
In general, in this paper, we propose a general and easy-to-use RL method for fine-tuning flow matching generative models. Our algorithm is directly built upon the flow matching training code in TorchCFM (Tong et al., 2024a). To facilitate reproducibility and highlight the effectiveness of our proposed RL tinettuning algorithm, we maintain the same flow matching training parameters and the calculation of in OT paths as in the TorchCFM examples. We have conducted detailed ablation experiments in the main text and appendix on the two most important hyperparameters our method: the entropy control parameter and the W2 distance regularization parameter
DRW-CFM- method involves only few hyperparameters, and we have carefully ablated those that affect the convergence behavior in our experiments of Sec. [5] namely and a. Specifically, controls the convergence speed of the fine-tuning method, while a has significant impact on the convergence point.
Massive experiments demonstrate that the parameter r effectively controls the convergence speed of the policy model. In the absence of W2 regularization, directly influences the rate at which the policy collapses to a delta distribution, as clearly demonstrated in the theoretical analysis of our exponential form weighting form
Interestingly, previous methods such as DDPO (Black et al. 2024) did not include KL regularization terms other explicit distance constraints, making it difficult control the divergence between the optimized policy and the pre-trained reference model. As a result, these methods often struggle to avoid over-optimization and policy collapse. Therefore, when employing our method, we strongly recommend setting a > 0 to provide sufficient W2 regularization and prevent policy collapse.
Remark. practice, we encourage users/readers to first use a small a to fine fine-tune the model with ORW-CFM-WF W2 method. along with relatively large Then, different values of a be tried to customize the optimal generative models that fit the user's 's needs for reward and diversity.
For experiments, we use batch size of 64 and employ separate test process to evaluate our performance.
In general, the proposed BRW-CFM-WE W2 method effectively fine-tunes flow matching models. We can control the convergence speed through and adjust the diversity of convergence points and the final reward through a. This provides sufficient flexibility for customizing the convergence policy to meet needs in various scenarios and enables us achieve stable and efficient RL fine-tuning using a simple reward- weighte method.
54

---

Published as a conference paper at ICLR 2025
F REWARD FUNCTIONS
In this section, we will describe the reward function we used in our experiments to demonstrate the effectiveness and generality of our method.
F.1 TARGET IMAGE GENERATION
In the MNIST target image generation task, our objective is to fine-tune the model to generate images of even numbers. We used pre-trained binary classifier C(x),, where x is the generated image, to determine whether the image corresponds to an even or odd number. The classifier outputs the probability of the image being even odd:
Ceven (z) p(even x)), Coda (z) = p( odd x))
(174)
The reward function is then defined as the difference between the classifier's probabilities of the image being even versus odd:
r(x) p( even
p(odd
(175)
Where p( even x²) Caven (x) the probability that the generated image is classified an even digit (i.e., digits (0, 2, 4, 6, 8} ), and pl odd x) = Claa (Ix) is the probability that the generated
image x is classified as an odd digit (i.e., digits {1,3, 5, 5,791
This reward encourages the model to maximize p(even x)) while minimizing p(odd x), guiding the model to generate even digits more frequently. The fine-tuning objective can be formalized as:
max (x) = Exzgg(x)) [Ceven (x) - Coat (a:)]
(176)
This reward function enables the model to be fine-tuned to produce numbers without needing explicit labels or dataset directly composed of even digits.
Furthermore, is worth noting that for this type of binary classifier-guidee reward, modify the task to generate the opposite class of images by simply using the negative of the reward function. For example, to generate odd digits, we define the reward r(x) = p (odd x) p( even | [x). The results for generating odd digits are provided in App. G.2
easily
F.2 IMAGE COMPRESSION
For the image compression task, we followed the reward function from DDPO (Black al.) 2024). where the goal was to either minimize the file size of the images after JPEG compression. The reward r(x) is proportional diversity using the regularization parameter a, which induces varying degrees of divergence between the fine-tuned model and the reference model. A lower a prioritizes compressibility, leading the model to generate images that occupy minimal storage space after compression, while higher increases diversity in the generated outputs.
the compression rate. We controlled the balance between compression and
F33 TEXT-IMAGE ALIGNMENT USING CLIP
For the text-image alignment task, utilized the CLIP model, which computes the similarity between text prompts and generated images. The reward function in this case was based on the probability score generated by the CLIP model. Given an image & and a text description i, the CLIP model outputs a probability Peup (z, t), which indicates the alignment between the image and the text. The reward function for this task was defined as:
r(x, t) Petp (I,
(177)
This reward allows the model to fine-tune its output images to better match the provided text prompts, optimizing for higher similarity as measured by the CLIP model. In practice, will use
55

---

Published as a conference paper at ICLR 2025
difference form in probability to promote training like "cat)) Peup (x,, "an image of cat") Pelip (z, "not an image of cat").
56

---

Published as a conference paper at ICLR 2025
G ADDITIONAL EXPERIMENTAL RESULTS
G.1 ALIGNING SMALL-SCALE FLOW MATCHING MODELS
We first evaluate by guiding unconditional generative model pre-trained on CIFAR-10 (Tog et al. 2024b) to maximize the text-image similarity score based on CLIP (Radford et al. 2021a). specifically aligning the generated images with given text prompt via reward (e-g.. "an image of cat"), As shown in Figure [1] without W2 regularization (i.e., tends to fall into a greedy policy that maximizes reward but sacrifices diversity. However, with W2 regularization, agent achieves more diverse convergent optimal policy while maintaining similar final performance.
0), the ORW-CFM method
1.0
0.07
0.9
0.05
0.7
ond pretsinc
0.
0.00
0.00
200
460
650
802
1060
5 260 450
650
800
1200
Teining Step
Thening Step
(a)
(b) a=
(c) Reward Curve
(d) W2 Curve
Figure 11: Diversity Control with different while = and reward is Text-Image Similarity Probability Score (e.g., Petip (x, "an image of cat") for a cat image). W2 distance is estimated by its upper bound (see App. C.4b. In practice, will use difference form in probability to promote training like r(z, "cat") = Pelip (z, "an image of cat") Pelip (z, "not an image of cat").
G.2 REWARD SHAPING
To demonstrate the adaptability of our model various reward objectives, can directly modify the reward function to change the fine-tuning target. For instance, we transformed the original even-pree oven-preference reward function, numbers instead. The results show that our method successfully fine-tunes the flow model according the new reward objective, leading to the generation of only odd numbers.
p(even) p( p(odd). into r p(odd) p(even) to generate odd
We set w(x) as w(x) = exp(T r(z)) with = 1 (except for the last one with r = 0 as control group), and experimented with different values of control the divergence between the fine-tuned model and the reference model, which induces diversity in the convergent fine-tuned model. As illustrated Fig. G.2 without W2 regularization, our method converges to the optimal generative policy, though with limited diversity. However, as the W2 regularization increases with higher values a, our model still performs the odd-number generation task effectively, while also showing increased diversity in the generated samples.
G.3 MORE TEXT-IMAGE ALIGNMENT RESULTT
In this section, we provide more text-image alignment results:
57

---

Published as a conference paper at ICLR 2025
5
3
3
3
5
3
3
3
5 5
3
3
5
5
3
3
3
3
5
3
3
3
5
3
7
3
3
3
3
5
7
5
3
3
3
3
9
7
3
3
3
3
9
3
7
7
3
5
7
9
7
3
7
7
9
7
3
3
5
7
9
7
5
3
5
9
3
3
0
3
3
0
3
3
3
3
3
7 7
3
3
7
2
7 3
3 3
3
3
3
3
2
9
0
3
7
3
7
2
9
7
3
3
3
7
2
9
3
F
3
7
3
5
9
3
3
2
3
7
7
5
7
3
9
7
7
3
9
7
7
3
3
5
2
3
3
5
3
3
3
5
3
3
3
3
3
5
3
3
3
5
3
3
3
5
7
7 9
7
3
5
3
7
7
2
?
3
5
7
0
3
3
7
7
O
3
5
7
7
O
(a)
0
(b)
0.1
(c)
0.2
(d)
0.3
3
3
3
9
3
5
3
5
3
3
0
S
5
3
9
S
3
3
0
3
S 9
0
3
3
9
3
7
5
3
3
3
9
3
9
3
3
3
9
7
3
3
3
3
9
3
3
5
7
9
7
3
7
7
9
7
3
5
7
9
1
3
7
7
9
5
3
0
3
1
3
0
3
3
3
7
2
7
3
3
7
3
7
7
3
7
9
5
9
3
5
F 3
7
3
5
9
7
5
F
3
7
5
9 7
7
5
F
3
7
3
5
9
3 1
3
3
7
5
7
3
7
5
7
7
7
3
7
7
5
3
5
3
5
3
3
3
3
3
5
3
3
5
S
5
5
3
3
3
3
5
3
3
3
7 9
?
5
7
5
9
5
7
7
0
7
3
5
7
7
3
3
7
7
0
(C)
0.4
(f)
0.5
(g)
0.6
(h)
0.7
3
5
5
3
5
3
3
7
3
S
6
3
L
>
3
3
3
9
3
7
3
3
3
3
9
3
7
3
3
3
3
9
3
7
0
2
8
/
4
8
1
3
5
7
9
3
5
?
1
3
5
1
9
3
5
?
1
3
5
7
3 7 3
5
?
1
3
4
7
2 7
3
9
7
7
3
1
1
0
7
2
1
2
0
2
6
3
7
3
7
9
5
3
3
5
7
5
F
3
5
/
6
6
2
9
7
4
7
7
7
3
7
7
7
3
3
7
3
7
5
7
6
4
4
7 2 0
3
5
3
3
3
5
3
3
3
5
5
3
3
7
8
C
9
3
5
3
7
7
0
5
3
5
7
7
0
?
3
5
3
7
0
8
8
8
/
7
9
(i)
0.8
(i) 0.9
(k)
1
(1)
T
Figure 12: Odd Number Generation in MNIST.
Q
(b) a=1
Figure 13: Diversity Control with different while and reward is Text-Image Similarity Probability Score (c.g., Peti (æ, "an image of dog") for a dog image). In practice, we will use difference form in probability to promote training like r(a, "dog") = Petp (æ, "an image of dog") Potip (z, "not an image of dog").
58

---

Published as a conference paper at ICLR 2025
(a) a
(b) a=1
Figure 14: Diversity Control with different a while and reward is Text-Image Similarity
Probability Score (e.g., Pelip (z, "an image of bird") for bird image). In practice, we will use difference form probability to promote training like r(x, "bird") = Pelip (z, "an image of bird") Palip (x, "not image of bird").
59

---

Published as a conference paper at ICLR 2025
H ALGORITHM PSEUDOCODE
In fact, we can implement the proposed fine-tuning algorithm by making minor changes based on the conditional flow matching implementation OT-CFM (Tong et al. 2024a). In this section, will demonstrate the Pseudocode of each fine-tuning algorithm proposed in this paper. For the origin CFM pipeline we refer the readers (Tong al. 2024a
Algorithm Reward Weighted Conditional Flow Matching (RW-CFM in Theorem
Require: Pre-trained model V8a and induced sample distribution ( efficiently samplable
Pe (a(x)) computable de (a((a)); initial network to, reward function r(æ1), ), weighting function
w(x1)= = =((((z)) where w(x1) o r((1), entropy hyper-parameter т, F can be exp or Softmax
< Opr
while Training do
21 ~ -0 [1);; t ~~ ((,1)) <~1((=)>1)
- r(11) ;
r)
CRW.CrM () < equation 3
0 <- Update(0, VoLRw
end while
return 08
Algorithm 2 Online Reward Weighted Conditional Flow Matching (ORW-CFM in Theorem 2
Require: Parameterized efficiently samplable go(x1), (((l)1)) computable u. (z/z1)) initial net- work ug. reward function r(z1), weighting function w(x1) F(r(z1)) where w(21) o r(z1). entropy hyper-parameter t, F can be exp or Softmax, pre-trained model yo and induced sample distribution ạ0,r (11)
O < Opeii (x( 49 (z1)
while Training do
21 qo (x1); t ~~ U(0 1);
x p((2)2))
r - r(11)()))) : w = F(T r
CORW (0) equation 5
- Update 0, V VoLORW.CIM())
end while
return
Algorithm 3 Online Reward Weighted Conditional Flow Matching with Wasserstein-2 Regularization (ORW-CFM-W2 in Theorem
Require: Parameterized efficiently samplable (ettt))ee(xx(())) computable u, (x|211, initial net- work Vg. reward function r(x1). weighting function w(x1) F(r(z1)) where w(x1) x r(z1). entropy hyper-parameter 7 and W2 hyper-paramett a, F can be exp or Softmax, pre-trained model Ugr.r and induced sample distribution (1.)
<- Pret (0(x1) <- 20 (21)
while Training do
21 qo 12)); t ~~ U(0 1);
x~ p<(l(x)
r - r(1))
w = +()
CORW-CHM-WZ + equation 10
- Updatee f, CoCoRW.CFMWW (8))
end while return
60

---

Published as a conference paper at ICLR 2025
1 DISCUSSION AND LIMITATION
Current RL fine-tuning methods, like DPO (Rafailov et al. 2023) or many reward- based RLHF methods for large language models (LLM) (Lee et al. 2023 Huguet et al. 2024 Gülçehree al.) 2023) normally rely on a well-collected or filtered offline dataset, requiring massive manual assistance and high costs, which far from automated fine-tuning. However, traditional online fine-tuning methods are usually too complicated, with too many hyper-parameters, and it is difficult for us to find the relationship between each hyperparameter and the convergent policy, making the overoptimization problem hard to handle.
In fact, the existing LLM model and large text-image generative models has a very strong data sampling capability, but traditional online RL fine-tuning methods may easily cause the model to forget old knowledge, completely collapse to a local greedy policy, and fail to jump out of the local optimum. However, if we can control the collapse speed of the policy (via 7) and the distance between the converged policy and the reference policy (via a)), then the online RL fine-tuning method can not only converge good and diverse policy without the need for additional manually labeled/filtered data, but also avoid over-optimization (Black et al. 2024
In general, one of the primary motivations of our paper is to propose an easy-to-use and controllable reinforcement learning algorithm for fine-tuning continuous flow matching generative models toward arbitrary reward objectives. Although flow matching has demonstrated powerful generative capa- bilities in models like Stable Diffusion 3 (Esser et al. 2024 SD3), an easy-to-use and theoretically comprehensive RL fine-tuning method for continuous flow matching models has not been widely studied. Previous methods are either offline (Lee et al. 2023), rely on filtered datasets and cannot be applied to arbitrary reward models (Rafallov et al., 2023). or lack an explicit divergence bound between the learned policy and the reference model, which may be incapable of addressing over- optimization and policy collapse issues (Black al.) 2024). Furthermore, most methods have not been extensively applied or studied in continuous flow matching fine-tuning tasks. To address these limitations, we propose an online reward- weiihted conditional flow matching with the Wasserstein-2 regularization BORW-CFM-VVE jet al.] 2019). Our goal is to achieve a controllable RL fine-tuning algorithm that allows us to control both the convergence speed and the convergence point (e.g., managing the trade-off between reward and diversity, see Fig. 4).
method based on reward- weighted RL (Peters & Schaal 2007 Peng
To this end, conducted detailed analysis of the optimal policy distribution induced by algorithm and experimentally verified the influence of the entropy coefficient r and the regularization coefficient a on the convergence speed and convergence point.
Interestingly, we established the similarity between our method and traditional RL methods with KL regularization in policy updates. We found that effectively controls the speed (or range) of the policy updates, thereby affecting the convergence speed, and when the W2 regularization is absent, influences the rate of policy collapse. The coefficient a directly affects the diversity of the final convergent policy by controlling the divergence from the reference moddelasssuming of course, that the pre-trained model is not purely greedy policy.
In practice, we recommend using the ORW-CFM-W W2 method with at least > 0 during fine-tuning to effectively avoid issues such as policy collapse.
Additionally, since our method does not make any assumptions about the reward function, it can readily adapt to new tasks by simply adjusting the reward, without the need to modify other hyper- pyper-parameters or collect new datasets manually. For example, in the MNIST dataset, we can encourage the generation of images of odd digits by defining the reward as - (p(even) - p(odd)). thereby penalizing even digits and rewarding odd ones (See App. c In other words, our method shifts the focus from men-tnnng tasks to designing more interesting reward functions, enabling the creation of various flow matching generative models that meet specific user needs without carefully designed offline dataset.
It is important mention that there are still many unresolved issues in RL for flow matching fine-tuning tasks, such as data efficiency, more general reward shaping methods, off-policy online fine-tuning methods, and online fine-tuning of large text-image generative models like SD3 (Esser et al. 2024). These topics warrant further investigation in future work.
61


Reviews:
[{"review_id": "cOIYwMxbnt", "reviewer": "Reviewer_dXDz", "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.", "questions": "See the weakness section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n\n- Combining the reward-weighted matching loss and the Wasserstein distance regulatization seems to be empirically effective. The experimental results look good.\n\n- 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)?\n\n- 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?\n\n- 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. \n\n- 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].\n\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q2y5iRWgFE", "reviewer": "Reviewer_tfNW", "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.", "questions": "Besides the points raised in the weakness section:\n\n1. It is probably better to also show quantitative metrics like diversity scores (e.g., feature pairwise distances) and FID scores.\n2. 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}.\n3. W2 distance is great, but I wonder if it makes a big difference if one instead uses KL divergence (both theoretically and empirically).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nCompared 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.\n\nThe 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.\n\nMany 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).\n\n[1] Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control\n. Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, Ricky T. Q. Chen. https://arxiv.org/abs/2409.08861", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bWVNa346I5", "reviewer": "Reviewer_S526", "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.", "questions": "1. What kind of reward functions can be fine-tuned without collapsing by the proposed W_2 regularization method?  \n2. Is this method capable of performing fine-grained fine-tuning tasks, such as controlling specific semantic parts of images?\n3. Why not use W_1 distance for regularizing?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n2. 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.\n3. 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:\nWhile 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.\nLack of Comparative Analysis with Other Regularization Techniques:\nThe 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.\nNarrow Empirical Validation:\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)