Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning
Authors: Riccardo De Santi
Marin Vlastelica
Ya-Ping Hsieh
Zebang Shen
Niao He
Andreas Krause
Venue: NeurIPS 2025 spotlight ()
Keywords: diffusion models
flow matching
generative optimization
fine-tuning
entropy-regularized control
Primary area: general_machine_learning

Abstract: Adapting large-scale foundational flow and diffusion generative models to optimize task-specific objectives while preserving prior information is crucial for real-world applications such as molecular design, protein docking, and creative image generation. 
Existing principled fine-tuning methods aim to maximize the expected reward of generated samples, while retaining knowledge from the pre-trained model via KL-divergence regularization. 
In this work, we tackle the significantly more general problem of optimizing general utilities beyond average rewards, including risk-averse and novelty-seeking reward maximization, diversity measures for exploration, and experiment design objectives among others. Likewise, we consider more general ways to preserve prior information beyond KL-divergence, such as optimal transport distances and Rényi divergences. 
To this end, we introduce Flow Density Control (FDC), a simple algorithm that reduces this complex problem to a specific sequence of simpler fine-tuning tasks, each solvable via scalable established methods. We derive convergence guarantees for the proposed scheme under realistic assumptions by leveraging recent understanding of mirror flows. Finally, we validate our method on illustrative settings, text-to-image, and molecular design tasks, showing that it can steer pre-trained generative models to optimize objectives and solve practically relevant tasks beyond the reach of current fine-tuning schemes.

Flow Density Control: Generative Optimization Beyond Entropy-Regularizn Fine-Tuning
Riccardo De Santi ETH Zurich
Marin Vlastelica ETH Zurich ETH AI Center
Ya-Ping Hsieh
ETH Zurich yaping hwweh@inff ethz.
ETH AI Center
rdesanti @ethz.tthhccc ch
m wwwstelicaiiinnn .ethz.c ch
Zebang Shen
Niao He ETH Zurich ETH AI Center niache@ethz.cm
Andreas Krause ETH Zurich ETH AI Center krauseaßethz, ch
ETH Zurich
zebang shen@inf ethz. ch
Abstract
Adapting large-scale foundational flow and diffusion generative models to optimize task-specific objectives while preserving prior information is crucial for real-worde world applications such as molecular design, protein docking, and creative image genera- tion. Existing principled fine-tuning methods aim to maximize the expected reward generated samples, while retaining knowledge from the pre-trained model via KL-divergence regularization. In this work, we tackle the significantly more gen- eral problem of optimizing general utilities beyond average rewards, including risk-averse and novelty seeking reward maximization, diversity measures for explo- ration, and experiment design objectives among others. Likewise, we consider more general ways preserve prior information beyond KI--divergence, such as optimal transport distances and Rényi divergences. To this end, we introduce Flow Density Control (FDC), simple algorithm that reduces this complex problem specific sequence of simpler fine-tuning tasks, each solvable via scalable established meth- ods. We derive convergence guarantees for the proposed scheme under realistic assumptions by leveraging recent understanding of mirror flows. Finally, we val- idate our method illustrative settings, texl-to-image, and molecular design tasks, showing that it can steer pre-trained generative models to optimize objectives and solve practically relevant tasks beyond the reach of current fine-tuning schemes.
Introduction
Current
This work
Large-scale generative modeling has recently seen markable advancements, with flow [30, 31] and diffu- sion models [51, 52, 23] standing out for their ability to produce high-fidelity samples across a wide range of applications, from chemistry [24] and biology [9] to robotics [8]. However, approximating the data distribution insufficient for real-world applications such as scientific discovery [6, 59], where one typi- cally wishes to generate samples optimizing specific utilities, e.g., molecular stability and diversity, while preserving certain information from a pre-trained model. This problem has recently been tackled via fine-tuning the case where the utility corresponds
Utility
Ex-ny [r(x))
F(p
Divergence
DRu(TTlle""
D(villP"")
We extend the capabilities of cur- rent fine-tuning schemes from KL-regularized expected reward maximization (left) to the optimization of arbitrary distributional util- ities F under general divergences 2 (right).
Figure 1:
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

to the expected reward of generated samples, and pre- pre-trained model information is retained via KL- divergence regularization, as shown in Fig. 1 (left). Crucially, this specific fine-tuning problem can be solved via entropy entropy-regularized control formulations [e.g.. 14, 55, 53] with successful applications in real-world domains such as image generation [14], molecular design [56], or protein engineering [56].
Unfortunately, many practically relevant tasks cannot be captured by this formulation. For instance, consider the tasks of risk-averse and novelty-seeking reward maximization. In the former case, one wishes to steer the generative model toward distributions with controlled worst-case rewards, thereby improving validity and safety. In the latter case, one aims to control the upper tail of the reward distribution to maximize the probability of generating exceptionally promising designs, e.g., for scientific discovery. Other applications that cannot be captured via maximization of simple expectations include manifold exploration [12]. model de-biasing [13], and optimal experimental design [38, 10] among others. Similarly, preserving prior information via KI. divergence has known drawbacks. For instance, it can lead to missing of low-probability yet valuable modes [29, 43], and it prevents from leveraging the geometry of the space even when this is known, e.g., in protein docking [9]. Replacing KL with alternative divergences can address these shortcomings. Driven by these motivations, in this work we aim to answer the following fundamental question (see Fig. I);
How can provably fine-tune a flow or diffusion model to optimize any user-specified utility while preserving prior information via an arbitrary divergence?
Answering this would contribute to the algorithmic-therrtne
foundations of generative optimization.
Our approach We tackle this challenge by first introducing the formal problem of generative optimization via fine-tuning. Then, we shed light on why this formulation is strictly more expressive than current fine-tuning problems [14, 53], and present a sample of novel practically relevant utilities and divergences (Sec. 3). Next, we introduce Flow Density Control (FDC), simple sequential scheme that can fine-tune models to optimize general objectives beyond the reach of entropy- regularized control methods. This is achieved by leveraging recent machinery from Convex [20] and General Utilities RL [60] (Sec. 4). We provide rigorous convergence guarantees for the proposed algorithm in both simplified scenario, via convex optimization analysis [42, 33], and in realistic setting, by building on recent understanding of mirror flows [25] (Sec. 5). Finally, we provide an experimental evaluation of the proposed method, demonstrating its practical relevance on both image and molecular generation tasks, showing how it can steer pre-trained models solve tasks beyond the inherent limits of current fine-tuning schemes (Sec. 6).
synthetic and high-dimensiin
Our contributions To sum up, in this work we contribute
A formalization of the generative optimization problem, which extends current fine-tuning formula- tions beyond linear utilities and general divergences (Sec. 3).
Flow Density Control (FDC), principled algorithm capable of optimizing functionals beyond the reach of current fine-tuming schemes based on cetrooy-regularizze control/RL (Sec. 4).
Convergence guarantees for the presented algorithm both under simplified and realistic assumptions
leveraging recent understanding of mirror flows (Sec. 5).
An experimental evaluation of FDC showcasing its practical relevance on both illustrative and nigh-dimensional text-to-image and molecular design tasks, showing how it can steer pre-trained models to solve tasks beyond the capabilities of current fine-tuning schemes. (Sec. 6).
2 Background and Notation
General Notation. We denote with X' C Rdd an arbitrary set. Then, we indicate the set of Borel probability measures on X' with P(2),, and the set of functionals over the set of probability measures F(X') as (X²). Given an integer N, we define [N] = NN..
Generative Flow Models. Generative models aim to approximately sample novel data points from a data distribution Pdata- Flow models tackle this problem by transforming samples X, Io from source distribution Po into samples X, = 21 from the target distribution Poatol 1, 17]. Formally, a flow is lime-dependem map e : [0. 1] x Rd - R such that th (t, (/)( A generative flow model is a continuous-time Markov process {Xe)o obtained by applying a flow dix to Xo ~ po as X, the (Xn), 1 [0, 1]. such that X, vh (Xo) Pdata- In particular, the flow v can be defined by velocity field [0, 1] Rd + IRd, which is a vector field related to &b via the following ordinary differential equation (ODE), typically referred to as ffow ODE:
d (x) (x)x
(1)
dt

---

pre
]
General
Flow Demilyy
distribationn
Bodse
Conwext GO
n
Pi
Linear GO
Regulattee Control
(a) Generative Optimization via Flow Model Fine-tuning.
(b) GO expressivity hierarchy
Figure 2: (2a) Pre-trained and fine-tuned policies inducing densities p and optimal density pi
w.r.t utility_ FF and divergence D. (2b) Expressivity and control hierarchy for generative optimization.
with initial condition vb((x) O. flow model X, bhy (Xo) induces a probability path of marginal densities p (oteosesi such that at time / we have that X₁ Pt. Given velocity field u and marginal densities p. we say that u generates the marginal densities p (pesosesi X, v/x(Xo) for all / [0, 1). This is the case if the pair (u.p) satisfy the Continuity Equation:
di. Pe(x) + div(p
In this case, we denote by p" the probability path of marginal densities induced by the velocity field Flow matching [30, 32, 31] can estimate velocity field u s.t. the induced marginal densities p" satisfy Pn Po and pi Pdatas where Pu denotes the source distribution, and Pdata the target data distribution. Interestingly, diffusion models [52] (DMs) admit an equivalent ODE-based formulation with identical marginal densities to their original SDE dynamics [31, Chapter 10]. Consequently, al- though in this work we adopt the notation of flow models, our contributions carry over directly to DMs.
Continuous-time Reinforcement Learning. We formulate finite-horizon continuous-time reinforce- ment learning (RI.) as a specific class of optimal control problems [57, 26, 54, 61]. Given a state space X' and an action space A, we consider the transition dynamics governed by the following ODE:
d vx(x) = (z)
dt
where de A is selected action. We consider vian) deterministic policy a function Ny (X,) == (X,, 1) E A mapping state (x, () E X' to an action E A such that a₂ (Xe; (), and denote with p the marginal density at time /. induced by policy
space X' Rd, (0, 1], and denote by (Marko-
Pre-trained Flow Models as an RL policy. A pre-trained flow model with velocity field wrre can be interpreted an action process af"e := uPre t), where of"e is determined by continuous- time (x t) [12]. Therefore, we can express the flow ODE induced by a
RL policy via at pre-trained flow model by replacing ae with a in Eq. (3), and denote the pre- pre-trained model by its (implicit) policy xpre, which induces marginal density ye" pf pre approximating Pdata.
Formal Problem: a General Framework for Generative Optimization
In this section, we aim to formally introduce the general problem of generative optimization (GO) via fine-tuning. Formally, we wish to adapt a pre-trained generative flow model q7Te to obtain new model лr" inducing an ODE:
dt [c(x) = (v)())) with x
(4)
such that instead of imitating the data distribution Pacta as typically in generative modeling,
induces a marginal density pi that maximizes a utility measure F IP() - R. while preserving information from the pre-trained model 77Te via regularization with an arbitrary divergence
D( pre). This algorithmic problem illustrated in Fig. 2a, and formalized in the following.
Generative Optimization via Flow Model Fine-Tuning
d arg max F(P)) aD [ 1|pp"e) s.t. px + ((x( 0 with a, T(x,t) (5)
3

---

NON- LINEAR GO
APPLICATION
FUNCTIONAL F/D
LINEAR GO
CONVEX GENERAL
REWARD OPTIMIZATION [14, 55]
Ba [[(o]]
MANIFOLD EXPLOR (12] GEN. MODEL DE- BE-BIISING
H(p") - B llogo (x))
supe
CVaR5(r))
(r(x)) in(ar) (btr")
RISK-WERRS
sup
B [[(]] Var(s")
x
NOVELTY-SEEKING OPTIMIZATION
SQS(p") = E ((x) In(e) 2 soo"")
(0)))((z) XI|
OPTIMAL EXPERIMENT DESIGN
af)) (log del() -T(()) (m)a
DIVERSE MODES DISCOVERY
x
LUG- CONSTRAINED GENERATION
(x-g)(((a)]] - 9 (W77
x
(z) Dir 11 = (2
KULLBACK LEIBLER DIVERGENCE 114, 551
pprt(s) di
/
RÉNYI DIVERGENCES
Doln" (1p"")) :- 3-1 log (n (sr)) ( Wyip" 11x!r))
x
/
OPTIMAL TRANSPORT DISTANCES
inf
[d(z, 1/
x
/
verv" gont) (saliin
MAXIMUM MEAN DISCREPANCY
(M 1p")) - Ipte Maore 1. Mip [e(or.
x
Table 1: Examples of practically relevant utilities F (blue) and divergences DD (orange). Apx. A provides mathematical details and practical applications for each functional. Notice that besides H. all non-linear functionals presented are novel in the context of fine-tuning of diffusion and flow models.
In this formulation, F and D are both functionals mapping the marginal density p induced by policy to a scalar real number, namely F, D P(X) - R. The constraint in Eq. (5) is the (controlled) Con- tinuity Equation (see Eq. (2)), which relates the control policy и to the induced marginal density pf-
3.1 The sub-case of KL-regularized reward maximization via entropy- regulaieee control Current fine-tuning schemes for flow generative models based on RL and control-theoretic formula- tions [e.g.. 14, 55] aim tackle the following problem, where we omit the flow constraint for clarity:
Linear Generative Optimization via Flow Model Fine-Ture Fine-Tuning
arg
E [=(=) aDx " p~p
(6)
Crucially, the common problem in Eq. (6), which we denote by Linear' GO, the specific sub-case of
the generative optimization problem in Eq. (5), where the utility F is a linear functional corresponding the expectation of (reward) function r X IR and D is the Kullback-Leibler divergence:
F(P)) (pi,r) r)
E [r())] and D(p "
DEL (rī l!
x-<p]
This specific fine-tuning problem can be solved via entropy-reeularizzed (or relaxed) control [14].
3.2 Beyond Linear Generative Optimization:
Expressivity Viewpoint
Let (pT)) = F(PT) - a D(pTI pr be the functional in Eq. (5). Then we denote by Convex GO the case where is concave in pF. and by General GO the case for arbitrary, possibly non-convex functionals 2, In terms of expressivity Linear GO C Convex GO General GO, as depicted in Fig. 2b (left). In Table we classify into these tree tiers sample of practically relevant utilities (F, blue) and divergences (D, orange). In Apx. we report complete definitions and applications. Except for entropy [12] and KL, all non-linear functionals in Table I first time in the flow and diffusion model teniin literature, while vastly employed in other areas. Moreover, the framework presented in this work for GO (Eq. 5) applies to any new choice of F D.
knowledge explicitly used for the
I For clarity, we adopt the term linear motivated by the linear utility even though the is non- non-linear. For clarity, we use the term convex GO, rather than concave GO, to denote the problem class where concave functionals are optimized.

---

Algorithm Flow Density Control (FDC)
1: input: gggeneall utility functional, K number of iterations, , pre- pre-trained flow generative model, {m 1K regularization coefficients
2: Init: == x.re
3: for 1,2,.... K do Estimate: Vxgs Vgdg(p) 1)
5: Compute Tr via first-order linear fine-tuning:
<- BATTOOYYBEELLLLNZZOCOOOOOOOSSSS (V.gk, (R, Tk-1)
end for
7: output: policy π == TK
Given the generality of generative optimization (5). natural question arises: how can it be solved algorithmically? In the next section, we answer this by leveraging recent machinery from Convex [20] and General-Utilities RL [60], to derive fine-tuning scheme that handles both convex and general GO, thus going beyond current entropy-regularizn control methods, as illustrated Fig. 2b (right).
Algorithm: Flow Density Control
In this section, we introduce Flow Density Control (FDC), see Alg. 1, which provably solves the generative optimization problem in Eq. (5) via sequential fine-tuning of the pre-trained model The To this end, we recall the notion of first variation of functional over a space of probability measures [25]. A functional G E ((),, where G P(()) + I, has first variation at μ E P(X) if there exists function (p) FX) such that for all μ² (()) it holds that:
G( + ex') G(x) + elu,, (g(()) + o(e).
where the inner product has to be interpreted at μ. namely (G()), can be interpreted as an infinite-dimense dimensional gradient the space of probability measures. Given this notion, and a pair of generative models represented via policies and n', can now state the following entropy-regudee regularizd first variation maximization fine-tuning problem.
expectation. Intuitively, the first variation of G
First Variation Maximization
(89 pT .PI - 1)Dri(P I!
arg max
Crucially, we can introduce function g X defined for all X such that:
g(x) = 8g
pT (xr) and E z-ppe [_(]))]] = (6g (pT) .P))
As a consequence, by rewriting Eq. (8) expressing the first term via an expectation as shown in Eq. (9), it corresponds to common Linear GO problem (see Eq. (6)), which can be optimized by utilizing established entropy-regularized control methods [e.g., 56, 14, 61].
We can finally present Flow Density Control (FDC), see Alg. 1, mirror descent (MD) scheme [42] specific sequence Linear GO problems. FDC takes three inputs: pre-trained flow or diffusion model "Fe the number of iterations K, and sequence of regularization weights { At each iteration, FDC first estimates the gradient of the functional first variation at the previous policy T i.e., ,666 (pi-1) (line 4). Then, updates the flow model тa by solving the fine-tuning problem in Eq. (8) via entropy-rege regulay-reullarized control solver such as Adjoint Matching [14], using V.9* dg (pf-1) ') in Eq. (9) (line 5). Ultimately, it returns a final policy π = -K. We report a detailed implementation FDC in Apx. D.
that reduces optimization of non-linear functionals
Gradient of first variation: computation and estimation. Surprisingly, estimating V. in Alg. 1 (line 4) rarely requires density estimation. Among the functionals in Table 1, only the Rényi divergence does, for which one can leverage the recent Itô density estimator [50]. All other functionals admit straightforward plug-in or sample-base approximations detailed in Apx. A. As an illustrative example, in the following we showcase three examples from Table 1:
Vz (o)
Entropy (H)
V v..0(p"))<)
V va(()) (r((x) S (1({²)}
CVaR
(x) where 0² arg m (00 pre) Wasserstein-1 (W1)
5

---

Here Q denotes either a utility F or divergence D, and 43 (p") is the 8-quantile of Z = rX)) with X ~ p" [47]. These gradients can be easily implemented. For entropy, the score term can be approx- imated via the score network in the case of diffusion models [12], and obtained via a known linear transformation of the learned velocity field in the case of flows [14, Eq) For CVaR, any standard sample-based estimator of (5(p") [47] can be used. For Wasserstein-1, o" actually corresponds the discriminator in Wasserstein-Gal GAN, which can be learned with established methods [2]. In Apx. A, we report the gradient of the first variation for all functionals in Table 1, explain their practical estimation, and present tutorial to derive the first variation of any new functionals not mentioned within Table 1.
Given the approximate gradient estimates and the generality of the objective functions, it is still unclear whether the proposed algorithm provably converges to the optimal flow model '², In the next section, we answer this question by developing a theoretical analysis via recent results on mirror flows [25].
5 Guarantees for Generative Optimization via Flow Density Control
In this section, recast (5) constrained optimization over stochastic processes, where the con- straint is given by the Continuity Equation (2). This formulation enables the application of mirror descent for constrained optimization and the notion of relative smoothness [3]. In our framework, convergence speed is governed by: 1. the structural complexity of the functional G (cf.. Section 4), the accuracy of the estimator g from (9), and the quality of the oracle ENTROPYREGULARIE COLVR in Alg. 1. To handle these cases, we will analyze two representative regimes: Idealized. G is concave, and both y and ENTROPY REGULARIIEDDCONTRLLSSS In this setting, classical results yield sharp step-size prescriptions and fast convergence rates.
exact.
General. is non-concave, with y and the oracle subject to noise and bias. While fast convergence is generally out of reach [34, 27], convergence to a stationary point remains attainable under mild assumptions.
Theoretical analysis: Idealized setting. We now present a framework leading to convergence guarantees for FDC (i.e., Alg. 1) for concave functionals G E F(X). We start by recalling the notion of Bregman divergence induced by a functional Q E ((X) between densities μ,, v E (P))), namely: Da( || v) == Q(n) - Q(v) (50(v), v)
Next, we introduce two structural properties for our analysis.
Definition (Relative smoothness and relative strong concavity [33]). Let G P(()) - IR concave
functional. We say that L-smooth relative to Q F(X) over P(X) if scalar s.t. for all H.v E (XX):
((v)= 2 G(H) + (8G()), v - μ4- - LD ()
(10)
and we say that G is I-strongly concave relative to Q E F(x) over P(X) if 2 0 scalar s.t. for all H.V (X)):
((v) < (()) + (5G()), n1) IDo(v
(11)
In the following, we interpret line (6) of FDC as a step of mirror ascent [42], and the KL divergence term as the Bregman divergence induced by an entropic mirror map Q = H. i.e., (KI(P,v) (x(HH)) v). We can finally state the following set of assumptions
well as the
convergence guarantee for an arbitrary functional G(-) = F(.) aD(, p! E (FX)..
Assumption 5.1 (Exact estimation and optimization). We consider the following assumptions:
1. Exact estimation: (p) is estimated exactly Vk [K].
2. The optimization problem in Eq. (8) is solved exactly.
Theorem 5.1 (Convergence guarantee of Flow Density Control with concave functionals). Given
Assumptions 5.1, fine-tuning a pre-trained model pre via FDC (Algorithm 1) with 77k = L vk E [K), leads to a policy n inducing a marginal distribution př such that:
G(pi) - ((pT)
--1 K
DKI
(12)
where pi = pi is the marginal distribution induced by the optimal policy x arg max, G(p)) F(r)) Dlpp "!
E
Theorem 5.1 provides fast convergence rate under specific step-size choice (7)- L). However, it critically depends on Assumption 1. which typically does not hold in practice. To address this limitation, now consider more general scenario where this key assumption is relaxed.

---

Theoretical analysis: General setting, Recall that p == Pi represents the (stochastic) density produced by the REAOOO/RBGGLLARIIEOCOOIIIIIIOO the following mirror ascent iterates, where 1 1/AE = "/e in Algorithm 1:
oracle at the k-th step of FDC, and consider
P
arg max
(8g (PT Th 1) ,) K
D (p PT
(MD²)
pEP(Opre)
In realistic settings, where only noisy and biased approximations of (MD;) are available, it is essential to quantify the deviations from the idealized iterates in (MD+). To this end, denote by Tk the filtration up to step k, and consider the decomposition of the oracle into its noise and bias parts:
by = E [89(pT)) ) - (v() [Tk]
Uk = (9(77) Th
(g by
(13)
Conditioned on Te. Us has zero mean, while Py captures the systematic error. We then impose: Assumption 5.2 (Noise and Bias). The following events happen almost surely:
|6 oooo + O, Σ IE 112 (Ibbllis
+ +1U
00, Enh lIbulloo < oo (14)
The first condition is a necessary requirement for convergence since when violated, it is easy to construct scenarios where no practical algorithm can solve the generative optimization problem. The second and third inequalities manage the trade-off between accuracy of the approximate oracle ENTROPY REGULARIZEUUULUU CONIEEDDOOOOOOOOOLL SOLVER and aggressiveness of the step sizes, "ik- Intuitively, lower noise and bias in the oracle enable the use of larger step sizes. To this end, Assumption 5.2 provides a concrete criterion that guarantees the success of finding the optimal policy with probability one.
Theorem 5.2 (Convergence guarantee of Flow Density Control for general functionals). Given the Robbins- Monro step-size rule: Σx 7hk oo, Ex th < xi, under Assumption 5.2 and technical assumptions (see Appendix C), the sequence of marginal densities pf induced by the iterates Trk of Algorithmm converges weakly to stationary point p1 of G almost surely, formally: p - pi a.s..
6 Experimental Evaluation
We analyze the ability of Flow Density Control (FDC) to induce policies optimizing complex non-linear objectives, and compare its performance with Adjoint Matching (AM) [14), a classic fine-tuning method. We present two types of experiments: (i) Illustrative settings to provide insights via visual interpretability, and (ii) High-dimensional real-world applications, namely (a) novelty- seeking molecular design for single-point energy minimization [18], and (b) manifold exploration for lext-to-image creative bridge design generation. Additional details are provided in Apx. E.
Risk-averse reward maxinization for better worst-case validity or safety. We fine-tune a pre- trained policy xrre (see Fig. 3a) by optimizing the CVaR, utility i.e., expected outcome in the B w-orsst-- case (see Tab. 1) with KL regularization, and costs interpreted as negative rewards. The cost has three regions: a high-cost plateau (dark orange), where the initial density lies; a moderate-cost left area (light orange); and predominantly low-cost right zone (yellow) punctuated by narrow, but catastrophic red-stripes. As shown in Fig. 3b, AM moves the model density into the yellow region, lowering average cost but exposing to rare extreme costs. contrast, FDC, run with K = 2 iterations and 3 = 0.01, successfully steers density into the safer, moderate-cost area, cutting the wor-orrtecase cost from 288.2 achieved by AM to 90.0, well below the initial 262.5, as shown in Fig. 3c and 3d.
Novelty-seeking reward maximization for discovery. We fine-tune a pre-trained policy pre to max- imize the SQ, utility, i.e., expected outcome in the 3-best-case (see Tab. 1). The reward shown in Fig. Be has a moderately high-reward left area (light gray). a medium-reward central plateau (darker gray) where the initial density lies, and a low-reward right region (black) with sparse, extreme-reware spikes depicted by thin white lines. As shown in Fig. 3f, AM drifts the density into the safer left basin - im- proving the average reward but only reaching a best-1% % expected reward of 55.5, as shown in Fig. 3g and Fig. 3h. In contrast, FDC, run for elevating the top- % reward to 596.1 (see Fig. 3h) far above both AM and the initial 66.6.
2 iterations and 3 = 0.99, pushes the density rightwards,
Reward maximization regularized via optimal transport distance. We fine-tune the pre-trained model with density in Fig. 3i to maximize a reward function that increases moving top right. We con- sider two W, distances induced by two ground metrics: d,,, which makes vertical moves more costly than horizontal ones, and dB, which does the opposite. Under dA, both AM and the OT-regularized model reach an expected reward of 35.0, O, but FDC-A incurs only A 1.95 versus 4.67 for AM,
7

---

250
120
0.06
0.59
100
Frrq
0.04
700
30
FDC
0.00
200
300
Cost
Iterationn
(a) Pre-trained samples
AM vs FDC samples
(c) 3-worst-case costs
(d) CVaR evaluation
400
329
roc
Feequeas
50C
RDC
40c
169
0.01
200
1DC
30
0.00
100
200
500
600
5
10
Reware
(e) Pre-trained samples
AM vs FDC samples
(g) 3-best-case rewards
(h) SQ evaluation
50
25
FOC
7%
FDG
Elr(x:)
W71
A%
20
Pre
29.5
0
AM
35.0
4.67
100
FDC-A
35.4
1.95
280
10
10
15
20 25
20
(i) Pre- trained samples
(j) AM vs FDC-A samples(k) AM vs FDC-B samples
(I) W/ evaluation
Figure 3: Illustrative settings with visually interpretable results. (top) Risk-averse reward maximiza- tion for valid or safe generation, (mid) Novelty-seeking reward maximization for discovery, (bottom) Expected rewards maximization under optimal transport distance regularization. Crucially, FDC optimize well these complex objectives, while AM [14], a classic fine-tuning scheme, fails at this.
and achieves a mean shift that 280% larger in the horizontal than in the vertical direction (Fig. 3j and Tab. 3l). By contrast, FDC-B under d preferentially shifts the density upward (Fig. 3k).
Conservative manifold exploration. We tackle manifold exploration [12] by fine-tuning pre-trained model TPP maximize the entropy utility (H in Tab. 1) under a KL regularization of strength a, a capability not possible with prior methods [12]. As in previous work, we consider the common setting where the pre-trained model density pf concentrates most of its mass a specific region as shown in Fig. 4a, where N 10000 samples are shown. By fine-tuning "hr via FDC, the density of the fine-tuned model shifts into low-coverage areas (see Fig. 4b and 4c). In particular, Fig. 4d demonstrates that reducing from 0.5 to 0.0 yields progressively higher Monte Carlo entropy estimates (7.00 at a = 0.5, 7.14 = 0), thus enabling control of the trade-off between preserving the original distribution and exploring novel regions, capability not supported by prior methods [12].
Molecular design for single-point energy minimization. We fine-tune FlowMol [15], pre- pre-trained on QM9 [46]. to discover molecules minimizing the single-point total energy computed via extended tight- at the GFNI-XTB level of theory [18]. Concretely, we maximize the negative energy. We do not aim maximize the average sample reward, but rather that of the top 0.2% samples. We employ FDC with novelty-seekin SQ utility (see Tab. 1) with 3 0.998, and make 2 gradient steps per K 10 iterations. We compare it with AM run for 240 steps. Fig. 4j shows that while AM generates better samples in average (namely 29.1 over 27.5 of FDC), the average quality of the top 0.2% molecules, indicated by SQ, higher for FDC than for AM (namely 41.8 over 39.7 of AM). This confirms (see Fig. 4i and 4h) that FDC can sacrifice the average reward to generate a few truly high-reward designs.
Text-to-image bridge designs conservative exploration. We perform manifold exploration by fine-tuning Stable Diffusion (SD) 1.4 [49] with prompt "A creative bridge design.", ". To this end, we maximize the KL-regularized entropy (see Tab. 1) with = 0.001 via FDC for K = 2 steps. As diversity metric, we utilize the Vendi score [19] with cosine similarity kernel on the extracted CLIP [21] features from sample of 100 images and compared it to the baseline pre-trained model in Fig. 4g. Beyond increasing the Vendi score, FDC also increases the CLIP score of the initial model.
7 Related Works
Flow and diffusion models fine-tuning via optimal control. Recent works have framed fine-tuning of diffusion and flow models to maximize expected reward under KI. regularization as an entropy-regularized optimal control problem [e.g., 55, 53, 56, 14]. Crucially, as shown in Sec. 3, the problem tackled by these studies is the specific sub-case of generative optimization (Eq. (5)),

---

H(pp"
Pre-trained
6.78
FDC 0.5
7.00
FDC 0.0
7.14
Pre-trained samples
(b) FDC 0.5 samples (c) FDC 0.0 samples
(d) Entropy evaluation
Vendi
CL.IP
Pre-trained
2.36 0.19
FDC a 0.001
2.47 0.22
(e) Pre-trained samples
(1) FDC samples
(g) Images evaluation
10
Eir(z)]
SQ,
FDC
FDC
00
Pre
15.4
24.2
AM FDC-A
29.1
39.7
27.5
41.8
24 30
40
10.8 41.4 42.0 42.6 43.2
(Hartree)
Negative (Hartreel
(h) Energy distributions
(i) Top 0.2% Energies
(i SQ Negative Energy (k) FDC-gDnnratee design
Figure (top) Illustrative manifold exploration experiment via KL-regularized entropy maximization, (mid) High-dimensional manifold exploration via text-to-image model fine-tuning for prompt "A creative bridge design". Left: images from pre-trained model, Right: images from model fine-tuned via FDC, with higher diversity indicated by higher Vendi score. (bottom) fovelty-seeking molec- ular design for Energy (kcal/mol) maximization by fine-tuning FlowMol [15]. FDC shows enhanced control capabilities for optimizing such complex objectives than AM, classic fine-tuning scheme.
where the utility is linear, and D DKL- In this work, we propose a principled method with guarantees for the far more general class of non-linear utilities and divergences beyond KL, including the ones listed in Tab. The framework introduced has strictly higher expressive power and control capabilities for fine-tuning generative model (see Sec. 3). This renders possible to tackle relevant tasks e.g., scientific discovery, beyond the capabilities of the aforementioned fine-tuning schemes.
Convex and General Utilities Reinforcement Learning. Convex and General (Utliitie) RL [20, 58, 60] generalizes RL the case where one wishes maximize concave [20, 58], or general [60, 4] functional of the state distribution induced by a policy over dynamical system's 's state space. The introduced generative optimization problem (in Eq. (5)) is related, with pf represent- ing the state distribution induced by policy π over subset of the state space. Recent works tackled the finite samples budget setting [e.g., 41, 39, 40, 44, 11]. Ultimately, to our knowledge, this the first work leveraging an algorithmic scheme resembling General RL for the practically relevant task of generative optimization of general non-linear functionals via fine-tuning diffusion and flow models.
Optimization over probability measures via mirror flows. Recently, there has been a growing interest building theoretical guarantees for optimization problems over spaces of probability measures in variety of applications. These include GANs [25], optimal transport [3, 28, 27], kernelized methods [16]. and manifold exploration [12]. We present the first use of this framework to establish guarantees for the generative optimization problem in Eq. (5). This novel link to probability-spe optimization sheds new light on large-scale flow and diffusion models fine-tuning.
8 Conclusion
This work tackles the fundamental challenge of fine-tuning pre-trained flow and diffusion generative models on arbitrary task-specific utilities and divergences while retaining prior knowledge. We introduce a unified generative optimization framework that strictly generalizes existing formulations and propose a rich class of new practically relevant objectives. We then propose Flow Density Control, a mirror-descent algorithm that reduces complex generative optimization to sequence of standard fine-tuning steps, each solvable by scalable off-the-shed methods. Leveraging convex analysis and recent advances in mirror flows theory, we prove convergence under general conditions. Empirical results on synthetic benchmarks, molecular design, and image generation, demonstrate that our approach can steer pre-trained models to optimize objectives beyond the reach of current fine-tuning techniques. for limitations, while our framework general, future work will need to assess to what extent the flexibility in selecting utilities and divergences yields concrete gains in specific applications.

---

Acknowledgements
This publication was made possible by the ETH AI Center doctoral fellowship Riccardo De Santi, and postdoctoral fellowship to Marin Vlastelica. The project has received funding from the Swiss National Science Foundation under NCCR Catalysis grant number 180544 and NCCR Automation grant agreement 51NF40 180545.
References
[1] Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209. 15571, 2022.
[2] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan, 2017.
[3] Pierre-Cyril Franko-Frankossi,, Anna Korba, and Flavien Léger. Mirror descent with relative smoothness in measure spaces, with application to sinkhorn and em. Advances in Neural Information Processing Systems, 35: 35:13323-17277 2022.
[4] Anas Barakat, Ilyas Fatkhullin, and Niao He. Reinforcement learning with general utilities: Simpler variance reduction and large state-action space. In International Conference on Machine Learning, pages 1753-1800. PMLR, 2023.
[5] Michel Benaïm. Dynamics of stochastic approximation algorithms. In Seminaire de probabilites XXXIII, pages 1-68. Springer, 2006.
[6] Camille Bilodeau, Wengong Jin, Tommi Jaakkola, Regina Barzilay, and Klavs F Jensen. Gener- ative models for molecular discovery: Recent advances and challenges. Wiley Interdisciplinary Reviews:
Computational Molecular Science, 12(5)=e1608 2022.
[7] Kathryn Chaloner and Isabella Verdinelli. Bayesian experimental design: review. Statistical Science, pages 273-304, 1995.
[8] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arX1v:2303.04137, 2023.
[9] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXiv:2210.01777 2022.
[10] Riccardo De Santi, Federico Arangath Joseph, Noah Liniger, Mirco Mutti, and Andreas Krause. Geometric active exploration in markov decision processes: the benefit of abstraction. arXiv preprint arXiv:2407.13360 2024.
[11] Riccardo De Santi, Manish Prajapat, and Andreas Krause. Global reinforcement learning: Beyond linear and convex rewards via submodular semi-gradient methods. arXiv preprint arXiv:2407.09900 09905, 2024.
[12] Riccardo De Santi, Marin Vlastelica, Ya-Ping Hsich, Zebang Shen, Niao He, and Andreas Krause. Provable maximum entropy manifold exploration via diffusion models. In ICLR 2025 Workshop on Deep Generative Model in Machine Learning: Theory, Principle and Efficacy.
[13] Alexander Decruyenaere, Heidelinde Dehaene, Paloma Rabaey, Johan Decruyenaere, Christiaan Polet, Thomas Demeester, and Stijn Vansteelandt. Debiasing synthetic data generated by deep generative models. Advances in Neural Information Processing Systems, 37:41539-41576, 2024.
[14] Carles Domingo-Eniinn Ennich, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.088.. 2024.
[15] Ian Dunn and David Ryan Koes. Mixed continuous and categorical flow matching for 3d de molecule generation. ArXiv, pages arXiv-2404, 2024.
10

---

[16] Pavel Dvurechensky and Jia-Jie Zhu. Analysis of kernel mirror prox for measure optimization. In International Conference on Artificial Intelligence and Statistics, pages 2350-2358. PMLR, 2024.
[17] Jesse Farebrother, Matteo Pirotta, Andrea Tirinzoni, Rémi Munos, Alessandro Lazaric, and Ahmed Touati. Temporal difference flows. arXiv preprint arXiv:2503.09817, 2025.
[18] Marvin Friede, Christian Hölzer, Sebastian Ehlert, and Stefan Grimme. dxtb-an efficient and fully differentiable framework for extended tight-binding. The Journal of Chemical Physics, 161(6), 2024.
[19] Dan Friedman and Adji Bousso Dieng, The vendi score: A diversity evaluation metric for machine learning. arXiv preprint arXiv:2210.024100 2022.
[20] Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. In International Conference on Machine Learning, 2019.
[21] Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A ofference-free evaluation metric for image captioning. In Proceedings of the 2021 Conference on Empirical Methods Natural Language Processing, pages 7514-7528, 2021.
[22] Saan-Baptiste Miriart-Urrutt and Claude Lemaréchal. Fundamentals of convex analysis. Springer Science & Business Media, 2004.
[23] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 1116840-8511 2020.
[24] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation 3d. In International conference on machine learning, pages 8867-8887. PMLR, 2022.
[25] Ya-Ping Hsieh, Chen Liu, and Volkan Cevher. Finding mixed nash equilibria of generative adversarial networks. In International Conference on Machine Learning, pages 2810-2819. PMLR, 2019.
[26] Yanwei Jia and Xun Yu Zhou, Policy evaluation and temporal-diffrrece learning continuous time and space: A martingale approach. Journal of Machine Learning Research, 23(154):1-55, 2022.
[27] Mohammad Reza Karimi, Ya-Ping Hsich, and Andreas Krause. Sinkhorn flow as mirror flow: A continuous-time framework for generalizing the sinkhorn algorithm. In International Conference on Artificial Intelligence and Statistics, pages 4196-4194. PMLR, 2024.
[28] Flavien Léger. A gradient descent perspective on sinkhorn. Applied Mathematics & Optimiza- tion, 84(2):1843-1888 2021.
[29] Yingzhen Li and Richard E Turner. Rényi divergence variational inference. Advances in neural information processing systems, 29, 2016.
[30] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02777, 2022.
[31] Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben- namu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:2412.06264 2024.
[32] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 030003,, 2022.
[33] Haihao Lu, Robert M Freund, and Yurii Nesterov. Relatively smooth convex optimization by first-order methods, and applications. SIAM Journal on Optimization, 28(1):333-34, 2018.
[34] Panayotis Mertikopoulos, Ya-Ping Hsieh, and Volkan Cevher. A unified stochastic approxima- tion framework for learning in games. Mathematical Programming, 203(():559-66) 609, 2024.
11

---

[35] Paul Milgrom and Ilya Segal. Envelope theorems for arbitrary choice sets. Econometrica. (0(2):583-601, 2002.
[36] Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, Bernhard Schölkopf, et al. Kernel mean embedding distributions: review and beyond. Foundations and Trends@ Machine Learning, 10(1-2):1-141, 2017.
[37] Mojmir Mutný. Modern Adaptive Experiment Design: Machine Learning Perspective. PhD thesis, ETH Zurich, 2024.
[38] Mojmir Mutny, Tadeusz Janik, and Andreas Krause. Active exploration via experiment design in Markov chains. In International Conference Artificial Intelligence and Statistics, 2023.
[39] Mirco Mutti, Riccardo De Santi, Piersilvio De Bartolomeis, and Marcello Restelli. Challenging common assumptions in convex reinforcement learning. Advances in Neural Information Processing Systems, 35:4489-4502, 2022.
[40] Mirco Mutti, Riccardo De Santi, Piersilvio De Bartolomeis, and Marcello Restelli. Convex reinforcement learning in finite trials. Journal of Machine Learning Research, 24(250):1-42, 2023.
[41] Mirco Mutti, Riccardo De Santi, and Marcello Restelli. The importance of non-markovianity in maximum state entropy exploration. In International Conference on Machine Learning, pages 16223-16239. PMLR, 2022.
[42] Arkadij Semenović Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983,
[43] Kushagra Pandey, Jaideep Pathak, Yilun Xu, Stephan Mandt, Michael Pritchard, Arash Vahdat, and Morteza Mardani. Heavy-tailed diffusion models. arXiv preprint arXiv:2410.1417,, 2024.
[44] Manish Prajapat, Mojmir Mutný, Melanie N Zeilinger, and Andreas Krause. Submodular reinforcement learning, arXiv preprint arXiv:2307.13372 2023.
[45] Friedrich Pukelsheim. Optimal design of experiments. SIAM, 2006.
[46] Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[47] R Tyrrell Rockafellar and Stanislav Uryasev. Conditional value-at-risk for general loss distribu- tions. Journal of banking & finance, 26(7):144311471, 2002.
[48] R Tyrrell Rockafellar, Stanislav Uryasev, et al. Optimization of conditional value-at-risk. Journal of risk, 2:21-42,10 2000.
[49] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models, 2021.
[50] Marta Skreta, Lazar Atanackovic, Avishek Joey Bose, Alexander Tong, and Kirill Neklyudov. The superposition of diffusion models using the it) o density estimator. arXiv preprint 27X1v:2412.17762, 2024.
[51] Jascha vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-2265. PMLR, 2015.
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper-
[52] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[53] Wenpin Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond. arXiv preprint arXivv2403.06227, 2024.
12

---

[54] Lenart Treven, Jonas Hübotter, Bhavya Sukhija, Florian Dorfler, and Andreas Krause. Efficient exploration in continuous-time model-based reinforcement learning. Advances in Neural 2023.
Information Processing Systems,
[55] Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine- tuning of continuous-time diffusion models 15X444240..55144 2024.
entropy entropy-regularie regularized control. arXiv preprint
[56] Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Sergey Levine, and Tommaso Biancalani. Feedback efficient online fine-tuning of diffusion models. arXiv preprint arXiv:2402. 16359, 2024.
[57] Haoran Wang, Thaleia Zariphopoulou, and Xun Yu Zhou. Reinforcement learning in continuous time and space: A stochastic control approach. Journal of Machine Learning Research, 21(198):1-34, 2020.
[58] Tom Zahavy, Brendan O'Donoghue, Guillaume Desjardins, and Satinder Singh. Reward is enough for convex mdps. Advances in Neural Information Processing Systems, 34:25746-25759, 2021.
[59] Claudio Zeni, Robert Pinsler, Daniel Zügner, Andrew Fowler, Matthew Horton, Xiang Fu, Sasha Shysheya, Jonathan Crabbé, Lixin Sun, Jake Smith, et al. Mattergen: a generative model for inorganic materials design. arXiv preprint 2rXiv:2312.03666 2023.
[60] Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4572 4583. Curran Associates, Inc., 2020.
[61] Hanyang Zhao, Haoxian Chen, Ji Zhang, David D Yao, and Wenpin Tang. Scores as actions: framework of fine-tuning diffusion models by continuous- time reinforcement learning. arXiv preprint arXiv:2409.00 084000 2024.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We carefully wrote the abstract a-posteriori so that it could represent fairly the content of the work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We present a in-depth theoretical investigation showing that depending on the oracle assumptions made, and that under general oracles (e.g., noisy and biased) the pre- sented problem cannot be solved with fast convergence rates, and we provide an asymptotic convergence guarantee.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
14

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions and derivations are presented within either the main paper or mentioned appendices.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We present simple algorithm that fundamentally based on sequential use of existing established methods. Therefore any proper implementation can lead empirical results analogous to the ones presented within this work and therefore in support of the main experimental claims.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
15

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Currently do not provide access data and code. We are although preparing the release of public version and are available if needed.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips.. cc/
httpii/guuidd//oeeuuuueeee ctssuummissoooooicyoy..oo. for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / ctttps.cccpubllt c/tui//uuieeeeeeee des/ Cooeeubmii wwwsssssiiiooooooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We believe to have stated all the necessary information to understand the results and claims made.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Within the appendix we report statistical analysis of several experimental results presented within the Experiments section, confirming statistical significance of the algorithmic scheme logic. We regard further statistical validation, especially in high
16

---

dimensional settings or concrete real-world applications as particularly challenging yet very important future work.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report I-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative eтror rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
We report our compute resources in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ett://nneui./eeeeee http:/ulll//thhcaaaiiieeeee ines?
Answer: [Yes]
Justification:
We have read the document and confirm conformity.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
17

---

Answer:
[NA]
Justification:
The work presented presents fundamental algorithmic innovations.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
The paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets g.g code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We reference, cite, and credit relevant contributions.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
18

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMS as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https: / //neruurpp.... wc/cooffeeeccee/225//11111 for what should or should not be described.
20

---

Contents
Functionals and Derivation of Gradients of First-order Variations
22
A.1 Overview utilities and divergences Table
22
A.2 A brief tutorial on first variation derivation
22
A.3 Derivation gradients of first-order variation for functionals in Table
23
Proof for Theorem 5.1
26
C Proof for Theorem 5.2
27
D Detailed Example of Algorithm Implementation
29
D.1 Implementation of EENIROYYEBLLLLLZZZOCOOOOOOOUSU ................................
29
D.2 Discussion: computational complexity and cost of FDC
29
E Experimental Details
31
E.1 Used computational resources
31
E.2 Experiments in Illustrative Settings
................................
E,3 Further Ablations
32
E,4 Real-World Experiments
33
21

---

A Functionals and Derivation of Gradients of First-order Variations
A.1 Overview of utilities and divergences in Table
In the following, we report the missing details for the functionals presented within Table 1, and discuss some possible applications.
Manifold Exploration and Generative Model De-biasing As mentioned within Sec. 3, maxi- mization of the entropy functional as been recently introduced as a principled objective for manifold exploration [12]. Moreover, we wish to point out that it can be interpreted also from the viewpoint of de-biasing a prior generative model to re-distribute more uniformly density while preserving certain notion of support, e.g., via sufficient KL--divergence regularization.
Risk-averse and Novelty-seeking reward maximization A definition of q3 can be found below, explanations of these utilities be found in Sec. 1, and experimental illustrative examples provided in Sec. 6.
Optimal Experiment Design The task of Optimal Experimental Design (OED) [7] involves choosing a sequence of experiments so as to minimize some uncertainty metric for an unknown quantity of interest X + R, where X' is the of all possible experiments. From a probabilistic standpoint, an optimal design may be viewed as probability distribution over X' prescribing how frequently each experiment should be performed to achieve maximal reduction in uncertainty about [45]. This problem has been recently studied in the case where f is an element of a reproducing kernel Hilbert space (RKHS), i.e., f E H&s induced by a known kernel k(r, x²) = ((x) ((x') where x, xx' E X' [37]. Given this setting, one might aim to acquire information about f according to different criteria captured by the scalarization function s(-) [38]. In particular, in Table 1, we report three illustrative choices for s:
D-design: log det(-) (Information)
A-design: -Tr(-) (Parameter error)
E-design: Amaa (Worst projection error)
as reported in previous work [Table 1 38].
Diverse Mode Discovery This objective corresponds re-interpretation of the Diverse Skill Discovery objective introduced in the context of Reinforcement Learning [58]. Consider the case where it is given discrete and finite set S of symbols interpretable as latent variables, which can be leveraged to (exactly approximately) perform conditional generation. This objective captures the task of assuring maximal diversity, in terms of KI. divergence between the different conditional components, represented as p. with k E S.
Log-barrier constrained generation This formulation can be found within the General Utilities RL literature [60]. In particular, here we show the case where constraints are enforced via log-barrier function, namely log(-). Nonetheless, the functional presented in Table remains meaningful for general penalty functions.
Optimal transport distances OT distances within Table and their relative notation are introduced below in the context of their first variation computation.
Maximum Mean Discrepancy Here & denotes a positive-definite kernel, which measures similarity between two points in sample space. Moreover, Mp denotes kernel mean embedding of distribution p [36]. In terms applications, choosing proper kernel could render possible to preserve specific structure of the initial pre-traines model that would be otherwise lost via KL regularization.
A.2 A brief tutorial on first variation derivation
In this work, we focus on the functionals that are Fréchet differentiable: Let V be normed spaces. Consider functional F V + R. There exists linear operator A V - R such that the following
22

---

limit holds
F(S+ + h) - F(F) - A 0.
lim |
(15)
|
We further assume that admits certain structure such that every element in its dual space (the space of bounded linear operator on V) admits some compact representation. For example, when V is the set of compact-supporte continuous bounded functions, there exists unique positive Borel measure with the same support, which can be identified as the linear functional. We denote this element as &F We will refer to (F as OF such that (6F[f), h) A[h). Sometimes we also denote it as 37 the first-order variation of F at f.
In this section, we briefty review strategies for deriving the first-order variation of two broad classes of functionals: those defined closed form with respect to the density (e.g., expectation and entropy) and those defined via variational formulations (e.g., CVaR, Wasserstein distance, and MMD).
Category 1: Functional defined in a closed form w.r.t. the density. For this class of functionals, the first-order variations can typically be computed using its definition and chain rule.
With definition (15) in mind, can try calculate the first-order variation of the mean functional. Consider a continuous and bounded function r : IR + and a probability measure μ on IR Consider the functional F(u) (x))) We have
[ + 6p) - F(u) 1r,8pp1 0.
(16)
We therefore obtain FFlu] = r for μ.. We will compute the first-order variations for other functionals in the next subsection.
Category 2: Functionals defined through a variational formulation. Another important subclass
of functionals considered in this paper is the ones defined via variational problem
F[J] sup G[f.g)
(17)
where f2 is set of functions vectors independent of the choice of f. and g is optimized over the set Q. We will assume that the maximizer g" (f( that reaches the optimal value for G[f, -] is unique (which is the case for the functionals considered in this project). It is known that one can use the Danskin's theorem (also known as the envelope theorem) to compute
8F[ ",G(U,9"()],
(18)
5f
under the assumption that differentiable [35].
A.3 Derivation of gradients of first-order variation for functionals in Table
Risk-Averse Optimization (Category 2) Recall that (p(p") = supfu E R|F2(U) s B}. where the random variable Z is defined 7 = r(z) with p= (x)). From [48], we have
s (x")) 3 inf {{ c 3 [min(r(x)) C, 00]
CVaR5(p") =
([(()))((x)) """"/("")>=(((()rrea""
Moreover, we have (* that solves the above optimization problem is exactly C* = (p"). By Danskin's theorem, one has (in a weak sense)
&CVaR,
33 minn(r(x)) (3'")),
(19)
5p
Risk-Seeking Optimization (Category 2) Recall that ()") = supfu E F2((()) S 3), where
the random variable Z is defined as Z = r(x) with x ~ n= (x). From [48], we have
SO E((()))((z)) 2 (6(p")
(1 inf
inf ( s+ 1- E (max(((x) Ç, 011
Moreover, we have Ç* that solves the above optimization problem is exactly C* = (p*).. By Danskin's 's theorem, one has (in weak sense)
(8Q)(p") (1- max
(20)
8p
23

---

PUNCTIONAL
VARIATION
OFTIMIZATION [14,55]
Eg-gr [r(x)(((>))
R(F") >- Range (Ha)")))
MODEL BLASING
(NARG))) R2-gr Inte) sfle")
Smin/r(x) - 06/171100
Feasy H(x)) Varly")
+()) (+(x)* 25. (+(]]-(())
RISK SEEKING OPTIMIZATION
"0%("* (x (2(x) (66%))
(1 Symax(rix) ",(p")
(R (9(=)0(2)"
OPTIMAL EXPERIMENT DESIGN
60 E flog des(-), - Tr(-)" 1-)
DIVERSE MODES DISCOVERY
Sun EQUATION (32)
x
LOG-BARREE BARRIER CONSTRAINED GENERATION
Bange (r()) Allo
See EQUATION (31)
KULLBACK LEIBLER DIVERGENCE
DeLip" liptt)) (xiilo piig) de
1 loggogp hagptre
RÉNYI
lym) J (e))
(4) ((())))))))
OPTIMAL TRANSPORT DISTANCES
Wolt [yfT) == inf MM0
SER EQUATION (29)
MAXMUU MEAN DISCRUPANCY
AMDIGG (P"") -m lapge P.e Mem
Aug wrg.aaaaenui0.f ppin)
Table Examples of practically relevant utilities F (blue) and divergences DD (orange), and their first-order variations.
Rényi Divergence (Category 1) Recall the definition of Rényi Divergence
Da(pllg) = 3-1 log /
dg(x).
(21)
We ignore higher-order terms like ((()p)²).
dg(x)
Dp(p + 5pl1g) Po(pl(a)
log
(22)
dg(z)
2
P
69 dg(((x)
log
(23)
3
dq(x)
3
f
2
Ez dg
log +
(24)
8 8-1
2 9
dg(x)
SB (g)
8-1
69 dg(x))
(25)
3 3-1
8
(1)
dq(x)
3 (
Ø-1
Ro(p.
dg(x)
(26)
óp
3-1
Optimal transport and Wasserstein-p distance (Category 2) Consider the optimal transport problem
OT,(u, v) = inf { da. dd.r)))()),))) 7(2,))de = u(v), / 7(x v)d u(x) (27)
24

---

where
r= 1- / 7xx v)dz u(v), / 7(2,
1lz. uldy v(x))
It admits the following equívalent dual formulation
{ fdu + / gdv f(x) + g(v) & c(x, (2
OT u) = sup f.g
(28)
By taking c(2, y) | vll". recover OT (u, W,(u, v)P. Let f and g" be the solution to the above dual optimization problem. From the Danskin's 's theorem, we have
W o)p
(29)
6u
the special case of p 1, we know that 9" -f" (note that the constraint can be equivalently written as || fll $ 1), in which case f is typically known as the critic in the WGAN framework.
Optimal Experiment Design. (Category 1) We take s(M) log det(M) as example. By chain rule, have
E ((r)((x) - Ar))
(#(x)((a)T An)
6F (F Tr
(30)
Log-Barrier Constrained Generation. (Category 1) By chain rule, we obtain
Bc
8F[p"] = c) C
(31)
Diverse modes discovery. (Category 1) By chain rule, we obtain
8F
8
If
p loggp dox ff prr-logg log (BielD de
5pm..
8p
8
(o pT log do
8.
(J 1 log (= dx
Eg
op
8p
n
E, log logp t + log (Bxlp=itt
(32)
Eg (p
Entropy. (Category 1) As first example, consider the entropy functional F(p) = - - dx.
By the definition of the first-order variation, we have 4f (p) -1 log p. and therefore V 4 (p) V lo p. This gradient term can be effectively estimated using standard score approximations; see [12].
25

---

B Proof for Theorem 5.1
Theorem 5.1 (Convergence guarantee of Flow Density Control with concave functionals). Given Assumptions 5. J.fiee-tuning pre-trained model wrr via FDC (Algorithm 1) with Th = L Ak [K], leads a policy π inducing a marginal distribution pi such that:
L - DKI(pt (pi l
((vi) ((P))
(12)
K
where pi = pT is the marginal distribution induced by the optimal policy x" E arg max (PT)) F(P)) (pi )
=
Proof. We prove this result using the framework of relative smoothness and relative strong convexity introduced in Section 5.
The analysis is based on the classical mirror descent framework under relative properties [33]. For notational simplicity, we let μl PT and fix an arbitrary reference density ja E PP To better align the notation of our theory with existing literature, we will proceed with the convex functional == G below.
We begin by showing the following inequality:
((UN) (()x-1) + (Óā 1), ph Pk-1) + LDo Pk-1)
(33)
< (1()==1)= + (5 (mx 14 14-1) LD (k-1) - (p,
(34)
The first inequality follows from the L-smoothness of g relative to as defined in Definition The second inequality uses the three-point inequality of the Bregman divergence [33, Lemma 3.1] with ((x)= = + (89()x-1),
(4--)); = 1k and = = j.k-
Next, using the I-strong concavity of 9 relative to O, again from Definition 1, we obtain:
((ux) - (()) + (L - 1)
LDo(u, pR).
(35)
By recursively applying the above inequality and using the monotonicity of G(ux) along with the non-negativity of the Bregman divergence, we obtain [33]:
L K (0(uz) - - < IDD(u, I. (,,1,) DelLuKK S (.D(Us L
(36)
Letting
K Cx
(37)
and rearranging terms, we arrive at the convergence rate:
1D IDolH, 1o) K
δ(pk) ((m) s CkLDelu
(38)
(1 1 -1
K
Finally, the convergence rate stated in the theorem follows by observing that (1+
2 21
1-1
26

---

C Proof for Theorem 5.2
To establish our main convergence result, we introduce two additional technical assumptions that are satisfied in virtually all practical settings:
Assumption C.1 (Support Compatibility). We assume that the support of PT is contained in a fixed compact domain Q for all k, and that for some we have supp(p")) = Ω
Assumption C.2 (Precompactness)) The sequence {8H(p77)) is precompact in the topology induced by the Lee norm.
We are now ready present the full proof. For the reader's convenience, we restate the theorem:
Theorem 5.2 (Convergence guarantee of Flow Density Control for general functionals). Given the Robbins-Monro step-size rule: Σ 7h = =0, Ex Th co, under Assumption 5.2 and technical assumptions (see Appendix C), the sequence of marginal densities p induced by the iterates Ti of Algorithm I converges weakly a stationary point P1 of 9 almost surely, formally: p -pi a.s..
Proof. We divide the proof into several key steps for case of reading.
Continuous-Time Mirror Flow. The main idea of our proof is to relate the discrete iterates (pT)kes produced by Algorithm I to continuous-tims mirror flow.
Define the initial dual variable as
hg 8PH(P log Ppre
and consider the gradient flow
hee 89(pt), Pr = 8( 8(-H)) (h+),
(MF)
where (h) log Jq. ex is the Fenchel dual of the negative entropy functional [25, 22]. This defines the deterministic mirror flow associated with G.
Continuous-Time Interpolation of Iterates. To connect the discrete algorithm with (MF), construct a continuous -imiiuouu--imeeeee interpolation of the dual iterates h² = 6H((T)) T5), Define the effective time
Or:
and let the interpolated process (t) be
rk
(())
(r 1 1²).
(Int)
T + 1
Intuitively, our convergence result follows two conditions hold:
Informal Assumption 1 (Closeness Continuous-Time Flow). The interpolated process l(()) asymptotically fo follows the dynamics of (MF) k - oo.
Informal Assumption 2 (Convergence of the Flow). The trajectories of (MF) converge station- ary point of G.
To formalize this, we invoke the stochastic approximation framework of [5]. Let Z be the space of integrable functions on Ω and let denote the flow of (MF). We define:
Definition (Asymptotic Pseudotrajectory (APT)). We say h(t) is an asymptotic pseudotrajectory (APT) of (MF) for all T > 0,
lim sup [h(1 + h) (b()l))) = 0. 100 OShST
If (t) is precompact APT, then [5] show:
27

---

Theorem C.1 (APT Limit Set Theorem). Let (t) be precompact APT for the flow (MF). Then, almost surely, the limit set of (t) is contained in the set of internally baain-transitive (ICT) points of (MF).
The proof of result thus follows from claims:
1. The iterates {h*} generate precompact APT under Assumptions C.I and5.2.
2. The ICT of (MF) consists only of stationary points of G.
The remainder the proof is devoted to verifying these two claims.
Convergence to Stationary Points. The second claim holds since the mirror flow admits G as strict Lyapunov function, and thus Corollary 6.6 in [5] ensures convergence the APT to the set of stationary points of G, provided that the set of equilibria is countable.
For the first claim, Assumptions C.I and C.2 ensure that the interpolated process is well-defined and precompact, while Assumption 5.2 allows us to apply standard stochastic approximation arguments [27]. We conclude the proof by applying Theorem C.I.
Quantitative Approximation to the Mirror Flow. For the first claim, we invoke the stochastic approximation techniques applied to the dual variables (see, e.g., [5, 27]) to obtain the following bound:
sup [1(t + 9) O,(M(0)I S C(T) 1, b(T) (T)],
(39)
0S
where C(T) depends only on T, A(I - 1,7 + 1) captures cumulative noise fluctuations, and B(T); (T) are the bias and step-size terms over the interval. This explicitly bounds the deviation of the interpolated process from the deterministic mirror flow.
Under the noise and bias conditions Assumption 5.2, standard stochastic approximation results [5, 27] imply
lim ( 1) lim b(T) O. T-+00
T-000
Hence, ((t) is an APT of the mirror flow.
Assuming precompactness of the dual iterates (stated Assumption C.2), Theorem 5.7 in [5] implies that the limit set of (t))) is internally chain transitive (ICT) for the mirror flow. Combining the quantitative approximation (39), the APT argument, and the limit set characterization. we conclude that the discrete iterates converge to stationary points of G. completing the proof.
Conclusion.
28

---

D Detailed Example of Algorithm Implementation
D.1 Implementation of
To ensure completeness, below we provide pseudocode for one concrete realization of ENTROPY REG-- ULARIZZDDCONRROLSOLVEE as in Eq. (8) using first-order optimization routine. In particular, de- scribe exactly the version employed in Sec. 6, which builds on the Adjoint Matching framework [14], casting linear fine-tuning stochastic optimal control problem and tackling it via regression.
Let uree the initial, pre-trained vector field, and finetuned its fine-tuned counterpart. We also use refer the accumulated noise schedule from [23] effectively following the flow models notation introduced by Adjoint Mathing [14, Sec. 5.2]. The full procedure Algorithm 2.
Algorithm 2 ENTROPY ENTRROPPEEUULAAEEEECOOOOOOOEEEE (Adjoint Matching [14]) based implementation
Input: N number of iterations, w pre-trained flow vector field, 7 regularization coefficient
as in Eq. (8), h: step size, Vf: reward function gradient, m batch size
2: Init: w //inetneed 12" pre with parameter
3: for = 0, 1, ,, N do
4: Sample trajectories (XiT== via memoryless noise schedule [14], e.g.,
sample E, NO,, I), Xo 1), then: 0:t
20g finetuned (Xe.))
Ot xx) + Vha(t)c
X = X,
at
Use reward gradient:
ät --=-Lv((X) v((1))
For each trajectory, solve the lean adjoint ODE, see [14, Eq. 38-39], from t = 1 to e.g.:
āi X2 ac
āæー ã + hā, Vx, (Xe
Where X₁ and ā are computed without gradients, i.e., X₁ = stopgrad(Xi), â. = stopgrad(ā.). For each trajectory compute the Adjoint Matching objective [14, Eq. 37]:
1-h Lg = Σ t=
fimetuned -t,) - uxre(x,,))) + (t))all
u
ott)
Compute the gradient Vol(() and update 0.
5: end for
6: output: Fine-tuned noise predictor ug finetuned
D.2 Discussion: computational complexity and cost of FDC
Flow Density Control (see Algorithm 1) is sequential fine-tuning scheme, which performs K iterations of base fine-tuning oracle, as shown in Algorithm Typically, for the case of Adjoint Matching [14], which is contextualized in Algorithm 2, the inner oracle also performs N iterations to solve the classic fine-tuning problem. a consequence, at first glance, this lead to FDC having computational complexity scaling linearly in K the one of classic fine-tuning. Nonetheless, this does not seem to capture well the practical computational cost. In particular, we wish point out the two following observations:
As discussed for the molecular design experiment in Sec. 6 and further in Appendix E, the FDC scheme might work well even with very approximate oracle to solve the entropy-regularized control problem at each iteration.
For many real-world problems very small number of iterations might sufficient to approximate the non-linear functional sufficiently well and hence obtain useful fine-tuned

---

models. This is shown in text-to text-to-image bridge design experiment in Sec. 6 and in Appendix E. In this case, merely K iterations of FDC lead promising results.
30

---

E
Experimental Details
E.1 Used computational resources
We run all experiments on a single Nvidia H100 GPU.
E.2
Experiments in Illustrative Settings
Shared experimental setup. For all illustrative experiments we utilize Adjoint Matching (AM) [14] for the entropy-regullaizze fine-tuning solver in Algorithm 1. Moreover, the stochastic gradient steps within the AM scheme are performed via an Adam optimizer.
Riek-averse reward maximization for better worst-case validity or safety. In this experiment, we execute FDC for 2 iterations with a total 1000 gradient steps within each iteration, AM solver (within the FDC scheme) with learning rate of 2e - the AM baseline, is run for 1000 gradient steps with a = 0.2857, and learning rate of Le -5, The resulting CVaR is computed via the standard torch quantile method. The values of 3 reported in the main paper effectively refers to the value of 1 - B. In the following, we report mean and sample standard deviation of AM and FDC over 5 seeds.
10 and 10. Mcanwhile,
CVaR,
Pre-trained
256.8 + 8.15
AM
225.3 + 78.9
FDC (1 iteration)
221.1 + 73.2 FDC (2 iteration)
90.0 + 0.05
Figure 5: Statistical analysis for CVaR
Novelty-seeking reward maximization for discovery. We run FDC for K = 2 iterations with total of 1000 gradient steps within each iteration, AM solver (within the FDC scheme) with learning 10 and 1) 0.625, and 8000 samples used estimate the first variation
rate of 36-6. gradient as explained in Appendix A. Meanwhile, the AM baseline, is run for 1000 gradient steps with a 0.666, and learning rate of 1 The resulting SQ is computed via the standard torch quantile method. In the following, we report mean and sample standard deviation of AM and FDC over 5 seeds.
SQa
Pre-trained
59.6 + 7.5
AM
56.7 2.7
FDC (1 iteration)
55.0 -:00:000 0.04
FDC (2 iteration)
452.5 + 250.0
Figure 6: Statistical analysis for SQ2 utility.
Reward maximization regularized via optimal transport distance. Within this experiment, we present two runs of FDC, namely FDC-A and FDC-B, compared against AM. Both FDC-A and FDC-B have been run for K 6 iterations of FDC, with a 0.1, AM oracle learning rate of 16 6.666. Both their discriminators to solve the dual OT problem presented in Appendix A and mentioned within Sec. 4, have been learned via simple MLP architecture with 800 gradient steps, by enforcing the 1-Lip. condition via the standard gradient penalty technique with regularization strength of Agp 10.0 and learning rate of le In particular, FDC-A is based on the distance defined, for two 2-dimensional points (I1, (2) and (V1, y2) by:
d / Vaz 31 + (K y2))²
Analogously, FDC-B leverages dg defined as:
dz V(((a y1)) + (x2 22
31

---

Eq. 8 as a learning-rate parameter, by choosing Where K = 7 in both cases. On the other hand, the AM baseline is run for 1000 gradient steps with leamning rate of le" Le and 1.538.
Conservative manifold exploration. We ran FDC for K 50 iterations and 2500 gradient steps total with 1 = 10 and a = 0.0, 0.01, 0.5, 1.0. We set the AM learning rate to and sample trajectories of length 400 for computing the AM loss.
B[r(x)]
W
A%
Pre
29.5 + 0.0
0
AM
35.08 + 0.04
4.68 + 0.0
100
FDC-A
35.38 + 0.04
1.92 + 0.03
288 + 15.0
Figure 7: Statistical analysis for W divergence.
E.3 Further Ablations
Runtime The only input hyperparameter Algorithm the number of iterations Towards evaluating its effect on algorithm execution, in the following, we consider an experimental setup analogous "Risk-averse reward maximization for better worst-case validity or safety" experiment in Fig. 3 (top row), and evaluate the effect of different numbers of iterations (K) on run-time and solution quality. We report results in Fig. 8 showing the depending on hlyper-paraeett
K for v = 20. As one can expect, for small step-sizes 77 the runtime and solution quality scale nearly linearly in K given fixed number of iterations N of the eetropy-regularizze solver (see Apx. D for the definition of N). As one can expect by interpreting
K
Runtime (s)
CVaR estimate (via 1000 samples)
0
0.00
254.49
44.71
241.37
2
89.38
167.04
133.31
288.72 72
4
176.79
271.00
220.37
84.89
6
264.06
96.55
Figure 8: Runtime vs. CVaR estimate as function of K, 20.
smaller values of 1) convergence can be achieved with less iterations K. In Fig. we report the same evaluation with 7) 10. The above tables hint the fact that the FDC fine-tuning process can be
K
Runtime (s)
CVaR estimate (via 1000 samples)
0
0.00
254.49
44.23
249.09
2
87.78
90.0
Figure Runtime vs. CVaR estimate function of K, 17 10.
interpreted experimentally classic (convex non-convex) optimization, although the space of generative models, with learning rate (or step-size) controlled by
Approximate Oracle the following, we investigate the use of an approximate entropy entropy-regularize regularizee control solver oracle (i.e., performing approximate entropy reguoar-re fine-tuning at each iteration of FDC), showing that this can also lead to optimality via increasing the number of iterations K. In the following (see Fig. 10) we consider N = 100 instead of N 1000 (as in previous experiments) and use K = 5 showing that FDC can retrieve the same final fine-tuned model in Fig. using only one tenth gradient steps (i.e. N 100 instead of N 1000) for the inner oracle.
32

---

K
Runtime (s)
CVaR estimate (via 1000 samples)
0.00
254.49
5.11
221.1
10.01
194.69
3
14.73
90.01
19.60
91.1
5
24.60
90.0
Figure 10: Runtime vs. CVaR estimate as function of K, 17 = 10, N 100.
E.4 Real-World Experiments
Molecular design for single-point energy minimization. In this experiment FDC is run for K 10 iterations, with merely gradient steps at each iteration (i.e., the AM oracle is very approximate), AM learning rate of le -e- 0.01 and a = 0. Meanwhile, the AM baseline is run for 240 gradient steps with a 0.0045.
Text-to-image bridge designs conservative exploration. For this experiment we ran FDC on single Nvidia HI00 GPU, with K 2, 200, a 0.001 and a 100 gradient steps in total. Similarly to previous work, we tuned the vector field resulting from applying classifier-fre guidance with guidance scale w 8 in SD-1.5.
33


Reviews:
[{"review_id": "ZPJDnX5C0z", "reviewer": "Reviewer_LUT6", "summary": "This paper proposes Flow Density Control (FDC), a general framework for fine-tuning flow and diffusion generative models using a broad class of utilities and divergences. Extending existing works that mostly use KL-regularized rewards, FDC supports a broad class of objectives, including risk-averse, novelty-seeking, and various divergences.", "questions": "1. Given the motivation around molecular and biological design, expanding the biological experiments (e.g., multiple objectives, more evaluation metrics, realistic constraints) would significantly strengthen the practical contribution. Can the authors comment on this>\n\n2. Please discuss how FDC compares with more recent non-KL regularized fine-tuning approaches such as those using Wasserstein, heavy-tailed distributions, or exploration-based objectives. I am not asking the authors to prove that this work can outperform novelty-seeking papers in novelty-seeking and can outperform risk-aware papers in controlling risks. I understand this work's major advantage is the generality of the schemework; at least, I would expect authors to discuss relevant works.  \n\n3. Can some of the non-linear objectives (e.g., CVaR, entropy maximization) be reformulated as expected reward maximization under transformed reward functions? (from my experience, I know CVaR can often be modeled via conditional expectation, so it is not principally different from other reward functions) If so, how does FDC differ in sample efficiency or convergence?", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength:\n\n1. This work extends current methods by supporting general utilities (e.g., CVaR, entropy) and divergences (e.g., Rényi, OT).\n2. Solid theoretical grounding with convergence analysis is also provided.\n3. Experiments show benefits over prior methods in both synthetic and practical domains.\n\nWeakness:\n\n1. As a researcher with experience in AI for Science, I found the experiments, while extensive and strongly relevant, are a bit short in evaluation. First, molecular design on QM9 is quite limited as this dataset is rather small and is not representative enough nowadays. Have the authors consider experiments on large datasets such as zinc250k？ Second, the discrete generative modeling context is especially relevant in biological design, for example, DNA sequences and proteins. But any sort of biology experiment is not included in experiments?\n\n2. The writing of the experiment section is too concise. Many necessary details are not provided. I found that section hard to parse. Besides, it seems ablation studies and sensitivity analysis for hyperparameters are also missing.", "was_revised": true, "final_justification": "After reading the detailed response with abundant additional computational results, my confidence in the framework’s generality and rigor has improved a lot. I support the acceptance of this work.", "rebuttal": "We thank the Reviewer for recognizing our work as with solid theoretical grounding as well as with extensive and strongly relevant experiments. In the following, we address several points and questions mentioned within the review.\n\n**Biology experiments**\n\nIn this work, we propose a novel probability-space optimization framework that enables fine‑tuning of flow and diffusion models for a broader class of utilities beyond the reach of classic control or RL methods. While this foundational contribution is significantly motivated by scientific discovery applications, showing impact in such applications goes beyond the scope of this work. The presented sample application on QM9, a computational chemistry setup, although not capturing the full complexity of modern chemical and biological tasks, shows that the proposed theoretically-grounded method is indeed (very) promising for biological discovery tasks, where we hope to see it tested as a natural next step.\n\n**Discrete generative modeling**\n\nWhile the focus of this work is specific to continuous-state flow and diffusion models, which cover many practically relevant biological discovery applications (e.g. protein docking for drug discovery [5]), we believe the algorithm proposed can be easily extended to control discrete diffusion or flow models as well. This is due to the recent introduction of entropy-regularized fine-tuning oracles for such discrete models (see, e.g. Eq. (5) in [1] for discrete diffusion), which can be employed to implement our method for such discrete models.\n\n**Experiments section too concise** \n\nWe agree with the Reviewer: unfortunately this was due to page limits. We plan to use the extra page of the camera-ready version to extend the experimental section with further experimental details as well as further ablation studies as reported in the following point.\n\n**Further ablation and sensitivity analysis**\n\nWe have shown promising performances of the proposed method for six different utilities and/or tasks of practical relevance. Since the algorithm proposed is particularly principled, its dependency on hyperparameters can be inferred from its (convex) optimization interpretation and is partially discussed in Appendix D.2 under \"Discussion: computational complexity and cost of FDC\". Nonetheless, we agree with the Reviewer that the experimental section can be further strengthened via ablation studies and sensitivity analysis wrt the algorithm hyperparameters. We wish to point out that our method can leverage as oracle any entropy-regularized fine-tuning method, such as Adjoint Matching [2], and the practical performance of FDC significantly depends on the one of such oracles. Luckily, there exist several oracles with established evaluations in complex settings (see, e.g., [2] for images or [6] for molecules and proteins). Effectively, the only input hyperparameter shown in Algorithm 1 is the number of iterations $K$. Towards evaluating its effect on algorithm execution, in the following, we consider an experimental setup analogous to \"Risk-averse reward maximization for better worst-case validity or safety\" experiment in Fig. 3 (top row), and evaluate the effect of different numbers of iterations \\(K\\) on run-time and solution quality. Notice that in the following tasks the optimal CVaR value is $90.0$.\n- Dependency on hyper-parameter $K$ for $\\eta = 20$. As one can expect, for small step-sizes $1$ / $\\eta$, the runtime and solution quality improves nearly linearly in $K$ given a fixed number of iterations $N$ of the entropy-regularized solver (see Appendix D.1 for definition of $N$). \n\n|  **K** | **Runtime (s)** | **CVaR estimate (via 1000 samples)** |\n| :----: | :-------------: | :------------------------------: |\n|    0   |       0.00      |        254.48877624511718        |\n|    1   |      44.71      |        241.36963653564453        |\n|    2   |      89.38      |        167.03791809082030        |\n|    3   |      133.31     |        288.72082519531250        |\n|    4   |      176.79     |        271.00385437011720        |\n|    5   |      220.37     |         84.88514938354493        |\n|    6   |      264.06     |         96.55234909057617        |\n\n- As one can expect by interpreting $\\eta$ in Eq. (8) as a learning-rate parameter, by choosing smaller values of $\\eta$ convergence can be achieved with less iterations $K$. In the following, the same setup with $\\eta =10$.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 44.23           | 249.0867721557617                |\n| 2    | 87.78           | 90.0                             |\n\nThe above two tables hint at the fact that the FDC fine-tuning process can be interpreted experimentally as classic (convex or non-convex) optimization, although on the space of generative models, with learning rate (or step-size) controlled by $\\eta$.\n\n- Furthermore, using an approximate entropy-regularized control solver oracle (i.e., performing approximate entropy-regularized fine-tuning at each iteration of FDC) can also lead to optimality via increasing the number of iterations $K$. In the following we consider $N=100$ instead of $N=1000$ (as in previous experiments) and use $K=5$ showing that FDC can retrieve the same final fine-tuned model as in the previous table using only one tenth gradient steps (i.e. $N=100$ instead of $N=1000$) for the inner oracle.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 5.11            | 221.07437133789062               |\n| 2    | 10.01           | 194.68543167114257               |\n| 3    | 14.73           | 90.01436312185169                |\n| 4    | 19.60           | 91.10320648061439                |\n| 5    | 24.60           | 90.0                             |\n\nWe thank the Reviewer for the question and will certainly discuss these dependences on hyperparameters as well as oracle quality in an updated version of the work.\n\n\n\n**Comparison with related works for fine-tuning of general utilities**\n\nTo the best of our knowledge, the only work that directly tackles a specific instance of convex or general utilities as depicted within our work is [3], which tackles the sub-case of entropy maximization and we compare in our work. In particular, in the \"Conservative manifold exploration.\" experiment, we explain how our framework can extend entropy maximization introduced in [3] to a \"conservative\" approach where a divergence measure from the prior model can be enforced, which is not possible in prior work. To our knowledge, no other work that aims to control flow models via continuous-time RL, or stochastic control, shows how to tackle similar general objectives. Other approaches seem to tackle specific sub-cases via approximating such complex non-linear objectives and then use classic RL machinery (e.g.,  [4]). In comparison to such schemes, our work renders possible to directly optimize non-linear objectives without the need of developing explicit and objective-specific approximations, which do not necessarily exist and that can be arbitrarily loose. We thank again the Reviewer for the question and we will further clarify these connections in the revised version of the manuscript.\n\n**Objectives reformulation**\n\nWe are not aware of such cases. The CVaR is typically defined as a conditional expectation (as in Table 1 in our work) rendering it linear given the quantile. But the latter quantity depends non-linearly on the measure rendering the CVaR non-linear by definition. For this reason, it is typically treated as a convex (or non-convex, depending on its parametrization) functional in related works (see, e.g., Table 1 in [7] or [8]). If the Reviewer can provide further details about this point, we would be happy to further discuss. \n\n**References**\n\n[1] Wang et al., Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design. ICLR 2025.\n\n[2] Domingo-Enrich et al., Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. ICLR 2025.\n\n[3] De Santi et al., Provable maximum entropy manifold exploration via diffusion models. ICML 2025.\n\n[4] Fan et al., Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization. ICLR 2025.\n\n[5] Corso et al., Diffdock: Diffusion steps, twists, and turns for molecular docking. ICLR 2023.\n\n[6] Uehara et al., Feedback efficient online fine-tuning of diffusion models. ICML 2024.\n\n[7] Mutti et al., Convex Reinforcement Learning in Finite Trials. JMLR 2023.\n\n[8] Mutti et al., Challenging Common Assumptions in Convex Reinforcement Learning. NeurIPS 2022."}, {"review_id": "pGH1JGKtsl", "reviewer": "Reviewer_hbzy", "summary": "The authors present a unified framework for optimizing general utilities by finetuning pre-trained flow models, subject to a constraint with respect to the pre-trained model (expressed by an arbitrary divergence). This framework generalizes prior work in fine-tuning flow/diffusion models, e.g., those that maximize expected reward subject to KL regularization. The authors present an algorithm, flow density control (FDC), to solve this general optimization problem, and prove convergence to a stationary point under mild assumptions. Finally, the authors conduct experiments in novelty-seeking molecular design and creativity-optimizing text-to-image generation.", "questions": "The main points that would cause me to increase my score are extensions to the experimental results: first and foremost, increased statistical rigor, and secondarily, more experiments showing the ability of FDC to \"unlock\" new capabilities in popular domains like text-to-image generation.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Note: I do not have the theoretical background to check the derivations and proofs in detail.\n\n**Strengths**\n\nThe paper is very well-written and very thorough. I understood exactly what problem the authors are trying to solve and why it's important. I do have some background in the theory of diffusion/flow models, specifically in fine-tuning them to maximize rewards (a special case of FDC), and I think FDC is a nice generalization of this idea. I appreciate the examples of various functionals presented in Table 1 --- even though not all of them are used in the experiments --- as well as the expanded explanations of their practical applications given in Appendix A. While I do not follow the theoretical analysis in detail, the primary result (Theorem 5.2) does seem strong --- the assumptions are fairly mild, and the authors are able to prove weak convergence to a stationary point of the utility functional.\n\nOverall, the clarity and thoroughness of the paper provide a strong starting point for operationalizing FDC for many practical applications, and on this basis, I think this work would be a valuable contribution to the community.\n\n**Weaknesses**\n\nThough they are not the primary focus of the paper, the included experimental results are relatively weak. The authors do not include error bars or any measure of statistical certainty, citing lack of time and compute. I think this is a poor excuse, as sampling from a flow model is not so expensive compared to training. Tests for statistical certainty are important for scientific rigor. Also, the text-to-image experiment merely maximizes entropy subject to KL regularization, which is a capability already demonstrated in prior work.", "was_revised": true, "final_justification": "The authors have answered my questions and addressed my main concern about statistical rigor. I will maintain my recommendation to accept.", "rebuttal": "We thank the Reviewer for recognizing our work as very well-written and very thorough, the problem we are trying to solve as important, and our contribution as valuable to the community. In the following, we address several points and questions mentioned within the review.\n\n**Statistical analysis**\n\nWe wish to point out that in order to perform a statistical analysis of the proposed fine-tuning method it would be needed to run such fine-tuning scheme multiple times rather than sampling multiple times. The latter is already done in the work for which we reported Monte Carlo estimates of expended performances (typically via approximately $5000$ samples). Nonetheless, we agree with the Reviewer regarding the relevance of statistical analysis. To this end, in the following we report the estimated expectations and sample standard deviations for several experiments given by $5$ seeds. In the following, we consider experimental setups analogous to the first three experiment types shown in the paper. These results confirm the insights already contained in the submitted version of the paper and further validate the working mechanisms of the proposed method. \n\n- Risk-averse reward maximization for better worst-case validity or safety. \n| **Method**                                | **CVaR$_\\beta$**          |\n|:-----------------------------------------|:-------------------------:|\n| Pre-trained                               | $256.8 \\pm 8.15$         |\n| AM                                        | $225.3 \\pm 78.9$         |\n| FDC (1 iteration)                         | $221.1 \\pm 73.2$         |\n| FDC (2 iterations)                        | $90.0 \\pm 0.05$          |\n\n- Novelty-seeking reward maximization for discovery.\n| **Method**                                | **SQ$_\\beta$**           |\n|:-----------------------------------------|:------------------------:|\n| Pre-trained                               | $59.6 \\pm 7.5$          |\n| AM                                        | $56.7 \\pm 2.7$          |\n| FDC (1 iteration)                         | $55.0 \\pm 0.04$         |\n| FDC (2 iterations)                        | $452.5 \\pm 250.0$       |\n\n- Reward maximization regularized via optimal transport distance.\n| **Method**             | **$\\mathbb{E}[r(x)]$**     | **$W_1^A$**        | **$\\Delta\\%$**        |\n|:-----------------------|:-------------------:|:------------------:|:---------------------:|\n| Pre                    | $29.5 \\pm 0.0$      | $0$                | $-$                   |\n| AM                     | $35.08 \\pm 0.04$    | $4.68 \\pm 0.0$     | $100$                 |\n| FDC-A                  | $35.38 \\pm 0.04$    | $1.92 \\pm 0.03$    | $288 \\pm 15.0$        |\n\nWe thank the Reviewer for raising this point, and in an updated version of the work, we will make sure to add further statistical analysis of the proposed method as suggested by the Reviewer.\n\n**Further ablation studies**\nTowards strengthening the experimental evaluation as suggested by the Reviewer, we conducted further experimental evaluations of the proposed method to analyse the effect on solution quality and runtime of hyperparameter choice and approximate (entropy-regularized) fine-tuning oracles. In the following, we consider an experimental setup analogous to the \"Risk-averse reward maximization for better worst-case validity or safety\" experiment in Fig. 3 (top row), and evaluate the effect of different numbers of iterations $K$ on run-time and solution quality.\n- Dependency on hyper-parameter $K$ for $\\eta = 20$. As one can expect, for small step-sizes $1$ / $\\eta$, the runtime and solution quality improves nearly linearly in $K$ given a fixed number of iterations $N$ of the entropy-regularized solver (see Appendix D.1 for definition of $N$). \n\n|  **K** | **Runtime (s)** | **CVaR estimate (via 1000 samples)** |\n| :----: | :-------------: | :------------------------------: |\n|    0   |       0.00      |        254.48877624511718        |\n|    1   |      44.71      |        241.36963653564453        |\n|    2   |      89.38      |        167.03791809082030        |\n|    3   |      133.31     |        288.72082519531250        |\n|    4   |      176.79     |        271.00385437011720        |\n|    5   |      220.37     |         84.88514938354493        |\n|    6   |      264.06     |         96.55234909057617        |\n\n- As one can expect by interpreting $\\eta$ in Eq. (8) as a learning-rate parameter, by choosing smaller values of $\\eta$ convergence can be achieved with less iterations $K$. In the following, the same setup with $\\eta =10$.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 44.23           | 249.0867721557617                |\n| 2    | 87.78           | 90.0                             |\n\nThe above two tables hint at the fact that the FDC fine-tuning process can be interpreted experimentally as classic (convex or non-convex) optimization, although on the space of generative models, with learning rate (or step-size) controlled by $\\eta$.\n\n- Furthermore, using an approximate entropy-regularized control solver oracle (i.e., performing approximate entropy-regularized fine-tuning at each iteration of FDC) can also lead to optimality via increasing the number of iterations $K$. In the following we consider $N=100$ instead of $N=1000$ (as in previous experiments) and use $K=5$ showing that FDC can retrieve the same final fine-tuned model as in the previous table using only one tenth gradient steps (i.e. $N=100$ instead of $N=1000$) for the inner oracle.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 5.11            | 221.07437133789062               |\n| 2    | 10.01           | 194.68543167114257               |\n| 3    | 14.73           | 90.01436312185169                |\n| 4    | 19.60           | 91.10320648061439                |\n| 5    | 24.60           | 90.0                             |\n\n\n\n**Exploration in image experiments**\n\nTo our knowledge, this is the first work that can perform KL-regularized entropy maximization, which we denote by Conservative Manifold Exploration, in a principled manner. Previous work (e.g., see [1]) could maximize entropy on the pre-trained model support only without any form of regularization. Crucially, such an un-regularized exploration objective risks to lead to an excessively explorative fine-tuned model that induces positive density in regions of invalid designs that do not correspond to any meaningful mode of the pre-trained model. On the other hand, the KL-regularized entropy objective that we introduced can render possible to balance the probability of sampling designs form diverse modes of the pre-trained model (even low-probability ones), while preventing sampling from regions of excessively low probability according to the prior model - thus 'Conservative', thanks to the divergence regularization term. This comparison is showcased in the \"Conservative manifold exploration\" experiment in Fig. 4 (top), where the method in [1] is retrieved for the sub-case of $\\alpha = 0$. This also shows that the framework presented in our work is strictly more expressive than previous formulations for exploration via diffusion or flow models. If the Reviewer was referring to other works of which we are unaware, we are happy to further discuss once given more details.\n\n**'Unlocking new abilities' in text-to-image generation**\n\nThe presented work proposes theoretically-grounded machinery to unlock promising abilities for this type of problems. Examples are CVaR maximization to prevent a text-to-image model to generate poor quality images in the worst case, according to a reward or preference model, SQ to generate particularly optimal images within a batch, and divergence regularized entropy-maximization for diverse text-to-image generation. Experimentally, we showcased the last application in the context of generation of bridge designs conditioned on a textual prompt showcasing increased diversity metrics (higher Vendi score) while maintaining high semantic quality score (CLIP). We regard this ability to generate more diverse images (due to the entropy term) while preserving their semantic quality (due to the divergence term) to be a novel and practically relevant ability for this class of models unlocked by the proposed scheme. We believe further applications of the proposed method on real-world problems, although both promising and interesting, go beyond the scope of this work.\n\n**References:**\n\n[1] De Santi et al., Provable maximum entropy manifold exploration via diffusion models. ICML 2025."}, {"review_id": "jkEVdYUxtk", "reviewer": "Reviewer_Zmzc", "summary": "This paper presents Flow Density Control (FDC), a new approach to fine-tuning pre-trained flow and diffusion generative models for a much broader class of objectives than previously possible. Existing work typically uses a linear expected-reward utility and KL-divergence regularization; FDC enables fine-tuning with general non-linear utilities (such as risk-averse, novelty-seeking, or diversity objectives) and more general divergences (e.g., optimal transport, Rényi).", "questions": "1. In line 244/245, why does the baseline method result in a worse loss than the basic model before fine-tuning?\n2. In line 278, why can't we improve the rewards for the top 0.2% while maintaining the average reward?\n3. Are there scenarios that require both risk-averse rewards and novelty-seeking rewards, and how do the methods perform in such contexts?", "limitations": "Yes, but only a little", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strength\n1. Theoretical Generalization: The paper rigorously extends finetuning of generative models to general utilities and divergences, going beyond standard entropy-regularized control. Section 3 formalizes this new generative optimization problem and classifies utility/divergence choices of practical interest.\n2. The paper provides theoretical guarantees for the methodology.\n3. The paper features a clear organizational structure and the figures are well-integrated and informative.\n\n\n## Weakness\n1. Although the showcased experiments are illustrative and varied, the empirical evaluation is still somewhat limited in scope. For the high-dimensional molecular and image tasks, more extensive quantitative metrics (e.g., additional standard baselines, statistical significance, generalization to other data sets or tasks) would strengthen the case that the generality of FDC yields practical upstream benefits in real-world applications.\n2. The paper selects a very limited number of baseline comparison methods. However, as I am not familiar with the frontiers of this field, I need to refer to the opinions of other reviewers.\n3. The paper fails to provide sufficient implementation details.\n4. The discussion does not explicitly benchmark or analyze the computational overhead of sequentially running entropy-regularized control solvers versus standard finetuning. In large models or high-dimensional tasks, this could be a meaningful practical consideration.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for recognizing our work as rigorous, well organized, and with varied experiments, as well as appreciating its theoretical generalization of generative models fine-tuning. In the following, we address several points and questions mentioned within the review.\n\n**Further evaluation** We agree with the Reviewer regarding the relevance of statistical analysis. To this end, we report the estimated expectations and sample standard deviations for several experiments given by $5$ seeds. We consider experiments analogous to the first three types shown in the paper. These results confirm the insights within the submitted version of the paper and further validate the working mechanisms on the proposed method.\n\n- Risk-averse reward maximization for better worst-case validity or safety. \n| **Method**                                | **CVaR$_\\beta$**          |\n|:-----------------------------------------|:-------------------------:|\n| Pre-trained                               | $256.8 \\pm 8.15$         |\n| AM                                        | $225.3 \\pm 78.9$         |\n| FDC (1 iteration)                         | $221.1 \\pm 73.2$         |\n| FDC (2 iterations)                        | $90.0 \\pm 0.05$          |\n\n- Novelty-seeking reward maximization for discovery.\n| **Method**                                | **SQ$_\\beta$**           |\n|:-----------------------------------------|:------------------------:|\n| Pre-trained                               | $59.6 \\pm 7.5$          |\n| AM                                        | $56.7 \\pm 2.7$          |\n| FDC (1 iteration)                         | $55.0 \\pm 0.04$         |\n| FDC (2 iterations)                        | $452.5 \\pm 250.0$       |\n\n- Reward maximization regularized via optimal transport distance.\n| **Method**             | **$\\mathbb{E}[r(x)]$**     | **$W_1^A$**        | **$\\Delta\\%$**        |\n|:-----------------------|:-------------------:|:------------------:|:---------------------:|\n| Pre                    | $29.5 \\pm 0.0$      | $0$                | $-$                   |\n| AM                     | $35.08 \\pm 0.04$    | $4.68 \\pm 0.0$     | $100$                 |\n| FDC-A                  | $35.38 \\pm 0.04$    | $1.92 \\pm 0.03$    | $288 \\pm 15.0$        |\n\nWe thank the Reviewer for raising this point, and in an updated version of the work, we will add further statistical analysis of the proposed method as suggested by the Reviewer.\n\n**Baselines** We believe that the baseline method we considered within our experimental evaluation, namely Adjoint Matching (AM) [1] is actually regarded as the strongest, as well as widely adopted, available method to target the task at hand. More specifically, given a known reward function one whishes to maximize by fine-tuning a diffusion or flow model, there is significant literature regarding how to solve this problem via RL or Control theoretic techniques (e.g., see an overview in [2]). Early methods required to solve the problem approximately or with hidden value function bias (e.g., [3]) as discussed in [Section 4.2, 1], while recently AM [1] rendered possible to tackle the problem exactly and with the strongest empirical performances to our knowledge. Therefore, we regard it as a strong and expressive baseline, widely adopted, and already known to dominate alternative methods. If the Reviewer has in mind alternative specific baselines that are competitive with the one proposed we would be happy to hear and further discuss this point.\n\n**Implementation details** Alg. 1 is composed of two operations: (a) gradient estimation (line 4), and (b) fine-tuning via an entropy-regularized control solver (line 5). Gradient estimation is discussed in Sec. 4 (from line 166) for several common objectives, and in Appendix A.3 for the remaining ones. Meanwhile, a detailed implementation of the fine-tuning solver is reported in Appendix D.1 under \"Detailed example of algorithm implementation\". Nonetheless, we thank the Reviewer to mention this point and understand that it would be particularly helpful to find a specific complete implementation of FDC, which we will add in an updated version of the work. \n\n**Computational overhead** The computational overhead of FDC is discussed in Appendix D.2 under \"Discussion: computational complexity and cost of FDC\". Nonetheless, we agree with the Reviewer that its experimental analysis is relevant. To this end, in the following, we consider an experimental setup analogous to \"Risk-averse reward maximization for better worst-case validity or safety\" experiment in Fig. 3 (top row), and evaluate the effect of different numbers of iterations $K$ on run-time and solution quality, as well as evaluate the effect of a computationally cheaper yet approximate (entropy-regularized) fine-tuning oracle.\n\n- Dependency on hyper-parameter $K$ for $\\eta = 20$. As one can expect, for small step-sizes $1$ / $\\eta$, the runtime and solution quality improves nearly linearly in $K$ given a fixed number of iterations $N$ of the entropy-regularized solver (see Appendix D.1 for definition of $N$). \n\n|  **K** | **Runtime (s)** | **CVaR estimate (via 1000 samples)** |\n| :----: | :-------------: | :------------------------------: |\n|    0   |       0.00      |        254.48877624511718        |\n|    1   |      44.71      |        241.36963653564453        |\n|    2   |      89.38      |        167.03791809082030        |\n|    3   |      133.31     |        288.72082519531250        |\n|    4   |      176.79     |        271.00385437011720        |\n|    5   |      220.37     |         84.88514938354493        |\n|    6   |      264.06     |         96.55234909057617        |\n\n- As one can expect by interpreting $\\eta$ in Eq. (8) as a learning-rate parameter, by choosing smaller values of $\\eta$ convergence can be achieved with less iterations $K$. In the following, the same setup with $\\eta =10$.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 44.23           | 249.0867721557617                |\n| 2    | 87.78           | 90.0                             |\n\n- Furthermore, a cheaper entropy-regularized control solver can lead to optimality via increasing the number of iterations $K$. We consider $N=100$ instead of $N=1000$ (as in previous experiments) and use $K=5$. FDC can compute an optimal fine-tuned model using only one tenth solver gradient steps.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 5.11            | 221.07437133789062               |\n| 2    | 10.01           | 194.68543167114257               |\n| 3    | 14.73           | 90.01436312185169                |\n| 4    | 19.60           | 91.10320648061439                |\n| 5    | 24.60           | 90.0                             |\n\n**Question 1** This is due to the fact that classic fine-tuning methods such as AM [1] cannot properly express the risk-averse objective (CVaR). In this case, fine-tuning the pre-trained model with AM to maximize expected rewards leads to a fine-tuned model more prone to generate worse designs in the worst-case. This can happen because expected reward maximization does not necessarily lead to optimization of the worst-case performances. On the contrary, FDC renders possible to directly control for the worst-case performance and trade-off average and worst-case generation quality of the fine-tuned model.\n\n**Question 2** Fine-tuning a generative model corresponds to moving a density to another area in space (or changing its distribution). As illustrated in Fig. 3 (middle row), it can happen that certain regions in space have higher average reward value while not presenting extraordinarily high 'discovery' regions (i.e., reward is high but with low variance). While other areas can have a lower average value, with small areas representing extraordinarily high reward sites (e.g., activity peaks in biological applications). A fine-tuned model that moves its density to the former region will have high average reward, while one fine-tuned to move to the latter area will have lower average reward, but potentially extraordinarily high reward for those few samples, or small amount of density, on the extraordinarily high reward sites. This is the standard setting where risk-seeking (or novelty-seeking) behaviour is necessary to discover the highest reward samples. Effectively, this can require sacrificing the average reward quality for the quality of the best performing samples.\n\n**Question 3** Scenarios of this type can occur and our framework can sharply capture this problem simply by considering a utility $\\mathcal{F}$ defined as an arbitrarily weighted sum of the CVaR and SQ utilities in Table 1. This leads to a general GO problem (see Sec. 3, Fig. 2.b), alike SQ maximization, which can be tackled via Algorithm 1, and enjoying the theoretical guarantees presented in Sec. 5. In particular, the proposed method and theoretical guarantees are not tied to the sample of utilities and divergences presented in Table 1, but is general to any novel objective one wishes to maximize, either by composing known utilities or by novel mathematical expressions, as discussed in Sec. 3.\n\n**References**\n\n[1] Domingo-Enrich et al., Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. ICLR 2025.\n\n[2] Masatoshi Uehara et al., Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review.\n\n[3] Masatoshi Uehara et al., Feedback efficient online fine-tuning of diffusion models. ICML 2024."}, {"review_id": "8BWWGuPYeJ", "reviewer": "Reviewer_P6Lm", "summary": "This paper introduces Flow Density Control (FDC), a method to fine-tune pre-trained flow /diffusion models beyond the standard expected-reward + KL formulation. This is achieved by leveraging recent machinery from Convex  and General Utilities RL.The authors prove convergence guarantees under ideal settings and give a proof outline of general settings.Experiments on 2D toy tasks, text-to-image generation, and molecular design shows that FDC can steer pre-trained models to solve tasks beyond the capabilities of current fine-tuning schemes.", "questions": "1.Gradient noise sensitivity: How robust is FDC when first-variation estimates (e.g., tail-quantile gradients) are noisy or biased?\n2.Ablation on K and runtime: Can you report how performance and wall-clock time scale with the number of FDC iterations?\n3.Add broader baselines:for instance, include a simple unregularized fine-tuning scheme and a policy-gradient method that directly targets tail-risk or high-quantile objectives.\n4.Ablation on practical tasks: for instance,Inverse Imaging tasks and Editing & Conditional Generation tasks.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. Clear derivation of FDC algorithm based on entropy-regularized control/RL, with convergence proofs (Theorems 5.1/5.2).\n2. This framework supports a broad range of utility and divergence measures, addressing a gap in the generative model fine-tuning literature.\n3. Experiments convincingly show risk-averse and novelty-seeking behavior; high-dim tasks show practical value.\n\nWeaknesses:\n1. Proof of Theorems 5.2 is insufficiently detailed and omits many of its required conditions. In particular, the interpolation closeness to the mirror flow is assumed but not quantified; the precompactness of dual iterates is stated (Assumption C.2) but not substantiated with uniform norm bounds; the APT noise conditions are not explicitly invoked; and the regularity of the mirror flow (MF)—required to ensure its gradient-flow structure and apply Sard’s theorem—is not rigorously justified. Given that convergence theory is a central contribution of the paper, these elements should be made explicit and backed by quantitative estimates.\n2. Lack of enough ablation studies. Gradient of first variation estimate accuracy tolerance and hyperparameter sensitivity (e.g. iteraiton size K) are under discussed, raising questions about stability in real-world settings.", "was_revised": true, "final_justification": "The authors’ rebuttal has allayed my initial concerns regarding theoretical rigor and experimental validation. They provided a much more detailed outline of the convergence proof, and the inclusion of runtime and sensitivity experiments strengthens the practical credibility of their approach, demonstrating that the method behaves as expected and that its claims are empirically grounded beyond the original limited discussion. \n\nThere are no major weaknesses in this paper, except for some minor points. So I suggest a weak acceptance recommendation for this paper.", "rebuttal": "We thank the Reviewer for recognizing our work as addressing a gap in the generative model fine-tuning literature, clear derivations of the algorithm, and experiments showing practical value. In the following, we address several points and questions mentioned within the review.\n\n**Proof of Theorem 5.2**\n\nWe agree with the Reviewer that the current exposition does not present certain important details, which we will add for completeness and will revise the proof derivation accordingly. We provide here a sketch of the arguments and will elaborate further in the revised manuscript:\n\n- **Quantitative approximation to the mirror flow:**\n\n  Using stochastic approximation techniques applied to the dual variables, one can derive the following quantitative estimate (see, e.g., [1–3]):\n\n  $$\n  \\sup_{0 \\le s \\le T} \\|| h(t+s) - \\Theta_s(h_t) \\|| \\le C(T) \\bigl[\\Delta(t-1,\\,T+1) + b(T) + \\gamma(T)\\bigr],\n  $$\n\n  where $C(T)$ is a constant depending only on $T$, $\\Delta(t-1, T+1)$ captures the cumulative noise fluctuations, and $b(T)$, \\$\\gamma(T)$ denote the bias and step‐size terms evaluated over the interval. This bounds the deviation of the interpolated process from the deterministic mirror flow.\n\n- **APT approximation:**\n\n  Under the noise and bias conditions stated in Assumption 5.2, it follows from standard results in stochastic approximation (cf. [1–3]) that both the noise fluctuation and bias vanish asymptotically:\n\n  $$\n   \\lim_{T \\to \\infty} \\Delta(t-1,\\,T+1) = \\lim_{T \\to \\infty} b(T) = 0.\n  $$\n\n  This establishes that the interpolated process is an **asymptotic pseudotrajectory (APT)** of the mirror flow.\n\n- **Limit set characterization:**\n\n  Given the precompactness of the dual iterates (which we state as an assumption but cannot presently prove from first principles), Theorem 5.7 in [1] implies that the limit set of the iterates is **internally chain transitive (ICT)** for the mirror flow.\n\n- **Avoiding Sard’s theorem:**\n\n  We note that invoking Sard’s theorem is not necessary in our argument. Since the continuous‐time mirror flow admits $\\mathcal{G}$ as a strict Lyapunov function, Corollary 6.6 (which holds for an arbitrary metric space) in [1] applies and ensures convergence of the APT to the set of stationary points of $\\mathcal{G}$, provided that the set of equilibria is *countable*. We thank the reviewer for the pointer and we will clarify this technical condition in the revision.\n\nWe will integrate these details into the revised version, ensuring that all assumptions are clearly stated and each step in the convergence argument is substantiated with the appropriate estimates and references.\n\n**Ablation on $K$ and runtime**\nIn the following, we consider the experimental setup analogous to the \"Risk-averse reward maximization for better worst-case validity or safety\" experiment in Fig. 3 (top row), and evaluate the effect of different numbers of iterations $K$ on run-time and solution quality, as well as evaluate the effect of a computationally cheaper approximate (entropy-regularized) fine-tuning oracle.\n\n- Dependency on hyper-parameter $K$ for $\\eta = 20$. As one can expect, for small step-sizes $1$ / $\\eta$, the runtime and solution quality improves nearly linearly in $K$ given a fixed number of iterations $N$ of the entropy-regularized solver (see Appendix D.1 for definition of $N$). \n\n|  **K** | **Runtime (s)** | **CVaR estimate (via 1000 samples)** |\n| :----: | :-------------: | :------------------------------: |\n|    0   |       0.00      |        254.48877624511718        |\n|    1   |      44.71      |        241.36963653564453        |\n|    2   |      89.38      |        167.03791809082030        |\n|    3   |      133.31     |        288.72082519531250        |\n|    4   |      176.79     |        271.00385437011720        |\n|    5   |      220.37     |         84.88514938354493        |\n|    6   |      264.06     |         96.55234909057617        |\n\n- As one can expect by interpreting $\\eta$ in Eq. (8) as a learning-rate parameter, by choosing smaller values of $\\eta$ convergence can be achieved with less iterations $K$. In the following, the same setup with $\\eta =10$.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 44.23           | 249.0867721557617                |\n| 2    | 87.78           | 90.0                             |\n\nThe above two tables hint at the fact that the FDC fine-tuning process can be interpreted experimentally as classic (convex or non-convex) optimization, although on the space of generative models, with learning rate (or step-size) controlled by $\\eta$.\n\n- Furthermore, using an approximate entropy-regularized control solver oracle (i.e., performing approximate entropy-regularized fine-tuning at each iteration of FDC) can also lead to optimality via increasing the number of iterations $K$. In the following we consider $N=100$ instead of $N=1000$ (as in previous experiments) and use $K=5$ showing that FDC can retrieve the same final fine-tuned model as in the previous table using only one tenth gradient steps (i.e. $N=100$ instead of $N=1000$) for the inner oracle.\n\n| **K** | **Runtime (s)** | **CVaR estimate (1000 samples)** |\n|:----:|:---------------:|:---------------------------------:|\n| 0    | 0.00            | 254.48877624511718               |\n| 1    | 5.11            | 221.07437133789062               |\n| 2    | 10.01           | 194.68543167114257               |\n| 3    | 14.73           | 90.01436312185169                |\n| 4    | 19.60           | 91.10320648061439                |\n| 5    | 24.60           | 90.0                             |\n\n\n**Gradient noise sensitivity**\nIn the following, we consider an experimental setup analogous to \"Risk-averse reward maximization for better worst-case validity or safety\" experiment in Fig. 3 (top row), and evaluate the effect on solution quality of different numbers of samples for Monte Carlo estimation of the sample-based gradient of first variation (as presented in Sec. 4). Notice that $90.0$ corresponds to the optimal CVaR value in this experiment. \n\n| **Samples for MC estimate** | **CVaR estimate** |\n|:---------------------------:|:-----------------:|\n| 8000                        | 90.0              |\n| 4000                        | 90.0              |\n| 2000                        | 90.0              |\n| 1000                        | 124.0             |\n| 500                         | 90.0              |\n| 250                         | 212.2             |\n| 125                         | 90.0              |\n| 62                          | 90.0              |\n| 31                          | 154.9             |\n\nOne can notice that in this experimental setup the method still converges to good solutions even for significantly lower amounts of samples than the one used for the experiments in the paper (i.e., $8000$). Further sample-complexity evaluations in real-world domains are significantly task and reward function dependent.\n\n\n**Further baselines and practical tasks**\nWe believe that the baseline method we considered within our experimental evaluation, namely Adjoint Matching [4] is actually regarded by the community as the strongest, as well as widely adopted, available method to target the task at hand. More specifically, given a known reward function one whishes to maximize by fine-tuning a diffusion or flow model, there is significant literature regarding how to solve this problem (i.e. sample from the optimal reward-tilted distribution) via RL or Control theoretic techniques (e.g., see an overview in [5]). Early methods required to solve the problem approximately or with hidden value function bias (e.g. [6]) as discussed in [Section 4.2, 4], while recently Adjoint Matching [4] rendered possible to tackle the problem exactly and with the strongest empirical performances (see e.g. [4]) to our knowledge. Therefore, we regard it as a strong and expressive baseline, widely adopted, and already known to dominate alternative methods. If the Reviewer has in mind alternative specific baselines that are competitive with the one proposed we would be happy to hear and further discuss this point.\n\n**Inverse Imaging and Editing and Conditional Generation tasks**\nIn this work, we tackle the problem of generative optimization. It is not particularly clear to us how the tasks mentioned by the Reviewer fit into the proposed generative optimization framework (in Sec. 3). Although interesting, we believe that exploring the connections of our framework with such problems and further evaluating the proposed method in other tasks goes beyond the scope of this work, which rather aims to introduce  novel algorithmic machinery beyond classic RL/Control schemes for generative optimization via fine-tuning.\n\n**References**\n\n[1] Benaim, Dynamics of Stochastic Approximation Algorithms, 1999\n\n[2] Mertikopoulos et al., A Unified Stochastic Approximation\nFramework for Learning in Games, 2024\n\n[3] Hsieh et al., The Limits of Min-Max Optimization Algorithms: Convergence to Spurious Non-Critical Sets, 2021\n\n[4] Domingo-Enrich et al., Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. ICLR 2025.\n\n[5] Masatoshi Uehara et al., Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review.\n\n[6] Masatoshi Uehara et al., Feedback efficient online fine-tuning of diffusion models. ICML 2024."}]

Author rebuttal:

Meta-review: This paper considers the problem to fine-tune pre-trained flow models. The authors extended the mostly studied KL-regularized reward fine-tuning setting to a more general reward and regularization with theoretical analysis. The proposed approach was applied on a variety of settings and problems including toy problem, text-to-image generation and molecular design.

All the reviewers vote to accept this paper. Initially, the reviewers mostly questioned about the theoretical rigor and experiment section and the authors included more analysis and experimental results that convinced the reviewers. 

In particular, I find the generalization considered in this paper nice and the great effort showed in designing experiments/applications. I vote to accept this paper as a spotlight presentation because it is worth for more audience to consider this framework from the list of applications and reward/regularization design. I would also encourage the authors to consider more application settings to expand the broader impact of the work.
Decision: Accept (spotlight)