Training Free Guided Flow-Matching with Optimal Control

## Abstract

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 ) 𝑅 subscript 𝑥 1 R(x_{1}) italic_R ( italic_x start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT ) 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.

## 1 Introduction and related work

SDE and ODE-based generative models such as diffusion and continuous normalizing flow (CNF) have exhibited excellent performance on various domains such as images (Ho et al., ; Esser et al., ), audio (Zhang et al., ; Défossez et al., ), and discrete data (Lou et al., ; Cheng et al., ). Particularly, the simplicity of Riemannian Flow Matching on SO(3) manifold (Chen & Lipman, ) empowers de novo generation of small molecules (Song et al., ; Xu et al., ) and proteins (Yim et al., ; Bose et al., ; Li et al., ), leading to enormous advancement in biomedicine. Controlled generation from pre-trained diffusion and flow matching priors has gained growing interest in numerous fields, as it encompasses a wide range of practical tasks including constrained generation (Giannone et al., ), solving inverse problems (Liu et al., ; Ben-Hamu et al., ), and instruction alignment (Black et al., ; Esser et al., ).

There are several lines of work for guiding diffusion and flow models. Classifier-free guidance (CFG) (Ho & Salimans, ) trains conditional generative models that take conditions as input. Reward fine-tuning approaches update the generative model parameters to align with certain target objective functions (Black et al., ). Both methods require specialized training routines which are costly and not extendable to new tasks. Training-free guidance on diffusion (Kawar et al., ; Chung et al., ; Song et al., ) alters the scores in the SDE generation process with the gradients from the target function to achieve posterior sampling. These guidances often rely on strong assumptions of the denoising process and require estimating target function gradients w.r.t noised samples which are often intractable. Accurate posterior sampling is only guaranteed for a 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 et al., ; Yim et al., ).

Notably, two recent works showed the effectiveness of guiding pre-trained flow models by differentiating through the ODE sampling process, outperforming popular guidance-based approaches. Particularly, Ben-Hamu et al. () differentiates a loss $R(x)$ through the forward-ODE w.r.t the initial noise $x_{0}$, which induces implicit regularization by projecting the gradient onto the “data manifold” under Gaussian path assumption. This strong confinement to the prior might hinder optimization when the target reward function diverges from the prior distribution. Liu et al. () formulates an optimal control problem where a control term $u_{t}$ at each timestep is solved to guide the ODE trajectory. However, the gradient decomposition trick used in Liu et al. () 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 differentiate-through-ODE approaches to better guide algorithm design in this space. Furthermore, existing works predominantly focus on the Euclidean manifold due to its simplicity, and there is a compelling need for a theoretically grounded guided flow matching framework on more complex geometries such as SO(3) which is heavily used in scientific applications.

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 an optimal control problem with a control term $u_{t}$ 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 a co-state flow and control term, with theoretical guarantees under continuous-time formulation.

2. In Euclidean space, we show that running cost bounds the KL divergence between prior and OC-Flow-induced 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 a unified view of the problem.

3. We present one of the first guided flow-matching algorithm on the $\text{SO}(3)$ manifold with theoretical grounds. Our approach extends the Extended Method of Successive Approximations (E-MSA) to $\text{SO}(3)$ with a rigorous convergence analysis. Additionally, we propose approximation techniques to enable computationally efficient OC-Flow on $\text{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., ; Liu et al., ) provides an efficient framework for training a generative model by approximating the time-dependent vector field associated with the flow represented as $\psi_{t}:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$. This vector field $u_{t}:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ defines a probability path of the evolution of an initial noisy distribution, denoted by $p_{0}$, towards a target distribution, $p_{1}$, with the pushforward probability as $p_{t}:=(\psi_{t})_{*}p_{0}$. The dynamics of the vector field that governs this flow can be described by the flow ordinary differential equation (ODE) of the form $\dot{x}_{t}=u_{t}(x_{t})$, where we follow the convention to use Newton’s notation with respect to time $t$ and the state at time $t$ is given by $x_{t}:=\psi_{t}(x_{0})$.
Lipman et al. () demonstrates that a tractable flow matching objective can be obtained by conditioning on the target data $x_{1}$. The primary goal of conditional flow matching is to train a model, $f^{p}_{t}:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$, such that it minimizes the difference between its output and the ground truth conditional vector field as $\mathcal{L}_{\text{CFM}}=\mathbb{E}_{t,p({x}_{0},{x}_{1})}\|f^{p}_{t}({x}_{t})
-u_{t}({x}_{t}\mid{x}_{1},{x}_{0})\|^{2}$.
The trained model $f^{p}$ can be employed as the marginal vector field during the inference phase. In this context, once a noise sample $x_{0}$ is drawn, the system’s evolution can be described by the following differential equation:

$$ $\dot{x}_{t}=f^{p}_{t}(x_{t}),\quad x_{0}\sim p_{0}(x).$ (1) $$

##### Rotation Group SO(3).

The formulation of flow matching can naturally extend to Riemannian manifolds (Chen & Lipman, ).
In a Riemannian manifold, a flow is defined as a time-dependent diffeomorphisms $\varphi_{t}:G\to G$, which describe the continuous evolution of points on the manifold over time, generated by a vector field $V$, with $V(p)\in T_{p}G$ for each $p\in G$ where $T_{p}$ is the tangent space at point $p\in G$. The flow evolves according to:$\frac{d}{dt}x_{t}=V(x_{t})=L_{x_{t}}V(e)$. The CFM objective in Equation 1 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 $\text{SO}(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 a rotation that defines its orientation (Yim et al., ; Bose et al., ; Li et al., ). Guiding such pre-trained generative models toward the desirable protein properties can potentially have a profound impact on the pharmaceutical industry.

Figure: Figure 1: Visualization of different backprop-through-ODE guidance algorithms. D-Flow (left): update $x_{0}$ directly; FlowGrad (middle):update each control term by gradient information; OC-Flow (ours, right): incorporates the running cost and reward ratio $\alpha$ in terminal reward. we leverage the co-state flow $\mu_{t}$ to integrate the gradient information to update the control terms $\{\theta_{t}\}$.
Refer to caption: extracted/5949639/figs/Method_compare_final.png

## 3 Optimal Control Framework for Guided Flow Matching

Given a pre-trained flow model, $f^{p}_{t}(x)$, parameterized by a neural network, our goal is to determine the optimal control terms $\theta_{t}$ that maximize the reward $\mathcal{R}(x)$ while maintaining the proximity of the resulting vector field to the original vector field induced by $f^{p}_{t}(x)$. The reward can be customized for diverse tasks such as inverse problem $\mathcal{R}=\|H(x)-y\|^{2}$, conditional generation $\mathcal{R}=(f(x)-c)^{2}$, and constrained generation $\mathcal{R}=\|x-y\|^{2}$. To ensure proximity, we incorporate a penalty on the state trajectory or control terms $\int_{0}^{T}L(\theta_{t})$, also known as the running cost. Optionally, one may also introduce a metric $d(\cdot,\cdot)$ to penalize the deviation between the new terminal state $x_{1}^{\theta}$ and the prior terminal state $x_{1}^{p}$. The modified terminal reward function is then defined as:
$\Phi(x_{1}^{\theta})=\mathcal{R}(x_{1}^{\theta})-d(x_{1}^{\theta},x_{1}^{p})$. and scaling the terminal reward by a constant $\alpha$, we can formulate the problem as a standard optimal control task:

$$ $J(\theta):=\alpha\Phi(x_{T}^{\theta})+\int_{0}^{T}L(\theta_{t})\,dt\quad\text{ subject to}\quad\dot{x}_{t}^{\theta}=h_{t}(x_{t}^{\theta},\theta_{t}).$ (2) $$

A fundamental result in optimal control theory is Pontryagin’s Maximum Principle (PMP) (Pontryagin ()), 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 is defined in terms of the state of the system, the control, and a new entity called the co-state $\mu$, which resides in the cotangent space of the state manifold.

$$ $H(t,x,\mu,\theta)=\langle\mu,h_{t}(x_{t}^{\theta},\theta_{t})\rangle-L(\theta)$ (3) $$

The co-state variables can be thought of as shadow prices that represent the sensitivity of the optimal value function to changes in the state variables. Meanwhile, the terminal reward is also encoded to $\mu_{T}$. PMP states that for a control to be optimal, the Hamiltonian must be maximized by the evolution of the co-state $\mu_{t}$. The details of PMP conditions can be found in Appendix  B.1.

### 3.1 OC-Flow on Eulidean 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 ()), which directly perturbs the prior trajectory. In fact, with the linear expansion, the additive control could be seen as a general case and is widely used in optimal control. Specifically, with $\theta_{t}$ representing the control input, the new state dynamics and the corresponding running cost can be defined as:

$$ $\dot{x_{t}}=h_{t}(x_{t}^{\theta},\theta_{t})=f^{p}_{t}(x_{t})+\theta_{t}\quad L (\theta_{t})=-\frac{1}{2}||\theta_{t}||^{2}$ (4) $$

The running cost can be interpreted as a constraint on the trajectory and, equivalently, as a constraint on the terminal probability distribution along a Gaussian path. This leads to the following observation:

Proposition 1. Under the Affine Gaussian Probability Paths framework, the per sample $\theta_{t}$ can be interpreted as a function of the data target $x_{1}$. The expectation of the running cost $\int_{0}^{1}\frac{1}{2}\|\theta_{t}\|^{2}dt$ provides an upper bound for a constant $C$ times the Kullback-Leibler (KL) divergence between the joint distributions of the data point $x_{1}$ and the prior terminal point $x^{p}$, denoted as $p_{1}(x^{p},x_{1})=p_{1}(x^{p}|x_{1})p_{data}(x_{1})$, and that of the data point $x_{1}$ and the controlled terminal point $x^{\theta}$, denoted as $p_{1}(x^{\theta},x_{1})$:

$$ $\mathbb{E}_{x_{1}\sim p_{data}(x_{1})}\left[\frac{1}{2}\int_{0}^{1}\|\theta_{t }(x_{1})\|^{2}dt\right]\geq C\cdot\text{KL}(p_{1}(x^{\theta},x_{1})\|p_{1}(x^{ p},x_{1})).$ $$

Figure: Algorithm 1 OC-Flow on Euclidean Space

One effective approach for directly applying PMP to optimal control tasks is the Extended Method of Successive Approximations (E-MSA) Li et al. (). E-MSA builds upon the basic MSA algorithm (Chernousko & Lyubushin ()), 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 a limited class of linear quadratic regulators (Aleksandrov ()).

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 $x_{t}^{\theta}$ is bounded, which can be safely assumed provided that appropriate regularization techniques are applied.

When the E-MSA algorithm is applied to the guided controlled generation task as defined earlier, the trajectory of the co-states $\mu_{t}$ can be calculated in closed form. Further, they can be expressed as $\nabla_{x_{t}}\Phi(X_{1}^{k})$. Overall, the following update rule is derived:

Theorem 2: 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 $\gamma>2C$ with $C$ is a function of $L$, such that under the addictive control and the running cost defined in Equation 4, the optimal update is following:

$$ $\theta^{k+1}_{t}=\frac{\gamma}{1+\gamma}\theta^{k}_{t}+\frac{\alpha}{1+\gamma} \nabla_{x_{t}}\Phi(X_{1}^{k}).$ (5) $$

This update rule for the control term $\theta_{t}$ guarantees an increase in the objective function defined in equation 2.

$$ $J(\theta^{k+1})-J(\theta^{k})\geq\left(1-\frac{2C}{\gamma}\right)\epsilon^{k}_ {\gamma},\quad\ \epsilon^{k}_{\gamma}\geq 0$ (6) $$

In practice, continuous ODEs are solved through discretization. The discretized formulation of our algorithm is outlined in Algorithm 1. In this version, the weight decay term is parameterized as: $\beta=\frac{\gamma}{1+\gamma}$ and the learning rate is defined as: $\eta=\frac{\alpha}{1+\gamma}$.

### 3.2 Practical implementation and acceleration

#### 3.2.1 Adjoint method and choice of optimizer

A significant portion of the computational time in the OC-Flow algorithm is spent on evaluating the gradient $\nabla_{x_{t}}\Phi(x_{1})$. The computational cost of directly backpropagating through the model to compute the Jacobian is prohibitive due to the complexity of $f^{p}$, which is parameterized as a deep learning model. Inspired by Liu et al. (), the gradient $\nabla_{x_{t}}\Phi(x_{1})$ can be efficiently computed using a vector-Jacobian product by the double-backwards trick, which is expressed as:

$$ $\nabla_{x_{k/N}}\Phi=\left(\nabla_{x_{(k+1)/N}}\Phi\right)\cdot\Phi_{x_{(k+1)/ N}}\left(x_{k/N}\right).$ $$

Alternatively, since our algorithm does not require a fixed step size and follows a deterministic optimization process, it is compatible with sophisticated optimizers such as LBFGS with line search, which automatically determines a quasi-optimal step size. Note that LBFGS may result in faster convergence but with the cost of increasing memory consumption (Table 1) similar to Ben-Hamu et al. (). It is also not compatible with OC-Flow in SO(3). In our experiments, we observed good convergence with the adjoint approach and only used LBFGS in QM9 where the memory consumption is smaller, for a fair comparison with Ben-Hamu et al. ().

#### 3.2.2 Asynchronous setting for Flexible Update Scheduling

In practice, discretization techniques are employed to simulate the ODEs governing both the state trajectory $x_{t}$ and the co-states $\mu_{t}$ and operate in a synchronous setting, where the number of time steps for the state trajectory $x_{t}$ coincides with the number of control terms $\theta_{t}$.

Here we show that OC-Flow can be extended to an asynchronous framework, providing greater flexibility in scheduling. We subdivide the time interval $\Delta t$ into $N$ equally spaced subintervals. Let $\{x_{t}\}$ denote the state trajectory over the time interval $[t,t+\Delta t]$, and $\{x_{t}^{\theta}\}$ represent the trajectory when the control term $\theta_{t}$ is applied in the first subinterval. The trajectory can be approximated as follows:

$$ $\begin{split}x_{t+\Delta t}=x_{t}+\frac{\Delta t}{N}\sum_{l=1}^{N-1}f^{p}(x^{ \theta}_{t+\frac{l\Delta t}{N}})+\frac{\Delta t}{N}\theta_{t}\approx x_{t}+ \Delta t\left(\frac{1}{N}\sum_{l=1}^{N-1}f^{p}(x_{t+\frac{l\Delta t}{N}})+ \frac{1}{N}\theta_{t}\right)\end{split}$ (7) $$

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.3 for the proof and justification of the approximations in Equation 7). In our peptide design experiment, asynchronous setting is applied for efficient computing.

**Table 1: Comparison of closely related methods for backprop-through guided-ODE in Euclidean manifold. For time complexity, $N$ is the number of steps, $L$ is the number of the layers of the prior flow model. Additionally, our algorithm is the only one allowing a variable number of control terms.**
|  | Effective | Running | Effective Num. of | Memory | Convergence | Generalization |
| --- | --- | --- | --- | --- | --- | --- |
|  | Update | Cost | Control Terms | Complexity | Analysis | to SO(3) |
| OC-Flow (Ours) | $\{\theta_{t}\}_{0}^{N}$ | $\|u_{t}\|^{2}$ | $[1,N]$ | $O(L)$ | ✓ | ✓ |
| OC-Flow-LBFGS(Ours) | $\{\theta_{t}\}_{0}^{N}$ | $\|u_{t}\|^{2}$ | $[1,N]$ | $O(NL)$ | ✓ | ✗ |
| FlowGrad | $\{\theta_{t}\}_{0}^{N}$ | 0 | $N$ | $O(L)$ | ✗ | ✗ |
| D-Flow | $\theta_{0}$ | implicit | 1 | $O(NL)$ | ✗ | ✗ |

### 3.3 connection to other backprop-through guided-ODE approaches

Several previous works discussed backprop-through-ODE guidance in flow-matching models. Notable examples include D-Flow (Ben-Hamu et al., ) and FlowGrad (Liu et al., ). An illustration of their algorithms and ours is shown in Figure 1. 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 in a manner similar to our optimal control target in Equation 2. Specifically, it directly applies gradient descent to the control variables:

$$ $\theta^{k+1}_{t}=\theta^{k}_{t}+\alpha\nabla_{\theta_{t}}\Phi(X_{1}^{k})= \theta^{k}_{t}+\frac{\alpha}{N}(\nabla_{x_{t}}X_{t+1})^{-1}\nabla_{x_{t}}\Phi( X_{1}^{k}),$ $$

which can be interpreted as a limiting case of our algorithm in Equation 5, where $\gamma\to\infty$ and given with $dt\to 0$, $\nabla_{x_{t}}X_{t+1}\to\text{I}$. However, as shown in Equation 6, the convergence rate is a complex function of $\gamma$, so in practice, $\gamma$ is treated as a tunable parameter. FlowGrad’s setting $\gamma\to\infty$ may undermine the convergence speed.

D-Flow optimizes the reward by applying gradient descent to the initial noise $x_{0}$:

$$ $X_{0}^{k+1}=X_{0}^{k}+\text{LBFGS}(\nabla_{x_{0}}\Phi(X_{1}^{k})).$ $$

In fact, with the update rule $x_{\delta t}=x_{0}+f(x_{0})\delta t+\theta_{t}\delta t$, the update of $\theta_{0}$ can be seen as an increment to $x_{0}$. The LBFGS algorithm provides a dynamic learning rate, aligning with our framework, where $\gamma$ 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 is 1. A more detailed compare of the algorithms can be found in Table 1.

## 4 Optimal Control Framework for Guided Flow Matching in SO(3)

Most optimization algorithms are designed for Euclidean spaces and encounter difficulties when extended to non-Euclidean settings, such as the topologically and geometrically complex structure of $\text{SO}(3)$. In this section, we review the current results for optimal control on $\text{SO}(3)$ and extend the applicability of the E-MSA to this manifold, providing 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 $\theta_{t}\in\mathfrak{so}(3)$, evolves according to the following differential equation:

$$ $\dot{x_{t}}^{\theta}=x_{t}^{\theta}\left(f_{t}(x_{t}^{\theta})+\theta_{t}\right)$ (8) $$

Here we employ the left-invariant vector field, which aligns naturally with the structural characteristics of the Lie group $\mathfrak{so}(3)$. This alignment is reflected in the system dynamics described by equation 8, where the tangent space is expressed as $T_{x}\text{SO}(3)=\text{SO}(3)\times\mathfrak{so}(3)$. Under the left-invariant vector field, it can be shown that the Hamiltonian becomes a linear functional in $\mathfrak{so}(3)^{*}$ (Jurdjevic (), Colombo & Dimarogonas ()). With the co-state $\mu_{t}\in\mathfrak{so}(3)^{*}$, the Hamiltonian function $H:[0,T]\times\text{SO}(3)\times\mathfrak{so}(3)^{*}\times\mathfrak{so}(3)\to
\mathbb{R}$ can be redefined as:

$$ $H(t,x,\mu,\theta)=\langle\mu_{t},f_{t}(x_{t},\theta)\rangle-L(\theta)$ (9) $$

Inspired by the Method of Successive Approximations (MSA), a direct approach to apply PMP conditions on the $\text{SO}(3)$ manifold involves iteratively updating the cotangent vector and the state trajectory. At each iteration, we utilize trajectories of $\mu_{t}$ and $x_{t}$ as shown in Step 4 and Step 5 in Algorithm 2, and then apply an update rule to determine the control term $\theta_{t}$ for the subsequent iteration. Drawing inspiration from our algorithm in Euclidean space, with weight decay $\beta$ and learning rate $\eta$ the update for $\theta_{t}$ can be written as:

$$ $\theta_{t}^{k+1}=\beta\theta_{t}^{k}+\eta\tilde{\mu}_{t}^{\theta^{k}},$ $$

where $\tilde{\mu}_{t}^{\theta^{k}}$ is defined by $\langle\tilde{\mu}_{t}^{\theta^{k}},v\rangle=\mu_{t}^{\theta^{k}}(v)$ for all $v\in\mathfrak{so}(3)$. The existence of $\tilde{\mu}_{t}^{\theta^{k}}$ can be derived from the Riesz Representation Theorem (Goodrich ()). This formulation leads to the introduction of the OC-FLow algorithm on $\text{SO}(3)$, which provides an optimization procedure for the task of guided flow matching, as detailed in Algorithm 2.

Figure: Algorithm 2 OC-Flow on $\text{SO}(3)$

### 4.2 Convergence of OC-Flow on SO3

To derive the proof of the convergence of our Algorithm 2, we first establish that under the PMP conditions on $\text{SO}(3)$, the objective function $J(\theta)$, as defined in equation 2, can be bounded. This is formalized in the following proposition:

Proposition 3: 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 $C>0$ such that for any $\theta,\phi\in\mathfrak{so}(3)$, the following inequality holds:

$$ $J(\theta)+\int_{0}^{1}\Delta_{\phi,\theta}H(t)\,dt-C||\phi_{t}-\theta_{t}||^{2 }dt\leq J(\phi),$ (10) $$

where $X^{\theta}$ and $P^{\theta}$ satisfy the PMP conditions on SO(3) manifold, and $\Delta H_{\phi,\theta}$ denotes the change in the Hamiltonian, defined as:

$$ $\Delta H_{\phi,\theta}(t):=H(t,x_{t}^{\theta},\mu_{t}^{\theta},\phi_{t})-H(t,x _{t}^{\theta},\mu_{t}^{\theta},\theta_{t}).$ $$

Proposition 2 provides a 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 $\Delta H_{\phi,\theta}(t)$ is not inherently bounded. Second, the term $||\phi-\theta||^{2}$ is non-negative, which complicates the minimization process. To address these issues, inspired by the method of E-MSA, we introduce a positive constant $\gamma$ and defining an Extended-Hamiltonian:

$$ $\begin{split}\tilde{H}(t,x,\mu,\theta,\phi)&=H(t,x,\mu,\theta)-\frac{\gamma}{2 }||\theta-\phi||^{2}=\langle\mu,f_{t}(x)+\theta\rangle-\frac{1}{2}||\theta||^{ 2}-\frac{\gamma}{2}||\theta-\phi||^{2}.\end{split}$ (11) $$

The introduction of the Extended-Hamiltonian enables the combination of the original Hamiltonian with the penalty term $||\phi-\theta||^{2}$ into a unified expression that can be optimized jointly. A natural approach to achieve this is by updating $\theta$ to maximize the Extended-Hamiltonian. The resulting update rule is given by:

$$ $\theta^{k+1}_{t}=\arg\max_{\theta}\tilde{H}(t,x^{\theta^{k}},\mu^{\theta^{k}}, \theta,\theta^{k})=\frac{\gamma}{1+\gamma}\theta^{k}+\frac{1}{1+\gamma}\mu_{t} ^{\theta}=\beta\theta^{k}+\eta\mu_{t}^{\theta^{k}}$ (12) $$

By performing this maximization step at each iteration, we ensure that the change in the Extended-Hamiltonian, $\Delta\tilde{H}$, is non-negative, indicating that the algorithm progresses towards an optimal solution. Furthermore, we can show that when the update process converges, i.e., when $\Delta\tilde{H}=0$ or equivalently $\Delta H=0$, the algorithm has reached the optimal control point. These insights can be formalized in the following proposition:

Proposition 4: Let $X^{\theta}$ and $P^{\theta}$ satisfy the PMP conditions . If the update rule follows Algorithm 2, we define $\epsilon_{k}:=\int_{0}^{1}\Delta_{\theta^{k+1},\theta^{k}}H(t)\,dt$, and $\epsilon_{k}$ is bounded as:

$$ $\epsilon_{k}:=\int_{0}^{1}\Delta_{\theta^{k+1},\theta^{k}}H(t)\,dt\quad\quad \lim_{k\to\infty}\epsilon_{k}=0.$ (13) $$

Furthermore, when $\epsilon_{k}=0$, we have $\theta=\theta^{*}:=\operatorname*{arg\,max}_{\theta}J(\theta)$

With these results, we can now extend the result in E-MSA to the $\text{SO}(3)$ manifold and establish a bound for the optimization algorithm based on the derived theoretical properties:

Theorem 5: Assume that the reward function, the prior model, and their derivatives satisfy Lipschitz continuity, bounded by a Lipschitz constant $L$. Let $\theta^{0}\in\mathfrak{so}(3)$ be any initial measurable control with $J(\theta^{0})<+\infty$. Suppose also that $\inf_{\theta\in\mathfrak{so}(3)}J(\theta)>-\infty$. If the update of $\theta$ satisfies equation 12, for sufficiently large $\gamma$, the following inequality holds for some constant $D>0$:

$$ $D\epsilon_{k}\leq J(\theta^{k+1})-J(\theta^{k})$ (14) $$

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 $\text{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 $\dot{\mu}_{t}$, which is a linear functional in the dual space $\mathfrak{so}(3)^{*}$. Instead we can optimize $\tilde{\mu}_{t}$ as defined in Section 4.1. We can decompose $\dot{\tilde{\mu}}_{t}$ into its projections onto a set of orthogonal bases within the $\mathfrak{so}(3)$ group.

A frequently used choice for the basis in $\mathfrak{so}(3)$ is the canonical basis $\{E_{1},E_{2},E_{3}\}$ (McCann et al. ()), which satisfies the condition $\langle v,E_{i}\rangle=2$ for all $v\in\mathfrak{so}(3)$. This property allows us to express the time derivative of the adjoint variable $\dot{\tilde{\mu}}_{t}$ as: $\dot{\tilde{\mu}}_{i}=\langle\dot{\tilde{\mu}},E_{i}\rangle$ and $\dot{\tilde{\mu}}=\frac{1}{2}\sum_{i=1}^{3}\dot{\tilde{\mu}}_{i}E_{i}$.
Thus, with the closed forms for the partial derivatives relates to Hamiltonian, the practical update for $\mu_{t}$ in Algorithm 2 can be written as following. Meanwhile the method of asynchronous setting and vjp acceleration can also be applied. See Appendix  C.3.

$$ $\dot{\tilde{\mu}}_{t,i}^{k}=-\langle\tilde{\mu}_{t}^{k},[f_{t}^{p}+\theta_{t}, E_{i}]\rangle-\langle\tilde{\mu}_{t}^{k},\frac{\partial f^{p}_{t}}{\partial x} x_{t}^{k}E_{j}\rangle,$ $$

$$ $\tilde{\mu}_{t-\Delta t}^{k}=\tilde{\mu}_{t}^{k}-\frac{\Delta t}{2}\sum_{i=1}^ {3}\dot{\tilde{\mu}}_{t,i}^{k}E_{i},\quad\tilde{\mu}_{T,i}^{k}=\langle\nabla \Phi(x_{T}^{k}),x_{T}^{k}E_{i}\rangle.$ (15) $$

## 5 Experiments

### 5.1 Text-Guided Image Manipulation

We first validate our OC-Flow on 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 as the reward (Black et al., ; Esser et al., ).
In our text-guided image manipulation task, we want to guide the generative model pre-trained on the celebrity face dataset CelebA-HQ (Karras, ) to three text guidance {sad, angry, curly hair} showing different facial expressions or traits.
Following the same setup in Liu et al. (), given an input image $x_{g}$, the reward for alignment with the text prompt can be effectively evaluated by the CLIP model (Radford et al., ) pre-trained in a contrastive way to score the similarity between arbitrary image-text pairs.
Following (Liu et al., ), we adopt the pre-trained Rectified Flow (RF) (Liu et al., ) as the generative prior.

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

We choose two state-of-the-art text-guided image manipulation baselines, StyleCLIP (Patashnik et al., ) and FlowGrad (Liu et al., ). We run StyleCLIP and FlowGrad with their official implementation and default parameter configurations. For $ours$, we set time step $N=100$, step size $\eta=2.5$, weight decay $=0.995$, and the number of optimization steps $M=15$.

Figure: Figure 2: Visualization of text-guided generated faces with different expressions.
Refer to caption: extracted/5949639/figs/text2img2.png

For qualitative comparison, we show generated examples of different text-guided expressions in Figure 2. Note that the faces in the reference input are mostly smiling. Therefore, due to the large gap between reference and guided distributions, 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.

Figure: Figure 3: Visualization of OC-Flow generated molecules with various dipole moments ($\mu$ / D) condition.
Refer to caption: extracted/5949639/figs/qm9.png

### 5.2 Molecule Generation for QM9

We further instantiate our OC-Flow for controllable molecule generation on the QM9 dataset (Ruddigkeit et al., ; Ramakrishnan et al., ), a commonly used molecular dataset containing small molecules with up to 9 heavy atoms from C, O, N, F. Following Hoogeboom et al. (); Ben-Hamu et al. (), we target for conditional generation of molecules with specified quantum chemical property values including polarizability $\alpha$, orbital energies $\varepsilon_{\text{HOMO}},\varepsilon_{\text{LUMO}}$ and their gap $\Delta\varepsilon$, dipole moment $\mu$, and heat capacity $c_{v}$. 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 et al., ), and the loss can be then calculated as the mean absolute error (MAE) between the predicted and the reference property values. The pre-trained unconstrained generative model is taken from Song et al. () (EquiFM), a 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-based methods of D-Flow (Ben-Hamu et al., ) and FlowGrad (Liu et al., ) 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 3. The unconditional EquiFM is also included as an upper bound for the guided models.
It can be seen that our approach, as a generalization of both D-Flow and FlowGrad, consistently outperforms both of them with lower MAEs.
We provide guided generation samples in Figure 3 with respect to different target dipole moments. A clear trend from hydrocarbons with more symmetric structures to molecules with more high-electronegativity atoms of oxygen and nitrogen can be observed, indicating an increase in the dipole moment.

Figure: Table 3: MAE for guided generations on QM9 (lower is better).

As we have theoretically demonstrated the impact of the regularization strength from the optimal control perspective, we further experiment with a different $\gamma$ and examine the quality of the conditionally generated molecules by evaluating additional metrics following Song et al. (). Specifically, we calculate the atom stability percentage (ASP), molecule stability percentage (MSP), and valid & unique percentage (VUP). Ideally, these metrics should not be greatly lower than the pre-trained model, and a higher strength of regularization should lead to higher scores. Indeed, as demonstrated in Table 4, in which we provide these scores for two different settings of $\gamma=0.01$ and $10$, all scores are higher with a higher strength of regularization at the cost of also a higher MAE. In this way, our OC-Flow effectively prevents exploitation from direct gradient descent that may hack the loss function and provides more flexible and fine-grained control over the guided generation.

**Table 4: MAE and other evaluation metrics for our approach with $\gamma=0.01$ / $\gamma=10$.**
| Property | $\alpha$ | $\Delta\varepsilon$ | $\varepsilon_{\text{HOMO}}$ | $\varepsilon_{\text{LUMO}}$ | $\mu$ | $c_{v}$ |
| --- | --- | --- | --- | --- | --- | --- |
| MAE $\downarrow$ | 1.383 / 1.557 | 367 / 365 | 183 / 188 | 342 / 339 | 0.314 / 0.320 | 0.819 / 0.852 |
| ASP $\uparrow$ | 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 $\uparrow$ | 64.4 / 69.9 | 67.9 / 70.5 | 65.8 / 69.8 | 68.5 / 70.8 | 68.0 / 70.1 | 67.1 / 68.9 |
| VUP $\uparrow$ | 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 et al., ), which includes 158 complexes clustered based on 40% sequence identity using mmseqs2(Steinegger & Söding, ).
Our experiments focus on PepFlow w/Bb, a model designed to exclusively sample peptide backbones while optimizing translations in Euclidean space and rotations in SO(3) space. The model employs the MadraX force field (Orlando et al., ) 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.

Figure: Figure 4: Visualization of OC-Flow generated peptide and unconditional generated peptide (5djd_C).
Refer to caption: extracted/5949639/figs/peptide_result.jpeg

In addition, we employ an existing metric, denoted as IMP, which measures the percentage of generated peptides that exhibit lower energy than the original ground truth. Additionally, we use the root-mean-square deviation (RMSD) to evaluate structural accuracy by aligning the generated peptides to their native structures and calculating the $C_{\alpha}$ RMSD. To further analyze structural characteristics, we compute the secondary-structure similarity ratio (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, ) among the generated peptides, representing their dissimilarities.

We compare our method to the pre-trained unconditional PepFlow model (Li et al., ), 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 metrics 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 can be found in Appendix  B.1.

**Table 5: Evaluation of OC-Flow peptide design.**
|  | MadraX $\AA\downarrow$ | RMSD $\AA\downarrow$ | SSR $\%\uparrow$ | BSR $\%\uparrow$ | Stability $\downarrow$ | Affinity $\downarrow$ | Diversity $\uparrow$ | imp(%) $\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Ground-truth | -0.588 | - | - | - | -84.893 | -36.063 | - | - |
| PepFlow | -0.195 | 1.645 | 0.794 | 0.874 | -45.660 | -26.538 | 0.310 | 14.3 |
| OC-Flow(trans) | -0.229 | 1.774 | 0.797 | 0.876 | -48.380 | -27.328 | 0.323 | 14.4 |
| OC-Flow(rot) | -0.221 | 1.643 | 0.794 | 0.872 | -48.636 | -27.211 | 0.310 | 14.5 |
| OC-Flow(trans+rot) | -0.263 | 2.127 | 0.797 | 0.869 | -48.853 | -27.468 | 0.338 | 15.0 |

## 6 Conclusions and Discussion

In this paper, we propose OC-Flow, a general and theoretically grounded framework for training-free guided flow matching under optimal control formulation. Our framework provides a unified perspective on existing backprop-through-ODE 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 to 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. We also note that our practical algorithms utilize discretization which could lead to gaps in convergence analysis derived under continuous formulation. We leave the discussions of discretization analysis as a future direction for study. We hope that our findings can guide algorithm design and motivate further model improvement in guided flow matching.

## Appendix A Background of Riemannian Manifold and SO(3) Group

A 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 a $C^{\infty}$ differential structure. When $G$ is endowed with a left-invariant Riemannian metric, it becomes a Riemannian manifold, where the inner product of any two tangent vectors $v,w\in T_{h}G$ at a point $h\in G$ is preserved under left multiplication. This property is expressed as:

$$ $\langle L_{h}(v),L_{h}(w)\rangle=\langle v,w\rangle,$ $$

where $L_{h}:G\to G$ is the left multiplication map, and $\langle\cdot,\cdot\rangle:TG\times TG\to\mathbb{R}$ represents the inner product. Moreover, the tangent space at any point $x\in G$ is given by $T_{x}G=L_{x}T_{e}G$, where $T_{e}G$ is the tangent space at the identity element $e$, which is identified with the Lie algebra $\mathfrak{g}$. Consequently, the full tangent bundle $TG$ can be written as $G\times\mathfrak{g}$.

At each point $x\in G$, a tangent space $T_{x}G$ is attached, representing the space of tangent vectors at that point. The collection of these tangent spaces forms the tangent bundle $TG$, which itself is a smooth manifold. Additionally, for any point $h\in G$, the cotangent space $T_{h}^{*}G$ is defined as the dual space of $T_{h}G$, consisting of linear functionals (co-states) that act on the tangent vectors.

Rotation Group $\text{SO}(3)$: The special orthogonal group $\text{SO}(3)$, describing 3D rotations, is a compact 3-dimensional Lie group. Its Lie algebra $\mathfrak{so}(3)$ consists of skew-symmetric matrices. The group $\text{SO}(3)$ is defined as:

$$ $\text{SO}(3)=\{r\in\mathbb{R}^{3\times 3}:r^{\top}r=rr^{\top}=I,\det(r)=1\}$ $$

It is a matrix Lie group, and its Lie algebra is given by:

$$ $\mathfrak{so}(3)=\{r\in\mathbb{R}^{3\times 3}:r^{\top}=-r\}$ $$

Parametrizations of $\text{SO}(3)$: The skew-symmetric matrices $r\in\mathfrak{so}(3)$ can be uniquely represented by a vector $\omega\in\mathbb{R}^{3}$, such that for any $v\in\mathbb{R}^{3}$, $rv=\omega\times v$, where $\times$ denotes the cross product. This vector is known as the rotation vector, where its magnitude $\|\omega\|$ represents the angle of rotation, and its direction $e_{\omega}=\frac{\omega}{\|\omega\|}$ defines the axis of rotation. The mapping from $\mathbb{R}^{3}$ to the skew-symmetric matrix is referred to as the hat operation, $\hat{(\cdot)}$.

Another common parametrization of $\text{SO}(3)$ is through Euler angles, described using three angles $(\phi,\theta,\psi)$. In the x-convention, the rotation is expressed as a sequence of three rotations: a rotation about the $z$-axis by $\phi$, followed by a rotation about the updated $x$-axis by $\theta$, and finally, a rotation about the updated $z$-axis by $\psi$.

One common inner product on $\mathfrak{so}(3)$ is the induced Frobenius inner product, given by:

$$ $\langle A,B\rangle=\text{tr}(A^{\top}B),\quad\forall A,B\in\mathfrak{so}(3),$ $$

which equips $\text{SO}(3)$ with a Riemannian structure. The manifold $\text{SO}(3)$ has constant Gaussian curvature and is diffeomorphic to a solid ball with antipodal points identified. The exponential map $\exp:\mathfrak{so}(3)\to\text{SO}(3)$, originating from the identity element, is defined as matrix exponentiation:

$$ $\exp(A)=\sum_{k=0}^{\infty}\frac{A^{k}}{k!},\quad\forall A\in\mathfrak{so}(3)$ $$

and can be represented more compactly via Rodrigues’ rotation formula:

$$ $\exp(A)=I+\frac{\sin\theta}{\theta}A+\frac{1-\cos\theta}{\theta^{2}}A^{2}, \quad\forall A\in\mathfrak{so}(3)$ $$

where $\theta=\|A\|_{\mathfrak{so}(3)}=\frac{1}{2}\|A\|_{F}$ is the rotation angle. Similarly, the logarithm map $\log:\text{SO}(3)\to\mathfrak{so}(3)$, also originating from the identity, is the matrix logarithm:

$$ $\log(R)=\sum_{k=1}^{\infty}(-1)^{k+1}\frac{(R-I)^{k}}{k}$ $$

or more compactly as:

$$ $\log(R)=\frac{\theta}{\sin\theta}A,\quad\forall R\in\text{SO}(3)$ $$

where $A=\frac{(R-R^{\top})}{2}\in\mathfrak{so}(3)$ and $\theta=\|A\|_{\mathfrak{so}(3)}$ is the rotation angle. In spherical geometry, the geodesic distance between two rotations is given by $d(R_{1},R_{2})=\|\log(R_{1}^{\top}R_{2})\|_{F}$, and interpolation between rotations can be performed using $\exp(tA)$.

## Appendix 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 as well as 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 $\mu_{t}$ such that the Hamiltonian (Equation 3) is maximized (or minimized, depending on the problem) with respect to the control at every time step. Additionally, the state and co-state evolve according to a 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 Maximum Principle
Let $\theta^{*}\in\mathcal{U}$ be an essentially bounded optimal control, i.e. a solution to (2) with $\mathrm{ess}\sup_{t\in[0,T]}\|\theta^{*}_{t}\|_{\infty}\langle\infty$ (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^{*}:[0,T]\to\mathbb{R}^{d}$ such that the Hamilton’s equations

$$ $\dot{X}^{*}_{t}=\nabla_{p}H(t,X^{*}_{t},P^{*}_{t},\theta^{*}_{t}),\quad X_{0}^ {*}=x,$ (16) $$

$$ $\dot{P}^{*}_{t}=-\nabla_{x}H(t,X^{*}_{t},P^{*}_{t},\theta^{*}_{t}),\quad P_{T} ^{*}=\nabla\Phi(X_{T}^{*}),$ (17) $$

are satisfied. Moreover, for each $t\in[0,T]$, we have the Hamiltonian maximization condition

$$ $H(t,X^{*}_{t},P^{*}_{t},\theta^{*}_{t})\geq H(t,X^{*}_{t},P^{*}_{t},\theta) \quad\text{for all }\theta\in\Theta.$ (18) $$

The PMP conditions can be naturally generalised to the Lie Group. Pontryagin’s Maximum Principle (PMP) for Lie groups (Saccon et al. ()) provides the conditions that govern the flow of the state $x_{t}$ and its associated cotangent flow $\lambda$, describing the Hamiltonian equations that the optimal state-adjoint trajectory must satisfy. Specifically, the Hamiltonian equations are given by:

$$ $X_{t}^{-1}\dot{X}_{t}^{*}=\frac{\partial}{\partial p}H(t,X_{t}^{*},P_{t}^{*}, \theta_{t}^{*}),\quad X_{0}^{*}=x,$ $$

$$ $\dot{\mu}_{t}^{\theta}=-\mathrm{ad}^{*}_{\frac{\partial H}{\partial\mu}}\mu_{t }^{\theta}-(dL_{x_{t}^{\theta}})^{*}\frac{\partial}{\partial x}H^{*},\quad\mu_ {T}^{\theta}=(dL_{x_{T}^{\theta}})^{*}\nabla_{x}\Phi(x_{T}^{\theta}).$ (19) $$

The dual map $(dL_{g})^{\star}:T_{gh}^{\star}\mathrm{SO}(3)\to T_{h}^{\star}\mathrm{SO}(3)$ pulls back a cotangent vector at $gh$ to a cotangent vector at $h$. The coadjoint representation $\mathrm{ad}^{*}_{X}$ acts on $\mathfrak{so}(3)^{*}$ and is defined as:

$$ $\langle\mathrm{ad}^{*}_{X}\mu,Y\rangle=-\langle\mu,\mathrm{ad}_{X}Y\rangle,$ $$

where $\mathrm{ad}_{X}Y=[X,Y]=XY-YX$ for $\mu\in\mathfrak{so}(3)^{*}$ and $X,Y\in\mathfrak{so}(3)$. Additionally, for each $t\in[0,T]$, the Hamiltonian maximization condition is satisfied:

$$ $H(t,X_{t}^{*},\mu_{t}^{*},\theta_{t}^{*})\geq H(t,X_{t}^{*},\mu_{t}^{*},\theta )\quad\text{for all }\theta\in\Theta.$ (20) $$

### B.2 Extended-Method 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) Chernousko & Lyubushin (), 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:

$$ $\dot{X}_{t}^{*}=f(t,X_{t}^{*},\theta^{*}),$ $$

Given an initial guess $\theta^{0}\in\mathcal{U}$ for the optimal control, for each iteration $k=0,1,2,\dots$, we first solve the state dynamics:

$$ $\dot{X}_{t}^{\theta^{k}}=f(t,X_{t}^{\theta^{k}},\theta_{t}^{k}),\quad X_{0}^{ \theta^{k}}=x,$ $$

to obtain $X^{\theta^{k}}$, followed by solving the co-state equation:

$$ $\dot{\mu}_{t}^{\theta^{k}}=-\nabla_{x}H(t,X_{t}^{\theta^{k}},\mu_{t}^{\theta^{ k}},\theta_{t}^{k}),\quad\mu_{T}^{\theta^{k}}=-\nabla\Phi(X_{T}^{\theta^{k}}),$ $$

to determine $\mu^{\theta^{k}}$. Finally, the control is updated using the maximization condition:

$$ $\theta_{t}^{k+1}=\arg\max_{\theta\in\Theta}H(t,X_{t}^{\theta^{k}},\mu_{t}^{ \theta^{k}},\theta),$ $$

for $t\in[0,T]$. This process is summarized in Algorithm 3.

Figure: Algorithm 3 Basic MSA

#### B.2.2 E-MSA

E-MSA introduces the augmented Hamiltonian

$$ $\tilde{H}(t,x,\mu,\theta,v,q):=H(t,x,\mu,\theta)-\frac{1}{2}\rho\|v-f(t,x, \theta)\|^{2}-\frac{1}{2}\rho\|q+\nabla_{x}H(t,x,\mu,\theta)\|^{2}.$ $$

Then, they define the following set of alternative necessary conditions for optimality:

Proposition 3 (Extended PMP) Suppose that $\theta^{*}$ is an essentially bounded solution to the optimal control problem (2). Then, there exists an absolutely continuous co-state process $\mu^{*}$ such that the tuple $(X^{*},\mu^{*},\theta^{*})$ satisfies the necessary conditions

$$ $\dot{X}^{*}_{t}=\nabla_{\mu}\tilde{H}(t,X^{*}_{t},\mu^{*}_{t},\theta^{*}_{t}, \dot{X}^{*}_{t},\mu^{*}_{t}),\quad X^{*}_{0}=x$ $$

$$ $\dot{P}^{*}_{t}=-\nabla_{x}\tilde{H}(t,X^{*}_{t},\mu^{*}_{t},\theta^{*}_{t}, \dot{X}^{*}_{t},\mu^{*}_{t}),\quad\mu^{*}_{T}=-\nabla_{x}\Phi(X^{*}_{T})$ $$

$$ $\tilde{H}(t,X^{*}_{t},\mu^{*}_{t},\theta^{*}_{t},\dot{X}^{*}_{t},\mu^{*}_{t}) \geq\tilde{H}(t,X^{*}_{t},\mu^{*}_{t},0,\dot{X}^{*}_{t},\mu^{*}_{t}),\quad \theta\in\Theta,t\in[0,T]$ $$

The key contribution is that the control terms $\theta$ can be updated by iteration using the Extended-PMP as shown below. Meanwhile, it is proven that under this update rule, for each iteration, the target function $J(\theta)$ is non-deceasing.

Figure: Algorithm 4 Extended MSA

## Appendix C Proofs and Theorems

### C.1 Proof for Proposition 1

We restate Proposition 1 here:

Proposition 1. Under the Affine Gaussian Probability Paths framework, the per sample $\theta_{t}$ can be interpreted as a function of the data target $x_{1}$. The expectation of the running cost $\int_{0}^{1}\frac{1}{2}\|\theta_{t}\|^{2}dt$ provides an upper bound for a constant $C$ times the Kullback-Leibler (KL) divergence between the joint distributions of the data point $x_{1}$ and the prior terminal point $x^{p}$, denoted as $p_{1}(x^{p},x_{1})=p_{1}(x^{p}|x_{1})p_{data}(x_{1})$, and that of the data point $x_{1}$ and the controlled terminal point $x^{\theta}$, denoted as $p_{1}(x^{\theta},x_{1})$:

$$ $\mathbb{E}_{x_{1}\sim p_{data}(x_{1})}[\frac{1}{2}\int_{0}^{1}\|\theta_{t}(x_{ 1})\|^{2}dt]\geq C\cdot\text{KL}(p_{1}(x^{\theta},x_{1})\|p_{1}(x^{p},x_{1})).$ $$

Proof
For Affine Gaussian Probability Paths, the conditional flow distribution of the prior vector field can be written as: $P^{p}_{t}(x|x_{1})=N(\mu_{t}(x_{1}),\sigma_{t}(x_{1})^{2}I)$, the conditional vector field is written as:

$$ $\phi_{t}(x)=\mu_{t}(x_{1})+\sigma_{t}(x_{1})x$ $$

with its dynamics:

$$ $\dot{\phi}_{t}(x)=\dot{\mu_{t}}(x_{1})+\dot{\sigma_{t}}(x_{1})x$ $$

Or

$$ $u_{t}(x|x_{1})=\dot{\mu_{t}}(x_{1})+\frac{\dot{\sigma_{t}}}{\sigma_{t}}(x-\mu_ {t}(x_{1}))$ $$

Now we add the addictive control term and assume it only alters the mean of the distribution by denoting it as $\theta_{t}(x_{1})$ by noticing it is not proportional to $x$:

$$ $\dot{\phi}_{t}(x)=\theta_{t}(x_{1})+\dot{\mu_{t}}(x_{1})+\dot{\sigma_{t}}(x_{1 })x$ $$

The resulting pushing-forward prob:

$$ $P^{\theta}_{t}(x)=\int P^{\theta}_{t}(x|x_{1})q(x_{1})dx_{1}$ $$

$$ $P^{u}_{t}(x|x_{1})=N(\mu_{t}^{\theta}(x_{1}),\sigma_{t}(x_{1})^{2}I)$ $$

Usually, besides maximising the reward, we hope the new distribution not too much away from the original distribution. One common constraint on the terminal state distribution is the the Kullback-Leibler (KL) divergence.

Inspired by Variational Flow Matching (Eijkelboom et al. ()), we can (KL) divergence between the joint distributions of the data point $x_{1}$ and the prior terminal point $x^{p}$, denoted as $p_{1}(x^{p},x_{1})=p_{1}(x^{p}|x_{1})p_{data}(x_{1})$, and that of the data point $x_{1}$ and the controlled terminal point $x^{\theta}$, denoted as $p_{1}(x^{\theta},x_{1})$:$KL(P_{1}^{p}(x,x_{1})||P_{1}^{u}(x,x_{1}))$

The KL term can be simplified as:

$$ $\displaystyle KL(P_{1}^{p}(x,x_{1})||P_{1}^{\theta}(x,x_{1}))$ $\displaystyle=\int\int P^{p}_{1}(x|x_{1})q(x_{1})\text{log}(\frac{P^{p}_{1}(x| x_{1})p_{data}(x_{1})}{P^{\theta}_{1}(x|x_{1})p_{data}(x_{1})})dxdx_{1}$ $\displaystyle=E_{x_{1}\sim p_{data}(x_{1})}[KL(P_{1}^{p}(x|x_{1})||P_{1}^{ \theta}(x|x_{1}))]$ $$

Given $P_{1}^{p}(x|x_{1})$ and $P_{1}^{\theta}(x|x_{1})$ are two Gaussians with same variance but different mean, and when we consdier the constraint per sample, the constraint can be written as:

$$ $\frac{1}{2}\frac{||\mu_{1}^{\theta}(x_{1})-\mu_{1}^{p}(x_{1})||^{2}}{\sigma_{1 }(x_{1})^{2}}$ $$

Given:

$$ $||\mu_{1}^{\theta}(x_{1})-\mu_{1}^{p}(x_{1})||^{2}=||\int_{0}^{1}\theta_{t}dt| |^{2}\leq\int_{0}^{1}\|\theta_{t}\|^{2}dt$ $$

Therefore, we have the following inequality:

$$ $\mathbb{E}_{x_{1}\sim p_{data}(x_{1})}[\frac{1}{2}\int_{0}^{1}\|\theta_{t}(x_{ 1})\|^{2}dt]\geq\frac{1}{4\sigma_{1}(x_{1})^{2}}\cdot\text{KL}(p_{1}(x^{\theta },x_{1})\|p_{1}(x^{p},x_{1}))$ $$

### C.2 Proof for Theorem 2

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 $\tilde{H}$:

$$ $H(t,x,\mu,\theta)=\mu_{t}\cdot f(x,t)+\mu_{t}\cdot\theta_{t}-\frac{1}{2}|| \theta_{t}||^{2}$ $$

$$ $\tilde{H}(t,x,\mu,\theta,\dot{x},\dot{\mu})=\mu_{t}\cdot f(x,t)+\mu_{t}\cdot \theta_{t}^{k+1}-\frac{1}{2}||\theta_{t}^{k+1}||^{2}-\frac{\gamma}{2}||\theta_ {t}^{k+1}-\theta_{t}^{k}||^{2}$ $$

Apply Extended MSA:

$$ $\dot{x}^{k}_{t}=\theta^{k}_{t}+f(x_{t}^{k},t)$ $$

$$ $\dot{\mu_{t}^{k}}=-\nabla_{x}H(t,x^{k},\mu^{k},\theta^{k})=-\nabla_{x}f(x_{t}^ {k},t)\mu_{t}^{k}$ $$

$$ $\mu_{1}^{k}=\alpha\nabla_{x_{1}}\Phi(x^{k}_{1})$ $$

$$ $\theta^{k+1}_{t}=\text{argmax}_{\theta_{t}}\tilde{H}(t,x,\mu,\theta,\dot{x}, \dot{\mu})$ $$

$\mu_{t}^{k}$ can be calculated in closed form as below, with $\mathcal{T}\exp$ is the time-order exponential:

$$ $\mu_{t}^{k}=\mathcal{T}\exp\left(\int_{t}^{1}\nabla_{x}f(x^{k}_{s},s)ds\right) \cdot\alpha\nabla_{x_{1}}\Phi(X_{1}^{k})$ $$

Thus, the update rule of $\theta_{t}$ is:

$$ $\theta^{k+1}_{t}=\frac{\gamma}{1+\gamma}\theta^{k}_{t}+\frac{\alpha}{1+\gamma} \mathcal{T}\exp\left(\int_{t}^{1}\nabla_{x}f(x^{k}_{s},s)ds\right)\cdot\nabla_ {x_{1}}\Phi(X_{1}^{k})$ $$

Further, the time-order exponential term can be simplified as following:

Evans () provides method to calculate $p(t)=\nabla_{x_{t}}x(1)$ efficiently by defining the adjoint and using the following ODEs:

$$ $\dot{p}(t)=-\nabla_{x}f(x^{k}_{t},t)p(t)$ $$

$$ $p(1)=\nabla_{x(1)}x(1)=I$ $$

Sakurai & Napolitano () provides an alternative viewpoint of the ajoint ODE, they show there would be a closed form solution:

$$ $\dot{p}(t)=A(t)p(t)$ $$

$$ $p(t)=\mathcal{T}\exp\left[-\int_{t}^{1}A(s)ds\right]p(1)$ $$

Thus, combine the closed form solution and the adjoint representation, we find that:

$$ $\nabla_{x_{t}}x(1)=\mathcal{T}\exp\left[\int_{t}^{1}\nabla_{x}f(x^{k}_{s},s)ds\right]$ $$

Given the linearity of the ODEs, multiply a constant to the terminal by change $p(t)=\nabla_{x_{t}}x(1)$ into $p(t)=\nabla_{x_{t}}x(1)\nabla_{x_{1}}\Phi(x_{1})$, the conclusions would not change, we have:

$$ $\nabla_{x_{t}}x(1)\nabla_{x_{1}}\Phi(x_{1})=\mathcal{T}\exp\left[\int_{t}^{1} \nabla_{x}f(x^{k}_{s},s)ds\right]\nabla_{x_{1}}\Phi(x_{1})$ $$

Therefore, our update rule is in fact:

$$ $\theta^{k+1}_{t}=\frac{\gamma}{1+\gamma}\theta^{k}_{t}+\alpha\nabla_{x_{t}} \Phi(X_{1}^{k})$ $$

The solutions can be naturally generalised to the space of $R^{N}$ if we use $Tr(A^{T}B)$ to replace $A\cdot B$.

### C.3 Asynchronous setting

In practice, as described by equation 15, discretization techniques are employed to simulate the ordinary differential equations (ODEs) governing both the state trajectory $x_{t}$ and the corresponding cotangent vector $\mu_{t}$. Most existing methods, as well as the algorithm presented earlier, operate under a synchronous setting, where the number of time steps for the state trajectory $x_{t}$ matches the number of control terms $\theta_{t}$.

However, OC-Flow can be extended to an asynchronous framework by approximation to allow greater flexibility in update scheduling.

Rather than employing the standard update rule $x_{t+\Delta t}=x_{t}+f(t,x_{t},\theta_{t})\Delta t$, we subdivide the time interval $\Delta t$ into $N$ equally spaced subintervals, applying the control term $\theta_{t}$ only during the first subinterval. The update for $x_{t+\Delta t}$ is given by:

$$ $x_{t+\Delta t}=x_{t}+\frac{\Delta t}{N}f(x_{t},\theta_{t})+\frac{\Delta t}{N}f ^{p}(x^{\theta}_{t+\frac{\Delta t}{N}})+\cdots+\frac{\Delta t}{N}f^{p}(x^{ \theta}_{\frac{(N-1)\Delta t}{N}})$ (21) $$

Moreover, for intermediate steps, we define:

$$ $x_{t+\frac{i\Delta t}{N}}=x_{t}+\frac{\Delta t}{N}f(x_{t},\theta_{t})+\sum_{l= 1}^{i-1}\frac{\Delta t}{N}f^{p}(x^{\theta}_{t+\frac{l\Delta t}{N}}),$ $$

where $x_{t+\frac{i\Delta t}{N}}$ denotes the state at the $i$-th subinterval.

Recall that $f(x_{t},\theta_{t})=f^{p}(x_{t})+\theta_{t}$, when we have $\frac{\Delta t}{N}$ is small enough, the update rule in equation 7 can be approximated as a case in Equation 4 by considering $\nabla_{\theta}x_{t+\Delta t}$:

$$ $\begin{split}\nabla_{\theta}x_{t+\Delta t}=&\frac{\Delta t}{N}+\frac{\Delta t} {N}\nabla_{\theta}f^{p}(x^{\theta}_{t+\frac{\Delta t}{N}})+\cdots+\frac{\Delta t }{N}\nabla_{\theta}f^{p}(x^{\theta}_{t+\frac{(N-1)\Delta t}{N}})\\ =&\frac{\Delta t}{N}+(\frac{\Delta t}{N})^{2}\nabla_{x}f^{p}(x^{\theta}_{t+ \frac{\Delta t}{N}})+\cdots+(\frac{\Delta t}{N})^{N}\nabla_{x}f^{p}(x^{\theta} _{t+\frac{(N-1)\Delta t}{N}})\\ =&\frac{\Delta t}{N}+O\left((\frac{\Delta t}{N})^{2}\right)\end{split}$ $$

Therefore, if we denote $x_{t}$ as the trajectory of the state variable $x$ over the time interval $[t,t+\Delta t]$, and $x_{t}^{\theta}$ as the trajectory when the control term $\theta_{t}$ is applied in the first subinterval, it can be reasonably approximated as:

$$ $x_{t+\Delta t}\approx x_{t}+\frac{\Delta t}{N}\sum_{l=1}^{i-1}f^{p}(x_{t+\frac {l\Delta t}{N}})+\frac{\Delta t}{N}\theta_{t}.$ $$

As a result, for asynchronous setting, the step 4 in algorithm 1 should be modified as:

$$ $X_{t+\Delta t}^{\theta^{k}}=\left(\frac{1}{N}\sum_{l=1}^{i-1}f^{p}(x_{t+\frac{ l\Delta t}{N}})+\frac{1}{N}\theta^{k}\right)\Delta t$ $$

For the case on the $\text{SO}(3)$ manifold, the asynchronous setting can be deployed using the Taylor expansion of the matrix exponential $\exp(A)$, and noting that when $\frac{\Delta t}{N}$ and $\Delta t$ are sufficiently small, the terms become commutative. We can derive the approximation as follows:

$$ $\begin{split}x_{t+\Delta t}=&\,x_{t}\exp\left(\frac{\Delta t}{N}f(x_{t},\theta _{t})\right)\exp\left(\frac{\Delta t}{N}f^{p}(x^{\theta}_{t+\frac{\Delta t}{N} })\right)\cdots\exp\left(\frac{\Delta t}{N}f^{p}(x^{\theta}_{\frac{(N-1)\Delta t }{N}})\right)\\ \approx&\,x_{t}\exp\left(\frac{\Delta t}{N}\sum_{l=1}^{i-1}f^{p}(x_{t+\frac{l \Delta t}{N}})+\frac{\Delta t}{N}\theta_{t}\right).\\ \approx&\,x_{t}\exp\left(\Delta t(f^{p}_{e}(x_{t})+\frac{1}{N}\theta_{t}) \right).\end{split}$ $$

The vector-Jacobian method can also be applied to compute the term $\frac{\partial f^{p}_{t}}{\partial x}(x_{t}^{k}E_{j})$ in Algorithm 2:

$$ $\langle\tilde{\mu}_{t},\frac{\partial f^{p}_{t}}{\partial x}x_{t}^{k}E_{j} \rangle=\mathrm{Tr}\left(\tilde{\mu}_{t}^{T}\frac{\partial f^{p}_{t}}{\partial x }(x_{t}^{k}E_{j})\right).$ $$

### C.4 Proof for Proposition 3

The key inequality we used for the proof of proposition 3 is following:

$$ $\langle h,v\rangle\leq\|h\|\|v\|\quad\quad\|x\|=1$ $$

for $h,v\in\mathfrak{so}(3)$ and $x\in\text{SO(3)}$
Before prove the proposition 3, we firstly show that the cotangent vector $\{\mu_{t}\}$ is also bounded.

Lemma2 Assume all functions satisfy Lipschitz condition. Then, there exists a constant $K^{\prime}\rangle 0$ such that for any $\theta$,

$$ $\|\mu_{t}^{\theta}\|\leq K^{\prime},$ $$

for all $t\in[0,T]$.
Proof Using necessary condition and setting $\tau:=T-t$, $\tilde{\mu}_{\tau}^{\theta}:=\mu_{T-\tau}^{\theta}$ we get

$$ $\dot{\tilde{\mu}}_{\tau}^{\theta}=ad_{\frac{\partial H}{\partial\mu}}^{*} \tilde{\mu}_{\tau}^{\theta}+(dL_{x})^{*}(\nabla_{x}\langle\tilde{\mu}_{\tau}^{ \theta},f\rangle)\quad\tilde{\mu}_{0}^{\theta}=(dL_{x_{T}^{\theta}})^{*}\nabla \Phi(x_{T}^{\theta})$ $$

With Lipschitz condition, we have $\|\tilde{\mu}_{0}^{\theta}\|\leq\|\nabla\Phi(x_{T}^{\theta})\|\|x_{T}^{\theta}
\|\leq K$ and $\|\nabla_{x}f(t,x_{t}^{\theta},\theta_{t})\|\leq K$ and $\|ad_{\frac{\partial H}{\partial\mu}}^{*}\mu\|\leq\|\frac{\partial H}{\partial
\mu}\|\|\mu\|\leq K\|\mu\|$. Hence,

$$ $\|\dot{\tilde{\mu}}_{\tau}^{\theta}\|\leq K\|\tilde{\mu}_{\tau}^{\theta}\|,$ $$

and

$$ $\|\tilde{\mu}_{\tau}^{\theta}\|-\|\tilde{\mu}_{0}^{\theta}\|\leq\|\tilde{\mu}_ {\tau}^{\theta}-\tilde{\mu}_{0}^{\theta}\|\leq\int_{0}^{t}\|\dot{\tilde{\mu}}_ {s}^{\theta}\|ds\leq\int_{0}^{t}(K\|\tilde{\mu}_{\tau}^{\theta}\|)ds$ $$

and by Gronwall’s inequality,

$$ $\|\tilde{\mu}_{\tau}^{\theta}\|\leq\|\tilde{\mu}_{0}^{\theta}\|e^{KT}=:K^{ \prime}.$ $$

This proves the claim since it holds for any $\tau$.

Now given all related terms are bounded, we can prove the proposition 2.

Proposition 3 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 $C\rangle 0$ such that for any $\theta,\phi\in\mathfrak{so}(3)$, the following inequality holds:

$$ $J(\theta)+\int_{0}^{1}\Delta_{\phi,\theta}H(t)\,dt-C||\phi_{t}-\theta_{t}||^{2 }dt\leq J(\phi),$ (22) $$

where $X^{\theta}$ and $P^{\theta}$ satisfy the PMP conditions in equation 19, and $\Delta H_{\phi,\theta}$ denotes the change in the Hamiltonian, defined as:

$$ $\Delta H_{\phi,\theta}(t):=H(t,x_{t}^{\theta},\mu_{t}^{\theta},\phi_{t})-H(t,x _{t}^{\theta},\mu_{t}^{\theta},\theta_{t}).$ $$

Proof:
Firstly, by the definition of Hamiltonina and PMP conditions, we always have:

$$ $I(x^{\theta},\mu_{t}^{\theta},\theta):=\int_{0}^{T}\langle\mu_{t}^{\theta},f_{ t}^{\theta}\rangle-H(t,x_{t}^{\theta},\mu_{t}^{\theta},\theta)-L(\theta)dt\equiv 0$ $$

Define $\delta\mu_{t}=\mu_{t}^{\phi}-\mu_{t}^{\theta}$ and $\delta f_{t}=f_{t}^{\phi}-f_{t}^{\theta}$, the difference in $I$ can be decomposed as:

$$ $\displaystyle 0\equiv$ $\displaystyle I(x^{\phi},\mu_{t}^{\phi},\phi)-I(x^{\theta},\mu_{t}^{\theta},\theta)$ $\displaystyle=$ $\displaystyle\int_{0}^{T}\langle\mu_{t}^{\theta},\delta f_{t}\rangle+\langle \delta\mu_{t},f_{t}^{\theta}\rangle+\langle\delta\mu_{t},\delta f_{t}\rangle dt$ $\displaystyle-$ $\displaystyle\int_{0}^{T}H(t,x_{t}^{\phi},\mu_{t}^{\phi},\phi)-H(t,x_{t}^{ \theta},\mu_{t}^{\theta},\theta)dt$ $\displaystyle-$ $\displaystyle\int_{0}^{T}L(\phi_{t})-L(\theta_{t})dt$ $$

Now define $U(t)=\int_{0}^{t}f_{s}ds=log(x_{t}x_{0}^{-1})$ and by integrating by parts:

$$ $\int_{0}^{T}\langle\mu_{t}^{\theta},\delta f_{t}\rangle dt=\langle\mu_{t}^{ \theta},\delta U_{t}\rangle|_{0}^{T}-\int_{0}^{T}\langle\dot{\mu_{t}^{\theta}} ,\delta U_{t}\rangle dt$ $$

$$ $\int_{0}^{T}\langle\delta\mu_{t},\delta f_{t}\rangle dt=\langle\delta\mu_{t}, \delta U_{t}\rangle|_{0}^{T}-\int_{0}^{T}\langle\dot{\delta\mu_{t}},\delta U_{ t}\rangle dt$ $$

Now we have:

$$ $\begin{split}&\int_{0}^{T}\langle\mu_{t}^{\theta},\delta f_{t}\rangle+\langle \delta\mu_{t},f_{t}^{\theta}\rangle dt\\ =&\langle\mu_{t}^{\theta},\delta U_{t}\rangle|_{0}^{T}+\int_{0}^{T}\langle \delta\mu_{t},\frac{\partial H^{\theta}}{\partial\mu}\rangle+\langle\mu_{t}^{ \theta},ad_{\frac{\partial H^{\theta}}{\partial\mu}}\delta f_{t}\rangle+ \langle\frac{\partial H^{\theta}}{\partial x},\delta f_{t}x_{t}^{\theta} \rangle dt\\ \end{split}$ $$

Similarly, we get:

$$ $\begin{split}\int_{0}^{T}\langle\delta\mu_{t},\delta f_{t}\rangle dt=&\frac{1} {2}\int_{0}^{T}\langle\delta\mu_{t},\delta f_{t}\rangle dt+\frac{1}{2}\int_{0} ^{T}\langle\delta\mu_{t},\delta f_{t}\rangle dt\\ =&\frac{1}{2}\langle\delta\mu_{t},\delta U_{t}\rangle|_{0}^{T}-\frac{1}{2}\int _{0}^{T}\langle\dot{\delta\mu_{t}},\delta U_{t}\rangle dt+\frac{1}{2}\int_{0}^ {T}\langle\delta\mu_{t},\delta f_{t}\rangle dt\\ =&\frac{1}{2}\langle\delta\mu_{t},\delta U_{t}\rangle|_{0}^{T}+\frac{1}{2}\int _{0}^{T}\langle\mu_{t}^{\phi},ad_{\frac{\partial H^{\phi}}{\partial\mu}}\delta U _{t}\rangle-\langle\mu_{t}^{\theta},ad_{\frac{\partial H^{\theta}}{\partial\mu }}\delta U_{t}\rangle dt\\ &+\frac{1}{2}\int_{0}^{T}\langle(dL_{x_{t}^{\phi}})^{*}\frac{\partial H^{\phi} }{\partial x}-(dL_{x_{t}^{\theta}})^{*}\frac{\partial H^{\theta}}{\partial x}, \delta U_{t}\rangle dt\\ &+\frac{1}{2}\int_{0}^{T}\langle\delta\mu_{t},\frac{\partial H^{\phi}}{ \partial\mu}-\frac{\partial H^{\theta}}{\partial\mu}\rangle dt\\ \end{split}$ $$

With mean value theorem and $x,\mu$ are bounded by constant $L$, we can always find $x_{t}^{\gamma}$ between $x_{t}^{\phi}$ and $x_{t}^{\theta}$, $\mu_{t}^{\gamma}$ between $\mu_{t}^{\phi}$ and $\mu_{t}^{\theta}$, $\gamma$ between $\phi$ and $\theta$, so that :

$$ $\begin{split}&\int_{0}^{T}\langle(dL_{x_{t}^{\phi}})^{*}\frac{\partial}{ \partial x}H(t,x_{t}^{\phi},\mu_{t}^{\phi},\phi)-(dL_{x_{t}^{\theta}})^{*} \frac{\partial}{\partial x}H(t,x_{t}^{\theta},\mu_{t}^{\theta},\theta),\delta U _{t}\rangle dt\\ =&\int_{0}^{T}\langle(dL_{x_{t}^{\theta}})^{*}\frac{\partial}{\partial x}H(t,x _{t}^{\theta},\mu_{t}^{\theta},\phi)-(dL_{x_{t}^{\theta}})^{*}\frac{\partial}{ \partial x}H(t,x_{t}^{\theta},\mu_{t}^{\theta},\theta),\delta U_{t}\rangle dt \\ &+\int_{0}^{T}\langle\nabla_{x}((dL_{x_{t}^{\gamma}})^{*})x_{t}^{\gamma}\delta U _{t}\nabla_{x}H(t,x_{t}^{\phi},\mu_{t}^{\phi},\phi),\delta U_{t}\rangle dt\\ &+\int_{0}^{T}\langle(dL_{x_{t}^{\theta}})^{*}\nabla_{x}^{2}H(t,x_{t}^{\gamma} ,\mu_{t}^{\phi},\phi)x_{t}^{\gamma}\delta U_{t},\delta U_{t}\rangle dt\\ &+\int_{0}^{T}\langle(dL_{x_{t}^{\theta}})^{*}\nabla_{\mu}\nabla_{x}H(t,x_{t}^ {\theta},\mu_{t}^{\gamma},\phi)\delta\mu_{t},\delta U_{t}\rangle dt\\ \leq&\int_{0}^{T}\langle(dL_{x_{t}^{\theta}})^{*}\frac{\partial}{\partial x}H( t,x_{t}^{\theta},\mu_{t}^{\theta},\phi)-(dL_{x_{t}^{\theta}})^{*}\frac{ \partial}{\partial x}H(t,x_{t}^{\theta},\mu_{t}^{\theta},\theta),\delta U_{t} \rangle dt\\ &+C\int_{0}^{T}||\delta U_{t}||^{2}+||\delta\mu_{t}||||\delta U_{t}||dt\\ \end{split}$ $$

Using the same method we get:

$$ $\begin{split}\int_{0}^{T}\langle\delta\mu_{t},\delta f_{t}\rangle dt\leq&\frac {1}{2}\langle\delta\mu_{t},\delta U_{t}\rangle|_{0}^{T}+\frac{1}{2}\int_{0}^{T }\langle\mu_{t}^{\phi},ad_{\frac{\partial H^{\phi}}{\partial\mu}}\delta U_{t} \rangle-\langle\mu_{t}^{\theta},ad_{\frac{\partial H^{\theta}}{\partial\mu}} \delta U_{t}\rangle dt\\ &+\frac{1}{2}\int_{0}^{T}\langle\frac{\partial}{\partial x}H(t,x_{t}^{\theta}, \mu_{t}^{\theta},\phi)-\frac{\partial}{\partial x}H(t,x_{t}^{\theta},\mu_{t}^{ \theta},\theta),\delta U_{t}x_{t}^{\theta}\rangle dt\\ &+\frac{1}{2}\int_{0}^{T}\langle\frac{\partial}{\partial\mu}H(t,x_{t}^{\theta} ,\mu_{t}^{\theta},\phi)-\frac{\partial}{\partial\mu}H(t,x_{t}^{\theta},\mu_{t} ^{\theta},\theta),\delta\mu_{t}\rangle dt\\ &+C\int_{0}^{T}||\delta U_{t}||^{2}+||\delta\mu_{t}||||\delta U_{t}||dt\\ \end{split}$ $$

With boundary conditions:

$$ $\begin{split}&\langle\mu_{t}^{\theta}+\frac{1}{2}\delta\mu_{t},\delta U_{t} \rangle|_{0}^{T}=\langle\mu_{T}^{\theta}+\frac{1}{2}\delta\mu_{T},\delta U_{T} \rangle\\ =&\langle(dL_{x_{t}^{\theta}})^{*}\nabla\Phi(x_{T}^{\theta}),\delta U_{t} \rangle+\frac{1}{2}\langle(dL_{x_{t}^{\phi}})^{*}\nabla\Phi(x_{T}^{\phi})-(dL_ {x_{t}^{\theta}})^{*}\nabla\Phi(x_{T}^{\theta}),\delta U_{t}\rangle\\ \leq&\Phi(x_{T}^{\phi})-\Phi(x_{T}^{\theta})+K||\delta U_{T}||^{2}\\ \end{split}$ $$

Using same method to $H(t,x_{t}^{\phi},\mu_{t}^{\phi},\phi)-H(t,x_{t}^{\theta},\mu_{t}^{\theta},\theta)$, Thus we obtain:

$$ $\begin{split}&[\Phi(x_{T}^{\theta})+\int_{0}^{T}L(\theta_{t})]-[\Phi(x_{T}^{ \phi})+\int_{0}^{T}L(\phi_{t})]\\ \leq&K||\delta U_{T}||^{2}-\int_{0}^{T}\Delta H_{\phi,\theta}(t)dt+\frac{1}{2} \int_{0}^{T}\langle\mu_{t}^{\phi},ad_{\frac{\partial H^{\phi}}{\partial\mu}} \delta U_{t}\rangle-\langle\mu_{t}^{\theta},ad_{\frac{\partial H^{\theta}}{ \partial\mu}}\delta U_{t}\rangle dt\\ &\quad+\frac{1}{2}\int_{0}^{T}\langle\frac{\partial}{\partial x}H(t,x_{t}^{ \theta},\mu_{t}^{\theta},\phi)-\frac{\partial}{\partial x}H(t,x_{t}^{\theta}, \mu_{t}^{\theta},\theta),x_{t}^{\theta}\delta U_{t}\rangle dt\\ &\quad+\frac{1}{2}\int_{0}^{T}\langle\frac{\partial}{\partial\mu}H(t,x_{t}^{ \theta},\mu_{t}^{\theta},\phi)-\frac{\partial}{\partial\mu}H(t,x_{t}^{\theta}, \mu_{t}^{\theta},\theta),\delta\mu_{t}\rangle dt\\ &\quad+C\int_{0}^{T}||\delta U_{t}||^{2}+||\delta\mu_{t}||||\delta U_{t}||dt\\ \end{split}$ $$

By definition:

$$ $\delta U_{t}=\int_{0}^{T}f(t,x_{T}^{\phi},\phi)-f(t,x_{T}^{\theta},\theta)dt$ $$

and so

$$ $\begin{split}||\delta U_{t}||\leq&\int_{0}^{t}||f(s,x_{s}^{\phi},\phi)-f(s,x_{ s}^{\theta},\theta)||ds\\ \leq&\int_{0}^{t}||f(s,x_{s}^{\phi},\phi)-f(s,x_{s}^{\theta},\phi)||ds\\ &+\int_{0}^{t}||f(s,x_{s}^{\theta},\phi)-f(s,x_{s}^{\theta},\theta)||ds\\ \leq&\int_{0}^{T}||f(s,x_{s}^{\theta},\phi)-f(s,x_{s}^{\theta},\theta)||ds\\ &+K\int_{0}^{t}||\delta U_{s}||ds\end{split}$ $$

By Gronwall’s inequality:

$$ $||\delta U_{t}||\leq e^{KT}\int_{0}^{T}||f(s,x_{s}^{\theta},\phi)-f(s,x_{s}^{ \theta},\theta)||ds$ $$

To estimate $\delta\mu$, we use the same substitution as in Lemma 6 with $\tau=T-t$,we get:

$$ $\begin{split}\delta\tilde{\mu_{\tau}}&=\delta\tilde{\mu_{0}}+\int_{0}^{\tau}( dL_{\tilde{x}_{s}^{\phi}})^{*}\nabla_{x}H(s,\tilde{x}_{s}^{\phi},\tilde{\mu}_{ s}^{\phi},\phi)-(dL_{\tilde{x}_{s}^{\theta}})^{*}\nabla_{x}H(s,\tilde{x}_{s}^{ \theta},\tilde{\mu}_{s}^{\theta},\theta)ds\\ &\quad\quad\quad+\int_{0}^{\tau}ad_{\frac{\partial H}{\partial\mu}}^{*}\tilde{ \mu}_{\tau}^{\phi}-ad_{\frac{\partial H}{\partial\mu}}^{*}\tilde{\mu}_{\tau}^{ \theta}ds\end{split}$ $$

Using Lemma 1 and Liptichitz conditions:

$$ $\begin{split}||\delta\tilde{\mu_{\tau}}||&\leq||\delta\tilde{\mu_{0}}||+\int_{ 0}^{\tau}||(dL_{\tilde{x}_{s}^{\phi}})^{*}\nabla_{x}H(s,\tilde{x}_{s}^{\phi}, \tilde{\mu}_{s}^{\phi},\phi)-(dL_{\tilde{x}_{s}^{\theta}})^{*}\nabla_{x}H(s, \tilde{x}_{s}^{\theta},\tilde{\mu}_{s}^{\theta},\theta)||ds\\ &\quad\quad\quad+\int_{0}^{\tau}||ad_{\frac{\partial H}{\partial\mu}}^{*} \tilde{\mu}_{\tau}^{\phi}-ad_{\frac{\partial H}{\partial\mu}}^{*}\tilde{\mu}_{ \tau}^{\theta}||ds\\ &\leq K||\delta U_{T}||+KK^{\prime}\int_{0}^{T}||\delta U_{t}||dt+K\int_{0}^{ \tau}||\delta\tilde{\mu}_{s}||ds\\ &\quad\quad\quad+\int_{0}^{T}||(dL_{\tilde{x}_{s}^{\theta}})^{*}\nabla_{x}H(s, \tilde{x}_{s}^{\theta},\tilde{\mu}_{s}^{\theta},\phi)-(dL_{\tilde{x}_{s}^{ \theta}})^{*}\nabla_{x}H(s,\tilde{x}_{s}^{\theta},\tilde{\mu}_{s}^{\theta}, \theta)||ds\\ &\leq e^{KT}K(||\delta U_{T}||+K^{\prime}\int_{0}^{T}||\delta U_{t}||dt)\\ &\quad\quad+e^{KT}K\int_{0}^{T}||\nabla_{x}H(s,\tilde{x}_{s}^{\theta},\tilde{ \mu}_{s}^{\theta},\phi)-\nabla_{x}H(s,\tilde{x}_{s}^{\theta},\tilde{\mu}_{s}^{ \theta},\theta)||ds\end{split}$ $$

Using the bound of $U_{t}$, we obtain:

$$ $\begin{split}||\delta\tilde{\mu_{\tau}}||&\leq K^{\prime\prime}(\int_{0}^{T}|| f(s,x_{s}^{\theta},\phi)-f(s,x_{s}^{\theta},\theta)||ds)\\ &\quad\quad+e^{KT}K\int_{0}^{T}||\nabla_{x}H(s,\tilde{x}_{s}^{\theta},\tilde{ \mu}_{s}^{\theta},\phi)-\nabla_{x}H(s,\tilde{x}_{s}^{\theta},\tilde{\mu}_{s}^{ \theta},\theta)||ds\end{split}$ $$

Also we obtain:

$$ $\begin{split}&\frac{1}{2}\int_{0}^{T}\langle\mu_{t}^{\phi},ad_{\frac{\partial H ^{\phi}}{\partial\mu}}\delta U_{t}\rangle-\langle\mu_{t}^{\theta},ad_{\frac{ \partial H^{\theta}}{\partial\mu}}\delta U_{t}\rangle dt\\ =&\frac{1}{2}\int_{0}^{T}\langle ad_{\frac{\partial H^{\theta}}{\partial\mu}}^ {*}\mu_{t}^{\phi}-ad_{\frac{\partial H^{\theta}}{\partial\mu}}^{*}\mu_{t}^{ \theta},\delta U_{t}\rangle dt\\ \leq&C\int_{0}^{T}||\delta\mu_{t}||||\delta U_{t}||dt\end{split}$ $$

Finally we get:

$$ $\begin{split}J(\theta)-J(\phi)&\leq-\int_{0}^{T}\Delta H_{\phi,\theta}(t)dt\\ &\quad+\frac{1}{2}K^{\prime\prime}\|\delta U_{T}\|^{2}\\ &\quad+K^{\prime\prime}\int_{0}^{T}\left(\|\delta U_{t}\|^{2}+\|\delta\mu_{t} \|^{2}\right)dt\\ &\quad+\frac{1}{2}\int_{0}^{T}\|\delta\mu_{t}\|\|f(t,x_{t}^{\theta},\phi_{t})- f(t,x_{t}^{\theta},\theta_{t})\|dt\\ &\quad+\frac{1}{2}\int_{0}^{T}\|\delta U_{t}\|\|\nabla_{x}H(t,x_{t}^{\theta}, \mu_{t}^{\theta},\phi_{t})-\nabla_{x}H(t,x_{t}^{\theta},\mu_{t}^{\theta}, \theta_{t})\|dt\\ &\leq-\int_{0}^{T}\Delta H_{\phi,\theta}(t)dt\\ &\quad+C\left(\int_{0}^{T}\|f(t,x_{t}^{\theta},\phi_{t})-f(t,x_{t}^{\theta}, \theta_{t})\|dt\right)^{2}\\ &\quad+C\left(\int_{0}^{T}\|\nabla_{x}H(t,x_{t}^{\theta},\mu_{t}^{\theta},\phi _{t})-\nabla_{x}H(t,x_{t}^{\theta},\mu_{t}^{\theta},\theta_{t})\|^{2}dt\right) ^{2}\\ &\leq-\int_{0}^{T}\Delta H_{\phi,\theta}(t)dt\\ &\quad+C\int_{0}^{T}\|f(t,x_{t}^{\theta},\phi_{t})-f(t,x_{t}^{\theta},\theta_{ t})\|^{2}dt\\ &\quad+C\int_{0}^{T}\|\nabla_{x}H(t,x_{t}^{\theta},\mu_{t}^{\theta},\phi_{t})- \nabla_{x}H(t,x_{t}^{\theta},\mu_{t}^{\theta},\theta_{t})\|^{2}dt\end{split}$ $$

Therefore, given the form of the addictive control terms and the running cost, we can derive the final term of our claim:

$$ $J(\theta)+\int_{0}^{1}\Delta_{\phi,\theta}H(t)\,dt-C||\phi_{t}-\theta_{t}||^{2 }dt\leq J(\phi),$ (23) $$

### C.5 Proof for Proposition 4

Due to the similarity between the bound derived in Proposition 3 and the bound obtained from the E-MSA method, the proofs of Proposition 4 and Theorem 5 follow the same reasoning as outlined in Section 3.3 of Li et al. (). For the sake of completeness, we provide the full derivations here.

Proposition 4 Let $X^{\theta}$ and $P^{\theta}$ satisfy the PMP conditions in equation 19. If the update rule follows equation 12, we define $\epsilon_{k}:=\int_{0}^{1}\Delta_{\theta^{k+1},\theta^{k}}H(t)\,dt$, and $\epsilon_{k}$ is bounded as:

$$ $\epsilon_{k}:=\int_{0}^{1}\Delta_{\theta^{k+1},\theta^{k}}H(t)\,dt\quad\quad \lim_{k\to\infty}\epsilon_{k}=0.$ (24) $$

Furthermore, when $\epsilon_{k}=0$, we have $\theta=\theta^{*}:=\operatorname*{arg\,max}_{\theta}J(\theta)$To establish convergence, define

$$ $\epsilon_{k}:=\int_{0}^{T}\Delta H_{\theta^{k+1},\theta^{k}}(t)dt\geq 0.$ $$

Proof By definition, if $\epsilon_{k}=0$, then from the update rule which maximizes the Hamiltonian, we must have

$$ $0=-\epsilon_{k}\leq-\frac{\gamma}{2}\int_{0}^{1}\|\theta^{k+1}-\theta^{k}\|^{2 }dt\leq 0.$ $$

and so

$$ $\max_{\theta}\tilde{H}(x_{t}^{\theta^{k}},\mu_{t}^{\theta^{k}},\theta,\dot{x}_ {t}^{\theta^{k}},\dot{\mu}_{t}^{\theta^{k}})=\tilde{H}(x_{t}^{\theta^{k}},\mu_ {t}^{\theta^{k}},\theta_{t}^{k},\dot{x}_{t}^{\theta^{k}},\dot{\mu}_{t}^{\theta ^{k}}),$ $$

Therefore, we always have the quantity $\epsilon_{k}\geq 0$ and it measures
the distance from the optimal solution, and if it equals 0, then we reach the optimum.

### C.6 Proof for Theorem 5

Theorem 5: Assume that the reward function, the prior model, and their derivatives satisfy Lipschitz continuity, bounded by a Lipschitz constant $L$. Let $\theta^{0}\in\mathfrak{so}(3)$ be any initial measurable control with $J(\theta^{0})<+\infty$. Suppose also that $\inf_{\theta\in\mathfrak{so}(3)}J(\theta)>-\infty$. If the update of $\theta$ satisfies equation 12, for sufficiently large $\gamma$, the following inequality holds:

$$ $D\epsilon_{k}\leq J(\theta^{k+1})-J(\theta^{k})$ (25) $$

for some constant $D>0$

Proof Using Proposition 3 we have

$$ $J(\theta^{k})-J(\theta^{k+1})\leq-\epsilon_{k}+C\int_{0}^{T}\|\theta^{k+1}- \theta^{k}\|^{2}dt$ $$

From the Algorithm 2 maximazaing step, we know that

$$ $H(t,X_{t}^{\theta^{k}},P_{t}^{\theta^{k}},\theta_{t}^{k})\leq H(t,X_{t}^{ \theta^{k}},P_{t}^{\theta^{k}},\theta_{t}^{k+1})-\frac{\gamma}{2}\|\theta^{k+1 }-\theta^{k}\|^{2}$ $$

Hence, we have

$$ $J(\theta^{k})-J(\theta^{k+1})\leq-(1-\frac{2C}{\gamma})\epsilon_{k}.$ $$

Pick $\gamma>2C$, then we shall have $J(\theta^{k})-J(\theta^{k+1})\leq-D\epsilon_{k}$ with $D=(1-\frac{2C}{\gamma})>0$.

Moreover, we can rearrange and sum the above expression to get

$$ $\sum_{k=0}^{M}\epsilon_{k}\leq D^{-1}\left(J(\theta^{M+1})-J(\theta^{0})\right )\leq D^{-1}\left(\inf_{\theta\in\mathcal{U}}J(\theta)-J(\theta^{0})\right),$ $$

and hence $\sum_{k=0}^{\infty}\epsilon_{k}<+\infty$, which implies $\epsilon_{k}\to 0$ and the algorithm converges to the optimum.

## Appendix D Experimental Details

### D.1 Text-Guided Image Manipulation

In our text-to-image generation experiment, we adopted the pipeline presented in Liu et al. (), utilizing the generative prior from Liu et al. (). We employed standard evaluation metrics: LPIPS and ID (face identity similarity) as introduced in Kim & Ye () to assess the differences between the original image and the manipulated image. Additionally, the CLIP score was used to evaluate the alignment between the generated image and the provided text prompt.

To enforce consistency with the original image, we introduced a constraint term to the terminal reward function to penalize significant deviations from the original image:

$$ $\Phi(x_{1})=\lambda\text{CLIP}(x_{1},T)-(1-\eta)|x_{1}-x_{1}^{p}|$ $$

Here, the hyperparameter $\lambda$ 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 $\eta$ results in greater emphasis on the terminal reward, leading to a higher CLIP score but lower LPIPS and ID scores. The weight decay is a function of $\gamma$, which is tuned to maximize $\left(1-\frac{2C}{\gamma}\right)\epsilon^{k}_{\gamma}$ for iteration k. In this experiment, due to the limitation of storage, we set $\gamma$ the same for all k. In our implementation, the learning rate $\eta$ was set to 2.5, and the weight decay $\beta$ was set to 0.995.

Baseline configurations were aligned with those reported in Liu et al. (), and the results presented in Table 2 reflect the same experimental conditions. For quantitative comparison, we used the CelebA dataset, randomly sampling 1,000 images, which were manipulated based on text guidance: {old, sad, smiling, angry, curly hair}.

### D.2 Molecule Generation

In our QM9 generation experiment, we mostly followed the conditional generation pipeline in Hoogeboom et al. (). An equivariant geometric GNN was trained for each property on half of the QM9 data as the classifier, which was then frozen during our training-free controlled generation. The EquiFM (Song et al., ) checkpoint, trained on the whole QM9 training data, was loaded as the generative prior. The test time properties were sampled from the whole training dataset, making it slightly different from the settings in Ben-Hamu et al. (). Therefore, we reimplemented the D-Flow algorithm with 5 optimizer steps and 5 inner steps each with a 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.

For all properties, MAE was used as the optimization target and $\gamma$ is the regularization coefficient such that $\gamma\int_{0}^{1}\|\theta_{t}\|^{2}dt$ is the additional OC loss. For all optimization methods, we always used a fixed number of 50 Euler steps so $\theta$ can be indexed by discrete indices. As the integral is done with a step size of 1/50, any $\gamma$ is effectively $\tilde{\gamma}=2\gamma/50=0.04\gamma$ when taking the derivative with respect to $\theta$ or $x$. Therefore, $\gamma=10$ effectively corresponds to $\tilde{\gamma}=0.4$, which is still a valid optimization scheme.

### D.3 Peptide Design

In our peptide experiments, we adopted PepFlow (Li et al., ) as the baseline model, utilizing the pre-trained checkpoint provided in the original PepFlow paper. The test dataset split was also based on the one defined in the PepFlow framework. For hyperparameter tuning, we randomly selected 10 samples from the dataset. After tuning, the full set of 162 samples was used for training and evaluation to ensure comprehensive performance assessment.

To achieve flexible update scheduling, we employed an asynchronous setting in OC-Flow. This approach maintains the same time step as PepFlow while providing finer updates to both control terms and state trajectories. The total simulation was conducted with 10 intervals, each further divided into 20 subintervals. This design enables improved flexibility and finer granularity in updates without increasing the number of overall time steps.

We applied the pre-trained model as the initialization for our experiments, allowing us to build upon the pre-trained weights and achieve consistent performance improvements through hyperparameter adjustments. The final hyperparameters were selected based on optimal performance during tuning and were applied across the full dataset. In OC-Flow(rot), we used $\alpha=0.95$ and $\beta=0.8$; in OC-Flow(trans), $\alpha=0.9$ and $\beta=1.2$; and in OC-Flow(both), $\alpha=0.95$ and $\beta=2.0$.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: See weaknesses.
soundness: 2
strengths: - Well-motivated problem, overall nicely written paper with clear literature review.
- The methodological and theoretical parts of the paper are well-sounded. 
- Providing a framework that has convergence analysis is always welcomed.
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.
weaknesses: **Major: questionable and inconsistent baselines' results in empirical benchmarks** 

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

Ben-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.

Li 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.

---

confidence: 2
contribution: 4
initial_rating: 6
presentation: 3
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? 
- How does OC-Flow compare to recent works in Riemannian FM? What about SO(3) and SE(3)?
- 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?
- Since OC-Flow is designed to be computationally efficient, could the authors comment on real-time applications requiring 'immediate' guidance?
soundness: 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.
- 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. 
- Unlike existing approaches, OC-Flow allows training-free guidance, making it computationally efficient and more applicable in real-life settings. 
- The SO(3) results such as on improved molecular generation accuracy, validate the importance of using this geometric inductive bias for generative models. 
- 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.
- The model consistently demonstrates improvement over previous work.
- I really appreciate the figures included in the paper to illustrate and compare the methods against existing approaches.
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​.
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.
- 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.
- 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.
- 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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
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.
* 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.
* 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.
soundness: 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.
* The paper provides a legitimate contribution to formalizing and extending existing guided-flow matching techniques to complex geometries (such as SO(3))
* The authors compare the proposed methodology to similar existing methods to demonstrate competitive performance
* Table 1 gives a good comparison to understand the contribution of OC-flow vs D-Flow and Flow-Grad
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.
weaknesses: __Theoretical Concerns__
* 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.

__Experimental Concerns__
* 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.
* 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.

* 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.

__Presentation Concerns__
* 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.
  * 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.
  * 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.
  * 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.
* 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.
* 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.

__Contribution Concerns__:
* 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.
* 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.
* Additionally, the experimental reproducibility of the paper is quite poor, with no experimental parameters given, and no experimental source code provided.
*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.


__Citations__
* Fjelde, T., Mathieu, E., & Dutordoir, V.. (2024). An Introduction to Flow Matching.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
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$)?

GradFlow 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?

I’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.
soundness: 3
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. 

To 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.

The approach generalizes existing methods that optimize the trajectory of ODE-samplers (D-Flow, GradFlow), and outperforms them in multiple benchmarks.
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.
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.

While 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.

I 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.

Meta-review: Among other things, the reviewers have highlighted the nice framing of guided flow matching as an optimal control problem, the contribution to formalizing and extending the existing guided-flow matching techniques to SO(3), and the applicability of the proposed approach to various domains such as image and molecular data. Furthermore, it was highlighted that the framing of existing work as special cases under the optimal control formulation helps clarify connections between different methods. Finally, the results were found promising.

As important areas for improvement, both reviewer faUf and jxTp had questions regarding the experimental setup that required clarification, and reviewer faUf raised concerns about the results of baseline methods and how they compare against what is reported previously in the literature. Another important area for improvement suggested by multiple reviewers was to provide more details on runtime comparisons. Other points include questions around how the theoretical assumptions hold in practice (raised by two reviewers). In particular, reviewer jxTp raised questions on the use of discretization in practice while the theory provided relies on a continuous framework. Furthermore, reviewer jxTp suggested that an ablation study on optimization in SO(3) would be helpful to demonstrate the additional benefits over Euclidian optimization. Finally, reviewer jxTp also asked for a clearer exposition of the contributions of the paper, while other reviewers found this clearly illustrated. Finally, several reviewers provided suggestions for clarifications and explanations to improve the accessibility of the paper, they 

The two major concerns were addressed in the rebuttal and revised version of the paper. The concerns on reproducibility and details of the experimental setup were addressed by sharing anonymized code, and including additional details in the revised paper, as well as explanations in the rebuttal itself. The second major point regarding runtime comparison and scalability was addressed by including a discussion on time and memory complexity in appendix D, as well as practical runtime and memory usage. The majority of the other concerns were addressed as well. The authors have included a discussion and theoretical analysis of the discretization gap in the revised paper. The authors have also incorporated the suggestion to include an ablation study to demonstrate the benefit of optimization in SO(3).  Finally, the authors have also taken into account several suggestions by the reviewers to improve the clarity of the paper. While reviewer jxTp still finds there to be room for improvement on the exposition of the paper, they have indicated that several other concerns have been addressed and raised their score.

In summary, most of the concerns initially raised by the reviewers were addressed through fruitful discussions between the authors and reviewers, and revisions in the submitted paper. Given multiple raised scores, there is now unanimous agreement among reviewers that this paper should be accepted. Therefore I recommend to accept this paper.

Decision: Accept (Poster)