Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference
Authors: Denis Blessing
Julius Berner
Lorenz Richter
Carles Domingo-Enrich
Yuanqi Du
Arash Vahdat
Gerhard Neumann
Venue: NeurIPS 2025 spotlight ()
Keywords: Stochastic optimal control
sampling
fine-tuning of diffusion models
Primary area: probabilistic_methods
TL;DR: We develop trust region methods for stochastic optimal control to improve sampling from unnormalized densities, transition path sampling, and diffusion model finetuning.

Abstract: Solving stochastic optimal control problems with quadratic control costs can be viewed as approximating a target path space measure, e.g. via gradient-based optimization. In practice, however, this optimization is challenging in particular if the target measure differs substantially from the prior. In this work, we therefore approach the problem by iteratively solving constrained problems incorporating trust regions that aim for approaching the target measure gradually in a systematic way. It turns out that this trust region based strategy can be understood as a geometric annealing from the prior to the target measure, where, however, the incorporated trust regions lead to a principled and educated way of choosing the time steps in the annealing path. We demonstrate in multiple optimal control applications that our novel method can improve performance significantly, including tasks in diffusion-based sampling and fine-tuning of diffusion models.

Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference
Denis Blessing Julius Berner*. Lorenz Richter*::4 Carles Domingo-Enrich
Yuanqi Du Arash Vahdat²
Gerhard Neumann'
Karlsruhe Institute of Technology ?NVIDIA 3Zuse Institute Berlin +dida Microsoft Research New England "Cornell University
Abstract
Solving stochastic optimal control problems with quadratic control costs can be viewed as approximating target path space measure, e.g. via gradient-based optimization. In practice, however, this optimization is challenging the target measure differs substantially from the prior. In this work, we therefore approach the problem by iteratively solving constrained problems incorporating trust regions that aim for approaching the target measure gradually in a systematic geomet- way. It turns that this trust region based strategy can be understood ric annealing from the prior to the target measure, where, however, the incorporated trust regions lead to principled and educated way of choosing the time steps in the annealing path. We demonstrate in multiple optimal control applications that our novel method can improve performance significantly. including tasks in diffusion- based sampling, transition path sampling, and fine-tuning of diffusion models.
particular if
Introduction
Even though the theory of stochastic optimal control (SOC) dates back several decades [12, 49], has recently attracted renewed interest within the machine learning community. Building on novel formulations that well-suited for gradient-based optimization (see [40] for overview) and drawing connections to diffusion models [15, 36, 92], recent work has led to significant progress in the numerical approximation of high-dimensional control problems using neural networks [42, 87]. Related problems are crucial many practical applications, ranging from sampling problems (e.g.. in statistical physics [48, 68]. Bayesian statistics [54, 85], and reinforcement learning [24]) to fine- tuning of diffusion models [38, 41, 132]. In this work, we aim to further advance SOC approximation methods by taking inspiration from trust region methods used in optimization [1, 88, 93, 110, 123], resulting in principled framework from the perspective of measure transport in path space.
Stochastic optimal control. SOC problems (with quadratic control costs) describe optimization problems of the form
[ (+ s) ds + s(x?)
with fxx: = (b ou) (X", s)ds as)ddW. Xo ~
(I) where optimizes the control of the stochastic differential equation (SDE). Since the law of the SDE solution X induces a so-called path measure TP' on the space continuous trajectories (specifying how likely certain trajectory is), finding the optimal control is equivalent to finding an optimal target path space measure Q. From the SOC literature it is known that the likelihood of
min IE wEEl
CCoresponneecce to denis. www.ssiggkii...eee.ee.......eeee edu
Equal contribution.
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

Q w.r.t. P" can be expressed in closed-form (see [34] and (3) below), which allows to minimize divergences D(P", Q2) via mradient-basee optimization (also termed iterative diffisiin optimization). Trust region methods. However, if the target Q is rather different from the initialization IP (typically the uncontrolled process with 40 = 0), many algorithms face challenges with high variances or mode discovery when directly minimizing D(P", Q), especially in high dimensions. To this end, we propose to approach the target measure gradually by a sequence (PP%) )a- where in the i-th step we add the constraint (TI being the approximated optimal control from the previous iteration and & > O a chosen trust region bound. We prove that the intermediate measures IP', define geometric annealing between the prior I and target measure Q, where the annealing step-sizes are chosen optimally, in the sense of having an approximately constant change in Fisher-Rao distance (Props. 2.2 and 2.3). Finding an optimal annealing schedule is paramount for the convergence speed of many measure transport and sampling methods [119], and understanding physical processes [30, 106]. While the direct computation of Fisher-Rac distances can be challenging, we show that trust region methods lead sense. Moreover, we
to the cost functional (1), with 04-1
to simple way of obtaining equidistant steps in an information-geoee show that the Lagrangian of the constrained problem can be written as another SOC problem and that the optimal Lagrangian multiplier can be obtained via a dual optimization problem without additional computational overhead (Sec. 2.1). Finally, adapt successful approaches based on SOC matching [41, 42] and log- log-variance divergences [87] to the constrained SOC problem to get a practical algorithm (Sec. 2.2).
Applications. The resulting frust region stochastic optimal control method can be viewed as an extension of various existing algorithms, yielding significant improvements on a range of applica- tions (Sec. 3). In particular, we consider (i) deep learning approaches to classical SOC problems (extending [42, 87I) enabling the usage of cross-entropy losses in high dimensions, (ii) diffusion- based sampling from unnormalized densities (extending [97, 129]) enabling efficient sampling from high-dimensional, multimodal densities with substantially fewer target evaluations,
(iii) transition path sampling in molecular dynamics (extending [72, 113)) yielding notably higher transition hit rates, and (iv) reward fine- iin-tuning of text-to-image models (extending [41]) achieving comparable performance while requiring significantly fewer simulations,
Contributions. Our contributions can be summarized as follows:
We develop a general framework for solving measure transport with trust regions and apply it to SOC problems using iterative diffusion optimization.
We prove that our framework leads to a sequence of SOC problems whose solutions define an equispaced annealing between initialization and optimum w.r.t. the Fisher-Rao distance.
Relying on different loss functionals, we propose two practical instantiations of our framework and demonstrate state-of-the-art performance unnormalized densities to transition path sampling and reward fine-tuning of text-to-image models.
series of applications, ranging from sampling from
Notation. We denote by U C C(R x [0, TT; Id)) the set of admissible controls and by P the set of all probability measures on C((O, R We define the path space measure IP E P as the law of a - stochastic process X = (Xt)1e(o.21 and we denote by P,, the marginal distribution at time 8. We refer to App. A for further details on our notation and assumptions.
2 Trust region constrained measure transport for optimal control
The idea of iterative diffusion optimization in optimal control based on path space measures is to consider loss functionals of the form
((u) = D(P", Q)
(2)
and minimize them with gradient-descent algorithms [87]. The loss functional (2) yields imple- mentable algorithms for SOC problems since the optimal path measure Q of (1) can be stated explicitly via the Radon-Nikodym derivative
dQ
c -WC.0)
with t(x,t)= -. J." s) ds g(Xr).
(X)
(3)
AIP
2(Xo)
where 2 = E le- W(X00)
Xo) and IP is the path measure of the uncontrolled process X X see App. D. In this work, we extend this attempt by using trust regions that shall make sure that the
"Note that the cost functional (1) corresponds (up to the normalizing constant) to the reverse Bullback-Leible (KL) divergence D) DKL.

---

optimization is conducted in more "regulated" fashion, where the essential idea is to divide the global problem into smaller (reasonably chosen) chunks. We quantify this in Prop. 2.3 below. To this end, we consider the iterative optimization scheme defined by
M++1 arg min PKL (T" s.t. DKL(P"IP"))
(4)
for any E N, where E > 0 defines a trust region w.r.t to the previous control iterate and where we often set We 0 (and thus I = IP. This corresponds to dividing the overall optimization problem into parts according to their distance measured in the KL. divergence between the respective preceding and succeeding path measures. Due to the convexity of the KI. divergence, we can show that in all but the last step we actually have an equality constraint in (4); see App. E.1. Thus, there exists an I E N such that 2 = u² is the optimal control of the global control problem defined in (I).
Remark 2.1 (Controlling the variance of importance weights). The constraint DKL (Pt [P) S E can be motivated by the goal to control the variance of importance weights Varp", (dP /dP":.. which can be explained by the inequality Varp", (dP /dIP") 2 eDac (1*))
see, e.g., [60]. For small E (which is a common choice in practice) we typically observe /dP")) /dlP%
par 1d assuming that dIP": //P" x 1. Low variance of importance weights is directly related to effi- ciency of many measure transport methods and too high variance makes it practically impossible to obtain reliable results. Note also that the reverse KL divergence allows for explicit expressions for the resulting constrained problem (see Sec. 2. and we leave alternative divergences for future research.
= 2: (see App. I.),, which can be explained by Taylor expansion and
In practice, under suitable regularity assumptions, problem can ap- proach the above constrained optimization problem using a relaxed Lagrangian formalism. To this end, we consider the loss functionals
Daa (Pa (pr)) S
C (u, A) = DKL (P) + X (DKKL =),,
(5)
where A > 0 is a Lagrange multiplier, and solve the saddle point problems
max min (()) (u, A). X20 WEU
(6)
We note that CTA is convex in u by convexity of the KL diver- gence (see App. E. and concave in X since can be expressed as the pointwise minimum min,, c(0) (u, A) among a family of linear functions of A. Thus, (6) has unique optima which we denote by U4+1 and A; respectively. We can now show the following evolution of the optimal measures.
Figure 1: Illustration of a sequence of distributions (IP":); resulting from our trust region method (orange) and a mea- sure transport corresponding to non- equispaced geometric annealing (blue), leading to high variance in the impor- tance weights for the initial steps.
Proposition 2.2 (Optimal change of measure as geometric annealing). Let Q2 be the optimal path measure defined in (3). The intermediate optimal path measures corresponding to (4) then satisfy²
dPP
dQ
(7)
dlP
dP";
and the optimal change of measure w.r.t. the base measure is given by*
dP4
dQ
/dP"
/dIP
1-8.
x
dP
(x)
(x)
with
8=
1+x
(8)
Proof. The first statement follows by the definition of the Lagrangian and the second follows by induction; see App. B.
Note that the sequence (B,), is monotonically increasing with values in [0, 1]. where we have Bo 0 and B3 = (as \ = 0 due to optimality). Thus, the formula in (8) can be seen as geometric annealing from the prior to the target measure. Note that when = 0, the second factor vanishes.
Importantly, the step-size of the annealing is automatically chosen such that we obtain a well-behaved sequence of distributions; see also Fig. I.
Proposition 2.3 (Equidistant steps on statistical manifold). Up to higher order terms in s, the sequence of measures P i E IO...., 1). are equispaced in the Fisher-Rao distance.
'For notational convenience we assume an Xo- -g-independent normalizing constant here and hereafter, which is possible whenever the optimal tilting of the initial density po is known, cf. App. D.3. As usual, the empty product defined such that Bo 0.
3

---

Proof. By Prop. 2.2. we obtain (KL(PP (P IP') I(8.) ( +OAA), where A, 34+1 - 3, and [(3e) is the Fisher information. The Fisher-Rao distance between IP,, and IP is then given Bu+1 (I( dr = 1aar-vIBB8 [[ + 0 = V2 + 0(4 3/2 ); see App. for details.
by B.
n
Remark 2.4 (Trust regions for general measures). The observant reader has likely noticed that far all our arguments do not rely on the fact that we consider path space measures, but work for general probability measures. We could therefore as well write our trust region method stated in (4) as
TP i++1 = arg min DEL (TP)Q) S.L. (P IP:) S.
(9)
PEP
We refer to App. H for a treatment when the measures admit densities on R2, which can, e.g., be considered for variational inference with normalizing flows.
2.1 Constrained stochastic optimal control
While the above formulation in principle works for arbitrary measures, in this work we focus on path space measures corresponding to optimal control problems. In this setting we can compute some of the objectives more explicitly and recover helpful relations.
Lagrangian as SOC problem. First, note that, using the Girsanov theorem (see App. A.2), it turns out that, for a fixed Lagrange multiplier X, the Lagrangian in (5) defines another SOC problem, i.c.
C A) Cfhc (u,
dε,
(10)
where.
(i) A) AT) [[," (LAI
A 1 uil + 2(म ||us|] + +)) (xx s) ds 9(X7) + log (Xx)
(11)
and X is still defined as in (I); see App. E.4 for details. Note that this cost functional more general than the one stated in (1), which one recovers when setting À = 0. We can show that the corresponding SOC problem satisfies the following optimality conditions.
Proposition 2.5 (Optimality for trust region SOC problems). For fixed X, let us define by
J. (2 - ux + ュリーх ||u:ll +) (X, 5) ds g(X7) X. =
(x, 1) inf E MEN
the value function of the SOC problem infreu (8) (TRC(U, A) corresponding to (11) and by wiY1 its solution. Then it holds
(i) (Estimator for value function) V - A) T W((05,)))) X = =
where (I(X" t) =f J.¹ (l s)) S. s) dW. + ((X :t.
(ii) (Connection between solution and value function) It holds uit > I+X Us 1
Proof. The statements can be proven using the verification theorem; see App. E.4 for details.
We note that Prop. 2.2, the Girsanov theorem, and (3) relate the functional W, in Prop. 2.5 to the importance weights
dQ
(X*)) (X ay and dr dIP", (X":) xe I+X, W.(XH..0
1
(12)
dlp.,
Dual problem for Lagrange multiplier. Next, we will outline how to find the optimal Lagrange multiplier À in (6) in the SOC setting. Plugging the optimal control u²ta in the Lagrangian (10) yields the dual function (()
Dual E C(R, given by
c(l (A) := cit (14+1; c( (1 Ae
(13)
We note that evaluating the SOC problem in (11) at the optimal control can be expressed via the value function given in Prop. 2.5, which yields
Bua (A) EE V2_ (x,+0) 0)]
de -(1 + X)E
log le- r w. x
de,
(14)
'The SOC problem is slightly general than (1) due to the shift in the quadratic cost.

---

Algorithm 1 Trust Region SOC with buffer (see App. E.2 for details)
Require:
Initial path measure P"a, target path measure Q, divergence D, termination threshold å for i = O, 1....... do
Sample trajectories X ~ P" by integrating the SDE in (I) with Brownian motion W and control uy Compute importance weights u = ardh (X":) x exp( (X .0) as in (12)
dQ
Initialize bufler B (W.X, X, w)
Compute multiplier = arg mAX D (A) as in (14) using and 1-dim. non-linear solver Compute u uf= arg min, D(P", p using B and if X, S 8 then return control 4441 with [p=++1 Q
dp*+ dpw w T dpd in Sec. 2.2
where we note that the expression in the expectation is proportional to the importance weights in (12). Note that can obtain Monte Carlo estimate of the dual function using only simulations X"; from the previous iterations. As it turns out, these simulations are in most cases already required when learning the control u4++1 and we can thus store them in buffer. We can then obtain
X = arg max Duai (A) using any non-linear solver with minimal computational overhead.
In theory, we can define u = w using the representations in Prop. 2.5 and proceed with the next iteration of our trust region method in (4). However, computing the optimal control M4+1 using the representations in Prop. 2.5 requires gradients and Monte Carlo estimators of the value functions. This is problematic since it relies on a large amount of samples for each state z due to the (typically) very high variance of the estimator; see App. C for details. Thus, propose versions of iterative diffusion optimization to learn parametrized approximations to 044-1 in the next section.
2.2 Learning the constrained optimal control
In this section propose strategies to learn the optimal control for each iteration. As before, the general idea is to minimize loss functionals based divergences between path space measures, namely ((u) DIP" pa++1). Such divergences often rely on the Radon-Nikodym derivative
d 2211
d
dIP
dla (X")
(X"))
dIP (X*)
dPu:
(15)
I| (X" 1 s)ds+
/" (X."" 8) dw, WXXX
exp
where used Girsanov's theorem and (12). Note that the Radon-Nikodyse derivative in (15) depends only on samples of the process with the already learned uy, Let us now suggest two concrete divergences. Those divergences are desirable for high-dimensional problems since both do not rely on computing derivatives of the stochastic process and can be optimized "off-policy" using trajectories X" with the control of the previous iteration, which can be stored in buffer; see Algorithm 1.
Log-variance divergence. This divergence can be considered w.r.t. an arbitrary reference measure, where we choose p for convenience [87, 98]. We can then define the loss functional
ddP"
LLv(M) == Var log
(X (Xx))
(16)
dP"
where the Madon-Nikodyye derivative can be explicitly computed as in (15). Note that for X; = 0. this loss reduces to the on on-policy log-variance loss typically used in the literature [97]. While this loss has beneficial theoretical properties [87], it requires to keep the full trajectory in memory for the gradient computation.
Cross-entropy divergence and SOC matching. Alternatively, we can consider the cross-entropy loss (i.e., the forward KL divergence computed using reweighting)
(CR(u) (u) = Dkr, (P" = E [log dIP" (X"))
dp
dPe
(x")
(17)
dPP
where the Radon-Nikodyee derivative is again given by (15). Contrary to the log-variance loss, dP"i+1
the reweighting ap% in (12) induces exponential terms. Our trust region constraint makes sure, however, that the variance of those weights stays bounded, see Remark 2.
5

---

LV AM
SOCM
TR-SOCM
TR-LV
36
100 26
10 100 36
Figure 2: Performance criteria for a Gaussian mixture target density with varying dimension d, averaged across four seeds. We show the errors of estimating the optimal control, the log normalizing constant, as well as the Sinkhorn and total variation distances over different dimensions (from left to right). We observe that our trust region methods (TR-SOCM and TR-LV) are the only methods that perform well in high dimensions.
To efficiently compute this loss, define the so-called (lean") adjoint state as in [41] via
d Ges (X.,
(Vb(X...)) 8) (x41(X., + V a)]
(18)
ds with d;+1 (XT, = 3 Vg(X)), satisfying 1 8) = V Vx, 3,+ 1 (X, s); see [41, Lemma 5] and observe that it differs from the standard lean adjoint by the factor B3 defined in Prop. 2.2. Similar to [42], we can use the expression for the optimal control in Prop. 2.5 and the Girsanov theorem to arrive at the SOC matching loss', simple regression objective given by
LsocM (u) = TE / ||o I|oT "x"i("X""
"(X. [)]] dIP
dP", (x":)
(19)
App. G.5 for details. Contrary to the log- variinee divergence above, this objective does not require to keep the whole trajectory X", in memory for backpropagation but can be computed at times & ~ Uni 7)) using a Monte Carlo approximation. We summarize our algorithm in (1) and compare the different losses against existing approaches for SOC problems in the next section.
Applications
In this section, explore several applications of SOC, comparing novel trust-region based optimization algorithm against existing methods. Specifically, we consider the three tasks sampling from unnormalized densities, transition path sampling, and fine-tuning -ottte models. For background information, detailed experimental setups, and additional results, we refer Apps. I to K, respectively. We also include further experiments on classical SOC problems in App. I..
3.1 Diffusion-based sampling
Using (3), we can show that sampling problems can be reformulated as SOC problems. To this end, we leverage the following corollary showing that the terminal distributions QT and Pr of the optimally controlled and uncontrolled processes differ by a tilting.
Corollary 3.1 (Sampling from tilted distributions). Let us f = 0 and assume that the terminal distribution of the uncontrolled process X is independent of po and admits density denoted by Pr- Then holdy that Qr & Pre
Proof. Using (3) it holds that JQ (X) 2(X0) with Z(X)) = Ele-s(X)) The results
s(XT)
AP
follows from the independence of XT and No: [41] and App. for details.
Cor. 3.1 shows that the optimally controlled process X" samples from a given unnormalized density Pearget when using uncontrolled process with known terminal distribution Ir and setting Pr see [33, 95, 97, 125, 139-131, 144, 149] and App. I for details. Such sampling
g log Piazes. problems are of immense practical interest, with numerous applications the natural sciences [109, 151], in Bayesian statistics [54], and reinforcement learning [24].
Numerical experiments. Here, we compare existing methods for solving SOC problems with our trust region method on challenging multimodal sampling problems. We use the Denoising Diffusion Sampler (DDS) [129] method, which leverages an ergodic Omsteen-Uhlenbecc. process initialized at
Instead of the uncontrolled process we could also express the adjoint state w.r.t. the process X": however, this relies on more costly rector-Jacobiaa products; see App. G.4. 'The loss similar to the SOCM-Adjoint loss in [42], which, however, involves matrix-valuee functions.

---

d 2-2
d 50
d 100
d 200
155
13-00
RE
1.36430.00
3.403as..0s
3.077=9.am
2.508=0.079
13
CE
0.001au.000
0.20220.150 139
0.526
0.641
137
IN
diverged
1.3
1..08
AM
1.364400000 3.432ac mo
3.407=0001
3.322:0.30
SOCM
2,908 =.001
831
2:971=0.246
3.501=0.005
TR-LV
0.00000.000
0.00000.000
0.002:0.002
0.002:0.001
30
Figure 3: The left. table reports | log 2| values for the Many Well target across different dimensions d. The middle plot compares the log-varianee loss of our trust region method (TR-LV) with that of Langevin preconditioning on the GMM target in dimension d = 100. The rightmost figure presents an ablation analysis of key components in our method, highlighting the importance of trust regions in preventing mode collapse and achieving low control error. All results are averaged across four seeds.
its equilibrium measure as uncontrolled process X. We consider five baselines, specifically, reverse and (importance weighted) forward KL, also known as relative entropy (RE) and cross entropy (CE) method, respectively. Additionally, we consider the log- logvvariance loss [98], adjoint matching (AM) [41], and stochastic optimal control matching (SOCM) [42], for the unconstrained problem in (2); see [40] for comprehensive overview of SOC losses. In all experiments, we deliberately avoid using gradient guidance from the target density in the diffusion process, often referred to as Langevin preconditioning (LP) [66]. Prior work has shown that LP is essential for preventing mode collapse in neural samplers [18, 66]. However, LP is computationally expensive, as it requires querying the target distribution at every discretization step, making such approaches impractical for many problems where evaluating the target gradient is costly.
First, we consider Gaussian Mixture Model (GMM) comprising 10 components and randomized mixing weights GMMs are particularly compelling as they admit an analytical solution for the optimal control, which enables direct computation of the I,² error between the learned and optimal controls, a reliable metric for detecting mode collapse. In addition, we assess the Sinkhorn distance [31] between samples from the target and the model, and the absolute error in estimating the log- normalizing constant, denoted | log Finally, we evaluate the total variation distance between the true mixing weights and the model's estimated weights. The results, shown in Fig. 2, indicate that for d = 2, all methods closely approximate the optimal control. However, for dimensions beyond d = 10, most methods suffer from mode collapse, as reflected by increased control errors, except for those employing trust region updates. Trust region methods maintain robustness across a wide range of dimensions and only begin to show signs of mode collapse in high dimensions (d 2 150).
We additionally evaluate our method on the Many Well target [135] with 32 modes. For quantita- tive analysis, we report the log- normalizatioo error | log 21. as other ground-truth quantities are unavailable. Additionally, for the high-dimensional case d = 200, we visualize pairs of marginal distributions in App. I. The results, presented in Fig. demonstrate that our method significantly outperforms competing approaches in estimating the normalizing constant. Furthermore, the visual- izations in App. illustrate that trust region updates effectively prevent mode collapse, even in high dimensions. In contrast, baseline methods either suffer from mode collapse or fail to converge.
Finally, we perform an ablation study on the GMM target, analyzing key components of our proposed method. Specifically,, we investigate the effects of incorporating replay buffer and applying trust region optimization. To this end, we compare a variant using a fixed Lagrangian multiplier X, selected via hyperparameter tuning, with one in which À is dynamically optimized using our trust region approach. Additionally, we evaluate the log- variancee loss both with and without using a replay buffer. Moreover, we compare our method to LV with Langevin preconditioning on the GMM target with dimensionality = 100. The results, shown in Figure 3, demonstrate that trust region optimization significantly reduces control error and decreases the number of target evaluations by several orders of magnitude.
3.2 Transition path sampling
Transition path sampling is of great importance for studying phase transitions and chemical reactions. The key challenge comes from the energy barrier that connects two sets A and B along the energy landscape, which makes direct sampling of transition paths extremely unlikely. These problems can also be formulated SOC problems [59, 62, 115]. Specifically, we set = -U, where U : R N - IR is the potential function, and g = log 1 as well as Po x 1 which constraints the initial and target states in the sets A and B. As in (3), it holds that dQ Lx(Xr) Recent work
dF 2(Xo)
7

---

Table 1: Quantitative evaluation transition path sampling problems. i denotes that results are taken from [113]. The results for TPS-DPS and TR-LV are averaged across three seeds.
Method
RMSD (A. D)
THP (%.. t>
ETS (kJ//mol)
Method
RMSD (A. 4
THP (S. t>
ETS (kll/mol)
Alanine Dipeptide
Chignolin
UMD (3600K)
0.32
6.25
812.47
148,80 80
UMD (1200K))
0.93
388.17
SMD
0.56 0.27
54.69
78.40 1 12.76
SMDi
0.85 0.24
34,38
179.52
138.87
PIPS
0.15
43.75
28.17
PIPS
0.17
0.00
TPS-DPS
0.17 = 0.18
39.58
28,13
46.31 = 10.16
TPS-DPS
1.06 0.08
25,00 =
10.69
-189.91
= 23.01
TR-LV
0.29 0.03
61.25
4.005
49.11 5.84
TR-LV
0.00 0.01
43.95 + 5.64
-303.98
28.05
TPS TPS-DPS DPS
TR-IV
2-
TWP 11
10000
04
<aerg) volustion
(a) Alanine Dipeptidc
(b) Chignolin
Figure 4: We compare our trust region method (TR-LV) with Diffusion Path Sampler (TPS-DPS) [113] on Alanine Dipeptide and Chignolin. All results are averaged over three random seeds, with both the mean and standard deviation reported. Our method identifies transition paths more consistently and robustly, as evidenced by higher THP values and lower standard deviations.
has leveraged neural networks to parameterize a bias force to solve the corresponding SOC problem, employing objectives such as the KL [44, 72, 141], or log-variance divergence [113].
Numerical experiments. We evaluate the performance of the trust-region-uun based log variance loss (TR-LV) on two transition path sampling problems: Alanine Dipeptide isomerization and Chignolin folding, with 22 and 138 atoms, respectively.
Our evaluation includes three metries: trbssch-aligned root mean squared distance (RMSD) between the final states of the sampled paths and the target state, transition hit percentage (THP) measuring the proportion of final states hitting within the target region, and energy of transition state (ETS) identifying the highest energy values along paths that reach the target.
We compare our method to standard molecular dynamics (MD) with increased temperature (UMD), steered MD (SMD) [75] with force applied to collective variables, and PIPS [72] which uses the cross- entropy loss. We also include TPS-DPS [1 13] as a key haseline, which employs an (unconstrained) log- variance loss to formulate TPS as stochastic optimal control (SOC) problem. Further experimental details are provided in App. J.
Table shows that TR-LV achieves superior target state RMSD and transition hit percentage compared to the standard log-vart variance objective (TPS-DPS) for both molecular systems. Notably, SMD performs well due to its use of collective variables with biased force guiding the sampling process. Figure 4 illustrates that the trust region constraint leads to significantly more robust training compared to TPS-DPS as indicated by low standard deviations across different seeds. Moreover, on Alanine Dipeptide, the trust region constraint initially regularizes optimization and accelerates convergence thereafter. Across both systems, the trust region constraint significantly enhances training stability and performance.
3.3 Fine-tuning of diffusion models
Interpreting -g as reward and the uncontrolled process X as a pretrained diffusion model (i.e., b includes the pretrained neural network), Cor. 3.1 shows that we can perform reward fine-tuning by solving the SOC problem in (I); see also [38, 41, 132]. Reward fine-tuning has recently shown impressive results, e.g., in image [28, 41] and molecule generation [38], and SOC provides a principled framework. A special case is given by posterior sampling [38]. Setting g log p 2)))2)) where p(yla xx) is the likelihood and we interpret Pr as a learned (diffusion) prior p(r), Bayes' theorem shows that the optimally controlled process samples from the posterior p(z|y).
Numerical experiments. We perform reward fine-tuning on Stable Diffusion 1.5 [102]. using ImageReward [140], which is a reward model designed to capture prompt alignment and image quality according to human preferences. We take the adjoint matching (AM) method as baseline and
8

---

AM,
AM, T7
TR-SOCM, 77 0
TR-SOCM, 77
1.50
1.30
Ieapiteeat
06-
07
322
4.2
0.5
3.0
c4adee Eragenaccie x.27
Figure 5: Comparison of Adjoint Matching against Trust Region SOCM for Stable Diffusion 1.5 fine-tuning w.r.t. four quality metrics, where = 0 and 17 = refor to ODE (DDIM) and SDE (DDPM) inference, respectively.
masterpiece, best quality, realistic photograph. sk, high detailed vintage motorcycle parked street at dusk, ncon reflections, shallow depth of field
close up photo of authropomorphic fox animal dressed In white shirt, fox animal, glasses
cobblestone
Figure 6: Comparison between images generated by the base Stable Diffusion 1.5 model (left) and its version fine-tuned with TR-SOCM (right), using the prompts and random seeds. The fine-tuned model generates higher quality images (hike) with hetter prompt alignment (fox).
compare it against our TR-SOCM loss (19), keeping all other hyperparameters fixed. Our TR-SOCM allows the principled use of buffers, and we perform three passes on each buffer of size 500, leading three times fewer trajectories for fixed number of model updates. For faster convergence, a modified version of TR-SOCM with annealing factor 33 = 1. For each algorithm, we evaluate checkpoints during fine-tuning (with ODE and SDE inference)
ImageReward and three additional metrics: CLIP-Score [69], which measures prompt aligmment, Human Preference Score [137], which measures numan-perceiver image quality, and Dreamsim diversity [51], which measures per-prompt diversity. We observe that TR-SOCM achieves similar performance metrics to AM fraction of the cost, as sampling the trajectories and solving the lean adjoint ODE, which dominates the computational costs, is amortized over the buffer passes; Figs. 5 and well as App. for more details.
Related works
In this section, we discuss the most related works, comparing our approach to existing methods for solving SOC problems. We provide more extensive comparison in App. C.
Iterative diffusion optimization. Many recently developed methods approach SOC problems by simulating the (diffusion) process X", computing a suitable cost function, and optimizing the param- eters of the control function using variants of stochastic gradient methods. These techniques are collectively referred
iterative diffusion optimization (IDO) methods [87]. While the underlying theory dates back to [33, 91], combinations with deep learning in the context of SOC have been explored by [15, 87, 95, 97. 129, 131, 149, 152]. One can derive most of the related objectives derivative dP"
starting from the Radon-Nikodyy minimize loss based on suitable divergence in (2). Previous works have, e.g., proposed the log-variance divergence 197, 113] or the forward KL divergence (corresponding to the cross-entropy loss 72, 76, 104, 150]), for which develop corresponding trust region versions (16) and (17). The SOC matching loss [42], which we extended trust regions in (19), is equal to the cross entropy loss in expectation but exhibits lower variance empirically. We refer to [41] for more IDO losses. However, all existing methods have either directly tackled the target measure Q or relied on form of hand-tuned annealing.
aip (X") as in (12) (with
uz). One can then
Trust region methods. We show how IDO methods generally be extended to trust region methods, enabling (1) automatic control on the variance of the importance weights and (2) principled usage of buffers, leading to faster and more stable convergence, in particular avoiding mode collapse

---

in high dimensions. Trust region methods have long history as robust optimization algorithms that iteratively minimize an objective within an adaptively sized "trust region"; see [29] for an overview. These methods have also been extended to optimize over spaces of probability distributions, particularly in reinforcement learning [2-4, 7, 83, 88, 90, 93, 110, 111, 138, 139, 142], black-box optimization [1, 118, 134], variational inference [9, 10] and path integral control [123]. To the best of our knowledge, these methods have not yet been extended path measures inference problems. Moreover, the connection between trust-region iterates and geometric annealing has not previously been established.
5 Conclusion
In this work, we develop a novel framework for solving SOC problems using deep learning. Our framework builds the fact that space measure induced by a controlled SDE. Instead of finding this optimal measure at once, we divide the unconstrained problem into sequence of constrained optimization problems by bounding the KL divergence to the measure from the previous iteration. We show that this defines well- behaved geometric annealing between the prior and the target path measure, resulting in equidistant steps on the Fisher-Rao information manifold. Crucially, each intermediate problem turns out to be an altered SOC problem that can be efficiently solved without simulations by using buffer of trajectories with the control from the previous iteration. our experiments, we show that our method significantly improves the learning of the optimal control, including applications in diffusion- based sampling and transition path sampling in molecular dynamics. Further, we show that our method can be scaled to improve the efficiency of reward fine-tuning for text-to-image diffusion models. In the future, we expect our framework to improve even more applications of SOC, potentially including the use of divergences other than the KL divergence for the trust region constraint. Finally, our results for general measures motivate the use of trust region methods for other learned measure transports, e.g., normalizing flows.
reformulate specific problems as finding optimal path
Acknowledgements
D.B. acknowledges support by funding from the pilot program Core Informatics of the Helmholtz Association (HGF) and the state of Baden-Württembe Württemberg through bwHPC, as well as the HoreKa supercomputer funded by the Ministry of Science, Research and the Arts the German Federal Ministry of Education and Research. The research of L.R. was partially funded by Deutsche Forschungsgemminsccaaff (DFG) through the grant CRC 1 1144 "Scaling Cascades in Complex Systems" " (project A05, project number 235221301).
and by
10

---

References
[1] A. Abdolmaleki, R. Lioutikov, J. R. Peters, N. Lau, L. Pualo Reis, and G. Neumann. Model- based relative entropy stochastic search. Advances Neural Information Processing Systems, 28, 2015.
[2] A. Abdolmaleki, J. T. Springenberg. J. Degrave, S. Bohez, Y. Tassa, D. Belov, N. Heess, and M. Riedmiller. Relative entropy regularized policy iteration. arXiv preprint arXiv:1812.02256. 2018.
[3] A. Abdolmaleki, J. T. Springenberg, Y. Tassa, R. Munos, N. Heess, and M. Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint arXiv: 1806. 06920, 2018.
[4] J. Achiam, D. Held, A. Tamar, and P.. Abbeel. Constrained policy optimization. In International conference on machine learning, pages 22-31. PMLR, 2017.
[5] T. Akhound-Sadeghe J. Lee, J. J. Bose, V. De Bortoli, A. Doucet, M. M. Bronstein, D. Beaini, S. Ravanbakhsh, K. Neklyudov, and A. Tong. Progressive inference- time annealing of diffusion models for sampling from Boltzmann densities. arXiv preprint arXiv:2506. 16471, 2025.
[6] T. Akhound- -SaounddSadegh,... J. Rector-Brooks, A. J. Bose, S. Mittal, P. Lemos, C.-H. Liu, M. Sendera, S. Ravanbakhsh, G. Gidel, Y. Bengio, et al. Iterated denoising energy matching for sampling from Boltzmann densities. arXiv preprint (rXiv:2402.0612/1 2024.
[7] Akrour, J. Pajarinen, J. Peters, and G. Neumann. Projections for approximate policy iteration algorithms. In International Conference on Machine Learning, pages 181-190. PMLR, 2019.
[8] M. S. Albergo and E. Vanden- Vanden-Eijnden. NETS: A non-equilibrium transport sampler. arXiv preprint arXiv:2410.02711, 2024.
[9] O. Arenz, Dahlinger. Ye, M. Volpp, and G. Neumann. A unified perspective on natural gradient variational inference with Gaussian mixture models, arXiv preprint arXiv:2209.11533 2022.
[10] O. Arenz, M. Zhong, and G. Neumann. Trust-region variational inference with Gaussian mixture models. Journal of Machine Learning Research, 21 (163)):1160, 2020.
[11] C. Beck, S. Becker, P. Grohs, N. Jaafari, and A. Jentzen. Solving the Kolmogorov PDE by means of deep learning. Journal of Scientific Computing, 88:1-28, 2021.
[12] R. Bellman. Dynamic programmingg Princeton University Press, 1957.
[13] J. Berner, M. Dablander, and P. Grohs. Numerically solving parametric families of high- dimensional Kolmogorov partial differential equations via deep learning. Advances in Neural Information Processing Systems, 33:1661:: 18:16615--1666 2020.
[14] J. Berner, L. Richter, M. Sendera, J. Rector-Brooks, and N. Malkin. From discrete-time policies to continuous-timg diffusion samplers: Asymptotic equivalences and faster training. arXiv preprint 2rXiv:2501.06//1 2025.
[15] J. Berner, L. Richter, and K. Ullrich. An optimal control perspective on diffusion- -ased generative modeling. Transactions on Machine Learning Research, 2024.
[16] K. Black, M. Janner, Du. I. Kostrikov, and S. Levine. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024.
[17] D. Blessing, J. Berner, L. Richter, and G. Neumann. Underdamped diffusion bridges with appli- cations to sampling. In The Thirteenth International Conference on Learning Representations, 2025.
[18] D. Blessing, X. Jia, J. Esslinger, F. Vargas, and G. Neumann. Beyond ELBOs: A large-scale evaluation of variational methods for sampling. arXiv preprint arXiv:2406.00 07423, 2024.
[19] D. Blessing, X. Jia, and G. Neumann. End-to-end learning of Gaussian mixture priors for diffusion sampler. arXiv preprint arXiv:2503.00524, 2025.

---

[20] P. G. Bolhuis, D. Chandler, C. Dellago, and P. L. Geissler. Transition path sampling; Throwing ropes over rough mountain passes, in the dark. Annual review of physical chemistry, 53(1)2291- 318, 2002.
[21] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne et al. Jax: Autograd and xla. Astrophysics Source Code Library, pages ascl-2111, 2021.
[22] Brekelmans, Masrani, F. Wood, G. V. Steeg, and A. Galstyan. All in the exponential fam- ily: Bregman duality in thermodynamic variational inference. arXiv preprint arXiv:2007.00642, 2020.
[23] R. P. Brent. An algorithm with guaranteed convergence for finding a zero of function. The computer journal,, 14(4):422-425, 1971.
[24] O. Celik, Z. Li, D. Blessing, G. Li, D. Palanicek, J. Peters, G. Chalvatzaki, and G. Neu- mann. DIME: Diffusion-based maximum entropy reinforcement learning. arXiv preprint arXiv:2502.02316. 2025.
[25] J. Chen, I.. Richter, J. Berner, D. Blessing, G. Neumann, and Anandkumar. Sequential controlled Langevin diffusions. In The Thirteenth International Conference on Learning Representationss 2025.
[26] Chetrite and H. Touchette. Variational and optimal control representations of condi- tioned and driven processes. Journal of Statistical Mechanics: Theory and Experiment, (015(12)P1100))) 2015.
[27] J. Choi, Y. Chen, M. Tao, and G.-H. Liu. Non-equilibrium annealed adjoint sampler. arXiv preprint arXiv:2506. /8/65, 2025.
[28] K. Clark, Vicol, K. Swersky, and D. J. Fleet. Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Representations, 2024.
[29] A. R. Conn, N. I. Gould, and P. L. Toint. Trust region methods. SIAM, 2000.
[30] G. E. Crooks. Measuring thermodynamic length. Physical Review Letters, 99110)000000, 2007.
[31] M. Cuturi. Sinkhorn distances: Lightspeed computation optimal transport. Advances in neural information processing systems, 26, 2013.
[32] M. Cuturi, L. Meng-Papaxantton Y. Tian, C. Bunne, G. Davis, and O. Teboul. Optimal trans- port tools (OTT): A jax toolbox for all things Wasserstein. arXiv preprint 22Xiv:2001.13324 2022.
[33] P. Dai Pra. stochastic control approach to reciprocal diffusion processes. Applied mathemat- ics and Optimization, 23(1):3133329, 1991.
[34] P. Dai Pra, L. Meneghini, and W. J. Runggaldier. Connections between stochastic control and dynamic games. Mathematics of Control, Signals and Systems, 9:303-326, 1996.
[35] A. Das, D. C. Rose, Garrahan, and D. T. Limmer. Reinforcement learning of rare diffusive dynamics. The Journal of Chemical Physics, 155(13), 2021.
[36] De Bortoli, J. Thornton, J. Heng, and A. Doucet. Diffusion Schrödinger bridge with applications score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-1777 2021.
[37] C. Dellago, G. Bolhuis, and D. Chandler. Efficient transition path sampling: Application to lennard-jones cluster rearrangements. The Journal of chemical physics, 108(22);9236-9245, 1998.
[38] K. Didi, F. Vargas, S. V. Mathis, V. Dutordoir, E. Mathieu, U. J. Komorowska, and P. Lio. A framework for conditional diffusion modelling with applications in motif scaffolding for protein design. arXiv preprint arXiv:2312.09236. 2023.
12

---

[39] Z. Ding, Y. Jiao. X. Lu, Z. Yang, and C. Yuan. Sampling via Föllmer flow. arXiv preprint arXiv:2311.03660, 2023.
[40] C. Domingo-Enrich. A taxonomy of loss functions for stochastic optimal control. arXiv preprint 00Xi::2004:: 2024.
[41] C. Domingo--Erich M. Drozdzal, B. Karrer, and R. T. Q. Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In The Thirteenth International Conference on Learning Representations, 2025.
[42] C. --omingo-Eniinn Enrrgo,, Han, B. Amos, J. Bruna, and R. Q. Chen. Stochastic optimal control matching, In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[43] A. Doucet, W. Grathwohl, A. G. d. G. Matthews, and H. Strathmann. Score-based diffusion meets annealed importance sampling. In Advances in Neural Information Processing Systems, 2022.
[44] Y. Du, M. Plainer, R. Brekelmans, C. Duan, F. Noe, C. P. Gomes, A. Aspuru- Guzik, and K. Neklyudov. Doob's lagrangian: A sample-efficient variational approach to transition path sampling. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems.
[45] Y. Du, M. Plainer, R. Brekelmans, C. Duan, F. Noe, C. P. Gomes, A. Aspuru-Guzik, and K. Neklyudov. Doob's lagrangian: A sample-efficient variational approach to transition path sampling. In The Thirty-eigh eightt Annual Conference on Neural Information Processing Systems, 2024.
[46] E. Erives, B. Jing, P. Holderrieth, and T. Jaakkola. Continuously tempered diffusion samplers. In Frontiers Probabilistic Inference: Learning meets Sampling, 2025.
[47] Y. Fan. O. Watkins, Y. Du, H. Liu, M. Ryu, C. Boutilier, Abbeel, M. Ghavamzadeh, K. Lee, and K. Lee. Dpok: Reinforcement learning for fine-tuning exx-tto-image diffusion models. arXiv preprint ar0iv:2305. 16381, 2023.
[48] M. F. Faulkner and S. Livingstone. Sampling algorithms in statistical physics: guide for statistics and machine learning. Statistical Science, 39(1):137-164, 2024.
[49] W. Fleming and Rishel. Deterministic and Stochastic Optimal Control. Applications of mathematics. Springer, 1975.
[50] W. H. Fleming and H. M. Soner. Controlled Markov processes and viscosity solutions, volume 25. Springer Science & Business Media, 2006.
[51] S. Fu, N. Tamir, S. Sundaram, L. Chai, R. Zhang, T. Dekel, and P. Isola. Dreamsim: Learning new dimensions of human visual similarity using synthetic data. arXiv preprint arXiv:2306.0993 2023.
[52] T. Geffner and J. Domke. MCMC variational inference via uncorrected hamiltonian annealing. Advances in Neural Information Processing Systems, 34:639 651, 2021.
[53] T. Geffner and J. Domke. Langevin diffusion variational inference. In International Conference on Artificial Intelligence and Statistics, pages 576-593. PMLR, 2023.
[54] A. Gelman, J. Carlin, H. Stern, D. Dunson, A. Vehtari, and D. Rubin. Bayesian Data Analysis, Third Edition. Chapman & Hall/CRC Texts in Statistical Science. Taylor & Francis, 2013.
[55] L. Grenioux, M. Noble, and M. Gabrié. Improving the evaluation of samplers on multi-modal targets. arXiv preprint arXiv:2504.088 08916. 2025.
[56] T. Gritsaev, N. Morozov, K. Tamogashev, D. Tiapkin, S. Samsonov, A. Naumov, D. Vetrow, and N. Malkin. Adaptive destruction processes for diffusion samplers. arXiv preprint arXiv:2506.0154/, 2025.
[57] W. Guo, M. Tao, and Y. Chen. Complexity analysis of normalizing constant estimation: from Jarzynski equality to annealed importance sampling and beyond. arXiv preprint arXiv:2502.04555 2025.
13

---

[58] J. Han, A. Jentzen, and W. E. Solving high-ding diteeeimesiona partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34)8850--8110 2018.
[59] C. Hartmann, O. Kebiri, L. Neureither, and L. Richter. Variational approach to rare event simulation using least-squares regression. Chaos: An Interdisciplinary Journal of Nonlinear Science, 29(6), 2019.
[60] C. Hartmann and L. Richter. Nonasymptotic bounds for suboptimal importance sampling. SIAM/ASA Journal on Uncertainty Quantification, 12(2):309- 2024.
[61] C. Hartmann, L. Richter, C. Schütte, and W. Zhang. Variational characterization of free energy: Theory and algorithms. Entropy, 19(11), 2017.
[62] C. Hartmann and C. Schütte. Efficient rare event simulation by optimal nonequilibrium forcing. Journal of Statistical Mechanics; Theory and Experiment, 2012(1 1004, 2012.
[63] A. Havens, B. K. Miller, B. Yan, C. Comingo-EEnichhe A. Sriram, B. Wood, D. Levine, B. Hu, Amos, B. Karrer, et al. Adjoint sampling: Highly scalable diffusion samplers via adjoint matching. arXiv preprint wrXiv:2504.1171: 2025.
[64] J. He, W. Chen, M. Zhang, D. Barber, and M. Hemández-Lobato. Training neural samplers with reverse diffusive kl divergence. arXiv preprint arXiv:2410.124533 2024.
[65] J. He, Y. Du, F. Vargas, Y. Wang, C. P. Gomes, J. M. Hernández-Lobatt, and E. Vanden-Eijnden. FEAT: Free energy estimators with adaptive transport. arXiv preprint arXiv:2504. 11516, 2025.
[66] J. He, Y. Du, F. Vargas, D. Zhang, S. Padhy, R. Ou C. Gomes, and J. M. Hernández- Lobato. No trick, no treat: Pursuits and challenges towards simulation-free training of neural samplers. arXiv preprint (XK::2502.06688 2025.
[67] D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus). arXiv preprint ar arXiv:/606.088 08016.0415,0000 2016.
[68] J. Hénin, T. Lelièvre, M. R. Shirts, O. Valsson, and L. Delemotte, Enhanced sampling methods for molecular dynamics simulations. arXiv preprint arXivv2202.04/66 046444444. 2022.
[69] J. Hessel, A. Holtzman, M. Forbes, R. L. Bras, and Y. Choi. Clipscore: A reference-free evaluation metric for image captioning. arXiv preprint arXiv:1104.08718, 2021.
[70] P. Holderrieth, M. S. Albergo, and T. Jaakkola. Leaps: A discrete neural sampler via locally equivariant networks. arXiv preprint arXiv:2502. 10843, 2025.
[71] L. Holdijk, Y. Du, F. Hooft, P. Jaini, B. Ensing, and M. Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths. In shirty-seventh Conference Neural Information Processing Systems, 2023.
[72] L. Holdijk, Y. Du, P. Jaini, F. Hooft, B. Ensing, and M. Welling. Path integral stochastic optimal control for sampling transition paths. In ICMI. 2022 2nd Al for Science Workshop, 2022.
[73] J. Huang, Y. Jiao, L. Kang, X. Liao, J. Liu, and Y. Liu. Schrädinger-Föllmer sampler: sampling without ergodicity. arXiv preprint arXiv:2106. 10880. 2021.
[74] X. Huang, H. Dong, Y. Hao, Y. Ma. and T. Zhang, Monte Carlo sampling without isoperimetry:
A reverse diffusion approach. arXiv preprint wxXiv23377000033 2023.
[75] S. Izrailev, S. Stepaniants, B. Isralewitz, D. Kosztin, H. Lu, F. Molnar, W. Wriggers, and K. Schulten. Steered molecular dynamies. In Computational Molecular Dynamics: Chal- lenges, Methods, Ideas: Proceedings of the 2nd International Symposium on Algorithms for Macromolecular Modelling, Berlin, May 21-24 1997, pages 39-65. Springer, 1999.
[76] J. Kappen and H. C. Ruiz, Adaptive importance sampling for control and inference. Journal of Statistical Physics, 162(5):1244-1266 2016.
14

---

[77] M. Kim, S. Choi, T. Yun, E. Bengio, L. Feng, J. Sector-Brooks, S. Ahn, J. Park, N. Malkin, and Y. Bengio. Adaptive teachers for amortized samplers. arXiv preprint 4rX1::2410.01332. 2024.
[78] M. Kim, K. Seong, D. Woo, S. Ahn, and M. Kim. On scalable and efficient training of diffusion samplers. arXiv preprint arXiv:2505. 19552, 2025.
[79] D. P. Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv: 142v:141266980 2014.
[80] G.-H. Liu, J. Choi, Y. Chen, B. K. Miller, and R. T. Chen. Adjoint schrödiiger bridge sampler. arXiv preprint arXiv:2506. 22565, 2025.
[81] J. Liu, G. Liu, J. Liang, Y. Li, J. Liu, X. Wang. P. Wan, D. Zhang, and W. Ouyang, Flow-grpo: Training flow matching models via online rl, 2025.
[82] Z. Liu, T. Z. Xiao, W. Liu, Y. Bengio, and Zhang. Efficient alignment via gradient-informer informed GFlowNets, 2025.
diffusion
[83] W. Meng, Q. Zheng, Y. Shi, and G. Pan. An off-policy trust region policy optimization method with monotonic improvement guarantee for deep reinforcement learning. IEEE Transactions on Neural Networks and Learning Systems, 33(5):2223 2235,, 2021.
[84] L. I. Midgley, V. Stimper, G. N. Simm, B. Schölkopf, and J. M. Hernández-Lobato. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208. 01893, 2022.
[85] R. M. Neal. Probabilistic inference using Markov chain Monte Carlo methods. 1993.
[86] M. Noble, L. Grenioux, M. Gabrié, and A. O. Durmus. Learned reference-based diffusion sampling for multi-modal distributions. arXiv preprint drXiv:2410.19449, 2024.
[87] N. Nüsken and L. Richter. Solving tigh-dimensional Hamilton-Jacobi-Belll PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications, 2:1-48, 2021.
[88] F. Otto, P. Becker, N. A. Vien, H. C. Ziesche, and G. Neumann. Differentiable trust region layers for deep reinforcement learning. arXiv preprint arXiv:2101.09207 2021.
[89] R. Yu B. Qiang, and J. M. Hernández-LLobato BNEM: A Boltzmann sampler based bootstrapped noised energy matching. arXiv preprint arXiv:2449909787 2024.
[90] J. Pajarinen, H. L. Thai, R. Akrour, J. Peters, and G. Neumann. Compatible natural gradient policy search. Machine Learning, 108: 1443-1466, 2019.
[91] M. Pavon. Stochastic control and nonequilibrium thermodynamical systems. Applied Mathe- matics and Optimization, 19(1):187-200 1989.
[92] M. Pavon. On local entropy, stochastic control and deep neural networks. arXiv preprint arXiv:2204. 13049, 2022.
[93] J. Peters, K. Mulling, and Y. Altun. Relative entropy policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pages 1607-1612, 2010.
[94] H. Pham. Continuous-time Stochastic Control and Optimization with Financial Applications. Stochastic Modelling and Applied Probability. Springer Berlin Heidelberg, 2009.
[95] L. Richter. Solving high-dimensional PDEs, approximation of path space measures and importance sampling of diffusions. PhD thesis, BTU Cottbus-Senftenberg 2021.
[96] L. Richter and J. Berner. Robust SDE-based variational formulations for solving linear PDEs via deep learning. In International Conference on Machine Learning, pages 18649-18666. PMLR, 2022.
[97] L. Richter and J. Berner. Improved sampling via learned diffusions. In International Conference Learning Representations, 2024.
15

---

[98] L. Richter, A. Boustati, N. Nüsken, F. Ruiz, and O. D. Akyildiz. VarGrad: A low-variance gradient estimator for variational inference. Advances Neural Information Processing Systems, 33:13481- -33448 2020.
[99] L. Richter, L. Sallandt, and N. Nüsken. Solving high-dimensional parabolic PDEs using the tensor train format. In International Conference on Machine Learning, pages 8998-9009 PMLR, 2021.
[100] L. Richter, L. Sallandt, and N. Nüsken. From continuous-ting time formulations to discretization schemes: tensor trains and robust regression for bsdes and parabolic pdes. Journal of Machine Learning Reseurch, 25(248): 1-00, 2024.
[101] S. Rissanen, R. OuYang. J. He, W. Chen, M. Heinonen, A. Solin, and J. M. Hernández-Lobato. Progressive tempering sampler with diffusion. arXiv preprint arXiv:2506.05231, 2025.
[102] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[103] D. C. Rose, F. Mair, and J. P. Garrahan. reinforcement learning approach rare trajectory sampling. New Journal of Physics, 23(1):113013 2021.
[104] R. Y. Rubinstein and D. P. Kroese. The cross-entropy method: a unified approach to combi- natorial optimization, Monte- Coote--Carlo simulation and machine learning. Springer Science & Business Media, 2013.
[105] M. Sabate Vidales, D. Šiska, and L. Szpruch. Unbiased deep solvers for linear parametric PDEs. Applied Mathematical Finance, 28(4):299-329, 2021.
[106] P. Salamon and R. S. Berry. Thermodynamic length and dissipated availability. Physical Review Letters, 51(13):1127, 1983.
[107] S. Sanokowski, W. Berghammer, M. Ennemoser, H. Wang, S. Hochreiter, and S. Lehner. Scalable discrete diffusion samplers: Combinatorial optimization and statistical physics. arXiv preprint arXiv:2502.0886 08696, 2025.
[108] S. Sanokowski, L. Gruber, C. Bartmann, S. Hochreiter, and S. Lehner. Rethinking losses for diffusion bridge samplers. arXiv preprint arXiv:2506. 10982. 2025.
[109] H. Schopmans and P. Friederich. Temperature-annent boltzmann generators. arXiv preprint wxX1v:2501.19077, 2025.
[110] J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz. Trust region policy optimization. In International conference machine learning, pages 1889-1897. PMLR, 2015.
[111] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.0/777.06347 2017.
[112] M. Sendera, M. Kim, S. Mittal, P. Lemos, L. Scimeca, J. Rector-Brooks, Adam, Y. Bengio, and N. Malkin. Improved off-policy training of diffusion samplers. Advances in Neural Information Processing Systems, 37:81016-81045 2024.
[113] K. Seong, S. Park, S. Kim, W. Y. Kim, and S. Ahn. Transition path sampling with improved off-policy training of diffusion path samplers. arXiv preprint arXiv:2405. 19961, 2024.
[114] Z. Shi, L. Yu, T. Xie, and C. Zhang. Diffusion-Prie PINN sampler. arXiv preprint arXiv:2410. 15336, 2024.
[115] A. N. Singh, A. Das, and D. T. Limmer. Variational path sampling of rare dynamical events. Annual Review of Physical Chemistry, 76, 2025.
[116] A. N. Singh and D. T. Limmer. Variational deep learning of equilibrium transition path ensembles. The Journal of Chemical Physics, 159(2), 2023.
16

---

[117] J. Sun, J. Berner, L. Richter, M. Zeinhofer, J. Müller, K. Azizzadeneshelii and A. Anand- kumar. Dynamical measure transport and neural PDE solvers for sampling, arXiv preprint 17Xiv:2407007877 2024.
[118] Y. Sun, D. Wierstra, T. Schaul, and J. Schmidhuber. Efficient natural evolution strategies. In Proceedings of the 11th Annual conference on Genetic and evolutionary computation, pages 539-546, 2009.
[119] S. Syed, A. Bouchard-Côté, K. Chern, and A. Doucet. Optimised annealed sequential Monte Carlo samplers. arXiv preprint arXiv:2408. /2057, 2024.
[120] C. B. Tan, J. Bose, C. Lin, L. Klein, M. M. Bronstein, and A. Tong. Scalable equilibrium sampling with sequential Boltzmann generators. arXiv preprint wxXiv:250218846) 2025.
[121] H. Y. Tan, S. Osher, and W. Li. Noise-free sampling algorithms via regularized Wasserstein proximals. arXiv preprint arXiv:2308. /4945, 2023.
[122] M. Tancik, P. Srinivasan, B. Mildenhall, S. Tridovich-Keill N. Raghavan, Singhal, R. Ra- mamoorthi, J. Barron, and R. Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems, 33:7537- 7547, 2020.
[123] D. Thalmeier, H. J. Kappen, S. Totaro, and V. Gómez. Adaptive smoothing for path integral control. Journal of Machine Learning Research, 21(191):1 2020.
[124] A. Thin, N. Kotelevskii, A. Durmus, E. Moulines, M. Panov, and A. Doucet. Monte Carlo variational auto-encoders. In International Conference on Machine Learning, 2021.
[125] B. Tzen and M. Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pages 3084-3114, PMLR, 2019.
[126] M. Uehara, Y. Zhao, K. Black, E. Hajiramezanali, G. Scalia, N. L. Diamant, A. M. Tseng, T. Biancalani, and S. Levine. Fine-tuning of continuous- time diffusion models entropy- regularized control. arXiv preprint arXiv:2402. 15194, 2024.
[127] R. Van Handel. Stochastic calculus, filtering, and stochastic control. Course notes., URL http://www. princeton. edt//van/acm2117ACM2777 pdf. 14, 2007.
[128] E. Hunden-Eijndon et al. Transition-pali theory and fath-finding algorithms for the study of rare events. Annual review of physical chemistry, 61:391-420, 2010.
[129] F. Vargas, W. Grathwohl, and A. Doucet. Denoising diffusion samplers. arXiv preprint arXiv:2302. 13834, 2023.
[130] F. Vargas, A. Ovsianas, D. Fernandes, M. Girolami, N. D. Lawrence, and N. Nüsken. Bayesian learning via neural Schrödinggr-Fööller flows. Statistics and Computing, 33(1):3, 2023.
[131] F. Vargas, S. Padhy, D. Blessing, and N. Nüsken. Transport meets variational inference: Controlled Monte Carlo diffusions. In The Twelfth International Conference on Learning Representations, 2024.
[132] S. Venkatraman, M. Jain, L. Scimeca, M. Kim, M. Sendera, M. Hasan, L. Rowe, S. Mittal, P. Lemos, E. Bengio, et al. Amortizing intractable inference in diffusion models for vision, language, and control. arXiv preprint arXiv:2405.2097. 2024.
[133] C. Wang, X. Zhang, K. Cui, W. Zhao, Y. Guan, and T. Yu. Importance weighted score matching for diffusion samplers with enhanced mode coverage. arXiv preprint urXiv:2505.19433 2025.
[134] D. Wierstra, T. Schaul, T. Glasmachers. Y. Sun, J. Peters, and J. Schmidhuber. Natural evolution strategies. The Journal of Machine Learning Research, 15(1)9949-980, 2014.
[135] H. Wu. J. Köhler, and F. Noé. Stochastic normalizing flows. Advances in neural information processing systems, 33:5933-5944, 2020.
17

---

[136] L. Wu, Y. Han, C. A. Naesseth, and J. P. Cunningham. Reverse diffusion Sequential Monte Carlo samplers. arXiv preprint arXiv:2508. 059266, 2025.
[137] X. Wu, Hao, K. Sun, Y. Chen, F. Zhu, Zhao, and H. Li. Human preference score v2: solid benchmark for evaluating human preferences of text-to-image synthesis. arXiv preprint arXiv:2306. 0966.0934.1.4 2023.
[138] Y. Wu, E. Mansimov, deep reinforcement learning using Kronecker-factored approximation. Advances in neural information processing systems, 30, 2017.
Grosse, S. Liao, and Ba. Scalable trust-region method for
[139] H. Xu, J. Xuan, G. Zhang, and J. Lu. Trust region policy optimization via entropy regularization for kullback-leibler divergence constraint. Nearocomputing. 589:127716,, 2024.
[140] J. Xu, X. Liu, Y. Wu, Y. Tong, Q. Li, M. Ding, J. Tang, and Y. Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[141] J. Yan, H. Touchette, and G. M. Rotskoff. ize dynamical phase transitions. Physical Review E, 105(2):024115, 2022.
Learning nonequilibrium control forces to character-
[142] T.-Y. Yang, J. Rosca, K. Narasimhan, and P. J. Ramadge. Projection Projection-based constrained policy optimization. arXiv preprint arXiv:2010.03152, 2020.
[143] S. Yoon, H. Hwang, H. Jeong, D. K. Shin, C.-S. Park, S. Kweon, and F. Park. Value gradient sampler: Sampling sequential decision making. arXiv preprint 2rXhv:2502.13280. 2025.
[144] D. Zhang, R. T. Chen, C.-H. Liu, A. Courville, and Y. Bengio. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization. In The Twelfth International Conference on Learning Representations, 2024.
[145] D. Zhang, Y. Zhang. J. Gu, R. Zhang, J. Susskind, N. Jaitly, and S. Zhai. Improving GFlowNets for text-to-image diffusion alignment. arXiv preprint arXiv:2406.0m 00633, 2024.
[146] G. Zhang, K. Hsu, J. Li, C. Finn, and R. Grosse. Differentiable annealed importance sampling and the perils of gradient noise. In Advances in Neural Information Processing Systems, 2021.
[147] L. Zhang, P. Potaptchik, G. Deligiannidis, A. Doucet, H.-D. Dau, and S. Syed. Generalised parallel tempering: Flexible replica exchange via flows and diffusions. In Frontiers in Proba- bilistic Inference: Learning meets Sampling, 2025.
[148] L. Zhang, P. Potaptchik, J. He, Y. Du, A. Doucet, F. Vargas, H.-D. Dau, and S. Syed. Acceler- ated parallel tempering via neural transports. arXiv preprint 2rXiv22502.10322 2025.
[149] Q. Zhang and Y. Chen. Path Integral Sampler: a stochastic control approach for sampling. In International Conference on Learning Representations, 2022.
[150] W. Zhang, H. Wang, C. Hartmann, M. Weber, and C. Schütte. Applications of the gross-entrope method importance sampling and optimal control of diffusions. SIAM Journal on Scientific Computing, 36(6):A2654 A2672, 2014.
[151] X. Zhang. L. Wang, J. Helwig, Y. Luo, C. Fu, Y. Xie, M. Liu, Y. Lin, Z. Xu, K. Yan, et al. Artificial intelligence for science in quantum, atomistic, and continuum systems, arXiv preprint arXiv:2307.08-44 2023.
[152] M. Zhou, J. Han, and J. Lu. Actor-critic method for high dimensional static Hamilton-Jacobe Bellman partial differential equations based on neural networks. SIAM Journal on Scientific Computing, A344,,64043,, 2021.
[153] Y. Zhu, W. Guo, J. Choi, G.-H. Liu, Y. Chen, and M. Tao. Mdns: Masked diffusion neural sampler via stochastic optimal control, 2025.
18

---

Appendix
Assumptions and auxiliary results
21
A.1 Additional notation
A.2 Technical assumptions
A.3 Useful identities
B Proofs
22
Further related works, broader impact, and limitations
22
C.I Further related works
22
C,2 Limitations
24
C,3 Broader impact
24
D Background on SOC
24
D.I Stochastic optimal control
................................
24
D.2 Iterative diffusion optimization
................................
25
D.3 On the initial value dependence of the normalizing constant
25
Details on trust region SOC algorithms
26
E.1
Characterizing the solutions of the trust region optimization problem
26
E,2 Implementation
26
E.3 Variance of the importance weights and trust region bounds
26
E.4
Lagragian formulation
27
Trust region SOC sequences and Fisher-Rar geometry
29
F.1
Basics on information geometry
29
F.2 Fisher-Rant geometry of an exponential family
30
F3 Fisher-Rao geometry of an exponential family of path measures
30
G Trust region SOC losses
32
G.I
Log-variance loss
...............................
32
G.2
Moment loss
................................
33
G.3 Cross-entropy loss
34
G.4 Stochastic optimal control matching via adjoint method
................................
34
G.5 Stochastic optimal control matching via lean adjoint method
................................
35
H Trust regions for probability measures
36
I Diffusion-based sampling
36
I.1 Experimental setup
37
I.2 Evaluation criteria...
38
I.3 Additional experiments
38
J Transition path sampling
39
19

---

J.1 Experimental setup
39
J.2 Additional experimental result discussion
40
K Fine-tuning of diffusion models
40
Classical SOC problems
42
L.1 Experimental setup
42
L,2 Benchmark problem details
42
L.3 Results
43
20

---

A
Assumptions and auxiliary results
A.1
Additional notation
For vectors v1, U2 R2, we denote by ||a]| the Euclidean norm and by 01 22 the Euclidean inner product. For a real-valued matrix A, we denote by (r(A) and A T its trace and transpose.
For sufficiently smooth function f: R x [0, T] + R, we denote by Vf = v its gradient w.r.t. the spatial variables and by 0 and 2z, its partial derivatives w.r.t. the time coordinate and the spatial coordinate Vi, respectively.
We denote by (n, Σ) multivariate normal distribution with mean E IR and covariance matrix E gdxd Moreover, we denote by Unif TT) the uniform distribution on [0, []. For random variables X1, X2. we denote by [[X]]] and Var[X;] the expectation and variance of X1 and by XX] the conditional expectation of X given X2.
A.2 Technical assumptions
Throughout our work, we make the same assumptions as in [42, 87], which are needed for all the objects considered be well-defined. Namely, we assume that:
The set 1 of admissible controls is given by
4 = {u C' (IE" [0, > o, v(z,s) E I [0, u(x, s) - C(1 + [|=]])) (20)
(ii) The coefficients and a are continuously differentiable, has bounded first-order spatial derivatives, and (ao "xx s) is positive definite for all (x, s) E Rd x [0, Furthermore, there exist constants C, C1, > 0 such that
[1b(,,)) s)ll & C(1 + |=]])
(linear growth) (ellipticity)
(21)
ei |||1 S 8 (oo (ooT)(z,s) s)B cal B1? for all (z. Idd, x [0, and B Rd.
A.3 Useful identities
Definition A.1 (Controlled SDEs). Let u u be control function. controlled and uncontrolled stochastic processes defined via the SDEs
Throughout, we consider
dX,"" (b + ou) (X" s)ds a(s)dW,,
XØ po,
(22)
dX, = b(X., s)ds + (s))W.
Xo ~ po,
(23)
where X ~ P", X ~ P. with T and P denoting the respective path space measures, and W is standard Brownian motion.
Theorem A.2 (Girsanov's 'ss theorem for path measures). Let u, u, E U. Then the Radon-Nikodyn derivative between nou and IP", evaluated along X is given by:
dpd log d (X*) = {" o - (u=x(X".X" dx." 2 [./
{" (bo-'b+ ull - -lo-llb-oo + oll?) (x,", s)ds,
(24)
Proof. See, e.g., Lemma A.1 in [87] or Appendix E in [131].
Corollary A.3 (Change of measure identities). Let u, E U. From Thm. A.2. we obtain the following
useful identities:
(i) log dp"
(X 10" u(X", 8) dW, fe |lu(Xy ([(XX,s)][+ s)|s
(i) log dp
dp"
(X) = _ u((x,,)) 8) dW, 5 T [lu(X,, s) 1
(iii) log dr" (X") = fo (u (X ddW, + 1/2 Sot ||u vl|"(XX, s)ds
(iv) log dp" a (X*)= Sh (u (i-v)(X3) () 8) dW, - 1 T ||u = s)ds fe
Lemma A.4 (Itô's formula). Lei X, solve the SDE
dx, = b(X, s)ds + (s)dW.,, and let f: Rdd [0, + TR. be smooth function. Then df(,,, (0, ()f s)ds + of ()(x,,s) dw..
21

---

where L is the infinitesimal generator given by
d
L
(oo' )is0,, 0n; + Σ
b. (x, t)0s..
T==1
B Proofs
Proof of Prop. 2.2. Let P be the measure defined by
oF AP"
dQ a
/2, where 2 is the
normalizing constant. Then we have that
dP
dF"
dp"
dF"
dlP
(1 + Xx) log
(1 + lo)
(1 Ac) log
log
+ (1 + ) Z
dP", di
dP
dQ
(25a)
x log dP",
dIP"
+ log
dIp
+ (1+ log 3.
(25b)
dQ
Using the definition of the Lagrangian in (5), this implies that
(1+ (.),))Da((PP (P) A; DEL (P*))" + DKL (P* + (1 A. log ((x))
(26a)
C (u, A.) + (1 + A log [(x)) Xce,
(26b)
Since we defined the minimizer of the Lagrangian (with optimal multiplier A₂) in the last expression as U,+1, we have that u,-1 = arg min ( (P) This shows that P = 1P by the uniqueness of the Radon-Nikodyne derivative. For the second statement, we introduce the unnormalized path measure p such that
dPP
(x)
dp
(X)
with E (Xo)
ddP"441 dP
(x) Xo
(27)
dP
2 (Xo)
dIp
and
df
dQ dIP
dpu dIP
us
IFAt
(X)
(X)
(X)
(28)
dIP
Assuming Zo 1, we have Puo I to for uo
U. By induction, it follows that
d (X)
dQ
A1-1
dIP
1-8:41
(X)
(29)
dIP
dIP
dP
with 8141 defined in Prop. 2.2, which proves the second statement.
Remark B.1. Using the left side of (27), we can rewrite the normalized version of (28) as
dP "1+1
1
dQ dP (x)
di "
143.
(X)
(X)
(30)
dP
2,41 (Xo)
dP
1
dQQ
dIP":
142
dP (x)
dIP (x)
(31)
2.t((X))
with E,4
1
C
Further related works, broader impact, and limitations
C.1 Further related works
Monte Carlo estimator. In theory, one could directly compute the optimal control using the representations in Prop. 2.5 (for
and = 0; see Item Thm. D.I) combined with Monte Carlo estimates of the value function in Item 4 in Thm. D.I [39, 73, 74, 121, 130]. However, in practice this can be problematic since it requires a large amount of samples for each state & due to the (typically) very high variance of the estimator for V [130]. In particular, note that the variance translates to a bias in the control due to the logarithmic transformation. Morcover, for nonzero f or
*One can obtain derivative estimates using adjoint states (as defined in Sec. 2.2) or using reparametrization tricks if the uncontrolled process has suitable, known marginals. For Gaussian marginals, one can also use Stein's 's lemma [73]. We further note that control variates for such estimators have been analyzed in 196, 105].
22

---

general b (e.g., in the fine-tuning setting), one needs to simulate the uncontrolled process to obtain samples.
PDE solver. One can also leverage the representation of the value function as the solution of an HJB equation (see Item 3 in Thm. D.1). While solving PDEs in high dimensions is very challenging. there exist scalable approaches based on tensor trains and neural networks" that leverage backward stochastic differential equations or the Hopf-Cole transform in combination with the Feynman- Kac formula [6, 11, 13, 58, 95, 96, 99, 100J. However, in practice, we only need the value function in the domain where the optimal path measure has sufficiently large values, which is typically not considered for PDE solvers.
Iterative diffusion optimization. To focus more on promising regions of the path space, methods for iterative diffusion optimization simulate stajeetraessssrsss of the controlled SDE to compute suitable loss and update the control. Typically, the control is parametrized as a neural network and optimized using variants of stochastic gradient descent. While such methods have been explored for general SOC problems with quadratic control costs [40, 42, 87, 95], many recent works have focused on the special case of sampling from unnormalized densities as described in Sec. 3.1; see, e.g.. [8, 15, 86, 112, 130, 131, 144, 149]. From the perspective of path measures, all these works propose to minimize suitable divergences between measures induced by controlled SDEs. While we demonstrate the benefits of leveraging trust region methods for the Denoising Diffusion Sampler (DDS) [129], our method could also be extended to other samplers.
Transition path sampling. Transition path sampling has been a longstanding problem in physics and chemistry to understand phase transitions and chemical reactions, with applications in energy, catalysis, and drug discovery [20, 128]. Computationally, MCMC-based approaches have been extended to path space to mix the transition path distribution, pioneered by [37]. As discussed in Sec. 3.2, transition path sampling can be formulated as a stochastic optimal control problem and has been numerically solved using reverse KL divergence [141]. cross-entropy divergence [71], and log- log-variance divergence [1 [113] the optimal control is known to be the Doob's s h-function [26, 45, 116] (for review, we refer to [115]). To solve the Doob's s h-function, [116] proposes a shooting- based method which requires MD simulation to reach the target state, while [45] proposes a Gaussian approximation conditioned by design and provides a simulation free optimization algorithm. Similarly to SOC, transition path sampling can also naturally be formulated a reinforcement learning problem, as shown in [35, 103).
both the initial and target state which satisfies boundary conditions
Diffusion and flow matching reward fine-tuning. Several of the early works on diffusion fine- tuning focused on directly optimizing the reward model making use of its differentiability [28, 140), without any KL regularization, which can lead to "reward hacking". Some other works [16, 47] framed reward fine-tuning as a reinforcement learning problem. but did not make the probabilistic connection to tilted distributions. [126] provides probabilistic view of the problem, but proposes an algorithm that is hard to scale. [41] gives a comprehensive view of flow matching reward fine-tuning, introducing memoryless noise schedules as the right ones, as well as a new scalable SOC algorithm that we use and adapt, namely adjoint matching. Using the memoryless noise schedule, a recent work [81] considers GRPO for flow matching fine-tuning. [82, 145] consider alternative algorithms that learn the value functions.
Diffusion-based sampling from unnormalized densities. Early work on sampling from unnormal- ized densities based on a Schrödinger-Fölle Föllmerrrrrrrrrrrree diffusions dates back to [33] and was later implemented using Monte Carlo [39, 73] and deep learning approaches [95, 130, 149]. Another line of work is based on Langevin diffusions [43, 52, 53, 124, 146] and denoising diffusion models based on Ornstein- -uieennee processes [15, 74, 129]. A unifying perspective was proposed in [97, 131], which consider general diffusion bridges. An extension based on underdamped diffusion processes was later proposed by [17]. Recent developments have led improved loss functions and training schemes [6. 8, 14, 27, 46, 56, 63, 64, 80, 89, 108, 114, 117, 133, 143, 144], exploration capabilities [19, 77, 78], or normalizing constant estimation [57, 65]. Other studies focus on the combination of MCMC and diffusion-basse
sampling methods [5, 25, 101, 112, 136, 147, 148). Approaches for discrete state spaces have been proposed in [70, 107, 153). Combinations of diffusion-basel sampling with additional access to ground truth data have been studied in [86, 120]. Lastly, [18, 55] study improved evaluation techniques.
"Note that some of these approaches correspond to regressions of the Monte Carlo estimators mentioned above [130] the IDO methods mentioned below [117].
23

---

C.2 Limitations
While our method for solving stochastic optimal control problems exhibits strong sample efficiency, relies on storing entire trajectories in the replay buffer during training. In large-scale settings such as fine-tuning lext-to-image models this necessitates keeping the replay buffer in CPU memory while training occurs on the GPU. This separation introduces additional computational overhead due to data transfers between CPU and GPU; however, the buffer still significantly accelerates the fine-tuning since the main computational cost in such settings stems from the simulation of trajectories.
C.3 Broader impact
This paper proposes new methodologies and theories that find numerical solutions for stochastic optimal control problems ranging from equilibrium sampling, transition path sampling, fine-tuning text-to-image generative models. Equilibrium sampling and transition path sampling are important in Bayesian statisties, physics and chemistry where they can be used estimate free energy, understand phase transition and rare events, thus holding promises to accelerate drug and material discovery. More efficient fine-tuning of sext-to-image models democratizes the generation of specialized high- quality visual content for creative applications. However, these capabilities also introduce risks such as the potential for generating convincing misinformation or deepfakes.
D Background on SOC
D.1 Stochastic optimal control
In this work, we consider stochastic optimal control (SOC) problems of the form
L." (all((X," $) + f(X", s)) ds + 9(x
min (soc(u) min E WEM
(32)
WEM
with state costs f. terminal costs 9 and control function U. where U denotes set of admissible controls; see App. A.2 for further details. Here, X" is controlled SDE of the form
dx,"" (b + ou)(X." s)ds + (s)dW., Xo po,
(33)
with base drift b, base distribution Po (typically a Gaussian or dirac delta distribution), and diffusion coefficient a. We denote the path measure induced by (33) by Pa E P. Moreover, we simply write P' for the path measure corresponding to the uncontrolled process, i.e..
dX. b(X.. s)ds + (s)dW,,,
Xo po-
(34)
Given time and state v, the cost functional J(; x, t) the expected cost-to-go for a control on the time interval [t. T) and is defined as
(u; t) =T= L." (HIlv(X), +) + f(X), s)) ds +((
X =
(35)
The value function V, or, optimal cost-to-go is obtained by taking the infimum over all controls in U4 that is,
v(x,t) inf z(u; well
(36)
Then we have the following well-known results on representations of the value function V and solution to the SOC problem u"; see, e.g., [33, 50, 87, 91, 94] for details.
Theorem D.I (Optimality for SOC Problems). Let us define the work functional as
1(X,))= J." s(X..s) + 9(XT).
(37)
Then we have the following representations of the value function V in (36) and the solution u" to the SOC problem in (32);
1. (Connection between solution and value function) The solution can be written as u² = - VV.
2. (Optimal change of measure) The Radon-Nikodym derivative of the optimal path measure Q w.E.t. the uncontrolled path measure IP satisfies
dQ dP
(X)
c -W(X00)
with Z(Xo) = Bie -
(38)
2(Xo)
3. (PDE for value Sunction) The value function V is the solution the Homiilonn-accoii-Bllllnn
(HJB) equation
(0, + LLV 211(0 [VV) + f(x,t) =0, V(x,T) g(x),
(39)
24

---

where L :: 2 -لك==(0o uncontrolled SDE in (34).
(oos)no")) + Ei b,02, denotes the infinitesimal generator of the
(Estimator for value function) For every (x, E Rd x [0. T] the value function can be written as [x, where X is the solution of the uncontrolled SDE in (34).
((x,t) -og [ W(Xx))
Combining the expressions for u" and V in Thm. D.I, we directly obtain the path integral representa- tion of the optimal control, i.e.,
u (x, = o(1) V. log w(X
X, =
(40)
In practice, computing the optimal control (40) is typically impractical, as it requires running multiple simulations for each state to obtain Monte Carlo approximation of the expectation; see App. C.I. To address this challenge, many approaches instead learn a parameterized control function, optimized using stochastic gradient methods. These techniques are collectively referred to iterative diffusion optimization (IDO) methods and are further discussed in the next section.
D.2 Iterative diffusion optimization
An alternative view on problem (32) is obtained by considering loss functions on path measures [87].
By the Girsanov theorem (see App. A.3) we have
dP dPP (x")
f/" "(X" dw. 2// (((xX,s)]]'
exp
(41)
Combining this with the optimal change of measure dQ/ /dP from Thm. D.1, we obtain an expression for dQQ /diP", from which we can compute the relative entropy CRE, i.e., the reverse Kulbback-Leibler (KL) divergence
CRE(U) = (KL((")Q) = B [[ (: [lu(X" s)| + f(X"' s)) ds ((x)) + log 2(X;)
(42)
Note that minimizing the stochastic optimal control problem in (32) is equal to minimizing the KL. divergence, that is,
arg min Lsoc(u) = arg min (RE(U):
(43)
in the sense that both have the same unique optimal control u² as a minimizer. As such, we can consider an arbitrary divergence D P P - IR+ for which D(P (P2) O holds if and only if P1 = P2. to solve stochastic optimal control problems. More generally, we can consider any loss function for which the unique minimizer is the optimal control =². Iterative diffusion optimization builds on this perspective and can be seen as a common framework for solving (potentially high- dimensional) SOC problems by leveraging parameterized control functions and stochastic gradient methods to minimize different loss functions.
D.3
On the initial value dependence of the normalizing constant
In general, the normalizing constant (Xx) in the optimal change of measure (3) depends on the initial value Xo. Let us demonstrate in the following why this is the case. To this end, let us first assume generic normalization constant Z that may or may not depend on Xo- As in (3), it then holds
dQ dTP
(X) = e (((.)))
(44)
2
We can then compute
(Xp)
ddQ
W(X.0)
E
(X)
Xo EE
Xo
(45)
(X,)
dIP
Now, for a chosen Po = Po we want that (X = Po which requires
2=E
w(X.)))
[x] = V(Xo.9)
(46)
Clearly, the right-hand side depends on Xo. Hence, in general, Qx(x)) = Po(X)) can only hold E depends on Xo- Conversely, if we wanted to have a global normalizing constant 2, which is independent of Xn, we would need to tilt the initial marginal of Q as well, namely via
Qo(Xo) = (X)))) Ele w(x.)) [x)]
Po())) e V(X..0)
(47)
However, the function ((, 0) is typically not known in practice.
25

---

E Details on trust region SOC algorithms
E.1 Characterizing the solutions of the trust region optimization problem
Proposition E.1 (Characterizing the solutions of the trust region optimization problem). The solution puct of the problem (9) is unique and it satisfies the following:
If (DDL(QQ I E, then P Q2.
If ( (Q 2 E, then DKL (P (I
E, i.e. past is also the unique solution of the
problem
arg min DEL (P"]Q) s.f. (P
(48)
uEu
Proof. To prove the first case, observe that Q is the only solution of the unconstrained problem arg milppp DKL (P) which means that it also the unique solution of the problem (9) since it satisfies the constraint DKL [Q J S To prove the second case, by the Karush- Kuhn- Tucker We assume that À = 0 and
(KKT) conditions, we have that either À 0, or (KL(P%:)) [P",)
DEL (IP [P= < reach contradiction, which will imply that 1KL(II [P", s. The
first-order optimality condition for the problem is as follows: for any perturbation e of the control 1,+1, we have that
(Dir. (P"41t""(()) Q)) +
- +))
DRi. (P= drj
d
dry
(49)
which means that Us+1 satisfies the first-order optimality condition for the relative entropy loss u ト+ DKL (P")Q). By [41, Prop. 2], the only control that satisfies the first-order optimality condition for the relative entropy loss is the optimal control u", which implies that IP Q. which yields contradiction because & > DKL. (PP. [TP":) = DRL. (Q[P")) 2 E.
Hence, we conclude that DKL (IP*++ IP4) = 8. To show that the solution Puu-1 is unique, we use that IP + (KL(PPPPPP) is strictly convex, and that (PI IP"s) S &} is convex set because is the sublevel set of a convex mapping.
E.2 Implementation
We provide a detailed version of Algorithm in Algorithm 2. The hyperparameters and used repositories for the experiments unnormalized densities, transition path sampling, and fine-tuning can be found in the respective sections in Apps. I to K.
E.3 Variance of the importance weights and trust region bounds
As mentioned in Remark 2.1, one motivation of the trust region constrain DKL (P* S defined in (4) is keep the variance of the importance weights between two consecutive measures Po, and pu small. This can be motivated by the inequality
dP dIPu:
dPP
Varp
Ep,.
(50a)
dlu
dlp
dIP"111
exp
Eper+1
log
-1 = exp (DKL (P (P"H))P"))) 1,
(50b)
dP",
which follows by Jensen's inequality. While a lower bound on the variance is not straight forward for path space measures (cf. [60]), we can consider the following heuristics. Let us assume that
dP
(51)
dP
p and -lmostt surely, which reasonable DKL (P PP4) < e with < By a Taylor approximation it then holds
dP 4P
de dPw
exp
log
1
2 2og
(52)
dP
d
Now. taking expectations w.r.t. P 11, respectively, using computations similar (50), and assuming DKL (P" IP:) = E, as argued in App. E.1, yields
Varp., ddP"+
2s.
(53)
dIp":
26

---

Algorithm 2 Trust Region SOC with buffer
Require: Neural network Ue with parameters target path measure Q, buffer size K, time discretization = (s:) C [0. TT, number of gradient steps M per trust region iteration, termination threshold Initialize 1 0 and Äg xc
for 0, ... do
Define u, us (detached)
Simulate K trajectories (Xx*) ).Es of the SDE in (1) with Brownian motion W and control u; (X oxp(-W.(X* (k), .0)) as in (12)
Compute importance weights w (8)
dQ
ar
Initialize buffer {(W()) (*)), w )
)
Compute multiplier A = arg maXXER+ Dual (A) as in (14) using and 1-dim. non non-linear solver if X, S 8 then
return control u, with IP: 2 Q
if adjoint matching loss then
Compute annealing 3241
IT
Ag as in Prop. 2.2
1+ Compute lean adjoint states a² a:41 (X. s). E S, in (18) and store in B for I..... .M do
if adjoint matching loss then
Estimate (0) = (x, soUnif(S) |||o Ta.
a. - uo(X.., s)[l TixT as in (19)
if log-variance loss then
Estimate [()) (ar |> ( A,
Wa,- 1) NWW,,yWW,,. 1④x, log w)
with A, u; (X.,: 85) us (Xa,, sy 85) in (16)
Perform gradient-deseee step on [())
E.4
Lagragian formulation
Using the Girsanov theorem (see App. A.3), we first note that we can write the Lagrangian as
((0) (u, A) = E L' |I((",)) w(X", 0) log 2(x + (PKK(P" =)
(54)
F. Je (bl s)11 llx <x"x" d W(X",, log 2(X)))
Aa
(55)
1-' (24a (u(X".)) A f((x, ds g(X7) log 3(x))
Ae
(56)
Cik (u, 2) de,
(57)
where ( A) is defined as in (11), A E R+ is the Lagrangian multiplier for the trust region
constraint, and we abbreviate fi 2(++) | | f. For fixed X, optimizing the Lagrangian
clil (u, A) with respect to u is again an SOC problem. As such, for given u, and X, we can define the value function as
(Z, t) inf IE [."" (1L [uxXX, (X," s)ll f((X,, s)) ds o(x
(58)
well
The next proposition provides representations for the value function and the solution to the SOC problem.
Proposition E.2 (Optimality for trust region SOC problems). For fixed X let us define by
S." [ (
(=) inf E NEN
A ul + أيبي ||u 1|2 + 1) (X", ds 9(X7) X.
the value function of the SOC problem inf C , A) corresponding to (11) and by uby its solution. Then it holds that
27

---

(i) (Estimator for value function) V 1 = -11 A) log E e 1+2 m Xx" 2
where (I(X"*,11= J J," (X",, (.)x"" s)lpdss + J." ui(X?"" .s) dW, + W(X"),1)
(i) (Connection between solution and value function) holds uAt - o VV TVV
Moreover. for uo = 0 and the optimal Lagrange multiplier Xo let us define the value function
J." (allull B (X", ds B441
(an( t) = inf
EE
X1
wEM
of the SOC problem given by the optimal change of measure
dp
(X)
1
(QQ
(x)(
dP (x)
A6+1
e P(41
(59)
(xo)(Xo)
(i+1(Xo)
as in Prop. 2.2 and (3), and (Xo) defined in (27). Then it holds that
(iii) (Estimator for value function) (Ir, t) = - log IE [e 9i- P.-1W(Xx, [X+ = z).
Ar (iv) (Connection between solution and value function) u:+1 uft2
VV++1-
Proof. For notational convenience, we abbreviate V= = Vi in this proof. From the verification theorem (see, e.g., [94, Theorem 3.5.2)), we obtain that the value function is the solution to the HJB equation
(de + LJV = - inf
{
fi +
L
la uil aa V
(60a)
ag3d
1+ la Tेx uill + vv). V(,T) =
= = - inf
(60b)
aekd
where the infimum is pointwise for every (x, 1) E Rd [0, T and the optimal a* defines the solution w". Solving for a yields a² = - ", _ VV, which proves Item (ii).
Plugging this result back into the HJB equation, we obtain
(Ge LJVV- ff- शाイ |
शाम
V
(TAx u; o VV
(61a)
=-f 2, =-f x((
||ull? + शाナ
VVI2 ou:
ou; VV
(61b)
||u|| + 2 loo VV - ou, VV + + VV. .V.
(61c)
Now, we define the infinitesimal generator of the SDE
dx. = (b(X,' 5) + oxn ds + odW,
(62)
as
1 (ao )r Bzg > (b. + (ou))))2 = L (uu)
L
(63)
2 4,j=1
Using (63), we can rewrite (61) as
(0. + L)V = -f 2(422 uill" + マाス T V 1 VV
(64a)
= =-f - I1lua11² + शम ||u + o VVI
(64b)
By Itô's formula (see App. A.3), we have
dV 8) = (0, + LXV ,8dd + o vV s) dW.
(65)
Plugging (64) into (65) and defining Y. ((X": 5) and Z. (-u; - (Xyi s), we obtain the pair of forward- forward-backward backward SDES (FBSDEs)
_x.:t = (b(X,", 8) + (X." +)) ds (s)dWW,, X ~ ~po
(66)
dY (-f(X, 8) [X s)]l + x 1 ds (u+(X)* .8) + Z,) -dW,
(67)
with Yr = 9(x4) ). This shows that
((x#) = Yi - 1.¹ (I(X)),) s) + 2x, s)]l 2nttl [za112) ds 1. (mi(X)", ,)) Z.) dW.,
28

---

which can be rewritten as
W .1) Yr S." 2(+2) [2,11 dds - [ Zo dd.dW...
(68)
Using the definition of Y, we can now write
I4X wiXX 9) X =c = I+X V(X
el+A Si 2. (axaa 1i" 212,1/ xx IB
II
where we leveraged Novikov's theorem to show that the Doléans-Dade exponential a martingale with vanishing expectation. This concludes the proof of Item (i). The proof of Items (iii) and (iv) follows directly from Thm. D.I.
F
Trust region SOC sequences and Fisher-Rao geometry
For fixed 6, suppose that we construct the sequence of controls (1:+1)i20 as the solutions of the problem (4), As shown in Prop. 2.2, we have that
dp"
dQ
A:
/dIP"o
1-4
with
B
II 2=0
(69)
dP
dP
dP
142
If we define the family (Q(1))
[(o.) such that
dQ(")
dQ
dlP
(70)
dP
JP
dPP
we can write Ia; the family (Q())) Next, we characterize this discretization more precisely using tools from information geometry.
Q(8.), Hence, we regard the sequence (P"))220 a discretization of
F1 Basics on information geometry
Let [p(z; } eee be parametric family of probability densities (or mass functions) on the sample space X, and let X be random variable with distribution p(x; 0).
Definition F.1 (Fisher information matrix). The Fisher information matrix at is defined as
( :x- Vo log p(X;0) 0) (Vo logp(X;))) ]
-x v log (x:
where V denotes the column gradient with respect to 0, and V the Hessian.
As average of positive semi-definite matrices, I(0) is positive semi-definite, which makes it possible to define geometric structure:
Definition F.2 (Statistical manifold). Let {p(x; leee be smooth parametric family of probability
densities on X, with parameter space θ C Rd. Then e itself can be viewed as a d-dimensional differentiable manifold
M = {p(:;0) A E Q @ O called the statistical manifold of the model. Endow M with the Riemannian metric
(0) = (0) Ex-p(-0) a (XX 0) 0, log log p(X:))
where 0 2 This known as the Fisher-Rao metric, turning (M,g) into the canonical enformttion-geomerri manifold of the model.
30,
Next, we review the definition the length curve on Riemannian manifold. Definition F.3 (Length of curve Riemannian manifold, and let 7: [a, b] - M be a piecewise smooth curve. Choose local coordinates f (0). ..... (d)) ((1) = (0 (t), (t).... (t)). Then the length of 7 is
Riemannian manifold). Let (M,g) be d-dimensional
open set U C M containing the image of 7, so that y(t)
M(7) = 1: Vou (0(1) (B(t)) (1) dt,
d0'
where (t) i.j 1,........... .....,d. d.
(t) and employ the Einstein summation convention on repeated indices
dt
29

---

A geodesic between two points 01.02 E M is piecewise smooth curve 7: [a, bj - M such that 7(a) fI. (6) 8 that minimizes the length functional locally. Any time reparameterization of geodesic is also geodesic, because the geodesic distance between 0,, 02 is the infimum over the lengths of all geodesics (or all piecewise smooth curves) between 0,, 02.
Definition F.4 Frishe-RRao distance). The geodesic distance induced by the Fisher-Rac metric is known as the Fisher-Rao Rao distance.
Lastly, present another statement which connects the Kullback-Leible divergence and the Fisher information matrix using a local expansion of the KL divergence. Proposition F.5 (Second-order expansion of KL). Let {p(a; b)ee be smooth parametric family of densities, and fix 0 E O For a small increment E IR,, consider
kL(p(:; + +1]1[((00)) [[)(=)) = J p(=;0 + log
p(z; 0 + dr.
p( 0)
Then one has the Taylor expansion
KL(p(0 + ()])p(0)
+ 8" (0) 8 + 0(1819).
where
(0) = :x |0, logp(X;0)00 1ox((;
is the Fisher information matrix. Equivalently.
0KL
8KKL
O,
06
08
6=0
6=0
Sketch. Expand both p(z; 0 + 8) and log p(:; 0 + 6) to second order in 6, substitute into the integral, and use J O log p dr = 0 and f p O log pdr (0) to verify cancellation of constant and linear terms, leaving the stated quadratic form.
F.2 Fisher-Rao geometry an exponential family Definition F.6 (The exponeniial-family manifold). Let
plæ; 0) (xp(O'T.(x)) A())) fi(e), 0= 69
be a regular d-parameter exponential family on X. The parameter space (equipped with the atlas coming from the coordinates 0') is d-dimensional differentiable manifold, which we identify with the statistical model
M = { :0) e)
Its tangent space at 0 T&M Idd, with basis {0/00'}.
Definition F.7 (Fisher-Raal Rao metric). The Fisher-Rao metric on M is the Riemannian metric whose components in the natural coordinate chart 0 are
A(0)
(0) Ex-pp.:0) [0: log p(X; 0, log (xx)( Ex= [a, log ((;00]
00 a0i
Equivalently, g(0) 2 A()), the Hessian of the log-partition function.
For general exponential families, the Fisher-Rao distance and the geodesics do not admit closed form. Yet, one-dimensional families can be handled explicitly, because geodesics are trivial: Proposition F.8 (One-parameter exponential family). If then (a, b) C R, and g(f) A" Hence
-1/ V (a) de
FR(01. 02)
F33 Fisher-Rao geometry of an exponential family of path measures
We can view the family (Q(r))re(o.)) defined in (70) as a one-parameter exponential family [22] by rewriting ((()) as
dg(+)
dlPuo = exp ( log
dQ dlP"o
-(T))
(71)
where the log partition function A(T) is defined as
dQ dIPno
)
A(r) = log Epso
(72)
30

---

Equivalently, we can write it as an exponential family centered on arbitrary E [0, 1]:
dQ (r+r))))
dQ = exp AT log
Ar(A7)) dTPuo
dQ
(73)
where
dQ
Ar
Ar log B
(74)
dIPug
Deriving an expression for the Fisher information. Observe that by construction
(
dQ
d0(r)
dQ dlP
dQ
exp (-A(T))
A, (Ar) = log Ep
= log Bp= m0
dPP"
dlP
dIP
(75)
= A(T + AT) A(r),
which means that (0) (T) for r E (0, 1). Thus, by Prop. F.8, we conclude that the Fisher information matrix, which is a scalar because the manifold is one-dimensional, reads
I(r) (T) = (0).
(76)
Computing the first and second derivatives of A, is ftrairht-forward:
log
4Q ap
ap
Ar
AC(AT)
aQ Ar TPpl
Bager)
(77)
dQ dlPw
)"1
dQ dlPua
)}
A²(0)
Eg(r)
log
18 B30
log
and this implies that
dQ dlP
I(r) = Varg(=)
log
(78)
Connecting the trust region constraint to the Fisher information. Applying Proposition F.5, obtain that
KL (L(((r+A7)100))
A-3 I(7) + 0(AT*),
(79)
2
When we set AT = B141, T = B₁ we have that
A+
I(T) + O(Ar").
(80)
2
Thus,
AT
28
+ 0()r3/2).
(81)
I(r)
Moreover, the Fisher-Rao distance between th and p(i), or rather, between 0 and 8
FR( J. 02
V(())dr.
(82)
Then, the difference between Fisher-Rao distances FR(O, B141) and FR(0, 3:) which equal to the Fisher-Rao distance FR,,, 8(+1) is
FR(O, B(+1) (O, B.) FR(B,, 34+1)= 8:11
I(r)dr
= (VI(B.)) + 0(:++ - (8+)1 - 3i) VI(B.)A + O(A+3)
(83)
=V((B)) I(B.) 3(3)3/3) 228 3(3)313).
In continuous time, we have curve 3 R=00 [0. 1). and
d
2 FR(0, VI(3(1) = VI(3(T) I(3(t))
= V2
(84)
dt
Thus, we have shown the following result:
Proposition F.9. Up to high order terms, the elements of sequence (IPa: )osiss-1 are equispaced in the Fisher-Rao distance. The last term pus is equal the target distribution Q
31

---

A Monte Carlo estimate for the Fisher information. By equation (71), we have that log d0 A(T). Hence, we can rewrite (78) as
aP"6
T [log _Pno
1 Varger)
dQ dQ(
I(r)= = 7²
log
(85)
dIP"
which provides way to estimate I(r), leveraging the Girsanov theorem estimate log dP"6 log
(d((r)
ag log
d dP
Remark F.10 (Analytical computation of annealing sequence). Using T AT 1.+1, & paired with (81) and (78) we can analytically compute the annealing sequence (B:),, up to high order terms, as
Bit1 = 3 + V
2x I(3.
+ O((3++1 - 3,)
(86)
with
I(3,) Varp
log
(87)
dPup
where we used that ((7) I
G Trust region SOC losses
In this section, provide non-exhaustive list of losses that be readily applied solve SOC problems within our trust region framework. More specifically, we aim for minimizing a divergence D P It between the path measures induced by the control u and the learnable control u. For a comprehensive overview of SOC losses without trust regions, see [40].
G.1 Log-variance loss
Here, we provide further details the log- variance loss [87, 97, 98] within our trust region framework. The log- varvaariance loss is defined as
ddP".
(xr)) with X P" dPP
(u) Var log
(88)
where X is defined as in (1), with u replaced by w E U, referred to as the reference process. Although the choice of is arbitrary, discuss particularly suitable options that facilitate sample reuse with replay buffers in combination with trust regions.
Using w = ut as reference control. First, we replace the reference control w with the control function of the previous iteration ug. Thus, the log-variance loss becomes
(u) = Var
log
cIP 11
dlu (x*))
]
Var log
ddP"++1
(X
dP
(X
(89)
dis..
dPu
The Girsanov theorem (see App. A.3) shows that
dP" (X*) (x*)) = 2 /
1 /-' (X." u(X", ,)),, dds + I" (w, u)(X" dW,
log dpu
(90)
Combining this result for u = with Prop. 2.2, we obtain
dP dP ()
dQ dP
dTP dp (X"))
(91a)
= 14,, (J 1 .a) Pas+ fo' ".
(91b)
Noting that the variance is shift-invariant, (90) and (91) imply that
I+ (2//" (ux("""" d) + /" ux(x" -.. dw. w(x
(iv(u) = Var
(92)
L/" [|un(X?' u(x"' s) . + {
1- (m ()(x' ()('
dw:
which can be implemented by discretizing the integrals; see App. E.2.
Please note that the loss reduces to Var
Lix (u)
log
dQ p..
(x
dlP
(xx))
) Var log
dQ
(xx))
(93)
dp
dlP
32

---

for A = O, which is how the loss is mostly used in the literature, where the variance is computed using the most recent control, see e.g. [97].
Using w 1 as reference control. Alternatively, when setting the reference control to the
optimal control u4+1- the log-variance loss is given by
dP
(u)(
Var
log
(X (X
(94a)
dlP
dP'4+1 (X*it1)
)
27
dIP"+4 dlp
(x ))
E
log
EE'
log
(94b)
dPn
[dP
2 dP
(X
dDP":
dI dlp
E
log
(X*1)
E
log
(xx1)
(X
dPn
dl
dlp
(94c)
which can be computed using (90) and (91). Hence, contrast to using w = u, (94) additionally incorporates the smoothed importance weights dIP /dP"..
G.2
Moment loss
The moment loss was introduced in [59] and is defined as
dIP
#++1
Emoment (u)
E
log
(X
with X" ~ P",
(95)
dp
where w E U is again an arbitrary reference control; see G.I. We distinguish again between u = 04 and w U++1.
Using w ui as reference control. this case, Lmmomen becomes
dP dlp
(Xm)) (xxm:)
)
dPP dipu.
dp
dlP (XxL) )
)
mommen (u)
log
log
(x
(96)
with
T+A; ..
dIP"++1 dP
(X":)=
with
Z 1+A, W.(x"
(97)
2
Contrary to the log- lorvvariance loss, the moment loss is not shift-invariant, thus requiring 2 which is commonly not available. As such, [59] proposes treat Z,+1 learnable parameter. Using (90) and (91) imply that
Cmoment (u, Z++1) = TF. (T((((
[[[[[T]] - Ai (2//
[[11l1u [x,, d)] (' "x''"" 8)- -dd W(X .0)
1 + 2 / |lu(XX"', s) u(x"" d) + 1/" (m (x . dW, log Zi41
(98)
which is optimized as min Z Lmoment (2, 2,_1).
Using U441 reference control. Using Us+1 yields
Lmoment (u, 2(-1) E
log
(X"+1)
(99)
dPu
dlp
dlp
d
E
log
(X"))
dlp (X":)
(Xxr)
(100)
dpu
dpu;
where d[P..4 /dIP depends on 2,+1, see (97). Hence, the difference between using w u; and u.41 lies in the additional importance weights dp 1/5.-11/0 1/9.i1/dPPP /dP".
33

---

G.3
Cross-entropy loss
The cross entropy loss is defined the forward KL divergence between 24-1 and i.e.
Cce(u)=
DEL (I* [PP))
E
log dlp
(X
(101a)
dI dIP
dr"; (x(+)))
E
log
(X
(101b)
dPu;
dlp
dI
dIP";
di
= E
log
(X
d (X*))
dpc (X
(101c)
dp
Using (90) and (91) implies that
(CE(u)= = I (
2 [ [X,; /-" (x(X",,) 8) dW. W(X";
(102)
2 1 [x, 8) u(X." ,)) I (us u)) ,dW,,
dp dipa - (x"))
log Z,
with importance weights dIP /dP".
G.4 Stochastic optimal control matching via adjoint method
Here, we provide further details on the trust region version of the stochastic optimal control matching (SOCM) loss introduced in [42]. We start from the cross-entropy loss, i.e. the forward KL divergence between 1 and u, that is,
(c DKL. (P"+-1|PP) = log dIP
(x*+1)
(103)
dip
Using Girsanov's 's theorem (see App. A.3), the cross-entropy loss can be written as
1
(ue(u) E
2 [|u+1(X."*+ ,0) u(X.'*1, ,8) dss
(104a)
[| (8) - u(X.' s) dIP dP"m11
= E
(104b)
dlu
Using the expression for the optimal control, u4+1
di 1 V VV++ see Prop. 2.5, yields
I+a;
Ith
F. 2 (
dP"4+1
(CE(")
1 (X? s) I+A.
VVi+ (x .)) u [ 2
(105)
dip
with
V.E.
e 14h, Wi .1)
X =
=
(106)
E e 14a,
WiX
We use the adjoint method [41, see Lemma 5] to evaluate the conditional expectation (106) 10 giving
V. I+x, (.( X = B a(t. u, x 14x, w, x = æ
(107)
where the adjoint state a(l, us, X satisfies the ordinary differential equation (ODE)
Je ā ui x -[-)))-<<(((""' (V(((X.'
oui(X" s)) "(uu, x
(108)
T+x V(((X' 5) + (x })]
(109)
with İTT u, x (Xr). Using the argument from [42, Theorem 1]. replacing the path-wise reparameterization trick with the adjoint method, we arrive at the trust region version of the stochastic optimal control loss given by
2 f A ( s) - o ā(us X u(X", 3) ds di"
LsocM (w) = E
K
(110)
dIpa,
Note that there exist other methods for computing derivatives of functionals of stochastic processes. We refer the interested reader to [42].
34

---

for some K independent of u. However, the adjoint state contains the Jacobian Vu, and the derivative V/lu, which can be expensive in practice. what follows, we rewrite the objective such that can get rid of these terms.
G.5 Stochastic optimal control matching via lean adjoint method
Starting again from the cross-entropy loss, we now employ the alternative expression for the optimal control as stated Item (ii). This yields the objective
1 2 / || - Vini(X., u(X., (X.,,))) s)l
dP
(CE(u) E
(111)
dP
where the gradient of the smoothed value function is given by
V. B |X
V. VaVis
(112)
PLII |x,
We evaluate the conditional expectation using the adjoint method: 8.+ ((X.0) X, x TF [X s)e 1 -
V
=
(113)
where (X 8) denotes the lean adjoint state [41], which satisfies the backward differential equation
ds a.t1(X,, - s) a [(Vb(X.,
d
+ Bon V(X,,
(114)
with terminal condition a:t1 (XT, T) = Bit1 Vg((T). Following the derivations in [42], we arrive at the objective
(sOCM( = E f 11o (X,,,X,, u(X., s) dIP
(x)
(115)
dPf
Finally, performing change of measure the previous control uy gives the expression:
(s) = EE //"' |lo "x,1 "(x,", s)]l s))l ds d|pu+1
diw (x*))
(116)
We remark that the adjoint ODE in (114) can be solved as
(441X,,5 B++1 exp ( Vb(X.,)) Vg(XT)
(117)
f O and Vb(X,,t)b((X = Vb(X,, 8) V)(X,, t) for all 8, E [0, 7] (i.e., the matrices at different times commute). This allows us solve the adjoint ODE exactly for our applications of sampling from unnormalized densities; see App. I.
Extensions for diffusion-based sampling. Consider the case where f 0 and b(x. t) = b (f)r with by C(O. 7), IR), which holds in certain settings for diffusion-based sampling [129, 149]. In this case (117) becomes
a:+1 (X = ( with 7(8)= = exp ( ((
(118)
The SOCM loss in (115) therefore reads
(soCM(u) Tp [." 183))1()) Vo(X)+t1) u u s)]
(119)
From Prop. E.2 Item (ii) directly follows that dP
dP
dlpu, (X) dI
(x) dPp
(X) x -
(120)
d.P
for up 0. Thus, the SOCM loss can rewritten as
2/4 J." [Bi+17(())
( Vv(X#) u(X.' (x)xss))l=d ds d
(soCM(u) = Epu;
dp": (X
(121a)
x Epp;
f [BB4)7(5so T VV((X)) " s)] s (0++) 013o((())))
(121b)
35

---

Lastly, using that PP: = P. PT = P. P.TPPT we arrive at
(s) x E I [1) VerraaaaoVVaaa Vg((7) u(X,' (83" s)] (0
(0++1 -0)+((0)))
(122a)
-f B [118)t17(s)) vg(x7) - s)]| e A ds, (122b)
where we marginalized out all t [0, except for t s. Note that for certain b1. PAtT can be sampled directly, removing the necessity for storing intermediate samples in the buffer.
Trust regions for probability measures H
Our goal is to sample from probability density of the form
Prarget (x)) (x) 2
with Z z= Prarget (z)d2,
(123)
where we evaluate Ptarget but typically do not have access samples from Ptarget- To tackle this problem, one can again formulate this problem as variational problem by minimizing a divergence between some and the target density Ptarget- We can again incorporate an additional trust region constraint, that is, an upper bound on the change of the variational distribution q within a single update step. Formally, we trying to solve the following problem:
/dg =
G++1 arg min DKL Pallarae))) S.t.
(gl())) s,
(124)
where 9 the variational distribution from the previous iteration. We again tackle the constrained optimization problem in (124) using Lagrangian multipliers. The Lagrangian is given by
2 (q. A, w) DKL (g |Parge)) + (DKa(gll -) (ʃ dda-))
(125)
with Lagrangian multipliers A, W. Taking the functional derivative SC (g. A, w/ /g and setting it to zero admits closed-form solution for the optimal density 9--1 as the geometric average between the old distribution and the (unnormalized) optimal distribution, that is,
T#x T+x
qi-1 (A) = arg min CDR (g.
9.
Pearges
with Z(A) = / /dg
Ptarget
(126)
(2)
Plugging the optimal distribution back into the Lagrangian yields the dual function
C (A) = ci (A), = - 1 + ^) log Z(()) - >e.
(127)
Note that we can use any non-linear optimizer for solving for the optimal Lagrangian multiplier by maximizing the dual function, i.e.,
A = arg max ((0) (A).
(128)
XER+
I Diffusion-based sampling
We consider the task of sampling from densities of the form
Ptarget
Ptarges
with
Z / Prarges (z)d2,
(129)
Z
where Ptarget E C(Rd, Rgo) can be evaluated pointwise, but the normalizing constant 2 is typically intractable.
Here, approach the sampling problem by using denoising diffusion-basee sampling based on the work of [129] (see [15, 97] for generalization). To that end, we consider a controlled ergodic Crnstein-Uhlentee (OU) process X (X,)se(o,r) i.c.,
dX = (-((s)x, + u(X"", s)) ds + n (
Xo Po:
(130)
with noise schedule Ç C(10, T]. R), (z) = N (N 12 and corresponding path measure P The target path space measure Q is induced by an uncontrolled ergodic Ornstein- OrnstcinUUleebbec. (OU) process, starting from the target Prarget and running backward in time, that is,
dX.- ((s)
dW..
Xx
Prarget:
(131)
I1 Note the dependence of CfA vanishes as g++1 satisfies the normalization constraint.
36

---

which fulfills Qo = Po for suitable choice of Ç. For integration, we follow [129] and use an exponential integrator. Lastly, it can be shown that the optimal control fulfills
" (z. a) 7V 2(() log P. (z))
(132)
which is later used to analytically compute the optimal control for Gaussian mixture model target densities, see e.g. [129]. Please note that P. = (0, for all [0. T] as the uncontrolled SDE is initialized at its equilibrium distribution.
L1 Experimental setup
Here, we provide further details on our experimental setup.
General setting. The codebase used in this work was developed from scratch but is loosely inspired hwwprchhhcoaaaa s0c-matching. All experiments are conducted using the Jax library [21] and are run on single 40GB NVIDIA A40 GPU. Our default experimental setup, unless specified otherwise, is as follows: We use the Adam optimizer [79] with learning rate of and gradient clipping with a value of We utilized 50 discretization steps using exponential integrators. The control function is parameterized as a fully-connected 6-layer neural network with 256 neurons and GELU activations all experiments, we used time horizon of 1.
by github. com/ffccbookreee
x 10
[67]. Time embedding is achieved via Fourier features [122]. For
The control is parameterized as
u (x. t) = fa t) (2)) 172
(133)
and for experiments using Langevin preconditioning (LP),, it parameterized as ufr () = (x, (t)
n? + V log Prarges (x))
(134)
where j and fa are neural networks parameterized by 0.
For non-trust methods, we train for 60k gradient steps with batch size of 2000, amounting total of 120M target evaluations. In contrast, trust region methods use a buffer of length 50k refreshed 150 times during training, resulting in total of 60k x 150 7.5M target evaluations. To optimize for the next control u,+1, we perform 400 gradient steps on the replay buffer using randomly sampled batches of size 2000. All experiments optimized using a line search method.
trust region bound of 0.1. The dual function is
For the Many Well target, we set the standard deviation of the prior distribution to 1 and to 2.5 for the Gaussian mixture target. For the randomization of the mixing weights, we uniformly sample positive values that are normalized and rescaled such that the ratio between the maximum mixing weight and the minimum is 3. The diffusivity is scheduled according to ((l) = (Cmax Cminn) cos? Am Omin 27
with Gmin = 0.01 and Cmax 10.
Evaluation protocol and model selection. We follow the evaluation protocol of prior work [18] and evaluate all performance criteria 100 times during training, using 2000 samples for each evaluation. We apply running average with window of 5 evaluations smooth out short-term fluctuations and obtain more robust results within a single run. We conducted each experiment using four different random seeds and averaged the best results for each run.
Benchmark problem details. The Many Well target involves d-dimensional double well potential, corresponding to the (unnormalized) density
Prasget (2) exp >
with m E N representing the number of combined double wells (resulting in 2m modes), and a separation parameter 8 E (0, oo) (see also [135]). In our experiments, we set m = 5 leading to 2m = 32 modes. The separation parameter set to = 4. Since Ptarget factorizes across dimensions, can compute reference solution for log 2 via numerical integration, as described in [84].
Moreover, we consider a Gaussian mixture model (GMM) target of the form
Prarget (I)
Σ
(x (x: pu. Ex),
(135)
37

---

where jod E IR Es E Raxd, The 2 O, and E T = 1. To compute the optimal control u", we exploit the fact that the optimal marginal path measures Qx)x) can be derived analytically [86],
Σ T&N (a| 1x" (le)ds Ege 22 ((e)ds + I." 2((s)e 2212 c(u)duds
K
(136)
and used this for computing the optimal control u², Finally, to compute the total variation distance, we leverage the known true mixing weights Tą and define the mode partitions Sş. C R as
S = fo E arg max [aN(a)1 Σ,) = k).
(137)
1.2 Evaluation criteria
Here, we provide further information how our evaluation criteria are computed.
Control 1 error. Assuming access the optimal control u", can compute the 1,2 error between the optimal and the learned control, i.e..
1 L. - (x" ssds
control L error = E
(138)
where X"" is obtained by simulating the controlled process with w", and compute the error using a Monte Carlo estimate. Note that this quantity is equivalent to the forward Kullback-Leibe divergence
DKL (Q))")
EE
log
dIP (Xx")
dQQ
(139)
Via Girsanov's theorem (see App. A.3) have that dQ
dP )
f" (u (X"" .) aW, +
2 I ||u - ull' ,sdds.
1 1T
(140)
The desired equivalence follows from the fact that, under mild regularity assumptions, the stochastic integral in (140) is martingale and has vanishing expectation.
Log-normalizing constant. By definition, the log-normalizing constant is given by
(
E
w(X.))
Xa
(141)
Applying a change of measure to the controlled process yields
Z(Xo) = EE W(X"
diP (x*)
dIP
Xx
= I e So" 2 [1a((2
Xo
(142)
which can be estimated via Monte Carlo using samples from the current control u.
Sinkhorn distance. We estimate the Sinkhorn distance W [31], an entropy-regularizee optimal transport distance, between model and target samples using the JAX-based ott library [32].
Total variation distance. Inspired by recent work [18, 55J, we assume access to ground truth mixing weights Ti. k {1..... K}. along with a partition {S1,..., Sk) corresponds to the k-th mode of the target distribution. We estimate the empirical mixing weights using
} of Rdd where each region S, Rd
Ls (X+]
(143)
Σ* E [1s (X#]
Using these estimates, we compute the total variation distance (TVD) between the empirical and true mode weights as
TVD
E Fx
(144)
Details on how the ground truth mixing weights and the corresponding mode regions S₁. are defined can be found in the descriptions of the target densities.
I.3 Additional experiments
Here, we provide results for additional numerical experiments.
Gaussian Mixture (GMM40). We further evaluate the performance of trust-region-baset losses by comparing them to existing SOC losses on the well-established GMM40 benchmark [84]. In this task, the target distribution is Gaussian mixture model with 40 components, where the means are uniformly sampled from the interval [-40, 40). and each component has initial variance of 1. We
38

---

4-2
10
1311
10'2
334
RE
CE
COBDI 4011
DidAO am:
TR-LV
42
SOCM
TR-SOCM
AM
833
3.0
120
12
110
EXLOOT WTDEE atSCca 22.07
Figure 7: Control L error function of the number of target evaluations for the GMM larget across varying problem dimensionalities results are averaged across four random seeds.
set the prior's s standard deviation to 17 30. The results, presented in Fig. 8, show that only two losses, Cross-Entropy (CE) and trust region with log-va variance (TR-LV), can consistently learn all 40 modes. Notably, TR-LV achieves this with approximately ten times fewer target evaluations than CE.
Control L error vs. target evaluations. We extend the results presented in Sec. 3.1 for the GMM benchmark by providing a detailed analysis of the control 2 error as a function of the number of target evaluations across varying problem dimensionalities d. For d = 2, all SOC losses achieve low control error. However, at d 10, some methods begin to exhibit elevated control error due to mode collapse. As the dimensionality increases further, only trust-region-batt -asodd losses consistently maintain low control error. While these methods show partial mode collapse for d 2 150, we anticipate that this issue can be mitigated by refining the control function architecture by employing larger buffer and batch sizes. Importantly, trust region methods also require significantly fewer target evaluations key advantage in many weal-world applications where evaluations costly.
Influence of trust region bounds. We further investigate the effect of different trust region bound values on the GMM target using TR-LV. The results are presented Fig. 10. The left figure shows that smaller trust region bounds significantly improve performance: an order of magnitude lower control error compared to & = 1. Additionally, smaller & values help stabilize training, as evidenced by the reduced standard deviation across random seeds. In contrast, training with
0.01 yields up
becomes unstable. However, this improved stability comes at the cost of slower convergence smaller bounds require more training iterations to effectively anneal from the prior to the target path measure, as illustrated in the middle figure. Finally, the right figure shows that the empirically observed smoothed effective sample size (ESS) aligns well with its Taylor series
approximation, ESS = (Var further details.
dp((-1)
+1)
IF+I for small values of s; see see App. E.3 for
d(P(s)
Transition path sampling
J.1 Experimental setup
We build upon the codebase provided by TPS-DPS [113] (github. http//iioung8//pp--prre Our experimental setups also follow [113] to ensure a fair comparison. The dual function is optimized using Brent's method [23].
MD simulation setup. run molecular dynamics simulation on the OpenMM platform. Both simulations are run at temperature 300K. For Alanine Dipeptide, we use the "amber99sbildn.xmi forcefield with a VVVR integrator to simulate in vaccum. Fach timestep set as femtosecond. Each path sampled is of length 1,000. For Chignolin, use the 'protein.fi wtt.nff1444..... B1144S0ooosc.cmm forcefield with
39

---

GT
RE
CE
AM
SOCM
TR-IV
1.00 150
Teage
Figure 8: Qualitative and quantitative results for the GMM40 target. The qualitative plots demonstrate only the CE (orange) and TR-LV (purple) losses successfully capture all 40 modes of the ground truth (GT, grey) distribution. This is further supported by the low L control error observed for these two methods. Results are averaged across four random seeds and reported for the log-variance loss due numerical instabilities.
implicit solvant model 'gbn2.xml' with a VVVR integrator. Each timestep is set as femtosecond. Each path sampled is of length 5,000.
Target hit. For Alanine Dipeptide, target hit is defined over the two dihedral angles & and and distance radius within 0.75Á. For Chignolin, long MD simulation is pre-loaded with Time-lagged independent component analysis (TICA) to select the first two dimensions that capture most variance. The region is then defined over the two dimensions with a radius of 0.75.
Training process. Annealing is applied from 600K to 300K. replay buffer is used with buffer size 1,000 and 200 for Alanine Dipeptide and Chignolin, respectively, and training over buffer per iteration is 1,000 times.
Hyperparameters. The trust region constraint is set to с = 0.01 for Alanine Dipeptide and y = 0.2 for Chignolin. Batch size for both systems is set to 16, Alanine Dipeptide is trained for 2000 iterations, while Chignolin is trained for 50 iterations.
Computing resources. Each experiment is run on a single 80GB NVIDIA H100 GPU.
J.2 Additional experimental result discussion
We discuss our results in comparison to [113]. First of all, we evaluate three seed average as we notice the high variance nature of the transition path sampling problem-running several times can have huge variance in results (also evidenced in Fig. 4). We can also observe the trust region constraint helps to stabilize the training significantly and thus have much smaller variance across three runs. Notably, for Alanine Dipeptide, both methods start with zero hitting percentage, while in Chignolin, in the beginning both methods already have some trajectories that hit the target, trust region constraint is already effective in improving the efficiency. We use almost the exact same setup as in [113] with the only difference being the batch size for Chignolin is 16 instead of 4. We do not tune the model as goal is to show the trust region constraint improves the training stability and thus the efficiency and accuracy in terms of number of energy calls.
K Fine-tuning of diffusion models
We take the adjoint matching (AM) implementation in http://wwwwnnnniaaa...aa..arerne http:///ww//.oooooooooooooeeeeaa as our baseline, and we modify it to implement TR-SOCM.
Fine-tuning experimental details. We generate images using classifier-free guidance, with guidance scale 7.5. We 50 inference timesteps sample the trajectories during fine-tuning, and the evaluation samples are also generated at 50 inference timesteps.
We fine-tune using the default hyperparameters in the repo: we use Adam using learning rate 3 x 10-6, beta set to 0.9, beta set to 0.95, and weight decay We use an effective batch size of 500 trajectories and 4 model backpropagations per trajectory. For the TR-SOCM loss, we use
40

---

(a) Relative entropy (RE)
(b) Cross entropy (CE)
(c) Log-variance (LV)
(d) Adjoint matching (AM)
(e) SOC Matching (SOCM)
(I) Trust region log-variance (TR-LV)
Figure 9: Qualitative results for the Many Well target with d 200. Level plots depict the ground truth density for pairs of marginal distributions, while blue dots represent samples generated by models trained using the respective loss functions (indicated in the sub-captions), Among all methods, only the trust-region-baseee log- variance loss successfully avoids mode collapse and convergence issues. Interestingly, although the cross-entrope cons-en loss achieves second-lowest estimation error log (see Fig. 3), the qualitative results suggest that model fails to adequately capture the target distribution likely due to the high variance of the importance weights. All visualizations are generated using the same random seed for consistency.
a trust-region bound & = 0.1, a buffer size of 100, and 10 passes per buffer. The dual function optimized using Brent's method [23].
41

---

0.05
0.2
107
1S5;:.))
0.6
catted
10
004
0.2
20"44
0.2
01
10
100
1590
2000
Larget wediuciion fana
Barnslam
Figure 10: Influence of different trust region bound values on the GMM target for TR-LV. The left figure considers varying problem dimensionalities d whereas the middle and right figure report results for d 100. The figure on the right shows empirically observed smoothed effective sample size (ESS) and approximation via Taylor series approximation, i.e., ESS Vaa respectively. All results averaged across four random seeds.
<P" dF"
+
2e with solid and dashed lines,
We use the 10000 fine-tuning prompts taken from the repository for [140], and the 100 validation prompts from the same repository (see https://g thub. cww/THUDM/ com/THHDM/ImaggeReeei ImageReward). The two prompts used in Figure 6 are "masterpiece, best quality, realistic photograph, 8k, high detailed vintage motorcycle parked on wet cobblestone street at dusk, neon reflections, shallow depth of field" and "close up photo anthropomorphic fox animal dressed in white shirt, fox animal, glasses".
L Classical SOC problems
Here, we consider classical SOC problems, for which the optimal control can be computed analytically, These problems have been widely used in recent studies to compare different loss functions [40, 42, 87]. Here, we leverage them to showcase that importance sampling works in high dimensions when using trust-region-ber bbase losses. To that end, we consider the comparison between the SOCM loss and trust-region-bast based counterpart.
L.1 Experimental setup
The experimental setup follows the setup used for diffusion-based sampling, as explained in App. I.I, including control function architecture, hyperparameter evaluation protocol, and model selection.
For discretizing the SDE, we leverage the Euler- Maruyama scheme, i.e..
8n+1 X, + (b + ou) AAt) (n)V En NNO,
(145)
Since the considered benchmark problems admit analytical solutions for the optimal control w", consider the 2 error between the learned and the optimal control for evaluating the models as explained in App. I.1.
L.2 Benchmark problem details
We consider two problems taken from [87], the Quadratic -hattinn-hhteebckk (OU) easy and Quadratic Ornstein-Uhlenbeer (OU) hard. For convenience, we briefly introduce them again here.
Quadratic Ornstein-Uhlenbeck (OU) The choices for the functions the control problem are
b(x,t) Ax, f(z,t) x Px, g(x) Qa,
(tt) = do;
(146)
where Q is a positive definite matrix. Control problems of this form are better known linear quadratic regulator (LQR) and they admit a closed form solution [127]. The optimal control given by
(², = 20 F(t)x,
(147)
where F(t) is the solution of the Ricatti equation
dt((t) AT F(t) + F(t) 2F(t)oo F(t) + P
(148)
dt
with the final condition F(T) Q. Within the Quadratic OU class, we consider two settings:
Easy: We set A 0.21, P 0.21, Q 0.11, oa I. 1, T L, Ziii 0.5A
42

---

100
700
500M
10:3
Carpet
140m7
Exgget erabartico
(a) Quadratic Ornstein- Orastein-Uhlenteeee Uhlenbeck process (easy)
Serge www.ereee
lerput mnidertion -107
(b) Quadratic Ornstein- MonstinnUhleeneck process (hard)
Figure 11: Control L error function of the number of target evaluations for the quadratic OU problem across varying problem dimensionalities d, All results are averaged across four random seods.
Hard: We set = - Fiait 0.5W(O, I).
-,
L.3
Results
We compare the performance of SOCM and its trust-region-basee variant (TR-SOCM) on the quadratic OnstteinUUllenbeck (OU) problem across varying problem dimensionalities d. Both approaches rely on importance sampling, which is known be challenging in high- nigh-dimensional settings. This experiment highlights the role of trust regions in scaling to such regimes. Results are presented in Fig. 11.
In low-dimensional settings (d < 10), both methods perform comparably, although TR-SOCM exhibits significantly better sample efficiency. As the dimensionality increases (d 2 50), the performance of SOCM deteriorates markedly, while TR-SOCM continues achieve low control eиror. For the more challenging variant of the quadratic OU problem, SOCM fails to meaningfully improve upon initialization, whereas TR-SOCM demonstrates consistent error reduction.
These results suggest that trust regions are particularly beneficial in digeedimensio and difficult problem settings, where they provide stability and improved performance.
43

---

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

---

The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: We provide detailed instructions for how to replicate our results in the respective repositories.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided not.
the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also provided via detailed instructions for how to replicate the results, access to a hosted model (c.g.. in the case of large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
example, the contribution
While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer:
[No]
Justification: We will make our code publicly available upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (tttps://nips..ccpuubll
for more details.
45

---

While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code. unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https cctippiccc/bblcc////ccccc Poiide/Cooeeuube commmissions seionPolicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We reference the repositorics our baselines and specify all training and test details for our methods in App. E.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is
necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix. or as supplemental material.
Experiment statistical significance
Question: Does the paper report eгror bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide error bars and corresponding explanations for all experiments that support the main claims of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1 I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). If error bars are reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
46

---

Answer: [Yes]
Justification: We provide sufficient information on the computer resources for each experiments in App. E.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //tpuri////urppp cc/ppbicc/fhht.ccuuudll........
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics,
the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We discuss both potential positive societal impacts and negative societal impacts of our work in App. C.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
legitimate to point out that
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e-g., pretrained language models, image generators, or scraped datasets)?
47

---

Answer: [NA]
Justification: Our paper poses no such risks since do not release new data models. Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: For each asset that we use, we cite the paper and provide the URL (including the license).
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperawi wwwrwwithhdeee... thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
48

---

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does involve crowdsoureing nor research with human subject.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is important, original, non- standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, or non-standard components.
Please refer to our LLM policy (https: //neuriis http///rrss.cc/oooeeeeeeeeeeeeee cc/Coonferencee/2025////// for what should or should be described.
49


Reviews:
[{"review_id": "gr9AYSfjk4", "reviewer": "Reviewer_31NG", "summary": "Stochastic Optimal Control (SOC) problems have appeared within numerous sampling algorithms within recent years. As solution of SOC problems can be approached by minimizing a statistical divergence in path-measure space, when the starting path measure of the SOC problem (often the one corresponding to the reference process) differs greatly from the optimal path measure, solving the SOC problem in \"one shot\" can be difficult. This work proposes a trust region approach to SOC, wherein a sequence of SOC problems  \n$$\nu_{i+1} = \\underset{u \\in \\cal U}{\\text{argmin}} D_{\\rm KL} (\\mathbb P^u \\mid \\mathbb Q) \\quad \\text{s.t. } D_{\\rm KL}(\\mathbb P^u \\mid \\mathbb P^{u_i}) < \\varepsilon,\n$$\nare solved such that the optimal path measure $\\mathbb Q$ is approached gradually via a sequence of controls $u_1, \\dots, u_I$. The constraint $D_{\\rm KL}(\\mathbb P^u \\mid \\mathbb P^{u_i}) < \\varepsilon$ ensures that $u_{i+1}$ does not change too drastically from $u_i$, the idea being that it will be more stable to solve $I$ problems of the form above than to solve one problem that attempts to go directly from $\\mathbb P^0$ to $\\mathbb Q$. The authors show that the sequence of path measures induced by the trust region scheme is the geometric annealing path between $\\mathbb P^0$ and $\\mathbb Q$ and that the constraint $D_{\\rm KL}(\\mathbb P^u \\mid \\mathbb P^{u_i}) < \\varepsilon$ provides a way to automatically choose step-sizes such that the waypoints along the path are approximately equidistant in the Fisher-Rao metric. Starting from the constrained problem above, the authors use a Lagrangian approach to derive tractable a expression for the optimal Lagrange multiplier $\\lambda_i$, which parametrizes the optimal control $u_i$. As the optimal control is also parametrized by the gradient of the value function, which is a conditional expectation,  the authors propose computational approaches for learning $u_i$ given $\\lambda_i$ without computing the value function directly. The numerical examples demonstrate robust performance, particularly in comparison to SOC approaches without trust regions, on a variety of tasks, including sampling from unnormalized densities, transition path sampling, and fine-tuning of diffusion models.", "questions": "* The result in Proposition 2.2 is very nice. After reading it and thinking about it for a moment it made sense to me because I know that the geometric mixture can be characterized variationally. I'll use standard probability measures instead of path measures here, but if $\\pi_0, \\pi_1 \\in \\mathcal P(\\mathbb R^d)$, then $\\pi_t \\propto \\pi_0^{1-t}\\pi_1^t$ is the minimizer of \n$$\n\\pi_t = \\underset{\\eta \\in \\mathcal P(\\mathbb R^d)}{\\text{argmin}} \\quad (1 - t)\\mathrm D_{\\rm KL}(\\eta \\| \\pi_0) + t \\mathrm D_{\\rm KL}(\\eta \\| \\pi_1), \\quad t \\in [0,1]. \n$$\nThis result can be found in, e.g., [Amari 2016]. So, given this characterization, it makes sense that requiring equispaced movement in KL would land one on the geometric annealing path. This is more of a comment than a question, but perhaps a result like this holds for path measures as well and could be connected to your result in Proposition 2.2, or be used to arrive at the trust-region method from a different angle? \n\n* Proposition 2.3 reminds me of a similar result in [Chopin et al. 2024], where they show that change in any $f$-divergence controls the change in Fisher Information along the geometric annealing path and use this fact to propose time-stepping schemes based on Fisher Information (Section 3.1/Proposition 2). Your result is applicable to path measures as well, but it is probably worth citing [Chopin et al. 2024] to make the connection. \n\n* This is minor, but in Algorithm 1 should $i$ be incremented on every iteration? And should $u^{i+1}$ be returned instead of $u_i$? I'm a little bit confused by the indexing.\n\n### References \n* Amari, S. I. (2016). Information geometry and its applications (Vol. 194). Springer.\n* Chopin, N., Crucinio, F. &amp; Korba, A.. (2024). A connection between Tempering and Entropic Mirror Descent. <i>Proceedings of the 41st International Conference on Machine Learning</i>, in <i>Proceedings of Machine Learning Research</i> 235:8782-8800 Available from https://proceedings.mlr.press/v235/chopin24a.html.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths \n* The paper is clear and well-written. I enjoyed reading it and learned a lot in the process.\n\n* While geometric annealing is a common design choice in algorithms for sampling from unnormalized densities (SMC, recent transport-based approaches, etc.), this is the first time I've encountered it for path measures. Moreover, one of the most critical components of using geometric annealing is effectively picking the annealing schedule. The trust-region-derived objective that is optimized to pick each $\\lambda$ here is nice and I imagine it could inform other methods which use geometric annealing. \n\n* The derivations of the methods are rigorous and elegant. \n\n* The trust region method proposed can be applied to any algorithm which uses SOC, including many recent sampling SOC-based sampling methods. \n\n* The numerical examples are challenging and span three different problem types (sampling, transition-path sampling, and diffusion model fine-tuning). The benefit of using the trust region SOC approach over regular SOC is clear from the numerical examples, in that it often leads to better performance and requires fewer target evaluations/energy evaluations/trajectories. \n\n## Weaknesses\n* The explanation of how one goes from the cross-entropy loss to the SOCM loss is very brief and relies on prior knowledge of adjoint-matching. I know that the details are in the appendix, but building a little more intuition for readers who may be unfamiliar with adjoint-matching would help with readability.", "was_revised": true, "final_justification": "This paper presents a compelling new algorithmic approach for solving SOC problems grounded in rigorous, interesting mathematics and justified by impressive empirical performance across a range of challenging benchmark problems. The two clarity issues I flagged have been resolved, and I feel that the discussion which has arisen out of the rebuttal discussions will further strengthen the final version of the paper. I therefore recommend this paper for acceptance in NeurIPS 2025.", "rebuttal": "We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.\n\n---\n> [...] a result like this holds for path measures as well and could be connected to your result in Proposition 2.2, or be used to arrive at the trust-region method from a different angle?\n\nIndeed, Proposition 2.2 can also be seen as an instance of the variational characterization you mention with the key difference that $\\pi_0$ is replaced with $\\pi_i$. To make this more explicit, consider the following optimization objective, which shares the same critical points in $u$ as $\\mathcal{L}\\_{\\mathrm{TR}}^{(i)}(u, \\lambda_i)$ (Eq. (5)):\n$$\n\\mathcal{L}(u) = \\tfrac{1}{1+\\lambda_i}D\\_{\\text{KL}}\\left(\\mathbb{P}^u | \\mathbb{Q} \\right) + \\tfrac{\\lambda_i}{1+\\lambda_i} D\\_{\\text{KL}}(\\mathbb{P}^u | \\mathbb{P}^{u_i} ).\n$$\nNext, defining $\\alpha_i = \\tfrac{1}{1+\\lambda_i}$, we arrive at\n$$\n\\mathcal{L}(u) = \\alpha_i D\\_{\\text{KL}}\\left(\\mathbb{P}^u | \\mathbb{Q} \\right) + (1-\\alpha_i) D\\_{\\text{KL}}(\\mathbb{P}^u | \\mathbb{P}^{u_i} ).\n$$\nSolving $u_{i+1} = \\text{argmin}\\_u \\mathcal{L}(u)$ gives \n$$\n\\frac{\\mathrm d\\mathbb{P}^{u_{i+1}}}{\\mathrm  d\\mathbb{P}} \\propto \\left(\\frac{\\mathrm  d\\mathbb{P}^{u_{i}}}{\\mathrm  d\\mathbb{P}}\\right)^{1-\\alpha_i}\\left(\\frac{\\mathrm  d\\mathbb{Q}}{\\mathrm  d\\mathbb{P}}\\right)^{\\alpha_i} \\quad \\iff \\quad \\frac{\\mathrm  d\\mathbb{P}^{u_{i+1}}}{\\mathrm  d\\mathbb{P}^{u_i}} \\propto \\left(\\frac{\\mathrm  d\\mathbb{Q}}{\\mathrm  d\\mathbb{P}^{u_i}}\\right)^{\\alpha_i}\n$$\nwhich is equivalent to Eq.(7) in Proposition 2.2.\n\nWe thank the reviewer for highlighting this connection. We will add additional details and references to the manuscript.\n\n---\n> Proposition 2.3 reminds me of a similar result in [Chopin et al. 2024] [...]\n\nWe agree that the work of Chopin et al. (2024) is highly relevant to our approach. In particular, the connection between geometric annealing and entropic mirror descent is intriguing. We are currently investigating whether similar connections can be established in the context of path measures, which we believe could offer valuable theoretical insights.\n\n---\n> This is minor, but in Algorithm 1 should be incremented on every iteration? And should $u_{i+1}$ be returned instead of $u_i$?\n\nWe thank the reviewer for pointing this out. We have corrected the indexing and incrementing accordingly in the revised manuscript.\n\n---\n> The explanation of how one goes from the cross-entropy loss to the SOCM loss is very brief and relies on prior knowledge of adjoint-matching.\n\nWe agree with the reviewer that the SOCM loss would benefit from further clarification. Accordingly, we have expanded the explanation in the revised manuscript to improve clarity and provide additional intuition.\n\n--- \n\nOnce again, we sincerely thank the reviewer for the valuable feedback and would be happy to provide any further clarifications if needed."}, {"review_id": "DwEXvW22Fb", "reviewer": "Reviewer_r2Du", "summary": "In this paper, the authors propose a novel method to solve stochastic optimal control problems. The idea of the paper is to divide the problem into constrained sub-problems, where the constraints maintain the solution of the actual problem closer to the previous one (having a constraint on the KL divergence between the solutions of two adjacent sub-problems). Doing this, the authors can reuse the sampled trajectories of the previous problem to solve the following one, reducing the target evaluations.", "questions": "- The approach appears sound and is applied across multiple domains. How sensitive is the method to domain-specific hyperparameters or design choices?\n\n- The experimental results demonstrate sample efficiency on Stable Diffusion, but the performance is similar to the baseline. Do the authors expect their method to outperform the baseline with more samples? \n\n- The proposed approach relies on the use of constraints to improve robustness. Could the authors provide an ablation or analysis showing how different values of the $\\epsilon$ parameters influence performance and convergence? Is there a clear trade-off? How to set the parameter?", "limitations": "Yes", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strenghts\n\n- The paper is well written and clearly presented, even for readers who are not experts in SOC problems.\n\n- It proposes a sound and principled approach to solving SOC problems.\n\n- The authors demonstrate the versatility of the method by applying it to several tasks and comparing it against existing approaches.\n\n- I thoroughly enjoyed reading the paper. Although the topic lies outside my area of expertise, I believe the method has potential utility in many AI applications.\n\n## Weaknesses\n\n- In the experiments, it would be great to show the contribution of the constrained in the performance of the method. Since the main idea is that the constraints will help in finding a more robust solution, it would be nice to have a trade-off between the $\\epsilon$ parameters and the number of iterations.\n\n- In the experiment on Stable Diffusion, the algorithm achieves similar performance to the baseline but with fewer samples. It is not clear if with more samples the algorithm will achieve better performance than the baseline.", "was_revised": true, "final_justification": "After reading the rebuttal and discussion, I remain positive about this paper. The authors provided clear and detailed answers to my questions.", "rebuttal": "We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.\n\n--- \n> The approach appears sound and is applied across multiple domains. How sensitive is the method to domain-specific hyperparameters or design choices?\n\nThe domain-specific hyperparameters and design choices were largely adopted from the respective baselines. Specifically, for transition path sampling, we followed the settings used in [1]; for fine-tuning, we adopted the configuration from [2]; and for diffusion-based sampling, we partially followed the setup in [3]. We kindly refer the reviewer to Appendices E.1, F.1, and H.1 for further details on these design choices.\n\nRegarding the trust-region bound $\\varepsilon$, we find that our method is generally robust across different domains. For all diffusion-based sampling tasks, we consistently used $\\varepsilon = 0.1$, which performed well without additional tuning. In the case of transition path sampling and fine-tuning, we observed that smaller trust-region bounds can yield further improvements. We attribute this robustness to the theoretical relationship between $\\varepsilon$ and the variance of the importance weights, as discussed in Appendix E.3, “Influence of Trust Region Bounds”. This relationship holds independently of the specific application domain (see also our response to the final question for further details).\n\n---\n> The experimental results demonstrate sample efficiency on Stable Diffusion, but the performance is similar to the baseline. Do the authors expect their method to outperform the baseline with more samples?\n\nOur primary goal in the fine-tuning experiments was to demonstrate improved sample efficiency, i.e., achieving competitive performance with significantly fewer samples. As noted, our method matches the baseline’s performance while requiring substantially less computational effort.\n\nWe are currently uncertain whether additional samples would allow our method to consistently outperform the baseline in terms of final performance, particularly given the strength of the baseline itself. However, since we observe consistent improvements under equal computational budgets in other domains, we plan to further investigate this question in the context of fine-tuning experiments.\n\n---\n> The proposed approach relies on the use of constraints to improve robustness. Could the authors provide an ablation or analysis showing how different values of the $\\varepsilon$ parameters influence performance and convergence? Is there a clear trade-off? How to set the parameter?\n\nWe have added such an ablation in Appendix E.3, “Influence of Trust Region Bounds” in the context of diffusion-based sampling. Indeed, in that experiment there is a clear trade-off between computational cost and performance: Smaller trust-region bounds lead to improved performance as shown in the left plot of Figure 10, but take longer to approach the target (see middle plot of Figure 10). We agree that this trade-off deserves more visibility, and have therefore integrated the ablation study in the main part of the manuscript.\n\nTo guide the choice of $\\varepsilon$, we have further analyzed its effect in Appendix E.3 by examining its relationship to the variance of the importance weights and the corresponding effective sample size. This connection often provides a practical and intuitive way to set $\\varepsilon$. In many of our experiments, we used $\\varepsilon = 0.1$, which corresponds approximately to an effective sample size of 0.8.\n\n---\n\nWe hope that our response addresses the reviewer’s question, and we would be glad to provide further clarification if needed.\n\n--- \n\n### References\n\n[1] Seong, Kiyoung, et al. \"Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers. ICLR 2025\n\n[2] Domingo-Enrich, Carles, et al. \"Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control.\" ICLR 2025\n\n[3] Domingo-Enrich, Carles, et al. \"Stochastic optimal control matching.\" NeurIPS 2024\n\n---"}, {"review_id": "3r5wuUdpjQ", "reviewer": "Reviewer_TLot", "summary": "This paper proposes a trust region method for stochastic optimal control (SOC) problems by reframing SOC as a measure transport problem in path space. The core idea is to iteratively minimize the KL divergence between controlled path measures while constraining the KL divergence between consecutive iterates within a \"trust region\". The authors demonstrate that this approach yields a geometric annealing path with equidistant steps in Fisher-Rao distance, which contributes to stable optimization. The method is shown to be efficient across diverse applications, including diffusion-based sampling, molecular transition path sampling, and diffusion model fine-tuning, consistently outperforming baselines with significant gains in stability and sample efficiency.", "questions": "Apart from the comments in the weakness section, I have the following questions:\n1. Could the authors explain why their proposed methods TR-SOCM and TR-LV are scalable to high dimensions (e.g., in Figure 2)?\n2. Have the authors tried shrinking $\\epsilon$, which is commonly used in optimization?\n3. How is the length of sequence $I$ affect the algorithmic performance?", "limitations": "Yes.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1.  One interesting finding in this work is that it elegantly links trust regions to geometric annealing, resulting in equidistant steps in Fisher-Rao distance. This theoretically grounded approach directly addresses the instability often encountered when prior and target measures differ substantially, by systematically approaching the target measure gradually.\n2. The proof in this work appears correct. The authors provide a clear interpretation of each theory with easy-to-digest examples.\n3. The true region SOC algorithm (Algorithm 1) is flexible with two compatible loss functionals, log-variance divergence and SOC matching loss, which are designed to work with buffer reuse. This \"off-policy\" optimization, where the $\\lambda$-dual can be fitted from already-simulated trajectories, is a key enabler for enhanced efficiency.\n\nWeaknesses\n1. The authors claim \"minimal computational overhead\" for the dual optimization of Lagrange multipliers. However, empirical comparisons with baselines are only qualitatively sketched, not based on running time. It would be better to see the benefits under the same computational cost.\n2. Hyper-parameter guidance on the effect of $\\epsilon$ and buffer size is missing, which could critically affect algorithmic performance.", "was_revised": true, "final_justification": "The authors have cleared my doubts. I believe this work makes a meaningful contribution to the field of stochastic optimal control. Therefore, I raised my score to 5 and recommend acceptance of this paper.", "rebuttal": "We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.\n\n---\n> The authors claim \"minimal computational overhead\" for the dual optimization of Lagrange multipliers. However, empirical comparisons with baselines are only qualitatively sketched, not based on running time. It would be better to see the benefits under the same computational cost.\n\nWe thank the reviewer for the helpful suggestion and have accordingly added quantitative results on running time. Specifically, we evaluate two aspects: \n\n(1) Overhead from dual optimization:\nFor the dual optimization step, we employed a line search algorithm based on a bracketing method. In the diffusion-based sampling experiments, for example, computing $\\lambda^*$ from 50k importance weights took approximately 0.014 seconds per optimization (averaged over 10k trials on an RTX 3060 GPU). Performing 150 such updates during training, the total overhead from dual optimization amounts to only 2.1 seconds, which is negligible relative to the overall training time.\n\n(2) Impact of the replay buffer in fine-tuning:\nWe also compared TR-LV, which uses a buffer, against AM (which does not use a buffer). We observed a significant improvement in efficiency when using a buffer:\n– Without a buffer, fine-tuning experiments averaged 1 day and 23 minutes.\n– With a buffer, the same experiments required only 11 hours and 9 minutes on average.\nThis clearly demonstrates that incorporating a replay buffer can substantially reduce computational cost in practice.\n\nWe thank the reviewer for this helpful suggestion and will include the corresponding runtime comparisons in the revised manuscript.\n\n---\n> Hyper-parameter guidance on the effect of $\\varepsilon$ and buffer size is missing, which could critically affect algorithmic performance.\n\nWe appreciate the reviewer’s observation regarding the importance of hyperparameter guidance for both the trust-region bound $\\varepsilon$ and the buffer size.\n\nTrust-region bound $\\varepsilon$:\nWe discuss the role of $\\varepsilon$ in Appendix E.3, “Influence of Trust Region Bounds”, where we highlight its connection to the variance of the importance weights and the resulting effective sample size. This theoretical relationship often provides an intuitive basis for selecting $\\varepsilon$. In many of our experiments, we used $\\varepsilon = 0.1$, which corresponds approximately to an effective sample size of 0.8 and showed good empirical performance.\n\nBuffer size:\nFor transition path sampling, we adopted buffer sizes consistent with prior work. In the fine-tuning experiment, we evaluated several buffer sizes and found that a size of 100 offered a good balance between performance and computational efficiency, as larger buffer sizes did not yield significant improvements.\n\nIn contrast, for diffusion-based sampling, we observed that buffer size has a substantial impact on performance. To explore this further, we conducted an additional experiment on the GMM task using the TR-LV loss, varying both the buffer size and dimensionality. We report the control L2 error in the table below:\n| Buffer size    | $d=2$ | $d=50$ | $d=100$ | $d=200$ |\n|----------|----------|----------|----------|----------|\n| $10k$    | $0.001 \\scriptstyle \\pm 0.000$   | $7.721 \\scriptstyle \\pm 0.051$   | $21.656 \\scriptstyle \\pm 1.504$   | $70.606 \\scriptstyle \\pm 1.886$   |\n| $25k$    | $0.001 \\scriptstyle \\pm 0.000$   | $3.507 \\scriptstyle \\pm 0.995$   | $4.313 \\scriptstyle \\pm 0.634$   | $59.814 \\scriptstyle \\pm 3.83$   |\n| $50k$    | $0.000 \\scriptstyle \\pm 0.000$   | $0.014 \\scriptstyle \\pm 0.000$   | $0.026 \\scriptstyle \\pm 0.002$   | $24.178 \\scriptstyle \\pm 3.040$   |\n| $100k$    | $0.000 \\scriptstyle \\pm 0.000$   | $0.014 \\scriptstyle \\pm 0.000$   | $0.030 \\scriptstyle \\pm 0.001$   | $11.594 \\scriptstyle \\pm 4.157$   |\n\nThe results show that larger buffer sizes consistently lead to lower control error, particularly as dimensionality increases.\n\nWe thank the reviewer for this valuable suggestion. We will incorporate these results and practical guidelines into the revised manuscript.\n\n\n---\n> Could the authors explain why their proposed methods TR-SOCM and TR-LV are scalable to high dimensions (e.g., in Figure 2)?\n\nBoth, the SOCM and LV loss do not require differentiation through simulations which makes them more memory efficient and therefore more scalable. For the LV loss we refer to Proposition 5.7 in [1], which shows beneficial scaling properties of this loss in high dimensions. Further, note that the additional use of trust regions allows losses that require importance weighting - such as the SOCM loss - to scale to high-dimensions as well (see e.g. Fig. 11) since the trust region keeps the variance of the importance weights approximately constant (see Figure 10). \n\nHowever, we agree that the statement about scalability in Figure 2 could potentially be misleading. We, therfore, changed it to \"We observe that our trust region methods (TR-SOCM and TR-LV) are the only methods that perform well in high dimensions.\"\n\n---\n> Have the authors tried shrinking $\\varepsilon$, which is commonly used in optimization?\n\nShrinking the trust-region bound is indeed a promising idea that could potentially enhance performance. However, to maintain simplicity and avoid introducing additional hyperparameters, such as a shrinking schedule, we opted to use a fixed value for $\\varepsilon$ in our current implementation. We acknowledge this as a valuable direction and highlight it as a potential extension for future work. At the same time, however, we want to highlight a tradeoff between a small $\\varepsilon$ and a potentially long runtime: the smaller $\\varepsilon$, the larger is the number of the overall iterations $I$, which might result in better convergence to the optimal control $u^*$, but at the same time typically increases the overall runtime.\n\n\n---\n> How is the length of the sequence $I$ affect the algorithmic performance?\n\nThe length of the sequence is intimately connected with the choice of the trust-region bound $\\varepsilon$: Small trust-region bounds require more iterations to approach the target and therefore larger $I$. Conversely, chosing large values of $\\varepsilon$ requires fewer iterations and therefore a smaller $I$.\n\nAs such, $I$ represents a trade-off between computational cost and performance. To see this, we kindly refer the reviewer to the supplementary material (Figure 10), where the left subfigure shows consistent performance improvements with smaller $\\varepsilon$ values. However, the middle subfigure highlights that these settings require more target evaluations and thus larger $I$ to reach convergence. (For reference, $(\\beta_i)_i$ denotes the annealing sequence, where $1 - \\beta_i = 0$ indicates that the target distribution has been reached.)\n\n---\n\n### References\n\n[1] Nüsken, N., & Richter, L. (2021). Solving high-dimensional Hamilton–Jacobi–Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications, 2(4), 48.\n\n---"}]

Author rebuttal:

Meta-review: This paper proposes a method for solving stochastic optimal control problems though a sequence of constrained problems. The idea is to gradually approach the optimal solution though a series of constrained problems, to ensure stability. All reviewers agreed the work is novel, and of high quality and importance.
Decision: Accept (spotlight)