Training Free Guided Flow-Matching with Optimal Control
Authors: Luran Wang
Chaoran Cheng
Yizhen Liao
Yanru Qu
Ge Liu
Venue: ICLR 2025 Poster (2025)
Keywords: flow matching
controlled generation
inverse problem
Primary area: generative models
TL;DR: We develope a general and theoretically grounded framework for guided flow matching on both Euclidean and SO3 geometries taking inpiration from optimal control.

Abstract: Controlled generation with pre-trained Diffusion and Flow Matching models has vast applications. One strategy for guiding ODE-based generative models is through optimizing a target loss $R(x_1)$ while staying close to the prior distribution. Along this line, some recent work showed the effectiveness of guiding flow model by differentiating through its ODE sampling process. Despite the superior performance, the theoretical understanding of this line of methods is still preliminary, leaving space for algorithm improvement. Moreover, existing methods predominately focus on Euclidean data manifold, and there is a compelling need for guided flow methods on complex geometries such as SO(3), which prevails in high-stake scientific applications like protein design. We present OC-Flow, a general and theoretically grounded training-free framework for guided flow matching using optimal control. Building upon advances in optimal control theory, we develop effective and practical algorithms for solving optimal control in guided ODE-based generation and provide a systematic theoretical analysis of the convergence guarantee in both Euclidean and SO(3). We show that existing backprop-through-ODE methods can be interpreted as special cases of Euclidean OC-Flow. OC-Flow achieved superior performance in extensive experiments on text-guided image manipulation, conditional molecule generation, and all-atom peptide design.

Published as a conference paper at ICLR 2025
TRAINING FREE GUIDED FLOW MATCHING WITH OP- TIMAL CONTROL
Luran Wang', Chaoran Cheng", Yizhen Liao,, Yanru Qu Ge liu?
1 University of Cambridge
University of Illinois Urbana-Champaign
Tsinghua University
1w7030cam ac. uk, chaoran7@ilinois.edd
1iaoyz0711@gman com, yanrugu2@1limois.cm geliuêin 11 www.iniiiiii.om..omm....... edu
ABSTRACT
Controlled generation with pre-trained Diffusion and Flow Matching models has vast applications. One strategy for guiding ODE-based generative models is through optimizing a target reward R(x1) while staying close to the prior distri- bution. Along this line, some recent work showed the effectiveness flow model by differentiating through its ODE sampling process. Despite the superior performance, the theoretical understanding of this line of methods is still preliminary, leaving space for algorithm improvement. Moreover, existing methods predominately focus on Euclidean data manifold, and there is a compelling need for guided flow methods on complex geometries such as SO(3), which prevails in high- stake scientific applications like protein design. We present OC-Flow, general and theoretically grounded training-free free framework for guided flow matching using optimal control. Building upon advances in optimal control theory. we develop effective and practical algorithms for solving optimal control in guided ODE-based generation and provide systematic theoretical analysis of the convergence guaran- tee in both Euclidean and SO(3). We show that existing backprop-through-OOD methods be interpreted as special cases of Euclidean OC-Flow. OC-Flow achieved superior performance extensive experiments on text-guided image manipulation, conditional molecule generation, and all-atom peptide design.
guiding
INTRODUCTION AND RELATED WORK
SDE and ODE- generative models such as diffusion and continuous normalizing flow (CNF)
have exhibited excellent performance various domains such as images (Ho et al.) 2020
Esser
al. 2024), audio (Zhang et al.) 2023 Défossez et al. 2022). and discrete data (Lou et al. 2023
Cheng et al.) 2024). Particularly, the simplicity of Riemannian Flow Matching SO(3) manifold
(Chen Lipman 2023) empowers de novo generation of small molecules (Song et al. 2024 Xu
let | 2022) and proteins (Yim al. 2023 Bose et al. 2023 Li et al. 2024), leading enormous
advancement in biomedicine. Controlled generation from pre-trained diffusion and flow matching
priors has gained growing interest in numerous fields, as it encompasses wide range of practical
tasks including constrained generation (Giannone et al.] 2023). solving inverse problems (Liu et al.
2023
Ben-Hamu al. 2024), and instruction alignment (Black al. 2023 Esser al. 2024).
There are several lines of work for guiding diffusion and flow models. Classifier-free guidance (CFG)
(Ho & Salimans, 2022) trains conditional generative models that take conditions input. Reward fine-tuning approaches update the generative model parameters to align with certain target objective functions (Black et al., 2023). Both methods require specialized training routines, which are costly and not extendable to new tasks. Training-free guidance on diffusion (Kawar et al.) 2022 Chung ctt al. 2024, Song et al. 2023) alters the scores in the SDE generation process with the gradients from the target function to achieve posterior sampling. These guidances often rely strong assumptions of the denoising process and require estimating target function gradients w.r.t noised samples which often intractable. Accurate posterior sampling is only guaranteed for limited family of objective functions such as linear. Therefore, efforts that deploy such guidance to flow models by bridging the ODE path and SDE path share similar constraints (Pokle al. 2023 Yim et al. 2024).

---

Published as a conference paper at ICLR 2025
Notably, two recent works showed the effectivvness of guiding pre-trained flow models by differ- entiating through the ODE sampling process, outperforming popular guidance- based approaches. Particularly, BBmm- et al. 2024) differentiates loss R(x) through the forward-ODE w.r.t the ini- tial noise 2o. which induces implicit regularization by projecting the gradient onto the "data manifold" under Gaussian path assumption. This strong confinement the prior might hinder optimization when the target reward function diverges from the prior distribution. Liu et al. (2023) formulates an optimal control problem where control term Ut at each timestep solved to guide the ODE trajectory. However, the gradient decomposition trick used in Liu et al. 2023) ignores the running cost of control terms and thus could lead to suboptimal behavior. Despite the good performance, there is a lack of systematic theoretical analysis on the convergence behavior and explicit regularization of the lifferentiate-throuug existing works predominantly focus on the Euclidean manifold due to its simplicity, and there is compelling need for a thcoretically grounded guided flow matching framework on more complex geometries such SO(3) which is heavily used in scientific applications.
approaches to better guide algorithm design in this space. Furthermore,
To fill the gap between the practical applications of guided generation and theoretical grounds, we propose OC-Flow, a general, practical, and theoretically grounded framework for training free guided flow matching under optimal control formulation. Our key contributions are as follows:
1. We formulate "controlled generation with pre-trained ODE-based priors" as optimal control problem with a control ferm uy and a running cost that regulates the trajectory proximity to the prior model while optimizing for target loss. Building upon advances in optimal control theory, we develop effective optimization algorithms for both Euclidean and SO(3) space through iterative updates of co-state flow and control term, with theoretical guarantees under continuous-time formulation.
2. In Euclidean space, we show that running cost bounds the KI. divergence between prior and OC-Flow-ing Floueinddcedd joint distribution. We develop a simple algorithm for OC-Flow through iterative gradient update and provide convergence analysis. We further demonstrate that Dflow and Flow-grad can be interpreted as special cases of Euclidean OC-Flow, providing unified view of the problem.
3. We present one of the first guided flow-matching algorithm on the SO(3) manifold with theoretical grounds. Our approach extends the Extended Method of Successive Approximations (E-MSA) to SO(3) with rigorous convergence analysis. Additionally, we propose approximation techniques to enable computationally efficient OC-Flow on SO(3).
4. We provide an efficient and practical implementation of OC-Flow, by introducing the vector- Jacobian product and asynchronous update scheme. We show the effectiveness of our method with extensive empirical experiments, including text-guided image manipulation, controllable generation of small molecules on QM9, and energy optimization of flow-based all-atom peptide design.
2 PRELIMINARIES AND PERSPECTIVES ABOUT FLOW MATCHING
Euclidean Flow Matching. Flow matching (Lipman et al. 2022 Liu et al. 2022) provides an efficient framework for training a generative model by approximating the time-dependent vector field associated with the flow represented as v/y [0, 1] Rdd + Rd. This vector field u, [0, 1] x Rd - Rd defines a probability path of the evolution of an initial noise distribution, denoted by Po. towards target distribution, Pl. with the pushforward probability p. == (v1)=Po. The dynamics of the vector field that governs this flow can be described by the flow ordinary differential equation (ODE) follow the convention
of the form dt- u, (x.), where to time t and the state at time t is given by 2, := (zx)- Lipman et al. (2022) demonstrates that a tractable flow matching objective can be obtained by conditioning on the target data 21 The primary goal of conditional flow matching is to train a model, ff : [0, 1] x IR + Rd, such that it minimizes the difference between its output and the ground truth conditional vector field as LCFM = =x,p( f2 - (1 I1, 2o)l The trained model fr can be employed as the marginal vector field during the inference phase. In this context, once noise sample zo is drawn, the system's is evolution can be described by the following differential equation:
use Newton's 's notation with respect
d S (xt),
Zo (x)).
Rotation Group SO(3). The formulation of flow matching can naturally extend to Riemannian manifolds (Chen & Lipman| (023). On the Riemannian manifold, a flow is defined as a time-
dependent diffeomorphisms pp: G G, which describe the continuous evolution of points on the
2

---

Published as a conference paper at ICLR 2025
RRR
J-R
state flow
J: aR
0
costate flow
8f
(
Neya
of
for
0 Axg
01
X
Orip
Ote
X1o
E(/N
E2
F//N
b//
x Kitr
D-Flow
X Kittter
FlowGrad
K.et
OC-Flow(ours) in SOI31
Figure 1: Comparison of hhapopppaaaiioottouuooooooo algorithms. For D-Flow (left) and FlowGrad (middle), the black curves represent the state trajectory at the k-th iteration, while the red curves show the updated trajectories at the (k 1)-th iteration, using gradient updates (blue dashed arrows). D- Flow updates the state at to only, while FlowGrad propagates updates across all time steps. OC-Flow (right) incorporates the running cost and wewght-weighting factor into the terminal reward. The co-state flow le (green curves) combines gradient information and system dynamics to iteratively update the control terms (0z), which in turn updates the states.
manifold over time, generated by a vector field V. with V(p) E T,G for each p G where Tp is the tangent space point p E G. The flow evolves according to: de V(x:) L,,, V(e). The CFM objective in Equation D can also be adapted for Riemannian flow matching, in which the ground truth vector field can be calculated as the time derivative using the exponential and logarithm maps. Details about rotation group SO(3) can be found in Appendix A In this work, we focus specifically on the (3(3), the Riemannian manifold of all 3D rotations equipped with the canonical Frobenius inner product. In previous flow-based protein design models, each amino acid is associated with rotation that defines its orientation (Yim et al.) 2023] Bose et al. [023 Li al., 2024). Guiding such pre-trained generative models toward the desirable protein properties can potentially have profound impact on the pharmaceutical industry.
OPTIMAL CONTROL FRAMEWORK FOR GUIDED FLOW MATCHING
The key challenge in guided generation is balancing optimization and faithfulness to the prior distribution. To address such a need, we propose the following framework. Given pre-trained flow model, ff (x), parameterized by neural network, our goal is to determine the optimal control terms O, that maximize the reward Rxx) while maintaining the proximity of the resulting vector field to the original vector field induced by [())). The reward can be customized for diverse tasks such as inverse problem = [z) - yll conditional generation R = (f(x) - c)². and constrained incorporate penalty the state trajectory or To ensure proximity,
generation R = || y/1 control terms 5 L(Oc). also known as the running cost. Optionally, one may also introduce a metric d(-,-) ) to penalize the deviation between the new terminal state = and the prior terminal state xf. The modified terminal reward function is then defined as: ((xf) = RR(x)) - d(xf. x.²). and scaling the terminal reward by constant a. we formulate the problem as standard optimal control task:
+ ((Oe) dt subject to
1=) (x(( +
i = (e 0z).
A fundamental result in optimal control theory is Pontryagin's Maximum Principle (PMP) (Pon tryagin (201)). ) which provides the necessary conditions for optimal solutions in control problems. Specifically, at the core of PMP is the introduction of the Hamiltonian function, H. This Hamiltonian defined terms of the state of the system, the control, and new entity called the co-state μ4, which resides in the cotangent space of the state manifold:
H x. f) = (H, hy (2) 0.) L()).
The co-state μ4, also known the adjoint variable, encodes the influence of the terminal cost function. Their evolution captures how the sensitivity of the system impacts the cost function, ensuring that the state variables evolve in accordance with the system dynamics. Consequently, in optimal control, the Hamiltonian must be maximized by jointly evolving the states and costates according to system of coupled differential equations. The details of PMP conditions be found in Appendix B.
3

---

Published as a conference paper at ICLR 2025
3.1 OC-FLOW ON EUCLIDEAN MANIFOLD
We first develop the algorithm and theoretical analysis for OC-flow in Euclidean space. One of the simplest choices for the control term is an additive control (Kobilarov & Marsden (201T)). which directly perturbs the prior trajectory. In fact, with the linear expansion, the additive control could be seen as general case and widely used in optimal control. Specifically, with 6, representing the control input, the new state dynamics and the corresponding running cost can be defined as:
1 - 10.11.
+ = (. 0.) = ff() + 0. I(O.)= =
The running cost effectively acts constraint on the trajectory 2
encourage proximity to the original prior distribution. To better understand the effect of running costs on the guided distribution, we provide the following proposition to formally prove that it can control deviation from the prior distribution measured by KL-divergence.
Proposition 1. For Affine Gaussian Probability Path, the expectation of the running cost upper bounds the KL divergence between the joiot distribution p1 P1(x", , 21) = Pi (xx" (1 Paaaa(r) and
the joint distribution after guidance p1 pi(x", .11 = P1 (x" (x1) Paata (21), with 21 Pdata- x² induced
by prior conditional vector field ui (a) (1) and x² sampled by applying control 0g (x(21) on wh:
[2 (21)) dt 2 CK ( (1")))((?,,))))
Bui-Paa (m1)
(5)
Furthermore, for square-shapet data x with non-zero probability path, the expectation of the running cost, combined with the L1 disttnne between the prior sample xt and the corresponding guided sample dñ, upper bounds the KL. divergence between the marginal distributions of the prior model ph: and the guided model pf:
Esp-pp(ar) Al|=Y - x111 + B L. [10c(z))1 dt 2 KL
P.) P
(6)
One effective approach for directly apply- ing PMP to optimal control tasks is the Ex-
Algorithm 1 OC-Flow Euclidean Space
1: Given: Pre-trained model: f", initial state: mo
tended Method of Successive Approximations
2: Initialize: Control terms 60, learning rate
(E-MSA) (Li et al. 2018). E-MSA builds
weight decay 3
upon the basic MSA algorithm (Chernousko & Lyubushin 1982), ). which iteratively updates the terms in the PMP conditions (Appendix [B.2). The primary enhancement of E-MSA over the basic MSA is its ability to extend convergence guarantees beyond limited class of linear quadratic regulators (Aleksandrov) (1968).
3: for k 0 to MaxIterations do
4: Solve for the state trajectory:
X X
Xxº + (P'() 0*) At
5:
Update control:
0% = 30/ + ((++)
6: end for
A key assumption is the global Lipschitz condition for the functions involved. However, note that this assumption can be relaxed to a local Lipschitz condition if we can demonstrate that wf is bounded, which can be safely assumed provided that appropriate regularization techniques are applied. Furthermore several prior work has shown the Lipschitz continuity for the deep learning models. (Gouk et al. (221 Khromov & Singh| (2024))
When the E-MSA algorithm is applied to the guided controlled generation task on Euclidean space, the trajectory of the co-states μ4 can be calculated in closed form. Specifically, they can be expressed as V., x, Φb (x²)) which enables us to derive the following update rule with convergence guarantees:
Theorem Assume that the reward function, the prior model, and their derivatives satisfy Lipschitz continuity, bounded by a Lipschitz constant L. Utilizing the E-MSA, for each iteration k. for each constant 7 > 2C with is a function of L,, such that under the addictive control and the running cost defined in Equation 4 the optimal update is following:
7 0¹
a
(7)
1+7
1 1+7
This update rule for the control term O2 guarantees an increase in the functioon defined in Equation 2
J(Oh+1)-10 - -(0 2
21-20
20
k
2.0 0.
(8)
4

---

Published as a conference paper at ICLR 2025
In practice, solving continuous ODEs requires discretization. The discretized version of the proposed algorithm is outlined in Algorithm 1. In this formulation, the weight decay term is parameterized as B -े.. and the learning rate is defined as 77 - As demonstrated in Appendix C.3 the discretization error introduced by the Euler method is of the order O(ât). This error diminishes as the number of ODE steps increases, ensuring the algorithm converges to the global optimum.
3.2 PRACTICAL IMPLEMENTATION AND ACCELERATION
3.2.1 ADJOINT METHOD AND VECTOR-JACOBIAN PRODUCT
significant portion of the computational time in the OC-Flow algorithm spent on evaluating the gradient v., (2)). The computational cost of directly back-propagating through ODE with vanilla Autograd requires saving all intermediate computation values, which results in demanding memory complexity of O(N (Pan et al. 2023; Chen et al. 2018) where N is the ODE steps. We instead employ the adjoint method where the gradient V (xt) is computed using a vector-Jacobian product by the Kouble-backwarde trick, which reduces the memory cost (
V VENIN &
(VzuFD/N D)
& PEuTD/N (IK/N).
3.2.2 ASYNCHRONOUS SETTING FOR FLEXIBLE UPDATE SCHEDULING
In practice, discretization techniques are employed to simulate the ODEs governing both the state trajectory and the co-states fer and operate in a synchromous setting, where the number of time steps for the state trajectory 2e coincides with the number of control terms Or.
Here we show that OC-Flow can be extended to asynchronous framework, providing greater flexibility in scheduling. We subdivide the time interval Át into N equally spaced subintervals. Let { denote the state trajectory over the time interval it, + At). and {=/} represent the trajectory when the control term Og is applied in the first subinterval. The trajectory can be approximated as:
N-1
N-1
I++A = Iz + N
At Σ
fd x 14
LAE
+
At 0 xt N
At
N
Σ f (I++
+N=+ N (10)
Consequently, the asynchronous setting allows the algorithm to be applied without modification while enabling finer updates to both the control terms and state trajectories by adjusting the frequency N of control term updates relative to the state trajectory simulation (see Appendix C.4 the proof and justification of the approximations in Equation 10 In our peptide design experiment, the asynchronous setting is applied for efficient computing.
Table 1: Comparison of runtime and memory complexity of different methods used in backprop- through guided-ODE in Euclidean and SO(3) manifold. For complexity, N the number of ODE steps, n is the number of effective control terms with synchronized and in the range [1, .V] and D is the complexity of computing 1-step gradient (VJP or Autograd), D depends on data and model size.
Number of Control Terms
Running Cost
Memory Complexity
Runtime
Convergence to Optimal
Generalization to SO(3)
Complexity
OC-Flow Flow
f 1
||09||
O
O(nD2)
0
O
O
X
D-Flow
Implicit
O(ND2)
O(ND*)
X
3.3 CONNECTION TO OTHER BACKPROP-TIIROUGE GUIDED-ODE APPROACHES
Several previous works discussed backprop-throuug--ODE guiggggnnne-ggguuueeeeeeeee in flow-matt now-matching models. Notable examples include D-Flow (Ben-Hami Hamu et al. 2024) and FlowGrad (Liu et al. (023). An illustration of their algorithms and ours shown Figure In this section, we demonstrate that our framework is more general, and both of these methods can be viewed as special cases of our algorithm.
FlowGrad formulates the optimization task manner similar to our optimal control target in Equation 2 Specifically, it directly applies gradient descent to the control variables:
0++1 o% V (xt) of
N (Va, X.-A:)
(x
(11)
which can be interpreted as a limiting case of our algorithm in Equation 7 where y + oo and given with di 0, V., X At T. However, as shown Equation 8 the convergence rate is a complex
5

---

Published as a conference paper at ICLR 2025
function of 7, so in practice, is treated tunable parameter. FlowGrad's setting y + may undermine the convergence speed. D-Flow optimizes the reward by applying gradient descent to the initial noise 2n:
X X LLFGS (xx)))
(12)
In fact, with the update rule 21+dr To + f(ro) dt + 0g dt, the update of fu can be seen as an increment to 2g. The LBFGS algorithm provides dynamic learning rate, aligning with our framework, where 7 is allowed to vary across iterations. Hence, D-Flow can be viewed as a special case of our asynchronous algorithm when the number of control terms 1. more detailed comparison of the algorithms can be found in Table 6 and additional discussion computation efficiency is in Appendix D
OPTIMAL CONTROL FRAMEWORK FOR GUIDED FLOW MATCHING ON SO(3)
Most optimization algorithms are primarily designed for Euclidean spaces and face significant challenges when applied to an--uulidean settings, such the SO(3) manifold, which plays crucial role in drug discovery and peptide design (Huguet et al. (2024)). This section extends the E-MSA algorithm to the SO(3) manifold and presents rigorous proof of its convergence.
4.1 OC-FLOW FOR SO3
To begin, we define the vector field governing the system dynamics. The state trajectory, influenced by control terms 0 E so(3), evolves according to the following differential equation:
# = = (f( + 0.).
(13)
In this work, the left-invariant vector field is utilized, under which the Hamiltonian can be shown to reduce to a linear functional in so))) (Jurdjevic (996). Colombo & Dimarogonas (2020). Given the co-state j E so(3)", the Hamiltonian function redefined as:
H [0, T| x SO(3) x so(3)* 5o( + R. (t,z, + pe (f((z)) 0.) L(0).
(14)
direct approach to apply PMP conditions on the SO(3) manifold involves iteratively updating the cotangent vector (4, and the state trajectory I, shown in Step 4 and Step 5 in Algorithm 2 and then apply an update rule to determine the control term 02 for the subsequent iteration with weight decay and learning rate the update for 0g can be written as:
0 30 nit
(15)
where pi is defined by (78* -4 (v) with jux so(3) for all s3(3). The existence of pt can be derived from the Riesz Representation Theorem (Goodric) (177)). This formulation leads to the introduction of the OC-FLow algorithm on SO(3).a jas detailed in Algorithm 2.
Algorithm 2 OC-Flow on S3(3)
1: Given: Pre-trained model: ff, initial state: 20
2: Initialize: Control term 0 E so(3,, learning rate weight decay |
3: for k = MaxIterations do
4: Solve for the state trajectory: X- = x (IP(I, X²") + 0z)
x8 = xo
5: Solve for the adjoint variables: jig" -ad n (db.g)
1 )
OH do
(dL,g)) (((()))
6: Update control: 05 60, + natt
7:
end for
4.2 CONVERGENCE OF OC-F ON SO3
To derive the proof of the convergence of our Algorithm 2, first establish that under the PMP conditions on S3(3), the objective function (()), as defined in Equation 2 can bounded. This is formalized in the following proposition:

---

Published as a conference paper at ICLR 2025
Assume that the reward function, the prior model, and their derivatives satisfy Lipschitz continuity. bounded by a Lipschitz constant L. Then, there exists a constant >> 0 such that for any 0, & E so(3). the following inequality holds:
Proposition 3:
((0) + + (t. Cl 0.1l J(ő).
(16)
where X and pe satisfy the PMP conditions on SO(3) manifold, and AHl denotes the change in the Hamiltonian, defined as:
(t == H 12. H1,,, (1, 22 14g 0z)
(17)
Proposition 2 provides lower bound for the difference in the objective function values under two distinct control terms that satisfy the PMP conditions described by the Hamiltonian equations in Appendix B.1
However, applying this result directly as an optimization algorithm presents several challenges. First, the difference in the Hamiltonian (o (t) not inherently bounded. Second, the term || 0||
non-negative, which complicates the minimization process. To address these issues, inspired by the method of E-MSA, we introduce positive constant y and define an Extended Hamiltonian:
(1, 25, 0, = H1,, x, ,4 0) 11000 oll = fx(x) + -
- 1010 2110 -1.. (18)
The introduction of the extended Hamiltonian enables the combination with the penalty term | 0 into unified expression that can be optimized jointly. A natural approach to achieve this is by updating to maximize the Extended- Extended-Hammitonian The resulting update rule is given by:
original Hamiltonian
gt+1 arg max H(t, 8* ,,",00000
7 8 0,,00) 1+7 1+ 1 =38
mpp
(19)
By performing this maximization step at each iteration, we ensure that the change in the Extended- Hamiltonian, AH, non-negative, indicating that the algorithm progresses towards an optimal solution. Furthermore, we can show that when the update process converges, i.e., when AH = 0 or equivalently AH formalized in the following proposition:
the algorithm has reached the optimal control point. These insights can be
Proposition 4: Let X and pd satisfy the PMP conditions. If the update rule follows Algorithm we define E = fo Aget1 , ((t) dl, and ex is bounded as:
Ck = L.¹ Agrts o² H(t dt,
lim C 0. #+000
(20)
Furthermore, when = 0. we have = 0" == arg max, ((0)
With these results, we can now extend result in E-MSA to the SO(3) manifold and establish bound for the optimization algorithm based on the derived theoretical properties:
Theorem Assume that the reward function, the prior model, and their derivatives satisfy Lipschitz
continuity, bounded by Lipschitz constant L. Let 0 E so(3) be any initial measurable control with J(0") < too. Suppose also that info)))()) J(0) > -c. If the update of satisfies equation for sufficiently large % the following inequality holds for some constant D O:
Deg J(0 J(0R).
(21)
Therefore, by invoking Proposition 3, we can conclude that after each update, the target function is non-decreasing and when the update process terminates, the optimal solution has been attained. This establishes the convergence of the OC-Flow algorithm on the SO(3) manifold.
4.3 PRACTICAL IMPLEMENTATION
In practice, directly optimizing Algorithm 2 using existing ODE methods is challenging due to the nature of the adjoint variable fu, which is a linear functional the dual space so(3)* Instead, we can optimize ü. as defined Section 4.1 We can decompose 14, into its projections onto set of orthogonal bases within the s3) group.
7

---

Published as a conference paper at ICLR 2025
sad
angry
curly
reference
StyleCLIP
FlowCrrad
OC-Flow (Owss)
Figure 2: Visualization of text-guided generated faces with different expressions.
A frequently used choice for the basis in so is the canonical basis {Ei, E2, Ea} (McCann al. (2023) satisfying the condition (v, E.) = 2 for all E so( Thus we can decompose the time derivative of the adjoint variable üi as: a. (i, E,) and jr = D. Thus, with the closed forms for the partial derivatives related to Hamiltonian, the practical update for j4: in Algorithm can af.
be written as follows. The vector-Jacobian method can be applied to compute the term 8e (1) Fig) in Algorithm 2, which significantly reduces complexity from O(D*) to O(D²). Meanwhile, the method of asynchronous can also be applied. See Appendix
C.4
OsL Ft
OST jun
(12 OlL Of 1t Ox
cr. = (ith UFF + O2, Ea])
Pt:
Es
x Ej
TT
(a Eg)
dx
At
Pi-ar
c 1ig
E,
PT (v4 24
(22)
5 EXPERIMENTS
5.1 TEXT-GUIDED IMAGE MANIPULATION
We first validate OC-Flow the traditional text-to-image generation task. Previous work has demonstrated the importance of alignment with the given text prompt using either automatic metrics or human preference the reward (Black al. 2023 Esser et al. 2024). In our text-guided image manipulation task, we want to guide the generative model pre-trained on the celebrity face dataset CelebA-HH HQ (Krras 2017) to text guidance (sad, angry, curly hair) showing different facial expressions or traits. Following the same setup in Liu et al. (2023). given an input image gg. the reward for alignment with the text prompt can be effectively evaluated by the CLIP model (Radford) itt al., 2021) pre-trained in a contrastive way to score the similarity between arbitrary image-text pairs. Following (Liu et al. 2023). we adopt the pre-trained Rectified Flow (RF) (Liu al., 2022) as the generative prior. Inspired by Proposition 1, for this image task where the square- like assumption is satisfied, an extra terminal constraint of x is added as part of the terminal reward function.
We choose two fftte-of-the-art text-guided image manipulation baselines, StyleCLIP (Patashnik et al. 2021) and FlowGrad (Liu et al. (023)) We run Style- CLIP and FlowGrad with their official implementa- tion and default parameter configurations. For ours, we set time step of 100, step size 7) = 2.5, weight decay of 0.995, the weight of the extra constraint of 0.4, and the number of optimization steps of 15. For qualitative comparison, we show generated examples of different text-guided expressions in Figure Due to the large gap between reference and guided dis- tributions, StyleCLIP fails to manipulate with sad. Lacking in regularization, FlowGrad may change the content too much with curly hair, Our OC-Flow generally produces the best results with a good alignment with the text prompt while preserving the generative prior so as to produce reasonable faces that are not distorted much.
Table 2: Comparison of methods on LPIPS, ID, and CLIP metrics. Lower LPIPS and ID indicate better performance, while higher ID and CLIP values are preferred.
Method
LPIPS
ID 1
CLIP
CG + RF CG + LDM
0.346
0.643
0.292
0,383
0.513
0.298
DiffusionCLIP StyleCLIP + e4c
0.398
0.659
0.285
0.359
0.704
0.267
FlowGrad RF OC-Flow (Ours) 0.207
0.302
0.737
0.299
0.732
0.302
5.2 MOLECULE GENERATION FOR QM9
We further instantiate our OC-Flow for controllable molecule generation the QM9 dataset (Rud) digkeit et al., 2012, Ramakrishnan et al. 2014), a commonly used molecular dataset containing small molecules with up 9 heavy atoms from C, O, N, F. Following [Hoogehoom et al. (202);: Ben-Hamu
8

---

Published as a conference paper at ICLR 2025
0.114 1,575
2,438
3,354
4.550
5,536
Figure 3: Visualization of OC-Flow generated molecules with various dipole moments condition.
ct al. (2024), we target for conditional generation of molecules with specified quantum chemical property values including polarizability a, orbital energies éHOMO, éLUMO and their gap As, dipole moment μ,, and heat capacity Cp. Such a conditional generation setting of molecules with desired properties has profound practical applications in drug design and virtual screening. To define the loss function, separate classifiers for each property are first trained to predict the property value for the generated molecule (Hoogeboom mean absolute error (MAE) between the predicted and the reference property values. The pre-trained unconstrained generative model is taken from Song al. (2024) (EquiFM), flow-based generative model that uses an equivariant vector field parameterization for generating the atom coordinates and types via the learned flow dynamics. To demonstrate the zero-shot guidance performance on such a conditional generation task, we compare our approach with other gradient- gradient-based methods of D-Flow (Ben-Hamu et al. 2024) and FlowGrad (Liu et al., 2023) on the same pre-trained EquiFM. To be comparable to D-Flow, we follow its setting to use the L-BFGS optimizer with 5 optimization steps with linear search. We generate 1000 molecules for each property and report the MAE in Table The unconditional EquiFM is also included as upper bound for the guided models. It can be that our approach consistently outperforms both them with lower MAEs, which better balances the reward optimization and the faithfulness to the prior. We provide guided generation samples in Figure 3 with respect to different target dipole moments. A clear trend from hydrocarbons with more symmetrie structures to molecules with more high-electroneegatiitt atoms of oxygen and nitrogen can be observed, indicating an increase in the dipole moment.
al. 2022), and the loss be then calculated as the
As have theoretically demonstrated the impact of the regularization strength from the optimal control perspective. we further experiment with a different y and examine the quality of the conditionally generated molecules by evaluating additional met- rics following Song et al.] (2024). Specif- ically, we calculate the atom stability per- centage (ASP), molecule stability percent- age (MSP), and valid & unique percentage (VUP). Ideally, these metrics should not be greatly lower than the pre-trained model, and higher strength of regularization should lead higher scores. Indeed, demonstrated in Table in which we provide these scores for two different settings of 7 = 0.01 and 10, all scores are higher with higher strength of regularization effectively prevents exploitation from direct gradient descent that may hack the loss function and provides more flexible and fine-grate mne-graieed control over the guided generation.
Table 3: MAE for guided generations on QM9 (lower is better).
Property Unit
As
SHOMO meV
$LUMM meV
Bohr
meV
D
Kindl
(O-Frs)(Our))
1_383
367
183
342
0.314
0.819
D-Flow
1.566
355
205
346
0.330
0.893
FlowGrad
2.484
517
273
429
0,542
1,270
EquiFM
8.969
1439
622
1438
1.593
6.873
Classifier
0.095
64
40
35
0.046
0.041
the cost of also higher MAE. In this way, our OC-Flow
Table 4: MAE and other evaluation metrics for our approach with = 0.01 / 7 10.
Property
Ae
ÊHомO
ÉLUMO
Co
MAE 4
1.383 1.557
1.557
367 / 365
183 188
342 / 339
0.314/0.33
0.819 0.852
ASP 1
94.8 96.0
95.2 96.1
95.2 96.1
95.3 96.1
95.8 96.1
95.2/96.1
MSP T
64.4/ 69.9
67.9/70.5
65.8/ 69.8
68.5 / 70.8
68.0/
67.1 68.9
VUP T
86.2/ 88.6
88.2 89.8
86.2/ 87.7
87.6/ 88.7
88.2 89.0
89.4/88.5
5.3 PEPTIDE DESIGN
We evaluate our OC-Flow approach for peptide backbone design using a test set derived from (Li al. 2024), which includes 162 complexes clustered based on 40% sequence identity using mmseqs2 (Steinegger & Söding 2017). Our experiments focus on PepFlow w/Bb, a model designed to exclusively sample peptide backbones while optimizing translations in Fuclidean space and

---

Published as a conference paper at ICLR 2025
bind: -49.041
bind: -41.665
bind: -46.676
bind: -42.627
bind: -50.410
Groud-truth PepFlow
(C-Floo((rans) OC-Flow(rot)
OC-Flow(trans+rot)
Figure 4: Visualization of OC-Flow generated peptide and unconditional generated peptide (5djd_C).
rotations in SO(3) space. The model employs the MadraX force field (Orlando et al. 2024) for energy optimization, and performance is evaluated using several key metrics. These metrics include MadraX energy, which assesses the total energy of the generated peptide structures, along with Rosetta-based measures of stability and affinity. Currently, our stability and affinity metrics are represented by their respective means: stability quantifies the energy states of peptide-protein complexes, while affinity measures the binding energies. In addition, employ existing metric, denoted IMP, which measures the percentage of generated peptides that exhibit lower energy than the original ground truth. Additionally, accuracy by aligning the generated peptides to their native structures and calculating the C. RMSD. To further analyze structural characteristics, we compute the secondary-structte (SSR), which reflects the proportion of shared secondary structures, and the binding site ratio (BSR), which quantifies the overlap between the binding sites of the generated and native peptides on the target protein. Structural diversity is assessed using the average of one minus the pairwise TM-Score (Zhang & Skolnick 2005) among the generated peptides, representing their dissimilarities.
the noot-mean-suurr deviation (RMSD) evaluate structural
similarity ratio
We compare our method to the pre-trained unconditional PepFlow model (Li et al. 2024), serving as a baseline. We also include ablations where our model guides only translations (Euclidean) or rotations (SO(3)). As shown in Table 5 our OC-Flow method, applied to both Euclidean and SO(3) spaces, consistently outperforms the baseline, even though we only optimize for the Madrax target function. This indicates that our algorithm not only achieves higher target function scores but also captures more natural structural configurations. It generates peptide backbones that are more stable, energetically favorable, and diverse, while improving key metries such as stability, affinity, IMP, diversity, SSR, and BSR. In comparison, optimizing in Euclidean space alone yields only marginal improvements in IMP, while optimizing rotations alone achieves comparable performance. More experimental details and ablation can be found in Appendix
E.3
Table 5: Evaluation of OC-Flow peptide design.
MadraX
RMSD
SSR %
BSR %
Stability
Aftinity
Diversity
imp(%)
Ground-truth
-0.888
-84.933
-3.0663
PepFlow
-0.195
1.645
0.794
0.874
-5.660
-26,538
0,310
14.3
(trran)
.0.229
1.774
0.797
0.876
448.800
-27.328
0.323
OC-Flow(rot)
-0.221
1.643
0.794
0.872
-48 636
-27.211
0,310
14.5
OC- DC-FFlo(trans-mrt]
-0.263
2.127
0.797
0.869
48.853
-27.468
0.338
15.0
6 CONCLUSIONS AND DISCUSSION
In this paper, propose OC-Flow, general and theoretically grounded framework for training- free guided flow matching under optimal control formulation. Our framework provides a unified perspective existing backprop-through-OOD approaches and lays the foundation for systematic analysis of the optimization dynamics of this setting. Extensive empirical experiments demonstrate the effectiveness of OC-Flow. Future extensions of OC-Flow include generalizing beyond additive control terms and bridging connection with fine-tuning regimes where control terms can be solved as learning updates the model parameters. Another extension could be scaling up the SO(3) OC-Flow to guide generative tasks for larger molecular systems such as protein motif scaffolding. One potential limitation of backprop-through-odd approaches, despite its superior result, is the higher computation cost compared to posterior sampling approaches. Such tradeoff has been demonstrated in Dflow/FlowGrad as well. Our practical implementations of OC-Flow improve the time and memory complexity (see analysis in Appendix D), where sampling on the image takes 216s compared to 15 minutes D-Flow. We hope that our findings can guide algorithm design and motivate further model improvement in guided flow matching.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Vladimir V Aleksandrov. On the accumulation of perturbations in the linear systems with two coordinates. Vestnik MGU, 3:67-76, 1968.
Heli Ben HamuH Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. arXiv preprint arXiv:2402.14017, 2024.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint 2xXiv:3305113300 2023.
Avishek Joey Bose, Tara Akhound-Sadega Kilian Fatras, Guillaume Huguet, Jarrid Rector-Brooks. Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Se (3)-stochastic flow matching for protein backbone generation. arXiv preprint arXiv:2310.02391, 2023.
Ricky TQ Chen and Yaron Lipman. Riemannian flow matching general geometries. arXiv preprint arXiv:2302.03660 2023.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Chaoran Cheng, Jiahan Li, Jian Peng, and Ge Liu. Categorical flow matching statistical manifolds. arXiv preprint 2rXiv:2405.16444 2024.
Felix L Chernousko and AA Lyubushin. Method of successive approximations for solution of optimal control problems. Optimal Control Applications and Methods, 3(2):011-114 1982.
Hyungjin Chung, Jeongsol Kim, Geon Yeong Park, Hyelin Nam, and Jong Chul Ye. Cfg++: Manifold- constrained classifier free guidance for diffusion models. arXiv preprint arXiv:2406.. 08070, 2024.
Leonardo Jesus Colombo and Dimos V. Dimarogonas. Symmetry reduction in optimal control of multiagent systems on lie groups. IEEE Transactions on Automatic Control, 65(1 11:4973-4980 2020. doi: 10. 00.1100/AAC.22200300077777
Alexandre Défossez, Jade Copet, Gabriel Synnaeve, and Yossi Adi. High fidelity neural audio compression. arXiv preprint N..iv:2210.134388 2022.
Floor Eijkelboom, Grigory Bartosh, Christian Andersson Naesseth, Max Welling, and Jan- Willem van de Meent. Variational flow matching for graph generation. arXiv preprint arXiv:2406. 04843, 2024.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Lawrence C Evans. An introduction mathematical optimal control theory version 0.2. Lecture notes available at http://math. berkeley. edu/evans/control. course. pdf. 1983.
Giorgio Giannone, Akash Srivastava, Ole Winther, and Faez Ahmed. Aligning optimization trajec- tories with diffusion models for constrained design generation. Advances in Neural Information Processing Systems, (6:51830-518611 2023.
Robert Kent Goodrich. A riesz representation theorem. Proceedings of the American Mathematical Society, 24(3):629 636, 1970.
Henry Gouk, Eibe Frank, Bernhard Pfahringer, and Michael J Cree. Regularisation of neural networks by enforcing lipschitz continuity. Machine Learning, 110:393-416, 2021.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12599. 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
11

---

Published as a conference paper at ICLR 2025
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867-8887.
PMLR. 2022.
Guillaume Huguet, James Vuckovie, Kilian Fatras, Eric Thibodeau-Laufer Pablo Lemos, Riashat Islam, Cheng- Hao Liu, Jarrid -rooks, Tara Akhound -adeeh, Michael Bronstein, et al. Sequence-augmented augmented se (3)-flow matching for conditional protein backbone generation. arXiv preprint arXiv:2405.203111 2024.
Velimir Jurdjevic. Optimal problems on Lie groups, pp. 368-406. Cambridge Studies in Advanced Mathematics. Cambridge University Press, 1996.
Tero Karras. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint OrXivv1110. 10/96, 2017.
Bahjat Kawar, Roy Ganz, and Michael Elad. Enhancing diffusion-based image synthesis with robust classifier guidance. arXiv preprint arXiv:2208. 08664 2022.
Grigory Khromov and Sidak Pal Singh. Some fundamental aspects about lipschitz continuity of neural networks. In The Twelfth International Conference on Learning Representations, 2024.
Gwanghyun Kim and Jong Chul Ye. Diffusionclip: Text-guided image manipulation using diffusion models. 2021.
Marin B. Kobilarov and Jerrold E. Marsden. Discrete geometric optimal control lie groups. IEEE Transactions on Robotics, 27(4):641 2011. doi: 101100/TRO.2011.22333
Jiahan Li, Chaoran Cheng, Zuofan Wu, Ruihan Guo, Shitong Luo, Zhizhou Ren, Jian Peng, and Jianzhu Ma. Full-atom peptide design based multi-modal flow matching. In Proceedings the International Conference on Machine Learning (ICML), 2024.
Qianxiao Li, Long Chen, Cheng Tai, and E Weinan. Maximum principle based algorithms for deep learning. Journal of Machine Learning Research, 18(165):1-29, 2018.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.027477 2022.
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 2022.
Xingchao Liu, Lemeng Wu, Shujian Zhang, Chengyue Gong, Wei Ping, and Qiang Liu. Flowgrad: Controlling the output of generative odes with gradients. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 14335-24344 2023.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834, 2023.
Brennan McCann, Annika Anderson, Morad Nazari, David Canales, and Ryne Beeson. On-manifold (3) for spacecraft coverage maximization. In AAS/AIAA Astrodynamics
pose optimization Specialists Conference, 2023.
Gabriele Orlando, Luis Serrano, Joost Schymkowitz, and Frederic Rousseau. Integrating physics in deep learning algorithms: force field as pytorch module. Bioinformatics, 40(4):btae166 April 2024. doi: 10. 1093/bioinformaatis//teee
Jiachun Pan, Jun Hao Liew, Vincent YF Tan, Jiashi Feng, and Hanshu Yan. Adjointdpm: Adjoint sensitivity method for gradient backpropagation of diffusion probabilistic models. arXiv preprint arXiv:2307.10711 2023.
Or Patashnik, Zongze Wu, Eli Shechtman, Daniel Cohen Or, and Dani Lischinski. Styleclip: Text- driven manipulation of stylegan imagery. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 2085-2094 2021.
Ashwini Pokle, Matthew J Muckley, Ricky TQ Chen, and Brian Karrer. Training-fie free linear image inversion via flows. arXiv preprint 2.Xiv:23100044. 2023.
12

---

Published as a conference paper at ICLR 2025
Lev Semenovich Pontryagin. Mathematical theory of optimal processes. Routledge, 2018.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 2014.
Lars Ruddigkeit, Ruud Van Deursen, Lorenz Blum, and Jean-Louis Reymond. Enumeration 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52 2:28844 11:2864-2875 2012.
Alessandro Saccon, John Hauser, and A. Pedro Aguiar. Exploration of kinematic optimal control on 10111002-0000 1307. 2010. ISSN 1474-6670. doi: the lie group so(3)*, IFAC Proceedings Volumes, URL Int tps:: www. sciencedirect https://doi.org/ http://///.....././..//.00.00000 /ttp:///ccc.oneee.......eeea/a/a
8th IFAC Symposium on Nonlin-
ear Control Systems.
Jun John Sakurai and Jim Napolitano. Modern quantum mechanics. Cambridge University Press, 2020.
Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming- Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-guided diffusion models for plug-and-play controllable generation. In International Conference on Machine Learning, pp. 32483-32498, PMLR, 2023.
Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao. Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei- Wei-Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36, 2024.
Martin Steinegger and Johannes Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature biotechnology, 35(1 1):1626-1022 2017.
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. International Conference Machine Learning, pp. 38592-38610. PMLR, 2023.
Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan Veeling, Regina Barzilay, Tommi Jaakkola, et al. Fast protein backbone generation with se (3) flow matching. arXiv preprint arXiv:2310.0522 2023.
Jason Yim, Andrew Campbell, Emile Mathieu, Andrew YK Foong, Michael Gastegger, José Jiménez- Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan Veeling, Frank Noé, et al. Improved motif-scaffolding with (3) flow matching. arXiv preprint arXiv:2441.0444 40..04082 2024.
Chenshuang Zhang, Chaoning Zhang, Sheng Zheng, Mengchun Zhang, Maryam Qamar, Sung- Bae, and In So Kweon. A survey on audio diffusion models: Text to speech synthesis and enhancement in generative ai. arXiv preprint arXiv:2303. 2023.
Yang Zhang and Jeffrey Skolnick. Tm-align: a protein structure alignment algorithm based on the tm-score. Nucleic acids research, 33(7):2302-2308 2005.
13

---

Published as a conference paper at ICLR 2025
A BACKGROUND OF RIEMANNIAN MANIFOLD AND SO(3) GROUP
Lie group G is a smooth manifold equipped with group operations, such as multiplication and inversion, which are smooth maps. Specifically, G' is considered smooth when it possesses C differential structure. When is endowed with a left-invariant Riemannian metric, it becomes a Riemannian manifold, where the inner product of any two tangent vectors E T2,G at a point E G is preserved under left multiplication. This property is expressed as:
(23) where LA G + G' is the left multiplication map, and (.. TG TG + R represents the inner product. Moreover, the tangent space at any point x E is given by T,G = L,T.G, where T.G is the tangent space the identity element e, which is identified with the Lie algebra g. Consequently, the full tangent bundle TG can be written as G x g.
(D), (w)) = {v, w),
At each point G, tangent space G is attached, representing the space of tangent vectors that point. The collection of these tangent spaces forms the tangent bundle TG, which itself is a smooth manifold. Additionally, for any point h G. the cotangent space TV, G is defined as the dual space of G. consisting of linear functionals (co-states) that act on the tangent vectors.
Rotation Group SO(3): The special orthogonal group SO(3), describing 3D rotations, is compact 3-dimensional Lie group. Its Lie algebra so(3) consists of skew-symmetric matrices. The group (3(3) defined as:
SO(3) = (r E R TT rr I, (r) =}.
(24)
It is a matrix Lie group, and its Lie algebra given by:
sa( = {r E 3x33 :r = -r].
(25)
Parametrizations of SO(3): The skew-symmetric matrices E s3) can be uniquely represented by a vector w E R², such that for any v R re = w x v, where x denotes the cross product. This vector is known the rotation vector, where its magnitude ||w|] represents the angle of rotation, and
its direction e = w/ [|| defines the axis of rotation. The mapping from R$ to the skmm-symeetric matrix is referred to as the hat operation, (.).
Another common parametrization of SO(3) through Fuler angles, described using three angles (ò, 0, v) In the x-convention, the rotation is expressed as a sequence of three rotations: rotation about the z-axis by é, followed by a rotation about the updated z-axis by 0, and finally, a rotation about the updated z-axis by th.
One common inner product on 5o(3) is the induced Frobenius inner product, given by:
(A, B) Tr(A B), YA, B E so()),
(26)
which equips SO(3) with Riemannian structure. The manifold SO(3) has constant Gaussian curvature and is diffeomorphic to a solid ball with antipodal points identified. The exponential map exp: : s3) + SO()), originating from the identity element, is defined as matrix exponentiation:
cxp(A) =
AA E so(3),
(27)
K!!
k=0
and can be represented more compactly via Rodrigues' rotation formula:
sin 0
cos A²,
exp(A)
A
(28)
Ø2
where 0 = || 2 || is the rotation angle. Similarly, the logarithm map log SO(3)
-
so(3), also originating from the identity, is the matrix logarithm:
1114++ (R - 1)*
log(R)
(29)
k
or more compactly as:
0
log( R)) sin 0 A, YR E SO(3).
(30)
where (R-R") R²) E s3) and (o(3)))(3 is the rotation angle. In spherical geometry, the 2 geodesic distance between two rotations is given by di d(R1, (R,,,8, R2) = log [R] R2)ll and interpolation between rotations can be performed using exp(
14

---

Published as a conference paper at ICLR 2025
B BACKGROUND OF PMP AND E-MSA
In this section, we shall introduce the details of the Pontryagin's Maximum Principle on both Euclidean Space and SO(3) manifold well the algorithms of MSA and E-MSA.
B.1 PMP
As described in section 3, PMP offers a set of necessary conditions for an optimal control strategy. PMP states that for an optimal trajectory, there exists a co-state trajectory 14 such that the Hamiltonian (Equation 3) is maximized (or minimized, depending on the problem) with respect to the control every time step. Additionally, the state and co-state evolve according to system of coupled differential equations, where the co-state variables evolve in the cotangent space of the state variables.
The governing differential equations are shown below:
Pontryagin's 's Maximum Principle Let 0" E LI be an essentially bounded optimal control, i.e. solution to (2) with ess suPre(O,T] |10% < oo (ess sup denotes the essential supremum). Denote by X the corresponding optimally controlled state process. Then, there exists an absolutely continuous co-state process P* [) + RR such that the Hamilton's 's equations
Pr, pr = V PT, ,?) are satisfied. Moreover, for each E [0, we have the Hamiltonian maximization condition
x
X
(31)
P (x7),
(32)
H(t, P.". 02) > H(t, P2.00 for all A e θ.
(33)
The PMP conditions can be naturally generalised the Lie Group. Pontryagin's 's Maximum Principle (PMP) for Lie groups (Saccon et al. (2010) provides the conditions that govern the flow of the state x, and its associated cotangent flow X, describing the Hamiltonian equations that the optimal state-adjoint trajectory must satisfy. Specifically, the Hamiltonian equations are given by:
δ H PT,
X-
x;= x
0p
0
i -dd a / (d dx 11,
2 (dL (
(34)
The dual map (dLg)* Ty SO(3) (0)) pulls back a cotangent vector gh to a cotangent vector at h. The coadjoint representation ad' acts on (3 and is defined as:
(a,z Y) (w,, -w,,ad xdx Y), where adxY- [X,Y] -XX for μ E so(3)* and X.Y 5o(3). Additionally, for each t e (), the Hamiltonian maximization condition is satisfied:
(35)
(t,, x;, (G.07) 2 Xi, for all 0 e e.
(36)
B.2
EXTENDED-METHON OF SUCCESSIVE APPROXIMATIONS (E-MSA)
B.2.1
METHOD OF SUCCESSIVE APPROXIMATIONS (MSA)
One numerical method for solving the Pontryagin Maximum Principle (PMP) is the Method of Successive Approximations (MSA) Chermousko & Lyubushin (1982). an iterative approach that alternates between propagation and optimization steps based on the PMP conditions. We first present the simplest form of MSA.
Consider the general state dynamics:
x; = f(t
(37)
Given initial guess 0 21 for the optimal control, for each iteration i: = 0,1, ...,.... ..., first solve the state dynamics:
x-0 -(1 X²x ,00),
x x,
(38)
15

---

Published as a conference paper at ICLR 2025
to obtain X"", followed by solving the co-state equation:
1e = Vt, X 14 0, Pr = =
--
(39)
to determine = Finally, the control is updated using the maximization condition:
0
arg max H(I X u 0), PEB
(40)
fort E [), This process is summarized in Algorithm 3.
Algorithm 3 Basic MSA
1: Initialize: 0 E U;
2: for = to #Iterations do
3: Solve X. ft, X O
X8
4: Solve iig" -- X
.14 O0))
-
5: Set 0f+1 arg maxoge HI,, X00 .1 00 for each i [0, 7):
6: end for
B.2.2 E-MSA
E-MSA introduces the augmented Hamiltonian
1 H(,,,,,,0,,0,)) H1,, u,, Gollu - f(t,
[ 0)
- 4,H((,,,,0111 (41)
Then, they define the following set of alternative necessary conditions for optimality:
Proposition 3 (Extended PMP) Suppose that 0 is an essentially bounded solution to the optimal control problem (2). Then, there exists an absolutely continuous co-state process J such that the tuple (X", |r" (*) satisfies the necessary conditions
X V (t, Of xT,
Xi
- 14 07 Xi,
1 -
Í1, Xi 14 07,05,0000000 Xi (4) 2 H(, X7, (i,,,,, E θ, E [0,
(42)
The key contribution is that the control terms can be updated by iteration using the Extended-PMP shown below. Meanwhile, it proven that under this update rule, for each iteration, the target function (0) is non-decreasing.
Algorithm 4 Extended MSA
1: Initialize: ( E U. Hyper-parameter: Pi
2: for == #lterations do
3: Solve X [1, X 01),
X =r; r:
4: Solve pa
(t, X
Pd
0t),
Pf =VV((X");
5: Set 0/++1 arg max Ht, Xp Pf
0, X pp. for each E [0,
6: end for
C PROOFS AND THEOREMS
C.1 PROOF FOR PROPOSITION
C.I.1 PART
Proposition 1. For Affine Gaussian Probability Path, the expectation of the running cost upper
bounds the KI. divergence between the prior joint distribution p1 (x², x1) = P1 (x?) (1) Pa (x1) and
16

---

Published as a conference paper at ICLR 2025
the joint distribution after guidance (z", .21)= ) pı P (21). with D1 Pdata- x² induced
by prior conditional vector field un (2) and x. sampled by applying control (x(11) on uy
!
En (a))
100 (21))1 dt 2 C (L , (
(43)
Proof. For Affine Gaussian Probability Paths, the conditional flow distribution of the prior vector field can be written as: Pf (x] (x) the conditional vector field is written Φx(x) px(1)) + oz
(44)
with its dynamics:
(x))) = pa(())) + d2(2))
(45)
or
(e(x)x1) (x))))) + - at (x-
Although control terms ((x1) are specifically derived for each ODE trajectory and are conditioned on the target point æ1, we can define marginalized control term 0z(x) (z) ((<)) (x) dæ to assemble aggregated control from all possible controls applied to each target point. Adding that to the marginal vector field can be thought of as effectively adding the turget-conditioed controls into the conditional vector field, and then marginalizing:
u(x) + (x(x) (w)(ar) + 0g (g(1))))))))))))))( dz
(46)
Given the definition of the Gaussian path, we can derive that the additive control terms only alter the mean of the distribution by denoting it as (z())) by noticing it is not proportional to z:
d (z) u,(((r))zz) (x1) (z(z)) (z()))) ox(1))
(47)
The resulting pushing-foree forward prob is equivalently:
(x) = Pt (c)
(48) Usually, besides maximising the reward, we hope the new distribution is not much away from the original distribution. One common constraint on the terminal state distribution is the Kullback-Leibler (KI.) divergence.
p (x(21) = (N (((x,))))
Inspired by Variational Flow Matching (Hijkelboom et al. (2024)), ) we can the KI. divergence between the joint distributions of the data point æ and the prior terminal point x?, denoted as pi (I(P, x1) pı (x* (z) ( and that of the data point æ and the controlled terminal point x". denoted as pı (a" .1): (L(PP (I, 21 Pl (x.
The KL term can be simplified as:
// (f(((()))(())) log Pp (ot x )aama (x1)
Pf ( D (21) da ddr, (49)
(1.)
Exy Expon [KL (2) II (x[x1))]-
Given ( and (x are two Gaussians with the same variance but different mean, and when we consider the constraint per sample, the constraint can be written as:
1
| (x1) of (x1))1?
(50)
2 o (x1))
given:
111² Or di 2
$ 5//'' 1101112 dt.
[n((())) (m) )112 1=
(51)
Therefore, we have the following incquality:
1
2 /.// ((((22))122
Ey, (11)
>
KL (x ,aaf,,,))))
(52)
40 (2
17

---

Published as a conference paper at ICLR 2025
C.1.2
PART 2
Proposition Part For square-shaped data with non-zero probability path, the expectation of the running cost, combined with the Ll-distance between the prior sample x; and the corresponding guided sample vi upper bounds the KL divergence between the marginal distributions of the prior model ph and the guided model pi:
Eg (xf
A
2 + B (0((2)))1 2 KL p?)-
(53)
Proof. Firstly, we can build a bound for the distance of terminal points of pre-trained model of (a) and our controlled model (x) given the same starting point. Given do](z) f*(x) with fP(() is
di
the prior model and (x) [*( O., meanwhile the Os are defined per sample, we can see them as a function of Given the ODEs governing of (z) and of (x) and use Lipschitz condition:
(=0()))))(())) - S. (1P()((()))))((((()))))))))) dt S /
de+ L." 1 dr.
(54)
By Gronwall's inequality:
(x) - ofx)| S et 110.1.tt.
(55)
From the bound we build above, we can
- of(x) 2(0,,z)) with g(0,,z) set S/" 10011av
(56)
and then by definition:
pp. xx) g).
(57)
Now we consider the push- forward functions and assume p non-zero for all
0 p.l
pf(x) (=)))
(x)
Ox
ri(x) pa (
(3))
(58)
Ör
Given the starting distribution is standard Gaussian, we can get the KL divergence as:
det 0 (x)
KL 11 = f (x)) (x)
- -1 1 + log [1,,7110)11
dx.
(59)
00 det
(x)
3x
With the Mean Value Theorem, find a point 3 between x and x g that ()) o - g) = (o))())))
(60)
Assume the gradient is bounded so that ( < k:
222 & Allg| + 1 2 | | (2)
lg
(61)
2
For the log term, with MVT again, there exists a point & on the line segment between r and x - such that:
1(x) - h(x g) = vh(E) g.
(62)
Applying this to h(x)) log det f(x):
log det f(x) log det f(x V det ()( g.
(63)
Thus,
A A(x)) (Vlg det ()(()) gl (Vog det (
(64)
18

---

Published as a conference paper at ICLR 2025
To bound (log ()) the gradient of log det f(x) with respect to 2 is:
(log det f(x) (det f())) f(x)):
(65)
Given the derivative of the determinant of matrix-valued function assuming (()) is square:
Videt, [(x)) det -
(66)
Therefore:
V det f(x)) = Tr -f((s)))))))
(67)
Assuming that:
|1 < M for some constant M;
[f) S N for some constant N.
Then:
(log det f(x)) S (z) [I)" n n MNn K.
(68)
Therefore:
det 0 (x) det 0 (x)
log
< Klgg|--
(69)
Overall, we can bound the KL divergence using the integration of 0g and the integration of the square of 0,:
(k + 1+]111911+ + 2²
(L || S Buiopf())
11113
22 +
f. [100011 2
(k+ + K)e+1l11 [1)
(70)
C.2 PROOF FOR THEOREM
Since the Extended Method of Successive Approximations (E-MSA) is applied in this context, its convergence properties are directly inherited. In this section, we focus on how our update rule is derived from the E-MSA update rule, specifically under the framework of additive control terms and running cost.
Define Hamiltonian H and Extended Hamiltonian H:
H(t, = pt f(x, t) + ue 0, 111
2(,,2,,,0,2.110 pe f(x,t) t) + Me 0.+1 110tt1122 11105 1110+11-000 0..1..
(71)
Apply Extended MSA:
$ = 0g -(xt,1), A = - v. x pth 0h) = f faf t)p
/f a V.., $(xt).
0/+1 argmax,, (t, μl, 0 2 14).
(72)
uf can be calculated in closed form as below, with T exp is the time-order exponential:
Texp (( aV (x}))
(73)
Thus, the update rule of 0, is:
(.. ((Vves(sss)))))) Va, Vo, ((x)).
0.+
Y Yigh
gh +
a T exp
Texp
(74)
1 1+ 1
19

---

Published as a conference paper at ICLR 2025
Further, the time-order exponential term can be simplified as follows. Evans (1883) provides method to calculate p(t) = Vz, (1) efficiently by defining the adjoint and using the following ODEs:
p(t) = v f t))(t), p(1) V=(1)=(1)
(75)
Sakurai & Napolitano 2020) provides an alternative viewpoint of the adjoint ODE, they show there would be a closed form solution:
[-] A(s) ds p(1).
p(t) ((t) p(t) = T exp
(76)
Thus, combining the closed form solution and the adjoint representation, we find that:
V (1) = T exp [. Vef( ds
(77)
Given the linearity of the ODEs, multiply a constant to the terminal by change p(t) = Va,z(1) into p(t) V,, ((x)), the conclusions would not change, have:
exp [[ Vaf(z", 5) ds Vz, ((1)).
V v2, F(x1)
(78)
Therefore, our update rule is in fact:
0 7
1 141 + aV v 2(x))
(79)
The solutions can be naturally generalised the space of R.V we use Tr to replace A B.
C.3 DISCRETIZATION ERROR
Given the discretization method using is Euler step, states due to Euler method with turning continous setting i = he(29, 0t) into discrete setting It-1 It he (2f 0.)At is bounded. Consider the first update of the states, define the local truncation error as:
show the error in terminal
Th = 2b41 It -he(x] 0t) At.
(80)
By definition:
ha ds he(a 0z) At.
(81)
Given At is small, with Taylor expansion:
Oh
Oh (s + higher order terms.
ha 0.) he(x,, 0.) +
(x x²
(82)
dx
Similarly:
x/6 - 0 = -, 0s) 1) higher order terms.
(83)
We can then obtain
@hh
Oh
h1 0.) = 11 0,) +
h4 +
(s - 1) + higher order terms.
(84)
Zr
0t
Therefore:
Aat1
A12
@hh
Oh
he (2 0s) = AL A1. Ne(2g.
0,)
h +
+ higher order terms,
2
dr
0t
2
Oh, 0r
Oh
7
h+
+
+ higher order terms.
(85)
Øt
Now we accumulate Local Errors to Global Error. Define error in kh step:
6k x² (1x)
Fk:
20

---

Published as a conference paper at ICLR 2025
ex+1 x² (tx+1) 24+1 (x'(tx) + At. ff==f(==" (tx), a" (tx), + Th) (Ix + At- f(zk, ug tx)).
(86)
Simplifying:
6k+1 = ek + A (f(2" (te), " (te), (x) - (xx, u 1x)) + Th-
(87)
Assuming I' i' Lipschitz continuous in a and u:
Il (tx), (tx), tx) f ug tx))l S (s(l=((te)) - Ixll Lyllexll
(88)
Error Recurrence Inequality:
||e+I|| S ||eall + At T. + ||Tell = (1 + At. [e]]le= +||Tx|]
(89)
Solving the Error Inequality given || s C(At)) where C' is a constant depending on f and its derivatives. We will solve the inequality:
Ex+1 S CE + C(At),
(90)
where eu = =|eex| and 1 At Ly Thus:
& Ek+1 S a k+ co + C(tt)
Σ
(91)
1=
Since eo = 0 (assuming zo is exact), the first term drops out. The sum becomes:
a 1-+1
S =
k-j
(92)
1=0
For small At:
a Et
a k+1 S (k+11) = c
(93)
Compute the Sum Sy:
clirk111 -
c1 1
S S اSعمi -1
(94)
At Lig
Final Bound on 6
cK+ & C(A))2 c 1-1 CAt(e -1)
(95)
At.L.
L.g
At the Final Time by:
CN = 1x= - CAl(etrts-) -11
(96)
L
For multiple rounds of updates, the error would be accumulated, for the N+1 th update, an additional local error added Tk:
N ' JN,N + (*, - (*,,N, xx)) At.
(97)
As the magnitude of this additional term is still O O(Ar2), the order of magnitude of Tu does not change, thus our conclusion does not change. The case in the SO(3) manifold should be similar, noticing the discretization also uses the Euler step.
C.4 ASYNCHRONOUS SETTING AND VJP IN SO(3)
In practice, as described by equation 22 discretization techniques employed to simulate the
ordinary differential equations (ODEs) governing both the state trajectory z; and the corresponding cotangent vector μ4. Most existing methods, as well the algorithm presented earlier, operate under a synchronous setting, where the number of time steps for the state trajectory 2; matches the number of control terms 0,,
However, OC-Flow can be extended to an asynchronous framework by approximation to allow greater flexibility in update scheduling.
21

---

Published as a conference paper at ICLR 2025
Rather than employing the standard update rule 24+A + f(t, 2e, 0t) subdivide the time interval At into equally spaced subintervals, applying the control term d, only during the first subinterval. The update for Z++ At is given by:
N f 0b) + fr (1
At fo (alm-riae)
It+A =
+ +
(98)
N
N
Moreover, for intermediate steps, we define:
244 R
Zt N O1),
AL
0.)
Σ N 1 (2g (524
Al
(99)
where x4 thr denotes the state the i-th subinterval.
Recall that (zt, 0z) (xx + +0, when we have At is small enough, the update rule in Equation 10 N can be approximated as case in Equation 4 by considering VoI++Ar.
At
At
At N
(2
V Vol++Ar
Vaf
N
N
At
At
At
Vxft (=
N
N
N
(100)
At
+
N
Therefore, if denote Fe the trajectory of the state variable over the time interval [t, t + At), and x as the trajectory when the control term 0g is applied in the first subinterval, can be reasonably approximated as:
At
A 0
L+t Lt +
ft
(101)
1=1
As result, for the asynchronous setting, the step 4 in Algorithm should be modified as:
X 1+
1 No AL. Σ ((o) +
(102)
For the case on the SO(3) manifold, the asynchronous setting can be deployed using the Taylor expansion of the matrix exponential exp(A), and noting that when At and At are sufficiently small, the terms become commutative. We can derive the approximation as follows:
1++11 = Iz exp
f(21.0))
exp
At (oe (2(++))
fd
...ep
exx
At N
fd (1) (iiN-pise)
N
N
RC R3 exp
At
p
(1e4)gg)
At
0z
(103)
N
N
1=1
(f((x.) 1 0: ).
Dr exp
At
The vector-Jacobian method can also be applied to compute the term 8ff (a Ez) in Algorithm 2:
Ole x E, = Pui
Ose dr
(104)
dr
C.5
PROOP FOR PROPOSITION
The key inequality we used for the proof of Proposition is following:
(h, v) < ||]||||||
| 1.
(105)
22

---

Published as a conference paper at ICLR 2025
for h, E s3 and SO(3) Before proving the proposition 3, first show that the cotangent vector {14} is also bounded.
Lemma 2: Assume all functions satisfy Lipschitz condition. Then, there exists constant K' > 0 such that for any 0,
| S K",
for all E [0, T].
Proof. Using necessary condition and setting
jul, pT--, get f)), / = (dL ) (v
ad on Po + (V
(106)
Du
With Lipschitz condition, we have || - (+ S K, (t, 0,)11 S K, and aldd ou pll
||||||| K k|]l Hence,
| K|l
(107)
and
|| | S a1 rbll 1 L' | ds J: ('Klll))
(108)
and by Gronwall's inequality,
| s l =: K"
(109)
This proves the claim since it holds for any т.
Now given all related terms are bounded, we can prove Proposition 2.
Proposition 3: Assume that the reward function, the prior model, and their derivatives satisfy Lipschitz continuity, bounded by Lipschitz constant L. Then, there exists conslant C >> 0 such that for any 0, ò E 50(3). the following inequality holds:
J(0) + ('' (s.ol(() dt - -1 - Ocll2dtt - J()),
(110)
where X and 1a satisfy the PMP conditions in Equation 34 and AH denotes the change the Hamiltonian, defined as:
AH) H1, (6) H6 0c).
(111)
Proof. Firstly, by the definition of Hamiltonian and PMP conditions, always have:
I(x (1.0) == S- (( H(t, Htt, u4 L(0) dt = O.
(112)
Define 6F4 14 pt and of: fp fa¹ the difference in can be decomposed as:
0 [x* S0, (1, Φ) I(x", (1.00)
(" (1i,5f)) (5pt, + (8p,, 8ft) dt
(113)
H(1,01 :) H(I 0)
1'' (L(Ø.) L((,)) dt.
T
Now define U(t) S6 f. ds = log(z,251) and by integrating by parts:
[" (/ 5f) (n, OU.N0 L'' (0. dt,
23

---

Published as a conference paper at ICLR 2025
L. (f, (8 8U:)10 - - (ōjū,,ü; dl.
(114)
Combine the two terms above:
0H (e (.) f.') dt (1l 80000 $ //' (ōpl du
) Mi ad on öft
1H
dr .8ful dt.
(115)
Similarly, we get:
L. (6p ole) di
1
//' (81, 55e) di + 2 (ôpt; 6f)
2
2 (8pe, SU,)
1
pT
(opt, 8U1) dt + 2//
1 l (opz, 8f() dt
2 J
1i, ad ayo oU, ) de (116)
2 (8p SUA)0
+
14:
ad
aue SU.
2
1 T
0H
OH
(dL
)
(d
88, dl
2
o
dr
Or
1
0H
dH
+ 2
81mn
dt.
J
0p
With mean value theorem and æ, are bounded by constant L, we can always find x between 2. and ag ss between μn and pg. y between & and 9, that
0 att (dL )r
(dL
H 2t. 0),0), SU, dt
dr
dx
0
H(1,
of 14 0) (d
H at pi 0), 0,, 8U, dl
dr
(v (2 H _ , 0), SU, dt
+
+
x1 o) 8U, dt
(117)
+
(V, 1 47, 6) SU.) dt
0
H(t. 2t 1. x) (dI=ag) 2
<
H (t, ar ni, 0), ,), dt
dr
dz
(" ||, 6 |8.
+
Using the same method we get:
[ (ppr, 5f.) di
2 2 (
(8,4 (U,) 10- +
aye 5U.
mi: ad
ene òU, dt 6p
+ 2
H(I, (.
H(1, 1 Mr, 0), SU di
dr
dr
(118)
T
0
H(t,
H 2 p6 0) 544 dt
2
0p
Op
T
C
|| + |1 11 dt.
24

---

Published as a conference paper at ICLR 2025
With boundary conditions:
" our, SU,/ 18--
-0= [1r++ 1 +22ur SUT
((dLz)) (V 0U,) 1 {(aL VV - (V oU,)
(119)
S (af)) (x+)) K 1110
Using same method to H ui , H(t, at pl, 0), we obtain:
(2g)) + L/ L(0,) dt
- + S'' de
((xx))
50 1// [4 ad
$ K K|0T 12
AH AH),o(t) dt
ouc ōU,
Mg ad
an1
oU, dt
cu
0
0
(1t, o 0)
(t, vi 1h 0), x. oU. dt
(120)
2
Or
Oz
1 T
H(, 2t 2
H(, 2 14 0).8 dt
2
J
0p
0p
+ C
|18U,112 + |1 SU: dt.
By definition:
8U, = S s(t,a),0) f(t, dt,
(121)
and
||0U%|| 1)( f(s, 0)| ds
J
$ 1 (1f(s,0g.)) f(s, o) ds f (1((s f(s.f( 22,
(122)
/ 11((, f( 0)1 ds ( ||U..lds.
0))
By Gronwall's inequality:
|| KT S' 110(s, 18 0)1 ds.
(123)
To estimate 814, we use the same substitution as in Lemma 6 with T we get:
ôjū, 8pā0 + ( (dLz))' V. 26 p. Φ) (dLgg)' 29, p.. 0) ds
(124)
/- (adz - ad و byf
/ ds.
25

---

Published as a conference paper at ICLR 2025
Using Lemma and Liptichitz conditions:
|ل6قū, || s 110 + ( [(ddيy)*V V (s, S. 0) (dLag) (s, i ngg 0)1 ds
+ S/" | 1i ad's 1 ds
Bsp
$ K||0UT|1 + kK" /-" ||0..]| K S/¹ |5 ds
(125)
+ ( ((d 1 Tg. po Φ)
R2 0) ds
$ c KTT K [10Ur|| K² / ||oU||| dt
+ t KT {-² (s, HsH 0)] ds.
Using the bound of Us obtain:
| |l S K" (L" (111F(1,2200)) f(s, 0) ds
(126)
+ KT K 1'' (s, pe 6) 2 2o n2, 0) ds.
Also we obtain:
14 ad 8U,) - U.g ad ono SU, dt
2
T
20
add"
an Dia
pi
ad
ane ue 8U,, 5m
Nl,,UU,) )dl
(127)
s C L/"
T
118 |1OUU/ dt.
Finally get:
26

---

Published as a conference paper at ICLR 2025
['' AHo,
J(0) - J(0)
1
+ K
1887/1
2
K"
(1100.11 + -116pell3) de
1
+ 2
16tu|l1[[(,, ob) - -(1,,2,,01) di
1
+ 2
Jo
V. -
V Mt, pZ 0.
T
AH),((!dl
(128)
(J² (1,(t, 01) ((a))-((tt,0 02)1
+
2
1'"
+C
$ [ Allo,((tdt
S/" 11f(t, 24, 01) - f(t, 2(t,2 0.)l
+C
+c //" IV, 2 p.) V, at, 0.)l dr.
Therefore, given the form of the addictive control terms and the running cost, we can derive the final term of our claim:
( / Avell(() Cllo: 0,11 dt S J()).
(129)
C.6 PROOF FOR PROPOSITION
Due to the similarity between the bound derived Proposition 3 and the bound obtained from the E-MSA method, the proofs of Proposition and Theorem follow the same reasoning outlined in Section 3.3 of Li et al. (2018). For the sake of completeness, we provide the full derivations here.
Proposition 4: Let X and p6 satisfy the PMP conditions in Equation 34 the update rule follows Equation! 19 we define Eg := So Agn12,0r (t) dt, and Eg is bounded as:
Ek = f A0tt1,6tH( dt,
lim E/ 0.
(130)
Furthermore, when = 0, we have = := arg max, J(O)To establish convergence, define
6 = 1" T
AHguts ot (t) dt 20 0.
(131)
Proof. By definition, if E 0, then from the update rule which maximizes the Hamiltonian, we
must have
-Ek S - 1i0+1 0² S 0,
(132)
and so
max (x," 8h .h 0, g0t .16.) (n .14 .0 18th 149h)
(133)
27

---

Published as a conference paper at ICLR 2025
Therefore, we always have the quantity 6ą 0 and it measures the distance from the optimal solution, and if equals 0, then we reach the optimum.
C.7 PROOF FOR THEOREM 5
Theorem Assume that the reward function, the prior model, and their derivatives satisfy Lipschitz
continuity, bounded by a Lipschitz constant L. Let 6 E so(3) be any initial measurable control with J(0") < t00. Suppose also that geso3)) J(0) > -C. If the update of satisfies equation 19 for
sufficiently large % the following inequality holds:
De# S ((B++1) J(t
(134)
for some constant D >> 0
Proof. Using Proposition we have
((ott)) -Ex +C 1/" |10 -0-11 dt.
(135)
From the Algorithm 2 maximizing step, we know that
Ht, Xf" pee ,0%) S Htt, X8" Pp 1 - -88122
-
(136)
Hence, we have
J(Oh) 1(0h+t)) - 20
7 Ck.
(137)
Pick y > 2C', then we shall have J(0h) 1(8h+1) S DEg with (1- -()-2)=
Moreover, we can rearrange and sum the above expression to get
M 4 S D- (J(ON+1)) J((O) < Tー1
(
inf J(0) - -(0)-J(0))
(138)
BEll
and hence Enzo < to0, which implies E + 0 and the algorithm converges to the optimum.
D COMPUTATIONAL EFFICIENCY
Table 6: Comparison of runtime and memory complexity of different methods used in backprop- through guided-ODE in Euclidean and SO(3) manifold. For complexity, N the number of ODE steps, n is the number of effective control terms with synchronized and in the range [1, N] and D is the complexity of computing 1-step gradient (VJP or Autograd), depends on data and model size. is the deficiency introduced by L-BFGS optimizer.
Number Of Control Terms
Memory Complexity
Runtime
Convergence to Optimal
Generalization to SO(3)
Complexity
OC-Flow
n
O(D2)
O(nD²)
FlowGrad
O(D2)
O(nD2
D-Flow
N
O(ND²)
O(cND2)
X
Red-Diff
N/A
O(D) O(LND)
In terms of memory complexity, both implementation and FlowGrad utilize the vector-Jacobian approach, which allows solving differentiation through ODE-integral with the adjoint method, significantly reducing memory from O(ND²) (D-Flow) to O(D²). In comparison, D-Flow does not employ the adjoint method, significantly increasing the memory complexity to O(ND?). as detailed in Table 6.
In terms runtime, both FlowGrad and OC-Flow are predominantly influenced by the number of control terms. Various strategies implemented reduce either the number of control terms (or
28

---

Published as a conference paper at ICLR 2025
Table 7: Illustration how different methods decrease the runtime and memory complexity and enhance model capability
Effective Timestep
Memory Complexity
Runtime Complexity
Generalization to SO(3)
OC-Flow
I
O(D*)
O
X
w/o-asyychromous
N
O(D2)
O(ND
X
w/o-VJP (adjoint)
N
O(ND2)
O(ND2)
X
OC-Flow-SO(3)
O(D*)
O(nDa)
woo-asynchronous
N
O(D2)
O(ND2)
w/o-VJP (adjoint)
N
O(ND*)
O(ND*)
the effective time steps requiring back- prock--ropggtiio)) Specifically, FlowGrad applies a straightening technique, and in oc-flow, this concept is extended to an asynchronous setting. Consequently. the runtime complexity is expressed (nnD2 for both FlowGrad and OC-Flow. In contrast, D-Flow does not imply an asynchronous setting, necessitating N steps of Autograd, resulting in O(N D²) complexity. addition, D-Flow is heavily relying on L-BFGS optimizer (Table 4), which also adds significant increase in runtime due to e.g., uncontrollable additional iteration of linear search, which we estimate with factor of constant
Table 7 provides a more direct comparison of the impact of the asynchronous setting and the vector- Jacobian product (VJP) method.
Compared to optimization based algorithms,, in some methods such as Red-Diff and DPS, the gradient- guidance can be approximated directly. These methods demonstrate significantly lower memory complexity and faster runtime. However, their capability is notably constrained. As reported in the Dflow paper, Red-Diff encounters difficulties in handling noise, even in simple linear inverse problems involving images.
When oc-flow is adapted to the SO(3) manifold, an additional calculation introduced for each control term. Specifically, this involves computing the trace of the product of two D D matrices, which adds a computational cost of O(D²). Despite this additional burden, the overall complexity remains in the order of OD2). The VJP method is also applied for oc-flow on the SO(3) manifold, as shown in Equation 15, for the term (it,, aff dh x F,), where as. "Be 4t is computed using VJP. This additional complexity is justified by oc-flow's improved convergence optimal solutions the SO(3) manifold. Notably, the OC-Flow-SO3 involves computing full Jacobian which could induce complexity of O(D*) if directing computing it. Our Jacobian- Vector Product derivation signifi- cantly reduces the complexity of OC-Flow-SOn from O(D*) to O(D²), which enabled our efficient implementation.
Table Memory Usage and Runtime on Text-guided Image (256*256) Manipulation. Note: ODE steps 100. optimization steps 15.
FlowGrad
DFlow
OC-Flow
Fast Simulation Peak GPU Mem (GB)
off
off
5.2
5.2
00M
5.4
5.4
Runtime per Sample (s)
114.7
206.2
115.7
216.7
In Table we show memory usage and runtime on digendimmnnionall images, evaluated on a single A100 with 40G memory. We compare FlowGrad, DFlow, and OC-Flow on images. (Liu et al. 2023) proposed fast simulation through skipping some Euler steps if the relative velocity change smaller than a threshold (1e-3 in (Liu et al. 2023)). We compare the 3 models with or without fast simulation. As for DFlow, we encountered Out-Of-Memory error even on 40G GPU. (Ben-Hamu et al. 2024) used gradient checkpoint to bypass the OOM error, however, due to the lack of code and implementation details of DFlow, we are unable to fully reproduce their implementation. For reference, DFlow takes 15 minutes per image (128*1 128) on single 32GB V V00 GPU, according to their paper. Generally speaking, FlowGrad and OC-Flow have similar memory usage and runtime except for DFlow.
29

---

Published as a conference paper at ICLR 2025
Table 9: Memory Usage and Runtime Comparison of OC-Flow Euclidean and SO(3).
SO(3)
Euclidcan
Peak GPU Mem (GB)
1.6
1.2
Runtime per Sample (s)
296.6
188.2
For peptides since FlowGrad and DFlow are not applicable in SO(3), we only report results for OC-Flow using our proposed asynchronous algorithm for efficient simulation, evaluated on a single A100 with 40G memory. Comparing the results, rotation is computationally 0.5x more expensive than translation due to the additional cost introduced by Eq. 22 However, as shown in Table 7 the costs of rotation and translation remain within the same order of magnitude. It is important to note that the extra cost of rotation is justified, as the additional computations required by Eq. [2 are inherent operations on the S3(3) manifold. Moreover, demonstrated in Table generation in SO(3) is crucial for the task.
Table 10: Runtime Comparison on Molecule. Note: ODE steps = 50, SGD steps = 20, L-BEGS steps 5 with inner steps =.
EquiFM
FlowGrad
DFlow
OC-Flow
Optimizer
N/A
SGD 37.6
SGD 31.3
L-BFGS
SGD 38.1
L-BFGS 103.7
Runtime per Sample (s)
2.4
102.8
In Table 10 we compare FlowGrad, DFlow, and OC-Flow on molecule, evaluated on single A100 with 40G memory. For reference, we also list EquiFM model as used our method. To align with DFlow, we also adopt L-BFGS optimizer and find it crucial in performance yet slows down the efficiency.
E EXPERIMENTAL DETAILS
E. TEXT-GUIDED IMAGE MANIPULATION
In our text-to-image generation experiment, we adopted the pipeline presented in Liu et al. (2023), utilizing the generative prior from Liu et al. (202)). We employed standard evaluation metrics: LPIPS and ID (face identity similarity) as introduced in Kim & Ye (202T) to assess the differences between the original image and the manipulated image. Additionally, the CLIP score used evaluate the alignment between the generated image and the provided text prompt.
To enforce consistency with the original image and inspired by Proposition introduced constraint term to the terminal reward function to penalize significant deviations from the original image:
(m)) XCLIP(#1, 7) (1 n)] afll-
(139)
Here, the hyperparameter was set to 0.7 across all experiments, and the Euler discretization step was set to N = 100 and the number of optimization iterations M = 15. As discussed in Theorem 2, increasing the learning rate 7) results in greater emphasis on the terminal reward, leading to higher CLIP score but lower LPIPS and ID scores. The weight decay is function of 7, which is tuned to limitation storage, we set 7 the same for all In our implementation, the learning rate r) was set to 2.5, and the weight decay 3 was set to 0.995.
maximize (1 20
7 ) for iteration k. this experiment, due
Baseline configurations were aligned with those reported in Liu et al. 2023), and the results presented in Table 2 reflect the same experimental conditions. For quantitative comparison, used the CelebA dataset, randomly sampling 1,000 images, which were manipulated based on text guidance: (old, sad, smi ling, angry, curly hair).
30

---

Published as a conference paper at ICLR 2025
E.2 MOLECULE GENERATION
In our QM9 generation experiment, we mostly followed the conditional generation pipeline in Hoogeboom et al. 2022)) An equivariant geometric GNN was trained for each property on half of the QM9 data the classifier, which was then frozen during our training-free controlled generation. The EquiFM Song et al. 2024) checkpoint, trained on the whole QM9 training data, was loaded as the generative prior. The test time properties slightly different from the settings in Ben--amu et al. (2024). Therefore, we reimplemented the D-Flow algorithm with 5 optimizer steps and 5 inner steps each with linear search using the L-BFGS optimizer. The results roughly matched those reported in the D-Flow paper with slightly worse MAES as we included the whole training dataset for property sampling. Our proposed OC-Flow also used the same optimization hyperparameters and also almost the same running time as the D-Flow. For FlowGrad, we followed the suggestion in the original paper to use 20 SGD steps to update the learnable parts, which ran slightly faster than OC-Flow and D-Flow.
sampled from the whole training dataset, making
For all properties, MAE was used as the optimization target and y is the regularization coefficient such that 7 fo 0,]] 1 di is the additional OC loss. For all optimization methods, we always used a can be indexed by discrete indices. As the integral is done with
fixed number of 50 Fuler steps a step size of 1/50, any is effectively 7 27/50 = 0.04 when taking the derivative with respect to 0 æ. Therefore, 7 10 effectively corresponds scheme.
0.4, which still a valid optimization
We noted the difference in the optimizer in these settings in order be consistent with the D-Flow setup. We provide additional ablation studies on the effect of the optimizer, in which all guided generation approaches used the SGD optimizer with 20 iterations and a learning rate of 1, following the FlowGrad setup. The results are summarized in Table [1 It can be clearly demonstrated that D-Flow performance is significantly worse than OC-Flow and even FlowGrad, the latter of which is special case of our OC-Flow. Indeed,
safely conclude that the advantage of D-Flow came solely from its optimizer using L-BFGS. Using the SGD optimizer caused to perform worse than the unconditional EquiFM baseline the dipole moment μ.. On the other hand, OC-Flow achieved consistent improvements, using either the L-BFGS or the simple SGD optimizer, demonstrating our superior performance.
Table 11: Ablation on optimizer. MAE for guided generations on QM9 (lower is better).
Property
Ct
As
EHOМO meV
LUMO MeV
Cy
Unit
Bohr
meV
K-mol
OC-Flow(OUrs)
1.907
346
187
300
0.362
0.972
D-Flow-SGD
5.753
1241
571
1195
1.639
2.982
FlowGrad
2.484
517
273
429
0.542
1.270
OC-Flow-LBFGSOOOus
1.383
367
183
342
0.314
0.819
D-Flow-LBFGS [Ben-Hamu et al. 2024
1,566
355
205
346
0.330
0.893
EquiFM
8.969
1439
622
1438
1.593
6.873
E.3
PEPTIDE DESIGN
In our peptide experiments, we adopted PepFlow ( al. 2024) the baseline model, utilizing the pre-trained checkpoint provided in the original PepFllow paper. The test dataset split was also based on the one defined in the PepFlow framework. For hyperparameter tuning, we randomly selected 10 complexes from the dataset. After tuning, the full set of 162 complexes was used for guided sampling and evaluation to ensure a comprehensive performance assessment.
To enable flexible update scheduling, we adopted an asynchronous setting in OC-Flow. This design maintains the same ODE time steps as PepFlow while utilizing fewer control terms. Specifically, 200 time steps used for ODE simulation, with 10 control terms, each controlling 20 time steps. As shown in Table 7 this approach reduces memory and runtime complexity without compromising the accuracy of the ODE simulation. Furthermore, for consistency and comparability across experiments,
31

---

Published as a conference paper at ICLR 2025
strictly controlled the initial noise during reruns to ensure consistency and comparability across experiments.
We used the pre- tre-trained model as the initialization for experiments, allowing to build upon the pre-trained weights and achieve consistent performance improvements through hyperparameter adjustments. In DC-Flow(rot) = 1.2; and in OC-Flow(both), a = 0.95 and 3 = 2.0. For all methods, we followed the same hyperparameter choices outlined in the Pepfiow paper ensure fairness.
used 0.95 and B 0.8; in OC-Flow(tans))
0.9 and
For evaluation, in addition to MadraX, the reward function used and optimized during training, we included several key metrics not used for training comprehensively assess the physical validity and overall performance of the generated structures:
Stability: Calculated using Rosetta over five independent runs, averaged to reduce high variance.
Affinity: Measured by Rosetta to determine the binding energy of designed peptides, also averaged over five runs.
IMP (Improvement Percentage); The percentage of peptides with improved affinities (lower binding energies) compared to the native peptides, aligning with the definition used in PepFlow.
Diversity: Calculated as the average of 1 TM-Score among generated peptides, reflecting structural dissimilarities.
SSR Seccondary-Structtiu Similarity Ratio):
The proportion of shared secondary struc- tures between the designed peptide and the native peptide.
BSR (Binding Site Ratio): The overlapping ratio between the binding site of the generated peptide and the native binding site on the target protein.
Due to the time-intensive nature of Rosetta evaluations, we drew 10 samples per pocket for our experiments, in contrast to PepFlow's 's use of 64 samples. This approach enabled us to achieve reproducible and fair comparisons across methods without excessive computational costs. Moreover, by employing OC-Flow with guidance, demonstrated that superior performance can be achieved with fewer samples while maintaining consistency in evaluation settings.
As shown in Table 5, demonstrated the importance of optimization on the SO(3) manifold for peptide design. To further evaluate the impact of using optimal control for updating rotations compared to standard gradient descent in Fuclidean space, implemented Naive-SO(3). In this implementation, the gradient of the terminal reward with respect to the control terms is computed directly and mapped to (3)),, and gradient descent is used update the control terms.
02- =30% + 17 vn[Vor (xx 1o(3)))
=Q-T1 (f gh+t 0 )a
0++1
(140)
For the experimental setup, we conducted a comparative study using a randomly selected subset of 30 pockets. Each peptide was tested under identical conditions, with the primary difference being the parameterization and update method for rotations.
Table 12: Comparison of OC-Flow-SO3 and naive SO3 gradient descent
MadraX
RMSD
SSR %
BSR %
Stubility 1
Affinity 1
Diversity f
imp(%) f
Ground-truth
00610
.91.107
-39,807
PepFlow
-0.157
1.932
0.788
0.882
-39,807
-28.080
0.322
11.6
Naive-SO13)
0.275
5.206
0.769
0.748
75.842
-21.901
0.635
OC-Flow-SO(3)
-0.191
1.943
0.794
0.874
-50.947
-29.027
0.332
14.0
As shown in Table 12 the results indicate that updating rotations on SO(3) using optimal control out- performs the naive method in terms of energy optimization and stability. The significant performance gap can attributed to the accuracy loss during the projection of the gradient onto so()), which may disrupt the delicate dynamics of the SO(3) space. Furthermore, due to the complexity of the SO(3)
32

---

Published as a conference paper at ICLR 2025
manifold, gradient- based methods are more prone to becoming trapped in local optima, whereas the optimal control-baser algorithm provides a pathway toward achieving the global optimum.
Our efforts to reproduce the PepFlow baseline involved extensive steps, including reaching out the original authors to address the absence of certain scripts and obtaining partial instructions for the evaluation pipelines. In the original PepFlow paper, affinity and stability reported as percentages. However, to facilitate more fine-grained comparisons, we opted to report absolute energy values instead. Additionally, ensure fairness in comparison, we included IMP (Interaction Metric for Peptides) as an evaluation metric, aligning it with the affinity measure used in PepFlow. Furthermore, our experiments were conducted using the latest version of MadraX, ensuring that our results are both robust and reproducible. These updates provide consistent and comprehensive framework for evaluating and comparing future methods in peptide design.
33


Reviews:
[{"review_id": "haMe8Qw8Jm", "reviewer": "Reviewer_faUf", "summary": "This paper proposed a new framework for controlled generation using pre-trained diffusion and flow matching models, dubbed OC-Flow. The method is based on sound theory in optimal control that offers additional convergence guarantees in Proposition 1 and Theorem 2 (under two key assumptions of affince Gaussian path and Lipschitz continuity of the gradient of guided loss). Several benchmarks on guided-image manipulation, molecular generation and protein design with generative models are performed to demonstrate the effectiveness of the method.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Well-motivated problem, overall nicely written paper with clear literature review.\n- The methodological and theoretical parts of the paper are well-sounded. \n- Providing a framework that has convergence analysis is always welcomed.", "weaknesses": "**Major: questionable and inconsistent baselines' results in empirical benchmarks** \n\n- While on the first task (section 5.1 text-guided image manipulation) the authors have report/insert exactly other baselines' results (originally in Table 2 of the FlowGrad paper); the results on two remaining tasks in section 5.2 (molecule generation) and section 5.3 (peptide design) do not match the results reported in their respective original paper. \n- More specifically, the results in Table 3 do not match those of Table 4 in D-Flow paper (Ben-Hamu et al. 2024); results in Table 5 do not match those of Table 1 in PepFlow paper (Li et al. 2024). In fact, if one instead takes into account the original results, the baseline D-Flow actually perform better in MAE metrics compared to OC-Flow in Table 3. For Table 5 the metrics reported are in different scale. \n- I am therefore request the authors to clarify this discrepancies between results reported in their paper and the results reported in the respective original works of compared baselines. Otherwise, I think the practical performance of OC-Flows remains questionable.\n\nBen-Hamu et al. (2024). D-Flow: Differentiating through Flows for Controlled Generation. Proceedings of the 41 st International Conference on Machine Learning, Vienna, Austria. PMLR 235, 2024.\n\nLi et al. (2024). Full-Atom Peptide Design based on Multi-modal Flow Matching. Proceedings of the 41 st International Conference on Machine Learning, Vienna, Austria. PMLR 235, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mIm1Rswdeq", "reviewer": "Reviewer_o479", "summary": "This paper provides a novel approach for guided generation using pre-trained diffusion and flow matching models. Traditional methods of guiding ODE-based generative models often require expensive retraining and work mainly on Euclidean manifolds, but OC-Flow uses an optimal control training-free framework beyond Euclidean spaces to the SO(3) manifold. Experiments on tasks like text-guided image manipulation, conditional molecule generation, and peptide design validate the method’s effectiveness​.", "questions": "- The paper shows promising results, but could the authors elaborate on potential ways to enhance scalability, especially when applied to e.g. large molecules or more complex target distributions in general? \n- How does OC-Flow compare to recent works in Riemannian FM? What about SO(3) and SE(3)?\n- Could OC-Flow be adapted for hybrid tasks where both Euclidean and Riemannian components are present? This might extend its applicability to even broader fields where you need this hybrid. Does the method allow for this directly or not?\n- Since OC-Flow is designed to be computationally efficient, could the authors comment on real-time applications requiring 'immediate' guidance?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- As far as I know, the approach is original in framing guided flow matching as an optimal control problem. The authors develop a general framework for non-Euclidean geometries with strong theoretical backing, which is fairly rare.\n- Another strength of this work is that it is a general approach, i.e. OC-Flow can be used effectively for a variety of applications such as image and molecular data. \n- Unlike existing approaches, OC-Flow allows training-free guidance, making it computationally efficient and more applicable in real-life settings. \n- The SO(3) results such as on improved molecular generation accuracy, validate the importance of using this geometric inductive bias for generative models. \n- Through the framing of existing approaches as special cases under their optimal control formulation, this paper helps clarify the connections between gradient-based techniques like FlowGrad and D-Flow.\n- The model consistently demonstrates improvement over previous work.\n- I really appreciate the figures included in the paper to illustrate and compare the methods against existing approaches.", "weaknesses": "- My main question is regarding scalability. While the model performs well on selected benchmarks, to me it is still unclear how OC-Flow scales to high-dimensional datasets such as large molecules. A discussion of potential scalability limits and memory efficiency in such cases would strengthen the paper.\n- Moreover, even though the formal contributions are great and well-formalized, to me the paper is still quite hard to read. Since the theoretical results are one of the main contributions, I think it would be valuable to add more intuitive explanations of the proofs and why they are there. For example, theorem one provides a bound based on VFM on KL between the model and a terminal point, but some intuition of why this bound is provided would make the paper more approachable.\n- Adding to this point, the theoretical assumptions made (e.g. Lipschitz continuity, boundedness) are clear and needed for the argument, but some reflection (perhaps on a high level) on whether these hold in practice would help to interpret the method's advantages.\n- While the focus is on continuous CNFs, a brief comparison with discrete flow techniques could contextualize OC-Flow's advantages or limitations more clearly, especially as discrete methods have shown promise in similar applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MVRZqETDsz", "reviewer": "Reviewer_jxTp", "summary": "This paper attempts to solve the problem of conditional generation using Flow-Matching models. In particular, they propose a unifying framework (OC-Flow) from which other approaches (such as D-Flow and FlowGrad) can be derived, and operates in Euclidean and SO(3) geometries. Subsequently, the authors provide extensive theoretical analysis of OC-flow to prove convergence and theoretical properties. Finally, the authors apply OC-Flow to text-guided image generation and peptide design tasks.", "questions": "* I would like to see a comparison of the runtime of OC-flow vs the other methods. While Table 1 suggests that the memory consumption of OC-flow is lower than D-flow and on par with Flow-Grad, I would be concerned that the additional complexity of solving in SO(3) adds significant computational costs.\n* I would like to have a better understanding of what parts of the paper are core to the methodology (i.e. actually implementing OC-Flow), versus what parts of the paper are necessary for proving convergence. I would then like to see separate sections/subsections for the proposal, and the subsequent analysis.\n* I would like to see a clearer presentation of a conventional flow-matching model, and how the proposed method extends this standard formulation, ideally in the form of before/after equations to get a clear and unambiguous idea of the elements being added/proposed.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "* The paper has a strong theoretical grounding, is well placed within the existing literature and goes to extensive efforts to prove theoretical properties of the proposed methodology. Additionally, the paper provides significant context and background, referencing existing works in Flow-Matching models. Finally, the authors provide a ton of detailed proofs in the appendix, and theoretical analysis in the main body of the text.\n* The paper provides a legitimate contribution to formalizing and extending existing guided-flow matching techniques to complex geometries (such as SO(3))\n* The authors compare the proposed methodology to similar existing methods to demonstrate competitive performance\n* Table 1 gives a good comparison to understand the contribution of OC-flow vs D-Flow and Flow-Grad", "weaknesses": "__Theoretical Concerns__\n* The main concern with the theoretical aspects of the paper is that the authors perform all the theoretical analysis in the continuous regime, but then implement the practical algorithms in a discrete regime. They even mention this limitation in the conclusion (“we also note that our practical algorithm…”). This presents a potentially significant hole in the paper, as the objects being proved, and the objects being validated empirically are not the same, and thus the theoretical content in the paper is not necessarily applicable to the experimental results (as noted by the author). To address this, we recommend that the authors 1) provide theoretical analysis of the discrete regime of the algorithm, 2) discuss in detail how the continuous approximates or bounds the behavior of the discrete implementation, or 3) implement a continuous version of the algorithm. These additions would help bridge the gap between theory and practice.\n\n__Experimental Concerns__\n* There is a significant lack of information regarding the experiments and reproducibility. There are no details as to how these models are trained, what the architectures are, the implementation, etc. This would need to be addressed before the paper could be accepted. To address this, we recommend that the authors give detailed descriptions of the model architectures and hyperparameters, training procedures and optimization details, data preprocessing steps, computational resources used, and code availability or plans for release. This would greatly improve the reproducibility of the work.\n* Similarly, no error bars/confidence intervals are reported on any of the experiments. In fact, it is unclear whether their model is better than existing baselines (i.e. in table 5, they claim 0.795 is better than 0.793, but no CI, similarly in tables 2, 3 and 4, they report outperforming existing methods but do not report CI.). Consequently, it is not possible to assess their claims that they outperform existing models, given the lack of confidence intervals and the close proximity of the performance values. To address this, we recommend that the authors 1) run multiple trials and report mean and standard deviation for all metrics, 2) perform appropriate statistical significance tests (e.g. t-tests) when comparing to baselines and 3) include error bars or confidence intervals in all tables and figures.\n\n* Finally, one of the major claims in the paper is that OC-flow can optimize in euclidean and SO(3) space, and that optimization in SO(3) provides benefits in tasks such as protein design. However, the authors do now present results split into Euclidean and SO(3) algorithms. It is unclear how/if the extension to SO(3) is even beneficial, and additional experimental details/results are needed to validate this claim as well. This is brought up by \"our OC-Flow method, fully optimized in both Euclidean and SO(3) space\" on page 10.\n\n__Presentation Concerns__\n* The paper has some issues with the presentation that make it quite difficult to asses which parts are novel contributions, and which parts are existing works that are being used. The authors/paper would benefit greatly from having a clear vision of what they are proposing and why, and subsequently moving large portions of the detailed proofs to the appendix to not muddy understanding with unnecessary detours. For example, what are co-state flow and E-MSA, and why do we care about these constructs? How do these constructs factor into the actual problem of performing guided matching flow generation? Are they purely used for proving convergence analysis? And if so, then it should be framed/explained as such. In fact, the structure of the paper would greatly benefit from having a section which clearly describes the proposed methodology in terms of implementation, and a separate section for the theoretical analysis of the proposed method, since the current structure makes it very difficult to separate the method as-such, from the additional theoretical concepts only necessary for proving convergence.\n  * To address these issues, we recommend that the authors clearly delineate novel contributions from existing work, as well as practical details from theoretical proofs. Adding sections such as \"contributions\", \"proposed methodology and implementation\", and \"theoretical results\" would greatly improve the structure of the paper.\n  * Additionally, we recommend that the authors provide clearer explanations of key concepts such as co-state flow and E-MSA, as well as highlighting their importance to the key contributions in the paper.\n  * Finally, by restructuring the paper to highlight the novel aspects, and moving detailed proofs to the appendix, the overall quality and clarity of the paper would be much improved.\n* Furthermore, several significant objects/theorems are introduced with very little explanation. For example, co-state variables are introduced as “shadow prices representing the sensitivity of the optimal value function to changes in the state variables”. But what are shadow prices? How does this analogy help when there is little thought/exposition given to the co-state/Hamiltonian introduced by the PMP? I would like to see a much more principled approach to writing, where each theorem introduced is clearly placed within the larger context of the work, and has a clear purpose in support of theoretical results.\n* Similarly, the lack of figures greatly hinders understanding. Additionally, figure 1 does not clearly articulate what it is presenting, and what the various sub-figures and equations represent.\n\n__Contribution Concerns__:\n* First of all, due to the presentation it is not clear what the contributions of the paper are, and what is preexisting work being leveraged for proving theoretical properties of the method. However, my understanding is that there are two main contributions of the paper: 1) they formulate conditional generation using flow-matching models as a control problem in equation 2, and 2) given that formulation, the authors demonstrate that OC-Flow is a generalization of D-Flow and Flow-Grad that can be optimized in SO(3), as well as proving various convergence properties.\n* In light of this understanding, it seems like the contributions of the paper are limited in scope. First of all, equation 2 seems to be fairly trivial extension of existing Flow-Matching/Continuous Normalizing Flow formulations (see Fjelde et al (2024)). Furthermore, given the lack of definitive experimental results, it is unclear whether this extension provides tangible benefit over existing methods, especially when considering the additional complexity. One of the major proposed benefits of the method is optimization in SO(3), but no ablation studies are given to demonstrate that SO(3) provides additional benefits over simple euclidean optimization.\n* Additionally, the experimental reproducibility of the paper is quite poor, with no experimental parameters given, and no experimental source code provided.\n*Finally, the scope of the contribution is somewhat niche. In particular, this paper focuses on classifier-guided generation using flow-matching models in SO(3). While useful for certain problems, it likely does not have wide-reaching implications outside of a few target applications.\n\n\n__Citations__\n* Fjelde, T., Mathieu, E., & Dutordoir, V.. (2024). An Introduction to Flow Matching.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "H1LR4KkHRV", "reviewer": "Reviewer_zuF2", "summary": "The paper proposes a novel method based on optimal control to optimize generations obtained by ODE-based generative models (e.g. flow matching). The paper proposes algorithms for generative models in Euclidean space and in SO(3), generalizes previously existing approaches, and studies the convergence of the proposed method.", "questions": "Proposition 1. What is the prior terminal point $x^p$? Is it $x^p_0$ or $x^p_1$ (that is, terminal as in time $t=1$ or $t=0$). If $t=0$ then the joint $p_1(x^p, x_1)$ is a delta distribution (since ODE is deterministic)? If $t=1$ then $x^p$ and $x_1$ are independent (since $x^p$ is generated from random noise independent of $x_1$)?\n\nGradFlow could in principle be used for manifolds too? The control terms would live in the tangent space of the manifold at the current point? They do not propose this in GradFlow so this is not something to compare against, and would even be consider a novelty and addition to the paper. But this work got me wondering if there’s any obvious reason why this would fail?\n\nI’m not sure I completely agree with the paper’s title “Training free…”. I understand training is the same, and that this method can be used for any pre-trained flow model. But it does require solving an optimization problem, albeit with few iterations (~20) but expensive ones. The difference is that this optimization happens at inference time.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper tackles the problem of changing the generation process of ODE-based generative models in order to produce samples that maximize a certain reward, while staying close to the original ODE trajectory (through a regularization term). This problem is relevant in multiple domains where additional signals/information are available at inference time. \n\nTo the best of my knowledge, this is the first paper to formalize this guidance and control framework in SO(3), which is a group used by many methods in structural biology.\n\nThe approach generalizes existing methods that optimize the trajectory of ODE-samplers (D-Flow, GradFlow), and outperforms them in multiple benchmarks.", "weaknesses": "**Computational cost.** The method proposed in the paper, as well as its predecessors (D-Flow, GradFlow) all require optimizing the sampling process *for each sample* produced by the generative model. In other words, producing a single sample requires solving an optimization problem, for which computing the loss requires simulating the full ODE. This has a high cost in memory and computation time.\n\nWhile GradFlow proposed a clever way of reducing the memory cost of this process, which is also adopted in this paper, this optimization process is inherently computationally expensive, significantly increasing the time cost of producing each sample. Previous work used some approaches to try to alleviate this (e.g. FlowGrad uses an adaptive solver to minimize the number of steps used during generation, by setting the step-size as a function of the estimated curvature of the flow at the current point), but simulation is still considerably slower than the baselines without this optimization process. While generation times are not discussed in this work, the D-Flow paper states that producing a single molecule takes around 3 minutes (they use 100 function evals for discretization of the ODE), and for images the time to generate a single output ranges from 4 to 15 minutes depending on the task. While the purpose of these works is not increasing generation efficiency, but generating better samples through guidance and optimization, they exacerbate the main limitation of diffusion models / flow models, which is their slow generation. In the paper I am unable to find generation time for the experiments. Given the method's nature, I think these should be reported and discussed.\n\nI think related to this point, experiments tend to be on the smaller end. Celeba-HQ for images, molecule generation with up to 9 heavy atoms, and peptides, which are short proteins (less than 50 residues). I understand these methods are able to produce better samples given the external guidance, while other approaches are unable to leverage such information, which is quite valuable. Still, I think providing values for computational cost / generation time, and comparing against plain approaches (baselines that do not require tuning) would be good. I would expect the cost of producing one sample is between 10x-100x more than baselines that do not optimize the sampling process (since there are ~20 optimization steps, and some gradient computation too), but happy to be shown otherwise. This does not account for the fact that lower memory requirements by the baselines would allow producing more samples in parallel too.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)