Transport meets Variational Inference: Controlled Monte Carlo Diffusions
Authors: Francisco Vargas
Shreyas Padhy
Denis Blessing
Nikolas Nüsken
Venue: ICLR 2024 poster (2024)
Keywords: SDEs
Diffusion Models
Optimal Transport
Annealed Importance Sampling
Schroedinger Bridges
Variational Inference
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Connecting optimal transport and variational inference, we develop a score-based annealing technique for Bayesian computation.

Abstract: Connecting optimal transport and variational inference, we present a principled and systematic framework for sampling and generative modelling centred around divergences on path space. Our work culminates in the development of the Controlled Monte Carlo Diffusion sampler (CMCD) for Bayesian computation, a score-based annealing technique that crucially adapts both forward and backward dynamics in a diffusion model. On the way, we clarify the relationship between the EM-algorithm and iterative proportional fitting (IPF) for Schroedinger bridges, deriving as well a regularised objective that bypasses the iterative bottleneck of standard IPF-updates. Finally, we show that CMCD has a strong foundation in the Jarzinsky and Crooks identities from statistical physics, and that it convincingly outperforms competing approaches across a wide array of experiments.

Published as a conference paper at ICLR 2024
TRANSPORT MEETS VARIATIONAL INFERENCE: CONTROLLED MONTE CARLO DIFFUSIONS
Francisco Vargas", Shreyas Padhy* University of Cambridge Cambridge, UK
Denis Blessing KIT Karlsruhe, Germany j181420kit.er
Nikolas Nüikken'
Kings College London London, UK nik. www.keueennnnnnn de
{fav25, 8p2058 cam. ac.uk .uk
ABSTRACT
Connecting optimal transport and variational inference, we present a principled and systematic framework for sampling and generative modelling centred around divergences on path space. Our work culminates in the development of Controlled Monte Carlo Diffusions for sampling and inference, score- booee annealing tech- nique that crucially adapts both forward and backward dynamics in a diffusion model. On the way, clarify the relationship between the EM-algorithh and iterative proportional fitting (IPF) for Schrödinger bridges, providing conceptual link between fields. Finally, we show that CMCD has a strong foundation the Jarzinsky and Crooks identities from statistical physics, and that it convincingly outperforms competing approaches across a wide array of experiments.
INTRODUCTION
Optimal transport Villani et al. 2009) and variational inference Blei et al. 2017) have for a long time been separate fields of research. In recent years, many fruitful connections have been established (Liu et al.] (201)), in particular based on dynamical formulations (Tzen & | (2019a). and in conjunction with time reversals (Huang al. 20 Song et al.) 2021). The goal of this paper is twofold: In the first part, we enhance those relationships based on forward and reverse time diffusions, and associated Girsanov transformations, arriving and sampling. In the second part, we build on this and develop novel score-based scheme for sampling from unnormalised densities. To set the stage, we recall classical approach (Kingma Welling 2014 Rezende & Mohamed| (2015) towards generating samples from target distribution (x)), which is the goal both in generative modelling and sampling:
unifying framework for generative modeling
Generative processes, encoders and decoders. We consider methodologies which can be imple- mented via the following generative process,
v(2), 21/2 <op>>p"nn p((z)z)) z),
(1)
transforming sample v(z) into sample ʃ p (x) (x)z))(dz). Traditionally, v(z) is simple auxiliary distribution, and the family of transitions p/ (zz 22) is parameterised flexibly and in such way that sampling according to (T) is tractable. Then we can frame the tasks generative modelling and sampling finding transition densities such that the marginal in matches the target distribution,
p(x) = D0 (d)))
To learn such transition, is helpful to introduce reversed process 2|æ (2)x)
x p(x:),
relying on an appropriately parameterised backward transition oe (=(x). We will say that I and ( reversals of each other in the case when their joint distributions coincide, that is, when
d (x(x))
0x(xxz)((()
To appreciate the significance of B notice that if (4p holds, then (2) is implied by integrating both sides with respect to z. Building this observation, it is natural to define the loss function
LD 0) D (9 (2 x))(())] (/) (x)))))(()),
(5)
"Equal contribution.

---

Published as a conference paper at ICLR 2024
where D is divergencd between distributions yet to be specified. Along the lines of Bengio et al. (202T); Sohl- Dicks-Dickttein et al. (2015); Wu et al. (2020); Liu et al. (b), we have now laid the foundations for algorithmic approaches that aim at sampling from (u(z) by minimising ( 0):
Framework 1. Let D be arbitrary divergence, and assume that (D 0) = 0. Then we have u(x) J (z)2))))(() and x(=)= (z)z))
(6)
that is, v(2) is transformed into u(xx) by pd (xx (z)]z), and ((x) is transformed into v(z) by q (z|z))
The sampling problem. Let V denote a probability density function on R of the form
p(=) Z = p(z)dz, where 1 : R - I+ can be differentiated and evaluated pointwise but the normalizing constant intractable. We are interested in both estimating Z and obtaining approximate samples from v/ given we can sample from a more tractable density μ.. Framework 1 provides us with an objective to tackle the sampling problem as once ( 0, we can generate samples from v(=) via the variational distribution (2|1z) Through variational inference and optimal transport, we discuss relationships to classical methods as well as shortcomings:
((z) = z:
Kl-ddivergecc, ELBO and variational inference. Choosing = DKL in (5), variational inference (VI) and latent variable model based approaches (Dempster et al., 1977 Blei et al.) [2017] Kingma & Welling 2014) can elegantly be placed within Framework Indeed, direct computation (see Appendix B shows that LDKL (Ø, 0) (aw) [ELBO,, (b, 0) + (e>> [In [(x)], so that minimising
(r (é, 0) is equivalent to maximising the expected evidence lower bound (ELBO), also known as the negative free energy (Blei et al., 2017). This derivation is alternative to the standard approach via maximum likelihood and convex duality (or Jensen's 's inequality) (Kingma et al. 2021 Section 2.2), and directly accomodates various modifications by replacing the -ii--divergence (see Appendix (B)
Couplings, (optimal) transport and nonuniqueness. Assuming 4 holds, is natural to define the joint distribution T ((x, z) := (2|(() (()())((=>==))) p" (xl x ())) which is coupling between μ(x)) and (()). Viewed from this angle, the set of minimisers of L(o, 0) stands in one-to-one correspondence with the set of couplings between sa(z) and v(2), provided that the parameterisations
chosen flexibly enough. Under the latter assumption, the objective in (5) admits an infinite number of minimisers, rendering algorithmic approaches solely based on Framework potentially unstable and their output hard to interpret. In the language of optimal transport (Villani 2003), minimising L(0 0) enforces the marginal ("transport") constraints in (6) without a selection principle based on an appropriate cost function ( 'ooptimal).
Methods such as VAFS (Kingma & Welling 2014) parameterise D (x/lz) and (2(xx) with a re- stricted family of distributions (such as Gaussians), thus restricting the set of couplings. Expectation maximisation (EM) minimises L(ó, 0) in a component
fashion, resolving nonquniqueness in a procedural manner (see Section 3.1.. Common diffusion models fix either p" (x) or ( [z|x)) and
thus select a coupling (Section 2.2). In this paper, we argue that the full potential of diffusion models can be unleashed by training the forward and backward processes at the same time, but appropriate modifications that resolve the nonuniqueness inherent in Framework [ need to be imposed. To develop principled approaches towards this, we proceed as follows:
Outline and contributions. In Section 2 we recall hierarchical VAEs (Rezende et al.) 2014) and, following Tzen & Raginsky (2019a). proceed to the infinite-depth limit described by the SDEs in 12.. Readers more familiar with VT and discrete time might want to take the development in Section explanation of (12 readers with background in stochastic analysis might take Framework as their starting point. In Proposition 2.2 we provide a generalised form of the Girsanov theorem for forward-reverse time SDEs, crucially incorporating the choice of reference process that allows us to reason about sampling and generation in a systematic and principled way. We demonstrate that a range of widely used approaches, such as score-based diffusions and path integral samplers, among others, are special cases of our unifying framework (Section 2.2). Similarly in Section [3. we unify optimal transport (OT) and VI under our framework by establishing a correspondence between inpetaationmmaiiissatiion (EM) and iterative proportional fitting (IPF). Going further, we show that this framework allows us to derive new methods:
2.1
In Section 3.2 we derive a novel score-based annealed flow technique, the Controlled Monte Carlo Diffusion (CMCD) sampler, and show that it may be viewed as an infinitesimal analogue of the
As usual, divergences are characterised by the requirement that D(ol [8) 2 O, with equality iff a 3.
2

---

Published as a conference paper at ICLR 2024
Finally, connect CMCD the foundational identities by Crooks and Jarzynki in statistical physics, and show that empirically outperforms a range of state-of satt--of-hhear inference methods in sampling and estimating normalizing constants (Section(4)
method from Section 3.1
2
FROM HIERARCHICAL VAES TO FORWARD-REVERSS TIME DIFFUSIONS
2.1 HIERARCHICAL VAES REZENDE ET AL., 2014
A particularly flexible choice of implicitly parameterising pa(x))) and (z|xz) can be achieved via
hierarchical model with intermediate latents: We identify x =: y/o and z =: yr with the 'endpoints'
of the layered augmentation (3/o Vl ................................ VL-1, 11) =: Yo.L. and define
7
gºt (YL..... (1o) = II h (1119(11), pyyo..... UL VL-IlU)) 1=1
II - (y--l))) (7) t=1
so that qº (z/x)) and p" (æ]|z) (2) can be obtained from 7 by marginalising over the auxiliary variables V1..... VL-1- Here, Φ (o0, Φ1-1) and 0 ....... 0z) refer to sets of parameters to be spec- ified in more detail below. Further introducing notation, we write qu (yo:)) := 9 //t:z//o)// (g)o) as well as p" (yo.L) :- p" (20. ( and think of those implied joint distributions em- anating from u(x) = (30) and v=)) ) v(yz), respectively, moving 'forwards' or 'backwards' according the specific choices for Φ and O. In the regime when L is large, the models in 7 are very expressive, even if the intermediate transition kernels are parameterised a simple manner. We hence proceed by assuming Gaussian distributions,
- (gylyrr =N (1 31-1+ da (3) 60 p'(y)-lly) = (Nt:NNyyylll 56 (3:), 60 (8)
where a > 0 controls the standard deviation, and ō > is small parameter, anticipating the limits 00, + to be taken Section 2.2 below. The vector fields of (3:) and by (3n) introduced in 8 should be thought of as parameterised by and 0, but we will henceforth suppress this for brevity.
The models (7-(8) could equivalently be defined via the Markov chains
yt++1 21 + Saz (y:) + VoaLi, Yo ~ P = Yo.r. ~ g (yo.1);
(9a)
y1-1 = y/ + (br())) + v VooEi yr. ~ V yo.r. ~ (yo.r.)
(9b)
where ( is an iid sequence of standard Gaussian random variables. As indicated, the forward process in 9al may serve to define the distribution ((t...)). whilst the backward process in (9b)
induces p (.o:L). Note that the transition densities p [</(z) and g (=(xx) obtained as the marginals
of will in general not be available in closed form. However, generalising slightly from Framework
we may set out to minimise the extended loss
C 0) = (yo.)))) (3)oz)),
(10)
where D refers to a divergence on the 'discrete path space' { Clearly, C (c, 6) still
implies (6), but is no longer equivalent. More specifically, in the case when D = DKI. the data processing inequality yields
D (yo (lp (Wo.L) 2 DKL (z(x))) (<<l==))))) (zlz)u(z)).
(11)
so that Cext (0,0) provides an upper bound for L (Ø) 0) as defined in (5)
2.2 DIFFUSION MODELS HIERARCHICAL VAES IN THE INFINITE DEPTH LIMIT
Here we take inspiration from Section and Tzen & Raginsky (2019a) Li et al. 2020); Huang et al. (2021a) to investigate the L, oc limit, using stochastic differential equations (SDEs). this end, we think of 1 = 0, ...
. step 8, that is, we set L as discrete instances in fixed time interval [o, T). equidistant with time T -1, The discrete paths you give rise to continuous paths (Yi)osesT e C(O0, T]; ( by setting Ys: = y: and linearly interpolating Ysi and Y8(+1)- To complete the set-up, think of a (o) ... ... a 1) and 6" (69 ... b%,) in (8) as arising from time-dependent vector fields b E ((, 7] x ded, IRd) via (y)) = a18-1 (Yr) and by (y)) = bi5-1 (Ya1).
Taking the limit + O, while keeping T > 0 fixed, transforms the Markov chains in (9) into continuous-time dynamics described by the SDEs (Tzen & Raginsky 2019a)
3

---

Published as a conference paper at ICLR 2024
dY, dr (Y)) Youp = (Y)osssT Qua Fra. dy, = by dt + Tw,, Y (X)osssT ~ pub Fub,
(12a)
(12b)
where d and a denote forward and backward Itō integration (see Appendix A for more details and remarks on the notation), and (Wi)osesT is standard Brownian motion. complete analogy with (9). the SDEs in (1) induce the distributions Q and Ipae on the path space C([0, 27; Relating back to the discussion in the introduction, note that we maintain the relations Yo = æ and Yr 2, and the transitions are parameterised by the vector fields a, b, in the sense that p=(x)) = 1 (x z) = P
.6 (x) and g0 (= [x) Q744 (z|Yo x) Q7 (z).
The following well-known result (Anderson 1982 Nelson 1967) allows us to relate forward and backward path measures via a local uscore-matching condition for the reversal relation in 4
Proposition 2.1 (Nelson's 's relation). For μs and a of sufficient regularity, denote the time time-marginals
of the corresponding path measure by P v.a =: p/. 0..a Then Fuaa- I if and only if
and
= a o" In pf for all E (0,
(13)
similarly clean characterisation of equality between forward and backward path measures is not available for the discrete- time setting as presented in (9.. In particular, Gaussianity of the intermediate transitions not preserved under time-reversal.
Remark 1.
A recurring theme in this work and related literature is the interplay between the score-matching condition in (33 and the global condition D IP p.a I 1 0, invoking Framework] | calculations involving the latter, we will rely on the following result:
To enable
Proposition 2.2 forwardd-bckwart Radon-Nikodym derivatives). Let FTon' = .prro be reference path measure (that is, To, IT and x define diffusions as in (2) and are related as in Proposition p. absolutely continuous with respect both T yu and pub. Then, Tuaa pla-stmost almmost surely, the corresponding NikoonNikodym derivative (RND) can be expressed as follows,
d P pia
du aro
dv
(X) = In
(Xr)
IIn
(Xr)
(14a)
Pub
drr
1
(at (at-))) (Yi). (dY; - (a: + xr)) (Xi) di)
(14b)
0
T (b. -0--777 vo (X:) (dx, 1 (be + vp (Y:) dt)
(14c)
Proof. The proof relies Girsanov's 's theorem Üstünel & Zakai 2013). using the reference relate the forward and backward processes. For details, see Appendix
Remark 2 (Role of the reference process). According Proposition 2.2 the Radon-Nikodym
derivative between P n.a and TP v.b can be decomposed into boundary terms (14a), as well as forward and backward path integrals (14b) and (14c). Since the left-hand side of (14a) does not depend on the reference I TO > the expressions in (14) are in principle equivalent for all choices of reference. The freedom in T and 7- allows us to "reweight" between (4a). (14b) and (14c), or even cancel terms. A canonical choice is the Lebesgue measure for Ho and T, and 7 0. see Appendix C.I Remark (Discretisation and conversion formulae). The distinction between forward and back- ward integration in (14) is related to the time points at which the integrands (at - (ex=)))) (Yi) and (be 7r (Yi) would be evaluated in discrete time approximations, e.g.,
/
T
a. d..(Y.). JrY Σ ar, (Y.:). Yr.), / a. (Y.) (x.())) T Yo Σ Earin (Y,, (Y. YY.)-
Alternatively, forward and backward integrals can be transformed into each other using the conversion
[. ar (Y.) dr. = J (x)) Tτ. - -o2/ (V ax)) dt.
(15)
We refer to Kunita (219)) and Appendix for further details. In passing, we note that (15) allows us eliminate the Hutchinson estimator (Hutchinson 1989 from a variety of common score-matching objectives, potentially reducing the variance of gradient estimators, see Appendix C.T
"The global condition P p.. Pub is captured by the local condition 13 due
12
Markovian nature.
4

---

Published as a conference paper at ICLR 2024
Framework can be translated into the setting of 12 noting that appropriate modifications:
continues to hold with ( p.a pub)
Framework 1'. For divergence D path space, minimise D
P Paa Fub). if IF
then 12a) transports μu. to v, and 12b transports v to j..
u. At optimality, 1 learned diffusion via V²V In pt de br. In this way, Framework I' is closely connected (and in some ways extends) score-mare mareeeit ideas (Song & Ermon 2019, Song et al., 2021). Indeed, recent approaches towards generative modeling and sampling can recovered from Framework I' by making specific choices for the divergence D, the parameterisations for and b, as well as for the
0, Proposition 2.1 allows us to obtain the scores associated to the
reference diffuuson Fro.." Prrr" in Proposition 2.2
Score-based generative modeling:
Letting μ be the target and fixing the forward drift ag and, motivated by Proposition [2.]] parameterising the backward drift as b Og 8t, we recover the SGM objectives in Hyvärinen
Dayan (2005); [Song & Ermon (2019); [Song et al. (202I) from
D = DKL: when IP v.a = prb, , the variable drift component will represent the score o²v In pt
Modifications can be obtained from the conversion formula (15.. see Appendix C.2
Score-based sampling- - ergodic drift: In this setting, becomes the target and we fix b, to be the drift of an ergodic (backward) process. Then choosing l'r = = 7 = b allows us to recover
the approaches Vargas et al. (2023a : Berner et al. (2022). Possible generalisations based on Framework 1
iicclude IWAE-type objectives, see Appendix C.3
Score-based sampling- - Föllmer drift: Finally choosing (x) = 2 we recover Föllmer sampling
(Appendix Follmer 1984
Vargas al. 2023b Zhang Chen 2022 Huang al. 2021b).
3 LEARNING FORWARD AND BACKWARD TRANSITIONS SIMULTANEOUSLYY
Recall from the introduction that complete flexibility in and b will render the minima of D P p.a (ub) highly nonunique. Furthermore, the approaches surveyed at the end of the pre- vious section circumvent this problem by fixing either f pu,bb However, to leverage the full power of diffusion models, both T n.a or Pubb should be adapted to the problem at hand. In this section, we explore models of this kind, by imposing additional constraints on a and b. We end this section by presenting our new CMCD sampler connecting it to prior methodology within VI (Doucet eet al. 2022b Geffner & Domke 2023 Papamakarios al. (017) and OT where we can view CMCD as an instance of entropy regularised OT the infinite constraint limit (Bernton et al. 2019).
3.1 CONNECTION TO ENTROPIC OPTIMAL TRANSPORT
One way of selecting particular transition between μ and is by imposing entropic penalty, en- couraging the dynamics to stay close to a prescribed, oftentimes physically or biologically motivated, reference process. Using the notation employed in Framework [ the static Schrödinger problem
(Schrödinger 1931 Léonard 2014a) is given by
π (ar, 2) E arg min =(0.2)
DKLL (Tr.((((a) (I( ( 2)) (m) u(x)) m=()) =z)==(z)) (())))) u(z)).
(16)
where r(x; 2) is the Schrödinger prior encoding additional domain-specific information. In analogous way, we can introduce a regulariser to the path-space approach of Framework 1' to obtain the dynamic Schrödinger problem
F" arg min Ey
lat - 100--1111 fx)1 dt
(17)
20
that is, the driving vector field ae determining pa should be chosen in such a way that (i), the corresponding diffusion transitions from μ to v, and (ii), among such diffusions, the vector field ag remains close to the prescribed vector field ft. in mean square sense. Under mild conditions, the solutions to (16) and (17) exist and unique. Further, the static and dynamic viewpoints are related through a foxxture-of-bridgee transitions induced by fi), see (Léonard 2014al Section 2).
construction (assuming that the conditionals r(z)2) a:) correspond to the
Coocurrently Richter & Berner 2024 propose an akin framework ours.
5

---

Published as a conference paper at ICLR 2024
Iterative proportional fitting (IPF) and the EM algorithm. is well known that approximate solutions for π² (x, x) and IP" can be obtained using alternating DKL- Projeproectsions keeping one of the marginals fixed in each iteration: Under mild conditions, the sequence defined by
n 2 (2, 2) arg min { (PKL(T(x, =)]|=> (xx ( = =-(x))
(18a)
n(w,²)
π 2m 2 (x, = arg min { ( 2)|l 2 (ar, z)) x=(=) v(=)) =(m,z)
n 122. 0,
(18b)
with initialisation x: (ar, 2) r(zr, z), converges r (æ. 2)
00 De Bortoli al. 2021 and
this procedure is commonly referred to as iterative proportional fitting (IPF) (Fortet 1940] Kullback]
1968 Ruschendori] 1995) Sinkhorn updates (Cuturi 2013). IPF can straightforwardly be modified to the path space setting of (17), and the resulting updates coincide with the Föllmer drift updates
discussed in Section see (Vargas et al.) 2021a) and Appendix [E.4]
To further demonstrate the coverage of our framework, we establish a connection between IPF and expectation maae-iaaaiiiiooooo (EM) (Dempster et al.] 1977). originally devised for finding maximum likelihood estimates in models with latent (or hidden) variables. According to Neal & Hinton (1998). the EM-algorithm can be described in the setting from the introduction, and written in the form
0.+1 arg min D (on, 0),
Dn+1 = arg min D (Ø, 9n41),
(19)
with LDxt. defined as in (5) If the initialisations are matched appropriately, the following result establishes exact correspondence between the IPF updates in (18) and the EM updates in :
Proposition 3.1 (EM + IPF). Assume that the transition densities p" (xx22 ())z) and (z|x) are parameterised with perfect flexiiill (1 and furthermore that the EM-scheme 19 is initialised at do
in such a way that n00 (z|xz) r(2) ar). Then the IPF iterations in (18) agree with the EM iterations
(79) for all n 2 1, in the sense that
x²" (x, z) qutn 13/2 (z/((/)))(z))
for n odd, T" 2) = ( (z)))))) for n evenn. (20)
is clear that flexibility of parameterisations is crucial, and thus EM - IPF fails for classical VAEs, but holds up negligle error for the SDDepparaettrrisation from Section 2.2, see also Liu et al, (b). Under this assumption, the key observation is that replacing forward- by reverse-DKt. in one or both of (18a) and (8b)) does not in theory change the sequence of minimisers.
From the proof (Appenix Ep.
In practice favoring the EM objectives over IPF can offer an advantage as optimizing with respect to forward-DKi. and backward- DK1. encourages moment-matching and mode-seeking behavior, respectively, and so an alternating scheme as defined in (19) might present a suitable compromise over optimizing single direction of DKL 's, empirical exploration is left for future work.
Whilst EM and IPF might seem appealing for learning sampler they both require sequentially solving series of minimization problems, which we can only solve approximately: this is not only slow but also causes sequential accumulation of errors arising from each iterate (Vargas al.) 2021a| Fernandes et al.] 2021). In order to address both issues we will present a novel approach (CMCD) that similarly to IPF learns both the forward and backward processes whilst preserving the desired uniqueness property. However, in contrast to IPF it does so in an end-to-end fashion and performs updates simultaneously. As an alternative in Appendix [..5 we also discuss regularised IPF objective and leave further empirical exploration for future work.
3.2 SCORE-BASED ANNEALING: THE CONTROLLED MONTE CARLO DIFFUSION SAMPLER
In this section, we fix prescribed curve of distributions (TT))c(o.T) whose scores V In ne (and unnormalised densities tt) are assumed to be available in tractable form; this is the scenario typically encountered in annealed importance sampling (IS) and related approaches towards computing poste- rior expectations (Neal) 2001 Reich, 2011] Heng et al. 2021 2020 Arbel et al. 2021 Doucet et al.) 2022a). The Controlled Monte Carlo Diffusion sampler (CMCD) learns the vector hield Vo,
dY (o== (Y)) V (Yi) dt + ovē d W.
Yo Yo~no.
(21)
'In precise terms, we assume that for any transition densities p(a) z)) and g(z) ((z]w),, there exist O, and d. such that p(z|2) 2)) = pd* 2) and q z) = q² (m)z).

---

Published as a conference paper at ICLR 2024
Algorithm 1 Controlled Monte Carlo Diffusions- Sampling and normalizing constant estimation
Require: To, TT, Tt, 0, K step-sizes Åls, network f trained via minimising Eq 24
Yo T0
In W = In To Yo) for k 0 to K do
Yuu+ N ( Yuu In + (X,) (tt) 20 Styy)
In W = In N Yop M6n Siz (Y 1 20 Dek
N Yuuss My to +h (08,, )ark 20 Aty
In W In W In πT (Yr)
return (Estimate of In Z = in W, Approximate sample Yr)
so that 21 produces the interpolation from the prior To to the posterior Tr, ie.. P o Vn
Te. for a E [0, T.. Note that ne were constant in time (T, = -o)), then Φ = O) would reduce (11
to equilibrium overdamped Langevin dynamics, preserving Vo. With n, varying in time, V be
thought of as a control enabling transitions between neighbouring densities n and TT+ót-
To obtain Vo, we invoke Framework I' but restrict F Frrb to retain uniqueness. Proposition motivates the choice b, (o² In T.
2.1
(b)) 20 V In ग. = - In n. V 5 leading
CD CCMC (d) = T
(22)
which is valid for any choice of divergence D. The additional score constraint b = - 22 V In Te restores uniqueness Framework I (see Appendix D for proof):
Proposition 3.2 (Existence and uniqueness). Under mild conditions on (TT)) (T.-..e.) unique minimiser d', up to additive constants, satisfying LCMCD(0") -.
22 admits
Given the optimal vector field Voi. we can produce samples from गт by simulating 21). Following Zhang & Chen (022); Vargas et al. 2023a) we can estimate Z in TT = tr /Zr unbiasedly via
d .r. Vinntv
a V d P (Y),
Z = E
(23)
VortVo
where the expectation is taken with respect to 21). and is valid for any (possibly suboptimal) Vel. The right-hand side, on the other hand, shows that optimality of VΦī yields a zero-variance estimator of Z, the statement holds almost surely in Y. without taking the expectation. To give a broader perspective, we give the following slight generalisation a well-known result from statistical physics:
Proposition 3.3 (Controlled Crooks' fluctuation theorem and Jarzynki's equality). Following Jarzyn Chen et al. 2019), define work and free energy as Wr(Y)= := - - o²0, In tt (Y) dt, In(Tt//Te) Then, we have the controlled Crooks' identity,
ski
1997);
Fr
In 2
d Fo. Vnn+V
(x)= exp (- (Fr- + + WX)
d
P FT. o
where CY := := 2 (Y,) . Vn dt So' Ao. (Yi)dtt By taking expectations and 0, this implies Jarzynski's 's equality B (exp(- 1 Wr)] -Zr/Zo.
The proof uses Proposition 2.2 to compute the RND de To. V / a o' In Vo followed by applying Itô's s formula to dn (Yi), see Appendix E.2 For ò = 0, we recover Crooks fluctuation
theorem (Crooks 1999). but the additional control allows CMCD to suppress said fluctuations by
adjusting the interaction term c+ (Y). Indeed, prior works (Neal 2001 Chopin 2002 Vaikuntanathan
& Jarzynski 2008: Hartmann et al. 2019 Zhang (02T) have used the Jarzynski equality to estimate
Z via importance sampling. but this approach might suffer from high variance, see (Del Moral et
Note the additional factor of 2 in Nelson's relation due to the noise scaling σ. 2 W, in
7

---

Published as a conference paper at ICLR 2024
2006), (Stoltz et al. 2010 Section 4.1.4). In contrast, the CMCD estimator version of (23) achieves
zero variance trained to optimality (see Appendix E.I.I for a convenient discretised version).
Finally, we would like to highlight that (Zong et al. (2023) concurrently proposes this generalisation of Crook's identity using different techniques in their sketch.
Our next result connects CMCD to Section 3.1 showing that minimising (22) can be viewed as jointly solving an infinite number of Schrödinger problems infinitesimal time intervals:
Proposition 3.4 (infinitesimal Schrödinger problems), The minimiser o" can be characterised as follows: For N N. partition the interval [0. T] into N subintervals of length 'T' I/N, and on each subinterval [( /N, iT /N], solve the Schrödinger problem {17 with marginals
( /N and prior drift J+ V In T Concatenate the solutions obtain the drift Vo(8), defined on [0, Then, V Vo as N o in the sense of (2 T] x 1R,,
π) (proof in Appendix D.4
Note the similarity of this interpretation to the sequential Schrödinger samplers of Bernton et al. (0T9)). Making specific choices for D in 22 we establish further connections other methods:
1. For D DKL, direct computation (see Appendix
D.I) based on Proposition shows that CCMCD (0)= = y 2 In dP Vo
DRI.
nr. 2 (X)
V In (Y)) ddl+ oV2 /
1
(0 In - JwW In
+const.
NYY tK+s NY + (o' In + V
Atx. 20 Atp)
Tol
(24)
VIn o
Atg, 20 At,)
the dineretscrttiitioonn in the third line is derived in Appendix D.5 Our goal is then to numerically
minimize CCMCD(
DKy (6) wrt to & (for a numerical minimisation scheme see Algorithm 2p. Note the first line in (4 is akin to the optimal control type objectives of Föllmer and DDS samplets recalled at the
end of Section 2.2 see also (Berner et al. 2022). Setting ø = in the third line recovers Unadjusted Langevin Annealing (ULA), see, e.g., eq. (14) in (Thin et al. 202T) or eq. (21) in (Geffner & Domke y;23); hence, we can view CMCD as a controlled version of ULA. Setting c=0 only in the numerator akin Monte Carlo Diffusion (MCD), see Algorithm 1 and eq. (34) in (Doucet et al. (202aa).
Finally, action matching (Neklyudov et al. 2023) can be recovered from D-DKı and Framework
Pr.7 in Proposition 2.2 appropriately, see Appendix C.4 by choosing the reference To.n== p
2. For the larvnnri divergence Dvar (2, P) -ar (In d9). see Appendix B we obtain
CCMD (O) Var (n TT (YT). To Yo)
J. (Y)) V In -2) / m. Pa))
Var
see Appendix D Here, Wd denotes Stratonovich integration, and the variance is taken with respect to samples from (1). In the limit - O, log-Va CMCD enforces integrated version of the instantancous change of density formula θ In T. (Y,) = - (Yr) for continuous-tinn normalising
flows of the form Y, = (Yi),, (Papamakarios et al. 2021 Section 4).
Remark 4 (Further related work). The task of learning the vector field Vol so that (2 reproduces (T.) 1e) Ti has been approached from various directions. Reich (2011 1011; Heng et al. (2021 Reich (202)); Vaikuntanathan & Jarzynski (2008) explore methodologies that exploit the characterisation of Ver in terms of the elliptic PDE (50) in Appendix [D.3] Arbel et al. (2021) propose leverage normalising flows sequentially to minimise KI. divergences between implied neighboring densities.
In an appropriate limiting regime, they recover the SDE (2). see Remark 9 These approaches
approximate Vó, sequentially in time, whilst CMCD learns (V Vor)eeoo 'all-at-once'.
EXPERIMENTS
We now empirically demonstrate the performance of the proposed CMCD sampler 24 in both underdamped (detailed in Appendix D) and overdamped (CMCD (OD)), Appendix D.6 formulations on a series of sampling benchmarks. We first replicate the benchmarks from Geffner & Domke (2023) on 6 standard target benchmark distributions. Following the experimental methodology in Geffner &
8

---

Published as a conference paper at ICLR 2024
ELBO (OD) ELBO (CTD)
ELBO (OD) ELBO (CD)
lb
W2
log. log. seg_ionospherr log So.oonar
brownian
funnel
30)
=
85)
lorenz
(00)
soods
26)
lgcp (600)
1gcp 1600)
(a)
D01
(b)
CMCD (UD)
Figure 1: Figure panes a) and b) report ELBOS across methods and targets following the experimental setup in Geffner & Domke (2023), the (OD) and (UD) columns group over and under-damped methods seperately. Figure reports IS In Z estimates and sample quality (where available) using
eOT. Higher ELBO and In denote better estimates, lower W2 signifies better sample quality.
Domke (2023), 3) we compare against two underdamped baselines, Unadjusted Langevin Annealing (ULA) (Wu al. 2020 Thin et al. 202I) and Monte Carlo Diffusion (MCD) (Doucet et al. 2022b) ;
and two overdamped baselines, Uncorrected Hamiltonian Annealing (UHA) (Geffner & Domke
2021 Zhang et al., (20T)) and Langevin Diffusion Variational Inference (LDVI) (Geffiner Domke
2023 Furthermore, we include comparisons of In Z estimation two datasets with known partition
function, the funnel and gmm, and compare against baselines from Vargas et al. (2023a). PIS (Barr
et al. 2020 Vargas et al. 2023b /Zhag & Chen 2022), DDS (Vargas et al. 2023a), and Sequential
Monte Carlo Sampler (SMC) (Del Moral et al. 2006; Zhou et al. (016))
We report the mean ELBO achieved by each method over 30 seeds of sampling, for Fuler discretisation steps K {8, 16, 32, 64, 128, 256}. comparing the underdamped and overdamped baselines their respective CMCD counterparts in Figure | We see that both overdamped and underdamped CMCD consistently outperform all baseline methods, especially at low K, and in fact, across most targets overdamped CMCD outperforms the underdamped baselines. Figure [ also reports In for two target distributions with known Z, comparing against PIS, DDS, and SMC. Again, CMCD recovers the log-partition more consistently, even at low K. Finally, as another measure of sample quality, we report the entropy-reeuularised OT distance (WV?) between obtained samples and samples from the target for funnel and gru. Hyperparameter tuming and other experimental details can be found in Appendix F and we provide GitHub repository to reproduce our results
5 DISCUSSION
Overall we have successfully introduced novel variational framework bridging VI and transport using modern advances in diffusion models and processes. In particular, we have shown that many existing diffusion-bar based methods for generative modelling and sampling can be viewed as special instances of proposed framework. Building dynamic entropy regularised transports (based on relationship between the EM and IPF algorithms) and annealed flows (with connections to fluctuation theorems due to Crook and Jarzynski, rooted in statistical physics). Finally, we have explored the CMCD inference scheme obtaining state-of-the-an results across suite of challenging inference benchmarks. We believe this experimental success is partly due to our approach striking a balance between parametrising a flexible family of distributions whilst being constrained enough such that learning the sampler is not overly expensive (Tzen & Raginsky 2019b annealed flow π, Goshtasbpour et 2023) and alternate divergences (Nüsken & Richter 2021
this, we have developed novel objectives for
Vargas et al. 2023c). Future directions can explore optimal schemes for the
Richter & Berner 2024 Midgley et al. 2022

---

Published as a conference paper at ICLR 2024
REFERENCES
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint arXiv:2303 08797, 2023.
Brian D.O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their
Applications, 12(3):11-3326, 1982.
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport Monte Carlo. In International Conference Machine Learning, pp. 318-330. PMLR, 2021.
Dominique Bakry, Ivan Gentil, Michel Ledoux, et al. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.
Ariel Barr, Willem Gispen, and Austen Lamacraft. Quantum ground states from reinforcement learning. In Mathematical and Scientific Machine Learning, pp. 635-653. PMLR, 2020.
Fabrice Baudoin. Conditioned stochastic differential equations: theory, examples and application to finance. Stochastic Processes and their Applications, 100(1-2):109-145, 2002.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. GflowNet foundations. arXiv preprint arXiv:2/ / 1.00266, 2021.
Julius Berner, Lorenz Richter, and Karen Ulirich. An optimal control perspective diffusion-bused generative modeling. In NeurlPS 2022 Workshop on Score-Based Methods, 2022.
Espen Bernton, Jeremy Heng, Arnaud Doucet, and Pierre Jacob. Schrödinger bridge samplers. arXiv preprint. 2019.
Patrick Billingsley.
Convergence of probability measures. John Wiley & Sons, 2013.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association. 112(518):859-8777 177899977777- 2017.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv preprint 1rXx::1509.00//9 2015.
Tianrong Chen, Guan-Horg Liu, and Evangelos Theodorou. Likelihood training of Schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Represen- tations, 2022.
net/ wtt.:eeee..ooommooooooocccc.cc.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. On the relation between optimal transport and Schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169(2):671-691 2016.
Yongxin Chen, Tryphon T Georgiou, and Allen Tannenbaum. Stochastic control and nonequilibrium thermodynamics: Fundamental limits. IEEE transactions on automatic control, 65(7):2979-2) 2991, 2019.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. Siam Review, 63(2):249-313, 2021.
Nicolas Chopin. A sequential particle filter method for static models. Biometrika, 89(3):539-552, 2002.
JMC Clark. A local characterization of reciprocal diffusions. Applied Stochastic Analysis, 5:45-59, 1991.
Gavin E Crooks. Entropy production fluctuation theorem and the nonequilibrium work relation for free energy differences. Physical Review E, 60(3):2721, 1999.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, 2013.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1):313-329 1991.
10

---

Published as a conference paper at ICLR 2024
Kamélia Daudel, Joe Benton, Yuyang Shi, and Amaud Doucet. Alpha-divergence variational inference meets importance weighted auto-encoders: Methodology and asymptotics. arXiv preprint TrXfi:2210006226, 2022.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the
Royal Statistical Society: Series B (Statistical Methodology), 68(3):4 1-436, 2006.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the EM algorithm. Journal of the royal statistical society: series B (methodological), 39(1): 1-22, 1977.
Arnaud Doucet, Will Grathwohl, Alexander G Matthews, and Heiko Strathmann. Score-based diffusion meets annealed importance sampling. Advances Neural Information Processing Systems, 35:21482-21494999 2022a.
Arnaud Doucet, Will Sussman Grathwohl, Alexander G de G Matthews, and Heiko Strathmann. Annealed importance sampling meets score matching. ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022b.
Paul Dupuis and Richard S Ellis. A weak convergence approach to the theory of large deviations. John Wiley & Sons, 2011.
Tarek A El Moselhy and Youssef M Marzouk. Bayesian inference with optimal maps. Journal of Computational Physics, 231 231(23)7711-7755 2012.
David Lopes Fernandes, Francisco Vargas, Carl Henrik Ek, and Neill DF Campbell. Shooting Schrödinger's s cat. In Fourth Symposium on Advances in Approximate Bayesian Inference, 2021.
Alessio Figalli and Federico Glaudo. An invitation to optimal transport, Wasserstein distances, and gradient flows. 2021.
H. Follmer. An entropy approach the time reversal of diffusion processes. Lecture Notes in Control and Information Sciences, 69:156-163, 1984.
Hans Föllmer. Calcul d' sans probabilités. In Séminaire de Probabilités XV 1979/80: Avec table générale des exposés de 1966/67 1978/79, pp. 143-150. Springer, 2006a.
Hans Föllmer. Time reversal on Wiener space. In Stochastic Processes Marirsstttt and Physies: Proceedings of the 1tt BiBoS- SympossSm held in Bielefeld, West Germany, September 10-15, 1984, pp. 119-129. Springer, 2006b.
Hans Föllmer and Philip Protter. On Itô's formula for multidimensional Brownian motion. Probability Theory and Related Fields, 1 16:1-20,, 2000.
Robert Fortet. Résolution d'un systeme d' ééquutttinn de M. Schrödinger. J. Math. Pure Appl. IX, 1: 83-105, 1940.
Peter K Friz and Martin Hairer. course on rough paths. Springer, 2020.
Tomas Geffner and Justin Domke. Meme variational inference via uncorrected hamiltonian annealing.
Advances in Neural Information Processing Systems, 34:639-655 661,, 2021.
Tomas Geffner and Justin Domke. Langevin diffusion variational inference. In International
Conference on Artificial Intelligence and Statistics, pp. 576-593. PMLR, 2023.
Shirin Goshtasbpour, Victor Cohen, and Fernando Perez-Cruz. Adaptive anncaled importance sampling with constant rate progress. 2023.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable reversible generative models with free-form continuous dynamics. In International Conference on Learning Representations, 2019. htttpps///ggareeeeeeeeee///rrrrr

---

Published as a conference paper at ICLR 2024
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. arXiv preprint arXiv:2211.0155, 2022.
Carsten Hartmann, Ralf Banisch, Marco Sarich, Tomasz Badowski, and Christof Schüttee. Characteri- zation of rare events molecular dynamics. Entropy, 16(1):300-37, 2013.
Carsten Hartmann, Christof Schütte, and Wei Zhang. Jarzynski's equality, fluctuation theorems, and variance reduction: Mathematical analysis and numerical algorithms. Journal of Statistical Physics, 175: 775:1214- 15::121-1261, 2019.
UG Haussmann and E Pardoux. Time reversal of diffusion processes. Stochastic Differential Systems Filtering and Control, pp. 176-182. Springer, 1985.
Jeremy Heng, Adrian Bishop, George Deligiannidis, and Arnaud Doucet. Controlled sequential
Monte Carlo. Annals of Statistics, 48(5), 2020.
Jeremy Heng, Arnaud Doucet, and Yvo Pokern. Gibbs flow for approximate transport with appli- cations to Bayesian computation. Journal of the Royal Statistical Society Series B: Statistical Methodology, 83(1):156-187, 2021.
Jose Hernandez-Lobato. Yingzhen Li, Mark Rowland, Thang Bui, Daniel Hernández-Lobate and Richard Turner. Black-box alpha divergence minimization. In International conference on machine learning, pp. 1511 1-500. PMLR, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6855 2020.
Chin-Was Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective diffusion-based generative models and score matching. Advances in Neural Information Processing Systems, 34: 2021a.
Jian Huang, Yuling Jiao, Lican Kang, Xu Liao, Jin Liu, and Yanyan Liu. Schrödinger-Föllmer sampler: Sampling without ergodicity. arXiv preprint arXiv:2106. 108.000000 2021b.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statisties-Simulation and Computation, 18(3):1059-1074
1989.
Aapo Hyvärinen and Peter Dayan, Estimation of non normalizel statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Christopher Jarzynski. Nonequilibrium equality for free energy differences. Physical Review Letters, 78(14):2690, 1997.
Ioannis Karatzas, Ioannis Karatzas, Steven Shreve, and Steven E Shreve. Brownian motion and stochastic calculus, volume 113. Springer Science & Business Media, 1991.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-2199 2021.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. ICLR 2015, 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. ICLR, 2014.
Diederik P Kingma, Max Welling, al. An introduction variational autoencoders. Foundations and Trends in Machine Learning, 1244)3307-392, 2019.
Peter E Kloeden, Eckhard Platen, Peter Kloeden, and Eckhard Platen. Stochastic disferential equations. Springer, 1992.
Takeshi Koshizuka and Issei Sato. Neural Lagrangian Schrödinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2023. URL htttps://opeettte ew. net. foltt/forume wttpee/eooommmmmmeemmmm.mmm
12

---

Published as a conference paper at ICLR 2024
Solomon Kullback. Probability densities with given marginals. The Annals Mathematiccal Statistics, 39(4):1236 1243, 1968.
Hiroshi Kunita. Stochastic flows and jump-diffusions, volume 92. Springer, 2019.
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems-Serice A, 34(4): 2014a.
Christian Léonard. Some properties of path measures. In Séminaire de Probabilités XIVI, pp- 207-230. Springer, 2014b.
Christian Léonard, Sylvie Rællyy, and Jean-Claude Zambrini. Reciprocal processes. a measure- theoretical point of view. 2014.
Xuechen Li, Ting- Kam Leonard Wong, Ricky T. Q. Chen, and David K. Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In Symposium on Advances in Approximate Bayesian Inference, pp. 1-28. PMLR, 2020.
Yingzhen and Richard E Turner. Rényi divergence variational inference. Advances neural information processing systems, 29, 2016.
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning, pp. 4082- 4092. PMLR, 2019.
Guan- Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Schrödinger bridge. In Advances in Neural Information Processing Systems, a.
Deep generalized
Xingchao Liu, Lemeng Wu, Mao Ye, et al. Let us build bridges: Understanding and extending diffusion generative models. In NeurlPS 2022 Workshop on Score-Based Methods, b.
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnaud Doucet. Continual repeated annealed flow transport monte carlo. International Conference on Machine Learning, pp- 15196-15219, PMLR, 2022.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-LLobato. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.. 01833, 2022.
Annie Millet, David Nualart, and Marta Sanz. Integration by parts and time reversal for diffusion processes. The Annals of Probability, pp. 208-238, 1989.
Jesper Møller, Anne Randi Syversveen, and Rasmus Plenge Waagepetersen. Log gaussian cox processes. Scandinavian journal of statistics, 25(3):451-42,, 1998.
Radford M Neal. Annealed importance sampling. Statistics and computing, 11:125-1333 2001.
Radford M Neal. Slice sampling. The annals of statistics, 31(3):705-767, 2003.
Radford M Neal and Geoffrey Hinton. view the EM algorithm that justifies incremental, sparse, and other variants. Learning in graphical models, pp. 355-368, 1998.
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learning stochastic dynamics from samples. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 25858 25889. PMLR, 23-29 Jul 2023. URL htt.tttp::. http:::://procecee ngs mm wwwpsseee2... nekl Bekkyuddv22a..tmm
Edward Nelson. Dynamical theories of Brownian motion, volume 3. Princeton university press, 1967.
Nikolas Nüsken and Lorenz Richter. Solving dimensimenii
Cacpttoo--Jcoob-BBelllle PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differential Equations and Applications, 2(4):1 2021.
13

---

Published as a conference paper at ICLR 2024
Nikolas Nüsken and Lorenz Richter. Interpolating between BSDEs and PINNS: Deep learning for elliptic and parabolic boundary value problems. Journal of Machine Learning, 2(1):31-64, 64, 2023.
Bernt Øksendal. Stochastic differential equations. In Stochastic differential equations, pp. 65-84. Springer, 2003.
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. OT-flow: Fast and accurate continuous normalizing flows via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9223-9232, 2021.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. arXiv preprint 1.Xiv1700.000555 2017.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1 2:66722660, 2021.
Stefano Peluchetti. Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling. arXiv preprint arXiv:2304.00017 2023.
Sebastian Reich. A dynamical systems framework for intermittent data assimilation. BIT Numerical Mathematics, 51(1):235-249, 2011.
Sebastian Reich. Data assimilation: dynamic homotopy-based coupling approach. arXiv preprint arXiv:2209.05279, 2022.
Daniel Revuz and Marc Yor. Continuous martingales and Brownian motion, volume 293. Springer Science & Business Media, 2013.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pp. 1530-1538. PMLR, 2015.
Danilo Jimenez. Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference machine learning, pp. 1278-1286. PMLR, 2014.
Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations, 2024. URL BRRLLtttp////peeeeeae 24
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: low-variance gradient estimator for variational inference. Advances in Neural Information Processing Systems, 33: 11431-1139922 2020.
Geoffrey Roeder, Yuhuai Wu, and David Duvenaud. Sticking the landing: Simple, lower-variance gradient estimators for variational inference. arXiv preprint arXiv:1703.09199 2017.
Sylvie Relly Reciprocal processes: stochastic analysis approach. In Modern stochastics and applications, pp. 53-67. Springer, 2013.
L Chris G Rogers and David Williams. Diffiusions, Markov processes and martingales: Volume 2, I1ō
calculus, volume Cambridge university press, 2000.
Ludger Ruschendorf. Convergence of the iterative proportional fitting procedure. The Annals of 1995.
Statistics,
Francesco Russo and Piemre Vallois. The generalized covariation process and Itô formula. Stochastic Processes and their applications, 59(1):11-144 1995.
Francesco Russo and Pierre Vallois. Itô formula for ff-fuciti of semimartingales. Probability theory and related fields, 104:27-41, 1996.
Erwin Schrödinger. Uber die Umkehrung der Naturgesetze. Sirzungsberichte der Preuss Akad. Wissen. Berlin, Phys. Math. Klasse, pp. 144-153, 1931.
14

---

Published as a conference paper at ICLR 2024
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. arXiv preprint arXiv:2303. 2023.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference Machine Learning, pp. 2256-2265. PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song. Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pp. 574-584. PMLR, 2020.
Yang Song, Jascha Sohl-Dickstein. Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. URL https: http::///ppeeeeeeeeeeeee net/ forum? (dd=xTTG122RH
Gabriel Stoltz, Mathias Rousset, et al. Free energy compulations: A mathematical perspective. World Scientific, 2010.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Michèle Thieullen. Reciprocal diffusions and symmetries of parabolic PDE: The nonflat case. Potential Analysis, 16:1-28, 2002.
Achille Thin, Nikita Kotelevskii, Arnaud Doucet, Alain Durmus, Eric Moulines, and Maxim Panov. Monte carlo variational auto-en.oders. In International Conference on Machine Learning, pp. 10247-10257. PMLR, 2021.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent Gaussian models in the diffusion limit. arXiv preprint arXiv:1905.09833 099889990....,,.... 2019a.
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference Learning Theory, pp. 3084-31144 14. PMLR, 2019b.
A Süleyman Üstünel and Moshe Zakai. Transformation of measure on Wiener space. Springer Science & Business Media, 2013.
Suriyanarayanan Vaikuntanathan and Christopher Jarzynski. Escorted free energy simulations: Improving convergence by reducing dissipation. Physical Review Letters, 100(19);190601, 2008.
Francisco Vargas. Machine-learning approaches for the empirical Schrödinger bridge problem.
Technical report, University of Cambridge, Computer Laboratory, 2021.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger bridges via maximum likelihood. Entopyy 23(9): 134, 2021a.
Francisco Vargas, Pierre Thodoroff, Austen Lamaeraft, and Neil Lawrence. Solving Schrödinger bridges via maximum likelihood. Entropy, 23(9), 2021b. ISSN 1099-4300. e23091134. URL htt.ts://///// //ttp:://www.... wwwpi. com 00000999--20022 /200//3300/////
doi: 10.3390/
Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers.
In The Eleventh International Conference Leurning Representations, 2023a. URL https
//ppenreview..r net. for
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil Lawrence, and Nikolas Nüsken. Bayesian learning via neural Setrödinger-Föllmme flows. Statistics and Computing. 33
Francisco Vargas, Teodora Reu, and Anna Kerekes. Expressiveness remarks for denoising diffusion based sampling. In Fifth Symposium on Advances in Approximate Bayesian Inference, 2023c.
Cédric Villani. Topics in optimal transportation. Number 58. American Mathematical Soc., 2003.
15

---

Published as a conference paper at ICLR 2024
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic normalizing flows. Advances in Neural Information Processing Systems, 33:5933-5944, 2020.
Winnie Xu, Ricky T. Q. Chen, Xuechen Li, and David Duvenaud. Infinitely deep Bayesian neural networks with stochastic differential equations. arXív preprint arXiv:2102.06559, 2021.
Mao Ye, Lemeng Wu, and Qiang Liu. First hitting diffusion models for generating manifold, graph and categorical data. In Advances in Neural Information Processing Systems, 2022.
Benjamin Zhang and Markos Katsoulakis. A mean-field games laboratory for generative modeling. arXiv preprini arXiv:2304. 13534, 2023.
Guodong Zhang, Kyle Hsu, Jianing Li, Chelsea Finn, and Roger Grosse. Differentiable annealed importance sampling and the perils of gradient noise. In Advances in Neural Information Processing Systems, 2021.
Jianfeng Zhang. Backward stochastic differential equations. Springer, 2017.
Qinsheng Zhang and Yongxin Chen. Diffusion normalizing flow. Advances in Neural Information Processing Systems, 34: 14280-1116, 2021.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling. In International Conference on Learning Representations, 2022. URL https://
openreview. net wtw./eewwwe///ouurnte.
uCb2ynRu7y
Wei Zhang, Some new results on relative entropy production, time reversal, and optimal control of time-inhomogeneous diffusion processes. Journal of Mathematical Physics, 62(4):043302, 2021.
Wei Zhang, Han Wang, Carsten Hartmann, Marcus Weber, and Christof Schütte. Applications the cross-entropy method to importance sampling and optimal control of diffusions. SIAM Journal on Scientific Computing, 3666) 2014.
Adrianne Zhong, Benjamin Kuznets-Speck, and Michael R DeWeese, Time-asymmetric protocol optimization for efficient free energy estimation. arXiv preprint arXiv:2304. /2287, 2023.
Yan Zhou, Adam M Johansen, and John AD Aston. Toward automatic model comparison: an adaptive sequential monte carlo approach. Journal of Computational and Graphical Statistics, 25 (3):701-726, 2016.
A STOCHASTIC ANALYSIS FOR BACKWARD PROCESSES
In this appendix, we briefly discuss background stochastic analysis relevant to the SDEs (2). here repeated for convenience:
dY, (Y,):, dt d W. Yo ~ μ. Yourn
(25a)
dY, = dY
(25b)
Recall that the forward Itō differential ले in 25a is far more commonly denoted simply by d, and theory for the forward SDE (25a) is widely known (Karatzas et al. 1991 Øksendal 2003). In contrast, reverse-time SDEs of the form 25b) are less common and there fewer textbook accounts of their interactions with forward SDEs. We highlight Kunita (2019) for an in-depth treatment, and alert the reader to the fact that bbackward stochastic differential equations' as discussed in Zhang (2017); Chen et al. 2022). for instance, are largely unrelated. We therefore refer to (25b) as a "reverse-time SDE.
bu. in this paper we stick the notation रे emphasise the symmetry of the setting.
16

---

Published as a conference paper at ICLR 2024
Remark 5 (Notation). We deliberately depart from some of the notation employed in the recent literature (see, for instance, Huang et al. (202la); Liu et al. (b)) by using Y, in both (25a) and (25b), and not introducing an auxiliary process capturing the reverse-time dynamies. From technical perspective, this is justified since (Y:)o merely represents a generic element in path space, and
full information is encoded in the path measures Qua P v.a and pa,4 Fobb, Importantly, placing (25a) and (25b) on an equal footing seems essential for a convenient formulation of Proposition 2.2 Slightly departing from the VAE-inspired notation from Section 2.1 we equivalently refer to these
path measures by Tua and T v.,,, highlighting the symmetry of the setting in 25
Intuitively. (25) can be viewed as continuous time limits of the Markov chains defined in 9 or, in
other words, the Markov chains ( are the Euler-Maruyama discretisations for (25 see Kloeden et al. 1992, Section 9.1). Throughout, we impose the following:
Assumption A.1 (Smoothness and linear growth of vector fields). All (time-dependent) vector fields in this paper belong to the set
u == C ((0, T) XR
there exists a constant C > 0
such that [|a,(a)) (v(y))l S Cl æ yll. for all / E [0, v,y E Rd
The preceding assumption guarantees existence and uniqueness for (25a) and (25b). and it allows us to use Girsanov's 's theorem in the proof of Proposition 2.2 (Novikov's 's condition can be shown to satisfied, see [Ōksendal (2003 Section 8.6)). Furthermore, Assumption A. sufficient to
conclude Nelson's 's relation (Proposition 2.1). see Haussmann & Pardoux (1985); Millet et al. (1989); Föllmer 2006b) and the discussion in Russo & Vallois 1996). Having said all that, it is possible to substantially weaken Assumption A.J with more technical effort. Moreover, we can replace the Raad throughout, assuming sufficient regularity, growth and
constant ø > 0 by a : [0, 7) x Rd
invertibility properties, and amending the formulas accordingly.
The precise meaning of (55) is given by the integrated formulations
Y = Yo + L (Y,)) ds + L o odw..
Yo ~ 1.
(27a)
Y- Yr J. b, b., ds [.
f Tow..
Yr
(27b)
where the forward and backward integrals need defining, Roughly speaking, we have
X. dz. = lim
Σ
(Zu11 Zee)
(28a)
'AC+0'
11
X. d. = lim
Σx Xtum (Zius -z.),
(28b)
'At-50'
Remark 3 for 'appropriate' processes (X, X.)osesT and Z )oSEST, and where the limit At of vanishing step sizes needs careful analysis (see Remark [6 below). The most salient difference
between (55) and (55b is the fact that Xx, is replaced by X in 28b
Remark 6 (Convergence of the limits in (8)).. If we only assume that X, E C((to, Rd), possibly pathwise, that is, deterministically. then the limits in (28) might not exist, or when they do, their values might depends on the specific sequence of mesh refinements. The following approaches are available to make the definitions
28 rigorous:
1. Itô calculus (see, for example, Revuz & Yor 2013 Chapter 9)) uses adaptedness and semimartingale properties for the forward integral in (28a). but note that the definition not pathwise (that is, the limit (28a) is defined up to a set of measure zero). For the backward integrals (28b) and, importantly for us, in (64) it can then be shown that the relevant processes are (continuous) reverse-time martingales (see Kunita (2019) for discussion of the corresponding filtration).. The latter property is guaranteed under Assumption A.I see the discussion around Theorem 2.3 in Russo & Vallois (1996).
17

---

Published as a conference paper at ICLR 2024
2. Föllmer's 'Itô calculus without probabilities' (Fölleer 2006a is convenient, since allows to us to perform calculations using (25) and Proposition 2.2 without introducing filtrations and related stochastic machinery. The caveat is that the results may in principle depend on the sequence of mesh refinements, but under Assumption A. those differences only appear
set of measure zero, see Russo & Vallois (1995); Föllmer & Protter (2000).
3. Similarly, the integrals in 28 can be defined in a pathwise fashion using rough path techniques, see Friz Hairer 2020 Section 5.4).
For the current paper, the following conversion formulas are crucial,
[ X d Z
[ x, d = -x,Z)t:
(29a)
Ju
L. X. d Z,
f" x, dz, = J/ X, Z.. 0
(29b)
where (X, Z) the quadratic variation process (if defined, see Russo & Vallois 1995); see in particular equations (3) and (4) therein), and denotes Stratonovich integration. For solutions 25
we obtain (15) from (29a). In particular, we can often trade backward integrals for divergence terms (see Appendix using the (backward) martingale properties
U 1(0) dw: = =0, if (Yt)oseST solves (5a)
(30a)
Jr f((Y)) E dw O, if (Yi)ogisT solves (55
(30b)
B VARIATIONAL INFERENCE AND DIVERGENCES
Various concepts well-known in the variational inference community have direct counterparts the diffusion setting. In this appendix review few that are directly relevant this paper.
Maximum likelihood. Framework 1 with D DKL leads via direct calculations
-EBOO.(..)
DDa (0,0) -E f In p (z (z)
+ / In (((z))(dd))
qº (dz|z)
(31)
go (z((z)
so that maximising (x))(())) [FLBO,, 0)] is equivalent to minimising L (1)
However, the traditional approach (Blei et al. 2017: Kingma et al. 2019) towards the evidence lower bound (ELBO) in (3Th is via maximum likelinood in latent variable models. Using the notation and set-up from the introduction, one can show using Jensen's inequality (or dual representations of the KL divergence), that
In (ʃ po(z, 2) dz In (x)) 2 (e,
(32)
with equality if and only if (=(2)²) (z(z)x)) The bound in (32 motivates maximising the (tractable) right-hand side, performing model selection (according to Bayesian evidence) and posterior approximation (in terms of the variational family (z)))) at the same time. The calculation in (31
shows that this objective can equivalently be derived from Framework I and connected to the KT. divergence between the joint distributions 46 (2r, 2) and (a z).
Reparameterisation trick (Kingma & Welling 2014: Rezende al., 2014). For optimising ELBO 0). it is crucial to select efficient low-variance gradient estimators. In this context, it has been observed that reparameterising z ~ go (=(xx) (x) in the form z = g(e, d; æ), see Kingma et al. (2019 Section 2.4.1), substantially stabilises the training procedure. Here, is an auxiliary random variable with tractable 'base distribution' that independent of á and æ, and g is deterministic function (transforming into 2), parameterised by ó and æ. We would like to point out that many (although not all, see below) objectives in diffusion modelling are already reparameterised, since the SDEs (12) transform the "auxiliary" variables (Wr)osesr into (Y))oSsTT. With this viewpoint, the
18

---

Published as a conference paper at ICLR 2024
vector field a, corresponds the parameter Φ, (Wi)sses corresponds to e, and corresponds to the solution map associated to the SDE 12a), sometimes referred to as the Itō map. In this sense, the objectives (74). 39) and 40) are reparameterised, but C from Section 3.2 is not if the gradients are detached as in (üüseen & Richter 2021] Richter et al. 2020 Richter & Berner, (024). We
mention in passing that sticking the landing (Roeder al. (017) offers further variance reduction close to optimality, and that the same method can be employed for diffusion objectives, see Vargas ct al. (2233b); Xu et al. (202I).
Reinforce gradient estimators. As an alternative the KL-divergence, Nüsken & Richter (2021) investigated the family of log-variance divergences
In dq dp
(gll)) = Varaaw
(z))
(33)
parameterised by an auxiliary distribution u. in order to connect variational inference to backward
stochastic differential equations (Zhang] 2017). The fact that gradients of (33) do not have to be
taken with respect to æ (see Remark (Nüsken & Richter 2021 Richter et al. (020)) reduces the computational cost and provides additional flexibility in the choice of u, but the gradient estimates potentially suffer from higher variance since the reparameterisation trick is not available. The latter drawback is alleviated somewhat by the fact that particular choices of u can be linked to control variate enhanced reinforce gradient estimators (Richter et al., 2020) that are particularly useful when reparameterisation is not available (such as in discrete models). We note that the same divergence has also been used as a variational inference objective in El Moselhy & Marzouk (2012).
Importance weighted autoencoders (IWAE). Burda et al. (2015) have developed multi-sample version of ELBO, (Ø. 8) that achieves tighter lower bound on the marginal bik-iiilddood in (32).
To develop similar objectives in a diffusion setting, we observe that for each K 2 1,
(a;)) ]
D
(K) KT. (gllp) E
In
(34)
ad
K
defines a generalised KI. divergence that reproduces the IWAE lower bound as per Framework 1 in the sense of equation (31). To the best of our knowledge. the precise formulation in (34) is new,
but similar to the previous works Hernandez-Lobatt et al. (2016); Li & Turner (2016): Daudel et al.
(2022..
We exhibit an example of (34) 14paplied in diffusion context in Section C.3 see Remark
C CONNECTIONS TO PREVIOUS WORK
C.1 DISCUSSION OF EQUIVALENT EXPRESSIONS FOR DKL. (P u ( FF
Notice that we can realise samples from b both via the reverse-time SDE in 12b or via its time reversal given by the following forward SDE (Nelson 1967
Anderson
1982
Haussmann Pardoux
1985):
- - (Y)) dt to dw,
Yo ~M
(35)
using Y, = Y--.. This allows to obtain an expression for DRL P u pub via Girsanov's theorem:
DKL F = DKL(Fo +E
1 2002 S. [|xx))) -
- o2 in (Y)) ||
dt
(36)
However there are several terms here that we cannot estimate or realise in a tractable manner, one being the score Vln In 2 and the other being sampling from the distribution 2
Indeed, by Jensen's inequality, we have that D D ((4t)))l))) implies D. Kg (gllp) > 0.
D (gllp), that in particular q p
9When Y, is an OU process and is Gaussian in the traditional DDPM setting (Sogg et al. 2021 and these quantities admit the classical tractable score matching approximations
19

---

Published as a conference paper at ICLR 2024
In order circumvent the score term, the authors Vargas al. (2021b) Chen al. (2022) use the Fokker-Plande (FPK) equation and integration by parts, respectively. trading of the score with a et al. (2021a)
divergence term, equivalent solution. From Proposition 2.2 we can avoid the divergence entirely and replace it by backwards integral (making use of the conversion formula (15) and the fact that the ensuing forward integral is zero in expectation). As hinted at in Remark 3 this replacement might have favourable variance-reducing properties, but numerical evidence would be necessary.
variant of the Feynman Kac formula to arrive an
C.2 SCORE-BASED GENERATIVE MODELING
Generative modeling is concerned with the scenario where (4(2) can be sampled from (but its density unknown), and the goal is learn a backward diffusion as in (126) that allows us to generate further samples from ((x)), see Song et al. (201)) We may fix a reference forward drift or, and, motivated by Proposition 2.1 parameterise the backward drift as b = Or 8t, so that in the case when P Ip v.b the variable drift component 8; will represent the score pdV In pf When the diffusion associated a ergodic and large, T p.a Fub requires that v=2) close to the corresponding invariant measure. Choosing Vt = Oe, and, for simplicity o 1. direct calculations using Proposition show that
show
(Y) dt J'' "(V. (t)
(ssM(() = DKL (F (. pu--o) =EywFrn
(V
const.
(37)
recovers the implicit score matching objective (Hyvärinen & Dayan 2005), up to a constant that does not depend on St-
Proof. We start by noticing that the contributions in [44a) and 14b) do not depend st, and can therefore be absorbed in the constant in (37) Notice that the precise forms of To. Pr and y are left unspecified or unknown, but this does not affect the argument. We find
DKL (F p = Ey Bx Fune [%%-
s(() (āv, - (2ar -s) si) (Xr) (t))tt)
+ const.
se())) (oāw, (xi)
+ const.
of (Y)) di f sx) di
+ const.,
2
where in the first line we use Proposition 2.2 together with b, =a 8, and 77- = a₂ and to proceed
to the second line substitute ' Y, using the SDE in 12ap. The last equality follows from the conversion formula between forward and backward Itô integrals, see (15). and the fact that forward integrals with respect to Brownian motion have zero (forward) expectation, see (30a).
Notice that the nonuniqueness Framework I' has been circumvented by fixing the forward drift a;: indeed LISM is convex in confirming Note that using integration by parts, LiSM is equivalent to denoising score matching (Song et al. 2020 2021):
f
DKL I pi p pa- =Ey BY
sY)) - V In Pejo (Y, Yo) df + const... (38)
2
Framework] 1' accommodates modifications of 37 in particular the divergence term in 337 can be replaced by backward integral, Appendix C. and Remark 3 Note that the settings discussed in this section are also akin to the formulations in Kingma et al. (2021); Huang et al. (2021a).
Finally, is worth highlighting that this setting is not limited to ergodic models and can in fact accommodate finite time models in the exact same fashion as the Föllmer drift is used for sampling
(Section C.3) by using a Doob's transform (Rogers & Williams
2000) based SDE for IP n.a as
opposed to the classical VP-SDE see Example 2.4 in Ye al. (2022).
20

---

Published as a conference paper at ICLR 2024
C.3 SCORE-BASELL SAMPLING
Consider the setting when v(2) is target distribution that be evaluated pointwise up normalisation constant. In order to construct a diffusion process that transports an appropriate auxiliary distribution p(x) p(z), one approach is fix drift b, in the backward diffusion
(1bb). and then learn the corresponding forward diffusion 12ab by minimising + D si vsb).
Tractability this objective requires that μ == p be known explicitly, at least approximately. In the following we review possible choices.
Choosing (xx = x/t, one show using Doob's transform (Rogers & Williams 2000 Theorem (ii), that Py = 6(x). for any terminal distribution vx)). Hence,
The Föllmer drift.
DRL( 60.0 bbb) leads tractable objective. In particular consider the choice To = 6n. y = O, corresponding to a standard Brownian motion, then it follows that y 2
minimising
I'- = _N(OTT0 and thus via Proposition 2.2
DKL 60.a [2n) = F
( 02
1
T (Y,) dt +hn
(d de
(Yr)
+ const.,
(39)
in accordance with Dai Pra 1991 Vargas et al. 2023b Zhang & Chen 2022). For further details, see Follmer 1984;; Vargas et al. (2023b); Zhang & Chen (2022 ; Huang et (2021b). As hinted at
in Appendix B replacing DKL in (39) by the log- variancecce divergence (33 leads to an objective that
directly links BSDEs, see (Nüsken Richter 2021 Section 3.2).
Ergodic diffusions. Vargas et al. (023a);; Berner et al. (2022 fix backward drift by that induces
an ergodic backward diffusion, that for large T,, the marginal at initial time I=bbb is close to the corresponding invariant distribution, and in particular (almost) independent of v(=) Defining == Vargas et al. (2023a); Berner al. (2022) set out to minimise the denoising diffusion sampler loss (fDs(f) := DKL u.b+o*) v.)). Choosing the reference process to be l'o,7 μ. y= (that is, the reference process at stationarity, with invariant measure uz))). direct calculation
f = r..b
based [4]
shows that
[ (()) di +In
dl' dv
LDDS (J) = E
(xT)
(40)
Remark 7 (IWAF-objective). In line with
34
we may also consider the multi-sample objective
FDD (f) = (()) ( Foo
R YI Y Protoot
In
1 K
Σ ==1
exp
(or P())) dt +In
o2
dr du
(Y7)
Proof. We start by noticing that the choice 7 b cancels the terms in 14ch. and the choice To cancels the first term in (4a)). Using a = be + o² fi. we therefore obtain
(fDs(f) DKI u.b+o" (pmn
(41a)
dr,
E 2
f((Y)) ((b. + fi)( dt ]26, + dt)
In
(Yr)
di
EE 2² ['" (2 d in
dT du
(Yr)
(41b)
As is implicit in Berner et al. 20022), it is also possible to choose 7 for the reference process, with To Fr- Leb, the Lebesgue measure on Rdd. , We notice in passing that although the Lebesgue
Vargas et al. 2023a) chose (time-inhomoegenous) backward Omstein- pracess, that n.b t=0 close to Gaussian, but generalisations are straightforward.
21

---

Published as a conference paper at ICLR 2024
measure is not normalisable, it is invariant under Brownian motion (the forward and backward drifts are both zero), and the arguments can be made rigorous by a limiting argument (take Gaussians
with diverging variances), or by using the techniques in Léonard (2014a Appendix A.1). By similar calculations as above, we obtain
(DDs( = o² { ((0)) dt di- =(()))))
( ((x)) Tw, + (n(())) - Inn((T))
(42a)
J0
o² [- f7 dt ["(v
(ac))ttt(((vv ( dt -m + const.,
E
(42b)
where we overload notation and denote the Lebesgue densities of μ and y with the same letters. In second line we have used the conversion formula (T3), together with the fact that the forward Itô integrals are forward martingales (Kunita 2019). and therefore have zero expectation. Comparing (40) and (42b), we notice the additional divergence term, due to the fact that the choice y = 0 does not cancel the terms in (64). See also the discussion in Appendix C.1
Finally we note that whilst the work in Berner et al. (2022) focuses on exploring a VP-SDE-base approach which is ergodic, their overarching framework generalises beyond ergodic settings, notice this objective is akin to the KL expressions in Vargas 2021 Proposition 1) and Liu et al. a Proposition 9).
C.4 ACTION MATCHING (NEKLYUDOV ET AT... 2023
Similar approach in Section 3.2 Neklyudov et al. 2023) fix curve of distributions (
In contrast to us, they assume that samples from unnormalised densities are not). Still, we can use Framework I to rederive their objective:
available, for each E [0, 7] (but scores and
Akin to the proof of Proposition 3.2 under mild conditions on (Te)eEjo.2)) there exists a unique vector field Vø, that satisfies the Fokker- Pokker-PPlnncl Planck equation
a V VT) of Am.
(43)
We can now use the reference process P no. Vo' (that is, To = To. 7pt Vol, DT = WTs Yt =
VØ; V In to compute the objective
DKL No. Wo
relying on the same calculational techniques as in Sections C.2 and (the particular choice of reference process cancels the terms in (14a)). Notice that the parameterisation in this objective constrains the target diffusion to have time-marginal. 7y, just as in Section 3.2 By direct calculation, obtain (up to factor of 2/0²) the action-gap in equation (5) in Neklyudov et al. (2023). Indeed, we see that
P. Ep In d V
DKL I TO. Ve vor
T
(Vrr voi)² (Yi) di - {' (Vr. voi)(() Tw.
E
1/2
VIn
(Y)) (Vih Vo.)
E
(V vo,)² (Y, dt
22

---

Published as a conference paper at ICLR 2024
where in the last line we have used the conversion formula
15 together with 30a
compute
% (V (V,(V,, (Vb. v dt
(Ve. vo)) Tw, EF
I
/-" Jzd (Vir vot) - (z))
T
( 1V-(WV
RR
= Vlnn((Y) (Vr (ot) dt
and cancel the two last terms in the penultimate line.
D CONTROLLED MONTE CARLO DIFFUSIONS (SECTION 3.2
D.1 DERIVATION OF DKL
CCMCD
The proof uses Proposition 2.2 choosing lo 'r to be the Lebesgue measure, with y+ - 0
(but notice that in (14 needs to be replaced by a 2 due to the scaling in (1)). We compute
dP Vlu
CMCD (Ø)
Ey - 2 In
(8)
EE In To (o) In nT (Yr)
In In (dy, Vn -n Vow())aa |
+ E
(āx, wn vn Vn - voo
E
=E
[In To (o)
In TT (Yr)
( In ग. + (Y)) हr.
/
T (-o²VIn + Vo.) T..
E
20
2
(n2 (Y.)
oz
=E [In Yo) - la ( (Y))
Vlln otz |' (6AVin Vol.) Vo.) (X;) 5 W.
where in the last equality we have inserted the dynamics (1 and used the martingale property 30ah. Notice that the expectation of the backward integral is not zero, see Appendix A.
D.2 DERIVATION OF CCMCD
In this section, we first verify the expression for CCCD in Section 3.2 using Proposition 2.2 and choosing To Fr to be the Lebesgue measure, Y Y 0. We recall that although the Lebesgue measure in not normalisable, the arguments can be made rigorous using the techniques in Léonard (2014a) Appendix A).
23

---

Published as a conference paper at ICLR 2024
The Radon- MidonnNikodymm derivative (RND) along 21 reads
d IP Vnnto"V
In
( (In mo (Y)) (lnr))(Y))
d
P 7o
+ V -ho at)
+ 2
- o -2 vi)dd)
20
(ln (In +0? 02 / (Y)
+ 2 ( (Y) dwo+ /' (lmm(()) Tw.
( (V())) dw, //" VolY)) Tw.) ov2
Using 29bb. obtain
Vin dw, + ('vinx(()) Tw.) v2o /' Vin odW,
v2
Furthermore, from (15) we see that
ov2 (L. ( dw, S' Vo(y)) āw.) - L" (Yi
(47)
from which the claim follows.
Remark (Estimating LCMCD without second derivatives). Using 47 we can equivalently write the RND as
dP VnntVe 1
In
(8) (NTTT(Yr) TYoo)
d -
ov2 Vo(Y)) dw. L'" Vo(Y.) Tw:
v (II (
/
that Va be estimated without the need evaluate Ad. Note that the identity 47 is similar to a finite difference approximation of Aø along the process Y.
D.3 EXISTENCE AND UNIQUENESS OF THE DRIFT
Before proceeding to the proof of Propostion 3.2 we state the following assumption on the curve of distributions (
Assumption D.1. Assume that क E C ((0,T) x R4 and that for all f = [0,2]
1. the time derivative 0, is square-interrable, that is, (t, L (R
2. satisfies Poincaré inequality, that is, there exists a constant C2 > 0 such that
Varz, ( - C. / gd
/
(49)
for all f (Rd)..
24

---

Published as a conference paper at ICLR 2024
Note that the boundary 0 7 {), agree to denote by dame the auward-ppintinn derivative' and interpret C (0. 7] x R) in that way. We remark that the Poincaré inequality (49) is satisfied under relatively mild conditions on the tails of (for instance, Gaussian tails) and
control of its derivatives, see, e.g., Bakry et al. 2014 Chapter 4). Under Assumption D.1 prove Proposition 3.2 as follows:
we can
Proof of Proposition
The Tokker-Plance equation associated 21 is given by
a (o)) =
(50)
The operator @ H -V. (Te Vo) is essentially self-adjoint in (RRR),, and, by (4)) coercive on CA == E (2(Rd)) fdr = 0}. Therefore, there exists a unique solution @; L ( to
(50). for any EE [0, This solution is smooth by elliptic regularity. By Proposition 2.1 and our general framework, any minimiser @ of (2) necessarily satisfies (50) well. We then obtain
(Te V x)) = 0.
Multiplying this equation by di de, integrating, and integrating by parts shows that J ([ 0) dbt = 0. proving the claim.
Remark 9 (Relation to previous work). Note we can earry out a change of variables to equation 50
Oe In Te (Vat Ver + T,)
Vln VΦ AΦ,
yielding the PDE
O Iu + II Vó+ = 0,
which when considered in terms of the unnormalised flow t = ZoT coincides with PDE in Vaikuntanathan & Jarzynski 2008); Arbel et al. (202])
0, ln V in V + Ao E,, 10, In T.] 0.
In particular, we note that the Markov chain proposed in Arbel et al. (202I) converges to our proposed parameterisation in equation 21 (see equation (12) in Arbel et al. (2021)
D.4 INFINITSSIMA SCHRÖDINGER BRIDGES (PROOF OF PROPOSITION 3.4
Throughout this proof, we assume that the Schrödinger problems on the intervals (iT/N, (i+ (1+1)T/N)
0...... N
admit unique solutions, with drifts regularity specified in Assumption A.I see
(éonard 20 Proposition 2.5) for sufficient conditions. We also work under Assumption D. so
that the drift Vø exists and is unique by Proposition 3.2
Given the interpolation (Te)sE(u,T)) we define the constraint sets
MM (T)
no. V ne.
at times
f, = N i= 0, i=.....,N
(51)
well
MT) = aeeu: Po.. = = for all
t tE 10
(52a)
In the set N is given by
u C((O, (T) a E 10 (
(++1T Rd. Rd),
for all O,...,N. N
N
3L > 0 such that [[a((x) - (x(')) S L[[ - yll, all [0, [0, a, y Rd
and recall that 2/ has been defined in Assumption A.1
25

---

Published as a conference paper at ICLR 2024
Algorithm 2 Controlled Monte Carlo Diffusions Traininn
Require:
no, иT, Tt, 0, step-sizes A network fo for i in epochs do
In WT, YT ~ Algorithm To. Ry, nt. o, {Ate)ko f9) Gradient descent step V In WT return f
By the construction in Proposition
3.4 the drift Vol(N) can be characterised by
VΦ((
E arg min E 70. aEM (x)
(1)
(54a)
20
arg min DRL no. No. Vlaxta Vo.T) (TM*(T)
(54b)
where the second line follows from Girsanov's theorem, see the proof of Proposition 2.2 We now claim that the CMCD drift Vo", by definition the minimiser in 22 can be characterised in a similar way by
V arg min E E P P Fo 20 / |a [a) (T)
(55a)
F Fo. Fo Vlatta p no. Vl.V..
arg min DKL (xM
(55b)
Indeed, the constraint Pro. ltnta for all t [o, T] implies that satisfies the Fokker-Plann equation dem, + V (TTa) = 0. By the Helmholtz decomposition (Figalli & Glaudo 2021 2.5.4), minimisers of az J dm, of gradient form, thus (55a) holds.
Section
Comparing (54h and (55)). it is plausible to infer the convergence VØ(N) - Vø" as the marginal constraints at the discrete time points 0, 1/T, .... I' become dense and approach the continuous- time constraint in (52).
To make this more precise, we note that since M (T) C C (T) for all N E N, we have that
DKL Fo. Vin Vn In < DKL To Vin Vo" F Fonn),
(56)
for all N E N. Since (-I.(-1 P no. In has weakly compact sublevel sets (Dupuis & Ellis 2011
Lemma 1.4.3), can extract subsequence _ To v+( that converges weakly towards path measure E P(C : R.).. To show that indeed P = P no. N Inn+VV"" it sufficient to note that by the constraints in (5T) those measures necessarily have the same finite-dimensional marginals. and to combine this observation with the continuity statement of Theorem 2.7.3 in Billingsley (2013). as well as the uniqueness from Proposition 3.2, The convergence of the drifts in the sense of [ T x Rd, Rd)) now follows from the lower semicontinuity of DxL. in combination with Girsanov's theorem.
D.5 DISCRETISATION AND OBJECTIVE
In the setting of CMCD with KL divergence we can use the EM approximations the RND presented in Proposition E. to express the objective as:
To Yo
N(Y ehen Yr, + (n n., + + ) Att 20 At&) Y (o m. In dr, (Y,,, Al,, 20 A1e)
CMCD (Φ) E
(57)
where the expectation is taken wrt the EM approximation of the SDE in 21). that is:
Yo NNYY,
V dr, (YY, A4,, 20 Al4).
(58)
26

---

Published as a conference paper at ICLR 2024
Algorithm 3 Forward transition Fri (Zi++1 Yu-- 2e Y..)
Require: Zie. Yu step-size Atk
Re-sample momentum Y N (Yis (1 oAts) + Von (Y 11, Atk, 20 1t
Update Y!" =Y, Afe V
In Tec (Zte)
2
Leapfrog step (Zir Yi:
Update Ztot1 Zin
+ + Yi
Update Y4111 Yuu + 2
Ale
VI (Z)z1))
return (Zee417 ((x+1)
Algorithm 4 Backward transition Bt, (Ziz: Yie Zeats: Y(o4s)
Require: Zena1 YIS step-size Al
Update Y!"" Y1++1 ALL V In In (Zie)
Inverse leapfrog
Update Zie Zint1 At&
x- 0- (1.1)
Update Yi -Y Atk
In Te, (Zie4)
2
Re-sample momentum Yo N (Y (1 o.tr) - Vol. (Y/ Zia Atk. 20
return (Zi, Yz,)
D.6 UNDERDAMPED LANGEVIN DYNAMICS
In this section, we motivate the underdamped generalisation of CMCD which is used across our experiments. This parameterisation is inspired by the underlying theory for overdamped approach, and we leave rigorous extension of those foundations for future work. However, have found this heuristic parameterisation to perform very well empirically.
Following Geffner Domke (2023) we parametrise as:
Yo, Zo -N(0, mo..
dz. Y,dl,
dY, (o2 In (Z,) ,²Y, VY, Z,)) dt + tov2 d.,
(59)
and it's time reversal as:
Yr, ZT ~ N(0, 1) STT,
dz. Yidt,
dY = ( In (Z.) o² Yr VV, Z,)) dt -o/2
(60)
D.6.1 TIME DISCRTEISATION AND ORJECTIVE
To discretise the above processes follow the exact discretisation scheme carried out in Geffner & Domke 22033,, however in this case we have to adapt the forward discretisation scheme include the non- non-linear drift when carrying out the momentum re-sample step, specific details for this scheme can be found in Algorithms 3 and 4 This discretisation in turn allows us to compute the discrete RND between these two processes which we require for Framework T'
Now via Propostion in GGefiner & Domke 2023) it follows that
Yo, Zn) Yr, Zrr
Fiz (Z Yomss Zee Yo,)
TT
Big (Zn, Yr, Zizzi M(etn)
k=
K-1 N (Y, Y, (1 At4 + Vol, (Y,,, Zo, At,, 20 At.
To Yo, Zo)
(61)
कт YT, ZT)
N (Y.
Y x Vol, (Y,, Zi, Atx 20 At&
k=
then we can use the above discrete time RND to approximate the KL divergence between SDEs (59 and (60) yielding objective for the under dampened setting:
27

---

Published as a conference paper at ICLR 2024
To o. Zo) WT Yr, Zr)
K
NY
Y1,
Vor (Y,, ) 20 A1a.)
DRL.
CCMDD-UD (0) (0) R
In
Vora Zi Atk: 20 At
k=0
(62)
Where the expectation is taken with respect to the discrete-time process Algorithm
E PROOFS
E.1 PROOF OF PROPOSITION 2.2
RADON-NIKODE DERIVATIVES)
Proof. We begin with the forward Radon-Nikodym derivative
M.a
dl dv
(Y) + 2/2 S (a: - (x.) dy, + 25 (" (b) - d?) (Y,) dl,
(Y)
= In
d dppb
(63)
following from Girsanov's theorem (see, for instance, Nüsken & Richter (2021) Lemma A.1) 1) and substitute au a b). To compute the backward Radon-Nikodym derivative, we temporarily introduce the time-reversal operator R, acting as (RY): := Yr-r on path and as (Ra) =
ar-((y) on vector fields. We then observe that
d F
R.
d p.o
In
(RY) = In
(Y);
d d
R.
for instance by comparing the discrete- time processes in 1a and 99 Equivalently,
d.
In
(Y)
In
(RY):
d P v.b
dF
since R² is the identity. Building on 663 the backward
derivative therefore reads
dy dv
((RY))) )o) 2/2 S/"" ((RR)). -(Bb).) dRY).
In
(8) In
dppu
T
+ 202 0
((Rb)) - -(Ra))) ((RY):)
dy
(Yr) + 1/2 ( (ar b.)(Y)) TY + 27
1 I" (br at) (Y,) de,
In
dv
64
where the integrals have been transformed using the substitution t + T - t. The result in 14 now follows by writing
n.a
dF n.a r
In
(Y) In
(Y) + In
(Y),
d d
d
using the assumption Toit
p pr.
and inserting 663
well 664..
E. DISCRETISATION AND CONNECTION TO DNFs (DIFFUSION NORMALISING FLOWS)
In this section we derive the main discretisation formula used in our implementations the forward- backwards Radon-Nikodie derivative (RND).
Athough pathwise definitions should be treated with care (hecause Itó integrals are defined only up set of measure zemo), the arguments made rigorous using the machinery referred to in Appendix A
28

---

Published as a conference paper at ICLR 2024
Proposition E.1. Letting To Fr = Leb and nx = O, we have that the RND in 14 is given by
JMm d P
= In - in 0² L'" oo(Y))
dr.
20² [ || (Gi)))))
In
2 f bY)) d aY, 2nd [ (Y)
dt.
Bun-almost almost surely.
and admits the following discrele-time approximation up to constant terms in a, and by (following Remark 3).
JPHo p.a d Ppb
K (Y)=-In
In
20 20-10-1-1 V1, -Y,, be (1t))) (421) const,
when using the Fuler-Maruyama discretisation:
Yuu = Yr, + at, (t,)( (2) + 11 g~N(0,1)
Proof. The first part follows by direct computation.
From here on, we will use the notation fr. = =ft.(Y) for brevity. Following Remark 3
we have that
d I n.a
(8) a (n - (xr)
d 1 2g
K
K-1
1
(Yi44 - .r.)) - 20
Σ
at, 11? (t++1 - -4)
6z
K
K |b |11 (t++1 tt).
blet (Yuu: - -Yi,) +
1=0
Adding and subtracting sum, resulting in:
Ytu) Y, (1/ - t2)) allows us complete the square in each
d
K-1
In
(Y) N (nu()) - Σ
| Y,, ac, (lit1 [:)l12
d P n6
20 20 4.l
1=0
Y, Yup b (1+11-14 14112
(67)
20 2o:4 -
Now notice that under the Euler-Maruyama discretisation ||me.m - Y, - a, (f++1 - t2)1 = (t++1 - -at--t))oll tt)o?||[112 where & ~ (O, [) does not depend on aj or b,; in particular when using DKL for the divergence have that E YYit Yt, ar, (1++1 1:)112 r² and thus:
dF n.a
K -
In
dP d (8) x (m - lnv( +
Σ
Dr. Yeur + beasa (ti+1 +t)112
20 10
(68)
Notice that in expectation (for computing DKI.), equation (68 matches equation (15) in Zhang
Chen (002) and thus provides a theoretical backing to the objective used in Zhang & Chen
(2021) Resolving the term Egaal | tcts Y,, ae (f++1 1l analytically may offer variance
reduction similar to the analytic calculations in noo-Blackwelizaiions of DKL in Ho al.) (2020).
et al. 2015 Equation 14) and the
Remark 10. The time discretised RND in equation 67 can be expressed as the ratio of the transition densities corresponding "(yo)q" (H1: (tl.x//oo//((/oox1111 (lonkklllk))(///) (y(K) with yo:K ~ discrete-time Markov chains sidering v(x) = D(x)/Z and the IS estimator Z pl(yox--lbx))( (w)))/()))/()))))) (yu.kyo)) it follows that (y (n))) (In 2] is ELBO of 2 (e.g. By 2] S In Z).
two
As a result con-
29

---

Published as a conference paper at ICLR 2024
Whilst superficially simple, Remark I0 guarantees that normalizing constant estimators arising from our discretisation do not overestimate the true normalizing constant. This result is beneficial in practice as it allows to compare estimators possessing this property by selecting the one with the largest value. As highlighted in Vargas et al. (2023a) many SDE discretisations can result in estimators that do not yield ELBO: for example, the estimators used in [Brmee al. (2022) can result in overestimating the normalising constant. Note similar remarks have been established in the context of free energy computation and the Jarzynski equality see Stoltz et al. (2010) Remark 4.5),
E.2 PROOF OF PROPOSITION 3.3 (CONTROLLED CROOKS' FLUCTUATION THEOREM AND THE JARZINKY EQUALITY)
Proof. Following the computations in Appendix D.I and using the formulae
we compute
In d PP I d d ) (x) - IIn(((o)) Inv(Y)))
+ % II dY, - oz / Vor(Y)) (Ln //' At(
Then via ItO's lemma applied to the unnormalised annealed log target In fr -lu In Ze we have
S/'' V Vn odY (X) dY, In Rt (T) In tYo) L/² ae (X)
thus arrive
d p dpn.Vaanvte in
-)) - +m (n mYo)
d P u. v o V
Oe ln tee (Y) dt / Vo)(Y.) (en,))))nn (Y,) dt 1/ (Yi
for arbitrary initial and final densities j and v. Crooks' generalised fluctuation theorem (Crooks) 1999) now follows from taking o O, and the controlled version in Proposition 3.3 follows from = To and v/ = TT. Finally notice that:
d dP In d d I
1
P Plar
'' d In on }
exp
- (()) + -n I (Yo)
which implies the Jarzynski equality when considering the boundaries μ au and mr. resulting in:
exp (%, dr in In dt ]
E
P in
e -(nn 20-
We want to highlight that in (Vaikuntanathan & Jarzynski 2008 Equations 10-14): we can see similar formulation to our proposed generalised Crooks' fluctuation theorem, that said Vaikuntanathan & Jarzynski (2008) seems to pose this as conjecture providing no rigorous proof. Furthermore, unlike our work, they do not formulate this result through SDEs, which we believe we are the first to do. In short, our work and concurrently Zhong et al. (2023) are the first to provide rigorous treatment in establishing the escorted version of Crooks' fluctuation theorem.
E.3 PROOF OF PROPOSITION 3.1 EM =
IPF
In applications, IPF is faced with the following challenges:
30

---

Published as a conference paper at ICLR 2024
1. The sequential nature of IPF, coupled with the need for each iteration undergo compre- hensive training as outlined in Section C.3 results in significant computational demands.
2. The reference distribution r(æ; 2) (or the reference vector field ft) enters the iterations in (8) only through the initialisation. As consequence, numerical errors accumulate, and it often observed that the Schrödinger prior is 'forgoteen' as IPF proceeds (Vargas et 2021a: la Fernandes et al. 2021 Shi et al. (023).
Thus to address these challenges this section will focus on establishing the connection between EM and IPF which in turn will provide us with family of algorithms that circumvent the sequential nature of IPF, further bridging variational inference and entropic optimal transport.
Proof. The proof proceeds by induction.
To begin with, the update formula in (8aa) implies that . (xx, z) = arg min ( (TKI, z) =)), z)) Tan (x=) = m(xr)) (œ.x)
recalling the initialisation T ((x, =) = r(x; To take account of the marginal constraint, we may write π (2, z) = (z()) and vary over the conditionals T<((x) By the chain rule for DRL. we
see that
(70) which is minimised at ग (() r(z(æ). From this, it follows that T' (ax x) = (zz)r(z((z) for the first IPF iterate. By assumption, the EM iteration is initialised such a way that q (2|1)) = r(=(m)) so that indeed nT' (x, z) = n00 (())))))))
DrL (x),
2)) = (xI(())))) + (Drr (T
The induction step is split (depending on whether n is odd or even):
1.) First assume that the first line of (20 holds for fixed odd n 2 1. Our aim is to show that this implies that
(x, 2) + (x) 2 (2)(
(71)
that is, the second line of (20) with n replaced by n + 1. From (18b), we see that
(x, 2) = arg min (DKI(T 2) 1| (x, 2)) ==(2) = ==)=)=
Again, we enforce the marginal constraint by setting <(x z) π (>(w))<(z) and proceed in (70
to obtain - (x. z) = x" (x)z))(z)) The statement in 771 is therefore equivalent to x (<<]z)
p/ (z]lz) To show this, we observe from the EM-scheme in 19) that
( 1)/2 = arg min C ( 1)/2,
In combination with the second line of 20) and the definition of (D 0) in (5) we obtain
(n+1)/22 arg min (xl(T"((: p (=(=)((())) = arg min (z) DKL (T (<<((<<<) llp"
9
where the second equality follows from the chain rule for DKL, as in (70). Since by assumption the
parameterisation of pd (æ=)=) >) flexible, we indeed conclude that n" (œ|z² = //t++///2 z)).
2.) Assume now that the second line of (20) holds for a fixed even n 2. We need to show that the first line holds with n replaced by + 1, that is,
n (xx, 2) = (2 (2) Using similar arguments before, see that nh* (æ, 2) n² (x))))(z)) that it is left to show that re" (x]l)) qdn/ (z)=:). Along the same lines as in 1.), we obtain
ón/2 2 = arg min LDxu (Ø, 0n/2) arg min DKL (g ( x (,, 2))
arg min E ( (2|z) || (=|=))
Again, this allows us to conclude, since the parameterisation in (2 (z)xz) is assumed to be flexible enough to allow for (2 = nr" (z/|z)
The proof for the path space IPF scheme verbatim the same after adjusting the notation. For completeness, consider drift-wise version below.
n1
31

---

Published as a conference paper at ICLR 2024
Remark 11 (Extension f-divergences). The proof does not make use of specific properties of DKL. other than that it satisfies the chain rule. As a consequence, the statement of Proposition straightforeardly extends to other divergences with this property, in particular to f-divergences, Proposition 6 in (Baudoin 2002).
E.4 DRIFT BASED EM
As remarked in the previous subsection, the proof of the equivalence between IPF and EM in path space follows the exact same lines, replacing the chain rule of DKL with the (slightly more general) disintegration theorem (Léonard 2014b). In this section, we provide a direct extension to the control setting, yielding yet another IPF- algorithm and motivating certain design choices for the family of methods we study.
Corollary F.2 (Path space EM). For the initialisation so the alternating scheme
n arg min DEL [ H. To puttoooo))
0
@at1 arg min DKL. ( (T (Tu.stoo puusto
(73)
agrees with the path space IPF iterations in (Brnnton al. 2019 et al. 2021).
Vargas al. 2021a De Bortoli
Proof. For brevity let (6,(0, = (K pustoPvo, Additionally, we parameterise
the forwards and backwards SDEs with respective path distributions v. T p as:
aY, = f((Y) dt V'V
dw Yoru,
dY, fY) dt + VV, dt + ď W. Yrov.
The proof will proceed quite similarly, so instead we will consider just the inductive step for the odd half bridge:
0₂ arg min CFB (F=(0).1,0)
We can show via the DKL. chain rule and the disintegration theorem (Léonard) 2014b) that the above
is minimised when 0 satisfies pp ppuffooVe = F p
dpf e.f lo which corresponds to
v0, V'V In p, ph Vo following Observation 1 in Vargas et al. 202 Similarly
as per Proposition 3.1 the results will follow for the even half bridges.
The above corollary provides us with convergence guarantees when performing coordinate descent DKL (I [urf+o'V0) subject initialising Φo practice, this indicates that the way of initialising & has a major impact on which bridge we converge to.
EM initialisation:
Thus as rule of thumb we propose initialising do such that we initialise at the Schrödinger prior: then one may carry out joint updates alternate heuristic, call this approach DNF (EM Init), as it is effectively clever initialisation of DNF inspired by the relationship between IPF and EMM
E.5 HUB-REGULARIZERS
As per Section 3.1 IPF resolves the nonquniqueness in minimising (0, 0) by performing the coordinate-wise updates (19) starting from an initialisation informed by the Schrödinger prior. On the basis of this observation, the joint updates (dn+ 0n+1) < (Φn, 0,,) A a. C (0, 0) suggest themselves, in the spirit of VAES (Kingma et al. 2019) and as already proposed in this setting by Neal & Hinton [1998). However, as is clear from the introduction, the limit lim (d., 0n). can merely be expected to respect the marginals in (6). and no optimality in the sense of (17) is expected. As remedy, we present the following result:
32

---

Published as a conference paper at ICLR 2024
Proposition E.3. For X divergence D on path space, and o, v (1., r R), ler
Schr (0.0)= D( puuto + AReg(o),
(74)
where Reg(E) = if and only if the HJB-equation deo + V + o/4 A Then Lsche (0, 0) 0 implies that the drift a; := oZ oe solves 17
2 0 holds.
The proof rests
optimal control reformulation of the Schrödinger problem (see Appendix Eb.
identifying the -uu-ttiunoon as the missing link that renders joint minimisation (74) theoretically sound for solving (17) The loss in (74) has two important benefits compared to standard IPF. First, it circumvents the need for the sequential updates used in IPF, thereby simplifying and speeding up the optimisation procedure. Second, it enforces the Schrödinger prior drift f directly, rather than recursively via eq. (8aa),
(18b). This prevents the prior from being forgotten, as is usually the case in regular IPF. In Appendix E.5 we detail possible constructions of Reg(Ø), discuss relationships to
previous work, and evaluate the performance of the suggested approach in numerical experiments.
This result can be found in Chen et al. (2021 Proposition 5.1), for instance, but since is relevant to the connections pointed out in Remark 13 below, we present independent proof:
Proposition E.4 (Mean-field game formulation). Assume that Φ E C (RR) satisfies the conditions:
1. The forward SDE
aY, = f((Y) dt o² (V)(()) dw. Youu
(75)
admits unique strong solution [0. satisfying moreover the terminal constraint Yr ~ v.
2. The Hamitton-lacobi--elllmae (HJB) equation
o² 0.6 f V +
Ao + V0
(76)
holds for all (1,x)) E [0, Rd.
Then V provides the unique solution to the dynamical Schrödinger problem as posed in
17
Proof. We denote the path measures associated the SDE
dY, f( + dW,
(77)
by and the SDE 75) by pd respectively. According to Girsanov's 's theorem, the Radon-Nil -ikkodmm derivative satisfies
dp
exp 0 / Voi(Y) -dW,
o² { Voll 2
(78)
dP
provided that the marginals agree at initial time, Po pd. Along solutions of 77 we have by Itô's formula
ff (o [( Vo.) of f ( (Y) of o Vor (o)(
(T(YY) (oo)))oo(Y))
/' (X) dt L' Vor( dw, = dIp
dpp
where we have used the HJB-equation 76 in the second line. Combining this with (78) we see that
dpp YY) = exp (- Yo)) exp (Yr()
(80)
The claim now follows, since the unique solution to the Schrödinger problem is characterised by the product-form expression in (80 see Léonard (2014a) Section 2), together with the marginal constraints Po yu and p v, which satisfied by assumption.
33

---

Published as a conference paper at ICLR 2024
Remark 12 (Summarised relationship to previous work), For X 0, coordinate wise updates of Csche (0. 8) recover the IPF updates from De Bortoli et al. (2021); Vargas et al. (022 according to Corollary 3.1 Note that Leehr is unconstrameed objective, in contrast to (17) previous works (Koshizuka & Sato 2023: Zhang & Katsoulakis, 2023) have suggested incorporating the marginal constraints softly by adding penalising terms the running cost in (17). Those approaches require limiting argument (from an algorithmic standpoint, adaptive tuning of weight parameter) to recover the solution (17). In contrast, the conclusion of Proposition E.3 holds for arbitrary À > Shi et al.
(2023); Peluchetti 2023) suggest an algorithm involving reciprocal projections onto the reciprocal class associated ft. From Clark 1991): Thieullen
(2002) Relly (2013). the equattin (76p is
local characteristic (Reg(Ø) = 0 forces (75) to be in the reciprocal class); hence Reg(ø) in 74) plays
a similar role as the reciprocal projection (Shi et al. 2023 Definition 3), see Remark [3 Liu et al.
(a) suggest an iterative IPF-like scheme involving temporal difference term (Sutton & Barto 2018
Chapter 6). As in Nüsken Richter
(2023). this is a an HUB-rcgulariser in the sense of Proposition
Remark 13 Finally, Albergo et al. (2023 Theorem 5.3) and Gushchin et al. 2022 develop
saddle-point objectives for
Remark 13 (Connection to reciprocal classes (Shi et al. 2023 Peluchetti 2023) and TD learning
(Liu et al., a)). The calculation in equation (79) makes the relationship between the HJB equation
(76) and reciprocal classes manifest (since reciprocal classes essentially be defined through the
relationship (0). see Léonard et al. (2014); RRllyy (013)). Moreover, equation (79) showcases the
relationship between TD learning (Sutton & Barto 2018 Chapter as suggested [il et al. a and
HJB regularisation. Indeed,
RegBSDE (Ø) Var (Tr( (Yo) + d² Vor
S/"' V Vo) dw,
(81)
where the variance is taken with respect to the path measure induced by 77 iss valid HJB-regulariser
the sense of Proposition E.3 The equivalence between RegaSDE (Ø) and the HJB equation (76)
follows from the theory of backward stochastic differential equations (BSDEs see, for example,
the proof of Proposition 3.4 in Nüsken & Richter (2023) and the discussion in Nüsken & Richter (222 i Section 3.2).
In the following, we present an analogue of Proposition E.4 involving the backward drift (Chen et al. 2019)
Proposition E.5. Assume that 0 E (2., R4 R) satisfies the following two conditions:
1. The backward SDE
aY f((Y) (V aw.
Yr N
(82)
admits a unique strong solution [0, T). satisfying moreover the initial constraint Yo
2. The Camilton-Jacobi-eelleee (HJB) equation
o²
0 + f V AO
V -
(83)
holds for all (t,x) [0, Rd
Assuming furthermore that the solution to 82 admits a smooth positive density p. have that
a, V0, + o²V In Ps provides the unique solution to the Schrödinger problem as posed in (7).
Remark 14. As opposed to Chen et al. (2016 equation (41)). the HJB-equation (83 does not involve the time reversal of the Schrödinger prior; the form of the HJB equations is not uniquely determined. On the other hand, 83) contains the divergence term V f. which discourages from enforcing this constraint in the same way as (6)). An akin result can be found in I Liu et al. (a) stated in terms of BSDEs.
be confused with reverse-time SDEs in 12
34

---

Published as a conference paper at ICLR 2024
Proof of Corollary E.5
Using the forward-backwant Radon- Nilionn- derivative in 14 we compute
In
dFms
MJ
(I) (atiee)
du (Tebb)
-ln
dv dLeb)
a [((Y)) dw, 0 ((Y)) Tw:
d v
[' VO,(Y.)
W. + o² / [00,
- In (1T(ee)))
dj (dLeb)
-In
dy
(V f.) di -o Tw, & vo.. (X)) dt.
dLeb 0
Here we have chosen 7 - O, and To- I' Leb. The initial measure for the Schrödinger prior is j., but the argument is unaffected by this choice (as the solution independent this). We now use the (backward) Itô formula along the Schrödinger prior,
0, (Y (Y J. 0,0,(Y) dt+ ( E Tw++/ / (0,00)) dt - (' (),(Y)) dt.
Using the HJB-equation (83). we see that
d (Y) - -in
dj
de cLLeb
In
(aLebb
In
0.(())) (o(o), (84)
d dupustovr
dLeb
and conclude in the proof of Proposition E.4
F CMCD EXPERIMENTS
In this section, we will cover further details pertaining to our experimental setup.
F.1 ELBO EXPERIMENTS AND COMPARISON TO GEFFNER & DOMKE (023
We compare our underdamped and overdamped CMCD variants against datasets from Geffner Domke 2023). which we describe in further detail below.
log_sonar (d = 61) and log_i onosphee (d =5) are Bayesian logistic regression models: x ~ NO0,, o.. 1), y ~ Bernoulli(iggmoid(""" u.)) with posteriors conditioned on the sonar and ionosphere datasets respectively.
brownian (d = 32) corresponds to the time discretisation of Brownian motion:
Oinn Log 2),
Olabs ~ Normalaaaa... 2),
I1 NO0 ainn);
24 ~ (I;-11, Ouun); i= 1....33.3
i 2, 2,..20
y/ Wx Onbss)
30.
inference is performed over the variables Qinn: alos and {x, given the observations
{ U 2
lorenz (d 90) is the discretisation of highly stiff 3-dimensional SDE that models atmospheric convection:
N loc
0, scale 1)
N loc = 0, scale =
21 N loc
0, scale 1)
N loc 10 (3i-1 24-1), scale = Qlinn
= 2, 30
V: N loc
Ii-1 (28 2-1) 91-1) scale Q'ion
2, ..........
2: N loc 2;-10h-1
8
3 2;-1; scale Qimn
2, 30.
0 N loc Ii, scale = 1)
2, 30
where Cim 0.1 (determined by the discretization step-size used for the original SDE).
The goal to do inference over I,, Vis 2; for = 30, given observed values o; for E (1,..., . 10) U (20....,30) 30).
35

---

Published as a conference paper at ICLR 2024
seeds (d 26) is random effect regression model trained on the seeds dataset:
T Gamma(0.01, 0.01)
ag NV 10)
01 ~ NO 10)
02 NO(0, 10)
41 (0)
b; ~ 0, 1
VF
1=1, 1=,...,,21 21
logits; op 01 + 0234 + a12241 + b
1,...,21 21
r Binomial (logits,, N,)
1 21.
The goal is to do inference over the variables 41, 02 012 and b; for = 1,..., 21, given observed values for Fs. ys and N..
For all target distributions, we follow the hyperparameter setup from Geffner & Domke (2023)
from their code repository [ for all baseline methods (ULA, MCD, UHA, and LDVI) well as our overdamped and underdamped variants. We first pretrain the source distribution to a mean-fiee meeldfieeld
Gaussian distribution trained for 150, 000 steps with ADAM and learning rate of 10-2, We then train for 150000 iterations with a batch size of 5, tuning learning rate between [10-5, 10-1, 10-3] picking the best one based on mean ELBO after training. For all methods, during training the mean-field source distribution is continued to be trained, as well as the discretisation step size and ôto. For the underdamped methods we also train the damping coefficient 7, and for methods involving a score network, i.e. MCD, LDVI, CMCD and CMCD (UD), we train the networks which are chosen to be fully-com fully-connected residual networks with layer sizes of [20, 20). In order to report the mean ELBO after training, we obtain 500 samples with 30 seeds and report averaged value over them.
F.2 In Z, SAMPLE QUALITY EXPERIMENTS AND COMPARISON TO ZHANG & CHEN 2022 VARGAS ET AL. 2023A)
Furthermore, we also include comparisons to a large-dimensional target distribution and two standard distributions with known In Z replicated from Vargas et (2023a), which we summarise below.
1gcp (d 1600) is a high-dimensional Log Gaussian Cox process popular in spatial statistics Moller et al. 1998). Using ad M x M 1600 grid. obtain the unnormalised target density N(;;; μ, K) K)ILeejMM exp (I.8) - a exp (zz)).
funnel (d N( ( O, (x) with 07 = (Neal 2003).
10) is challenging distribution given by (r(22110
gmm (d = 2) is two-dimensional Gaussian mixture model with three modes, given by the following target distribution
1 3.V
0.7
0
1
-2.
0.7
N xr (x)
x::
0
0.05
+
3
.
x
0
0.05
0.95
+ W
0.95
For these target distributions, we follow the hyperparameter setup from Vargas et al. (2023a) from their code repository [4 or the baseline methods of DDS and PIS, and replicate them as closely as
|tttaas://ttual.ooo//ooooo////// denoidnnniiooooooooo sing_di www.iiiffffuiinona_aaaa____
36

---

Published as a conference paper at ICLR 2024
Concatenate
MLP Block (2)
Output
Embedding Layer
Architecture from (Geffner & Domke 2023) used across experiments for CMCD drift network. Softplus activations are used.
Figure 2:
possible for CMCD. Unlike the previous, we don't 't pretrain the mean mhean-field Gaussian source distribution
(0,, m) [). We select the optimal learning rate in [10-3 10-0, 10-5], the optimal standard deviation of the source distribution iin in [1, 2, 5, and the optimal a [0.1, 0.5, 1, 1.5, 2]. Instead of training e ôto, we sweep batch size of 300 for 11000 steps, where we keep the source distribution parameters fixed, as well 6. For evaluation, over the seeds. DDS and PIS use a 1i8demensiooal positional embedding, along with an additional network for the time parameters, however MCMD uses regular exact comparisons, we select differing network architecture sizes that result in an equivalent number of parameters for funnel and gmm. For gep, due the high dimensionality of the dataset, we choose a small network for CMCD. We summarise these below. For gmm. and funnel, it is possible to sample from the target distribution, and we report OT-regularised distance (W22 with regularisation 7 10-2 Similar to the mean ELBO, we draw 2000 samples from the models and the targets, and average W over 30 seeds. We use the Python Optimal Transpor implementation of hetropy-regularised distance. Results for comparisons to DNF can be found in
optimal value in [10- 10110, ,, 11. The models are trained with
use 30 seeds with batch size of 2000, and report average performance
network. In order make
15 library's ss default
T 5
Table 1: Network Sizes for comparison. Note that CMCD has less parameters for the despite the Funnel target despite the larger drift due to the PIS and DDS networks having an additional grad network.
GMM
LGCP
FUNNEL
DDS
[10, 10
664 64]
[64, 64]
PIS
10, 10
64, 64
[64, 64
CMCD
38, 38
64, 64
[110, 110)
F.3 COMPARTSONS WITH THE LOG-VARIANCE LOSS MODE COLLAPSE FAILURE MODE
Here, we report performance using the log-variance divergence-based loss NNüseen & Richter 2021 introduced at the end of Section 3.2
N(Y YI + (V + V oI, (Y,, At.. 20 Atx) MY,, Y &e (Vln Mecti VIn @i ( 114 Atk. 20 Att))
CMCD Var
(1) Var In
(24)
careful reader will note this loss simply consists of replacing the expectation in the KL loss with a variance. A major computational advantage of this loss is that the measure that the expectations taken with respect can be any measure and
restricted the forward backward SDEs
like in KL (Richter & Berner 2024 Richter et al. 2020 Nüsken & Richter (021). I), this allows us to detach the samples and thus accommodating for much more computational objective.
which we find performs quite well compared to our default loss function, especially for multimodal target distributions. We consider the very multi-modal mixture of Gaussian target distribution from Midgley et (022). and report the ELBO and Z numbers the table helow. For this experiment, we use batch size of 2000 and train neural networks with size [130, 130] for 150k iterations.
htttps://pythhoot.t http://ppyyoono....t.......
37

---

Published as a conference paper at ICLR 2024
Table 2: ELBO and on 40-GMM
ELBO
In
W2
log-variance LOSS
-1.279 0.096
-0.065 + 0.101
0.0143 0.001
KL LOSS
-2.286 0.1109
-0.244 + 0.3309
0.0441 0.012
samples
samples
6
20
-00
8
Figure 3: (left) 2000 samples drawn from the CMCD algorithm trained with the default loss function, and (right) 2000 samples drawn from the algorithm trained with the log- divergence based loss. We can that the default loss function misses many modes in the target distribution, whereas the log-variance loss has not missed any modes. We report final results after sweeping over A., and learning rates for both methods, picking the one with the lowest training loss. We highlight that concurrent work by Richter Berner (2024) explores the log variance divergence more detail and proposes an akin general framework for diffusion-basse sampling.
10
mmm M
Train loss
200
600
8D0
1000
Train Steps
Default
Variance loss, A.
0.96
Default lows, A. 0.66
Figure 4: Plots showing training loss curves for the log-variance loss and the default loss for different values of At,- We find that a low value of A,, = 0.1 is needed in order to obtain a low training loss for the default loss, whereas the log- log-variance loss is much more robust to different values of A The x-axis reports an evaluation every 150 steps of training
F.4 SPECIFICATION AND TUNING: SMC, AFT, AND NF-VI
We adopted the implementations provided by the studies in Arbel et al. (2021) Matthews al. (2022 and initialized them with default hyperparameters before fine-tuning.
38

---

Published as a conference paper at ICLR 2024
Sequential Monte Carlo (SMC). For SMC, utilized 2000 particles sampled from zero-mie cero-mean unit Gaussian distribution, implementing re-sampling if the effective sample size (ESS) fell below 0.3. We employed Hamiltonian Monte-Carle (HMC) for particle mutation, executing one Markov Chain Monte Carlo (MCMC) step after each annealing step. The numher of leapfrog steps was fixed at 10, and an extensive grid search over different step sizes conducted, consistent with AAreel et al. (2021). This search spanned four different step sizes, contingent on the temperature, resulting grid search over 256 parameters. The finalized values are presented in Table 3 For SMC, K defined by the number of temperatures.
Furthermore, we report the results for SMC in Table 3 for the tuned hyperparameters used for each target. Please note that we were able obtain similar In Z values as in Vargas al. (2023a) suggesting SMC was well-tuned. Finally, for each result, we report the mean and standard deviations across 30 different seeds, results can be seen in Table
Annealed Flow Transport Monte Carlo (AFT). We maintained similar setup to SMC, with a few adjustments: using 500 particles for training and 2000 for evaluation to accommodate the added complexity from the normalizing flows. We also decreased the number of temperatures and increased the number of MCMC steps to mitigate memory requirements from the flows. K' is defined as the number of temperatures trained simultaneously. Inverse autoregressive flows (IAFs) were employed in all experiments except for lgep, using a neural network with one hidden layer whose dimension matches the problem's dimensionality. For lgep, diagonal affine flow was used due to memory constraints arising from the high dimensionality. AFT flows were trained for 300 iterations until convergence.
MCMC steps, with the latter fixed at 4, resulting in maximum of 64 flows
Variational Inference with Normalizing Flows (VI-NF). We utilized the same flows as for AFT. In this case, K denotes the number of flows to stack. The flows were trained over a total of 2000 iterations with batch size of 500. For some targets
Table 3: Tuned MCMC Step Sizes.
CMN
LGSP
LONDNZ
SHOWNIAN
LOG SONAR
LOC. 10NO0000S1EE
SCEDS
PUNNEL
0...6.0.....
0..00,002
0.0. 0.01)
1022002 0.0,000
0.2.0.00,,2 0.]
[0..0.2 02,022
p.L.O.L,,000 a0i)
K.06.0.2 0.20000
Table 4: SMC Results. ELBO and In Z values for different number of steps K and experiments.
LOCF
LOWENG
300WNIAN
LOG SOMAR
SU1D5
PUNNEL
0,842
7.797
S7510
416
2.784
2214
3,04
100,770
1.221
-1018 0.116
0,265 0.056
LS5
206-277-2
4.065
-22122 280
1047. 478
-24.744
115.891 d: 1,666
114..07 2 L088
88.008
0.465
-065:2:00
32
10.119 0.017
56 106
5.9
-19045. 207 K34.5
-1535-19.22
-130.557
0.61
-137.
1 00.394
-78
0-973
-X.08 0.17
64
0.000 2 0.015
203.506
7.42
9604,505 + 115:723
4 1,042
118885 2 0.307
118,812
= 0.152
26.255 0.159
IU550 = 0.067
0,023 0.000
5.427
04100075
1.76
110.001
0.001
0.307
74.774
0:250 0.108
ELBO
BU022 0,005
233 051 3 9.623
2012229917
- 1101177 A: 3:3.22
-LIV 4,009
- -123.772
3: 3.560
-15.183 1-347
-IU42T 2 0.296
-6.003 0,05
7.196
-2/90/22
-5.02 2 2.114
1110000
1825
-113.466
1 1.361
-13.524
0.845
-0.922 0.184
-12110
9205
0.778
-21.575
-01257
.28
0000 6,00
490.58 0.441
-8524.107
165.
1.008 0.27
108 108.296
0.087
111.,003 0.256
73,699
0/86
-0.222.0.12
X 256
0032 0.005
153
3 1.40
-2811 161 100.68
1.302 4 0.135
- -00 (108366:2
0.071
-111.61
0.1.7
-73.113 8 0.87
-0151 - +0.36
F.5 FURTHER ABLATION WITH NF-STYLE METHODS AND AFT
We further run both flow models (AFT and NFVI) on all possible target distributions (subject to OOM errors). Results can be found in T 6
F.6 WALLCLOCK TIMES FOR In Z CALCULATION
In order to calculate the average wall-clock time for In Z calculation, we calculate the time it takes to draw 30 seeds of 2000 samples each from the methods below, and use these samples calculate the mean and standard deviation of in Z across 30 seeds.
F.7 TRAINING TIME COMPARISONS TO SMC
In this section, we explore total time comparison between our approach CMCD and SMC.
39

---

Published as a conference paper at ICLR 2024
Table 5: In Z comparison. In Z values for different number of steps K, experiments and methods. Not all methods could be evaluated on every K/experiment combination due to numerical instabilities or out- memory (OOM) problems.
Dataset
Method
K
16
K 32
64
K 128
K' 2%
funnel
CMCD
-0.3037
1507
+ 0.0773
401085
0.1143
-00573
0,0444
0.0641
log Jososphere
VI-DNE
-0.3768
= 0.2157
-0.3817 + 0,1627
.02211 +: 0.0999
406941
# 0.6841
0.1947 = 0.1325
-02124
0.0637
VI-NF
-0.2062 0.079
+0,206 0.082
-0.2061 0.087
-0.1941 0.101
0.1824 0.097
-0.197. 0.099
AFT
-0.87512 0,543
-0.385L 0.351
-0.1881 0.192
-0.271 0.227
-02351 0.139
-0.1961 0.111
gmm
CMCD
-0.1358
0.0899
-0,01331
0,1292
0.0195
0.00366
0.0177
-0000
0.0368
-00081
0.0520
VI-DNE
-0.3676 1: 0.6314
-0.258
0.412
-04983 1 0.3878
40,4499
0.5379
-04652 1: 0.3223
-0.004
0.6381
VI-NF
-0.3555 0.698
-0,455+
0.258
-0064+ D.138
41054+ 0.15
-0066- 0.188
-0.45++ 0.177
AFT
-0.336-
0.372
-0,006+
0.082
0.02+ 0.068
-0,016+ 0.042
40000++
0.029
0.001 0.026
CMCD
491 019 3.553
147 2.634
502.705
2
506045
1,761
165 1,553
500. 43
VIDDNP
5.858
5.861
5,860
8.896
AFT
126.651 5.764
344.1454: 23.95
191.613 173.873
420.2594: 91.43
480, 126 33.059
491.028 8.057
Table 6: ELBO comparison. ELBO values for different number of steps K, experiments and methods. Not all methods could be evaluated on every K/experiment combination due to numerical instabilities or out-of- mf--of-memory (OOM) problems.
Dutaset
Method
K=
K= 16
K 32
64
K = 126
K 250
sceds
CMCD
74.501 L 0.049
-T4,527 1 1.006
-74,44 0.05
73.067 1 0.038
-73.8
-73.684 0,035
VI-FF
73.063
0.003
-74.547 0.011
-74.74
T8.58
0.014
00014
-786676
0.011
147.457
-116.134
8.157
-90.002
+-0419
776
0.188
CRAFT
1.531
-9L2 0.305
-50.95
-70.555
0.175
-74.99
-74.22
Tog, Jonosphre
CMCD
113.211
0.089
112.548
11126433
0.082
112.22 0.046
-11.98 0.04
0.045
VINF
111.003
0.022
-111.902
0.022
-111.802
0.017
-111.881 1 0.017
00M
OOM
AFT
168.174 + 21.224
138.788
+ 8:374
123.013
8.771
11886644
#: 0.801
116.07 = 0.496
114.90 = 0.781
Ing, Sonn
CMCD
112:274 + 0.124
110.04 + 0.111
110.460
0.106
109 503 0.076
109.08 + 0.066
100.25 0.052
VI-FF
100.553 0.035
-199.346
1 0.001
100, 441 0.035
100
0.041
100.711
0.000
00M
AFT
209.249 : 12,506
-168355
1 8,00
-22.772
3,057
-1211 82338833440 2.505
114.91
+ 0.101
1122.211 1.182
lgsp
CMCD
469.475 d: (1,250)
2412246255 0:287
486.789
0.249
492.745 à: 0.200
207.074 d: 0.267
4:0 708
11,230
AFT
62 808
340,087 126.851
30.455
424.715
12.416
Iarenz
CMUD
3180.997 0.184
-1180.797
6.188
11170.514
11.154
-1174.300
0.148
-1172.450
0.158
0.15
VI-NF
0.84
-1471.778
0.582
-1400.649
-1433.536
0.316
Inosiion
-0.753
0.075
-0,200
0.058
0.578 0.046
VI-N
0.733 0.019
0.797
0.001
0,816 0.018
00M
As both methods are quite inherently different it is not immediately obvious how carry out an insightful comparison. In order to do so we chose the LGCP which is our most numerically intense target and we phrase the following question:
"For how long do we have to train CMCD to outperform the best-run SMC"
For this, we look at our Figure pane c) and we can see that at K 8 CMCD already outperforms SMC at K = 256 with 2000 particles. So we choose these two approaches to compare to. In Table g is a brief comparison of total time calculations, note we have included tuning time for SMC which is akin to our training time as without tuning SMCs hyperparameters ELBOs and In Z estimations were much worse. We can observe that the total runtime for training and sampling CMCD to reach a better In Z value does not exceed the time required to tune SMC.
G REGULARISED IPF-TYPE EXPERIMENTS
For the purpose of completeness in this section, empirically explore the regularised IPF-type ob- jectives proposed in the main text. We explore a series of low-scale generative modelling experiments
Method
Average Time (s)
Min Time (s)
Max Time (s)
CMCD (OD)
9.665
5.592
21,475
ULA
9,204
4.673
20.721
UHA
9.427
5.588
20.263
MCD
9.204
4.673
20.721
Table 7: Wallelock times for evaluation in seconds.
40

---

Published as a conference paper at ICLR 2024
Mode
ULA
MCD
CMCD
DDS/ PIS O(K- O(K- G
Sampling
O(K (d G(d))
O(K (d G(d)))
OJK (d+ G(d) + N())) O(K (d G(d) + (()))
ELBO
O(K (d G(d))
O(K G(d) + M(d))
Table 8: Sampling and loss calculation complexity across SDE based methods, K represents the number of integration steps, Gi(d) represents the cost of evaluating the score of the target and N (d) for evaluating the drift/score networks both quantities are dimension dependant. PIS and DDS have an additional grad network cost N2 which is dimension independant.
Method
Train + Sample Time (min)
In Z
ELBO
CMCD
33.12 + 0.12
491.059 + 3.553
469.475 + 0.2589 453.395 4.4300
SMC
62.62 0.10
477.162 + 4.998
Table 9: Training+ Tuning Sampling time comparisons for CMCD and SMC at comparable In Z estimates.
where the goal is retain generative modelling performance whilst improving the quality of the bridge itself (i.e. solving the SBP problem better).
Across our experiments, we use DKL. and let Po = PT = Leb, which can be simplified to the heward-backwarrd KL objective used in DNF (Zhang et al. 2014). see Appendix E.I.I We use the Adam optimiser (Kingma & Ba 2015) trained on 50,000 samples and batches of size 5000 following Zhang & Chen 2021p. For the generative modelling tasks we use 30 time steps and train for 100 epochs whilst for the double well we train all experiments for 17 epochs (early stopping via the validation set) and 60 discretisation steps. Finally note we typically compare our approach with À to DNF (A 0), with DNF initialised at the reference process, which we call DNF (EM Init),
see Appendix E.4 for further details.
G.
2D TOY TARGETS GENERATIVE MODELLING
Here we consider the suite of standard 2D toy targets for generative modelling explored in Zhang & Chen (2021) In contrast to Zhang & Chen 2021) we consider the SDE dY, -o dt 2ddW as the Schrödinger prior across methods. We parametrise DNF and our proposed approach with the same architectures for fair comparison. Furthermore, we incorporate the drift of the above Schrödinger prior into DNF via parameterising the forward drift as in (75). partly motivated by Corollary E.2
In order to assess the quality of the bridge we consider three different eгтor metrics. Firstly we estimate DKL between the Schrödinger prior and the learned forward process (i.e. Fyn J ||a, (Y) dr)) Secondly, evaluate D +o $p -u. ob- tain a proxy error between the learned and target marginals. Finally, we estimate the cross entropy between F
a. and to assess how well the constraint at time T is met.
In Table 10 we observe that similar values of DKL attained across both approaches in the tree, sierpinski. and checkerboard datatsets whilst achieving significantly lower values the SBP loss across all training sets, and for tree, swirl and checkerboard validation datasets. At the same time, we can see that the cross-entropy errors are effectively the same across both approaches. Overall we can which train approach, obtain much better
conclude that on the empirical measures fit for the target Schrödinger bridge, and on the validation results we can see that we generalise to 3/5 datasets in improving the bridge quality whilst preserving the marginals to similar quality.
G.2 DOUBLE WELL RARE EVENT
In this task we consider the double well potential explored in (Vargas et al. 202 Hartmann et al. 2013) where the Schrödinger prior is specified via the following overdamped Langevin dynamics dy, - UY, dt + a dW, The potential Uty) typically models landscape for which is difficult transport y into
41

---

Published as a conference paper at ICLR 2024
(a) 2
(c) A = 0 (Random Init) proposed regularised objective.
(b) A = 0 (EM Init)
(d) DNF (No Prior)
Figure 5:
(b) A set O but using clever EM motivated initialisation, C À set to 0 with random initialisation the forward drift, ( for reference DNF with ft 0 (uninformative Schrödinger prior).
This is notably challenging task as we are trying to sample a rare event and as noted by Vargas et al. (2021a) many runs would result in collapsing into path rather than bifurcating. In Figure we can observe how our proposed regularised approach (5a) is able to successfully transport particles the well whilst respecting the potential, whilst both variants of DNF using the EM-Init for (5b) and random init (5c) fail to respect the prior as nicely and do not bifurcate, with the random init in particular sampling quite inconsistent trajectories. Finally for reference we train a DNF model with ft = 0 and Φ (dd) initialised at random to illustrate the significance of the initialisation of ø.
G.2.1 DOUBLE WELL POTENTIAL
We used the following potential (Vargas et al. 2021ap:
5 (x² - 1)2
1
x2
1²
+
x²
+ exp
(85)
2
with 0.35, furthermore, we used the boundary distributions:
a ~~N N ((o).
(0.0125
0.15) ).1s): -~N( (())
0.0125
0 0.15)
).
0
0
The Schrödinger prior is given by:
dY = - Vv, +od
(86)
with 0.4. The terminal time is T 1. Furthermore, employ the same exponential discretisation scheme as in the generative modelling experiments.
Target
Method
KL
SBP Loss
PINN Loss
Cross Ent
Val
Train
Val
Train
Val
Train
Val
Train
== 0.5
1.67#0.02
1.000.01
4..84
42.31=1152
0.06+0.00
0.05a0.00
2.87+0.01
2.80=0.01
tree
DNF (FM Init)
1.63+0.02
1.39M0.01
15.33+1.79
49.60+1.68
1.7480.04
1.6440.04
2.88+0.01
2.80=0.01
A= 0.5
2.9520.06
0.12+0.01
39.30s0.90
25.2440.62
0.2620.01
0.1020.00
2.4920.01
2.7720.02
olympics
DNF (EM Init)
2.70=0.05
0.02+0.01
40.20H+.77
1.533
1.6440004
2.0540.08
2.54+0.01
2.77+0.02
sierpinski
A= 0.5
2.3120.01
2.20x0.01
28.5421.49
26.6720.90
0.0420.00
0.0340.00
2.8240.01
2.8320.01
DNF (EM Init)
2.30:0.01
2.20+0.00
30.8721.93
53
1.25+0.14
7.22:0.14
2.80#0.01
2.82:0.02
A= 0.5
15.6740.29
1.95+0.03
121.81+1.4
40.24t:1.74
1.010.03
0.1420.00
2.97+0.01
2.6020.02
swirl
DNF (RM Tnit)
13.7740.38
1.9240.04
151.6743.68
07.55&1..86 1.86
5.89+0.15
2.6340.08
2.9540.02
2.74+0.03
== 0.5
4.7920.01
4.7010.01
34.4720.80
33.7020.91
0.0320.000
0.02:00.00
2.8220.00
2.8120.01
checkerboard
DNF (EM Init)
4.78:0.01
4.70+0.02
39.76t5.83
1.10
3.66#0.07
3.68:0.06
2.81
2.8120.02
Table 10: Generative Modelling Results comparing DNF (Zhang & Chen 2021) (A 0) to our PINN regualirsed approach with = 0.5. We observe that PINN regularisation obtains similar KL and Cross entropy losses to DNF whilst achieving lower distances to the prior.
G.3 IMPLEMENTATION DETAILS
G.3.1 NEURAL NETWORK PARAMETERISATIONS
Following Zhang & Chen 202T and the recent success in score generative modelling choose the following parameterisatios::
(x(x)) fi(x) + o vo(tzx),
(87a)
bx(x:) fi(x) + o Vo(t,z) o² (t,
(87b)
42

---

Published as a conference paper at ICLR 2024
Figure Generated samples trained by our approach (A 0.5) left and DNF (A 0) right. Qualitatively we can observe that both learned models have similarly matched marginals.
where 8g is a score network (Song et al. 2021 De Bortoli et al. 2021 Zhang & Chen 2021 and
a(),az) neural network potential. We adapt the architectures proposed in Onken et al. [202] )
Koshizuka & Sato (2023) to general activation functions. Note that these architectures allow for
fast computation of Aø comparable that of Hutchinson's trace estimator (Grathwohl al. 2019
Hutchinson 1989).
Finally, we remark that the parametrisation in (87b) allows us to learn the score of the learned SDE and thus seamlessly adapt our approach using the probability flow ODE (Song al. 2021 at inference time.
G.3.2 PINN LOSS
For the PINN loss across all tasks, we sample the trajectories from Y Fr.Vo and thus employ detach(() before
the same discretisation used in the KL loss. However, detach the trajectories 0.7 calculating the gradient updates in a similar fashion to Nüsken & Richter (202I).
43


Reviews:
[{"review_id": "De8pik4v9w", "reviewer": "Reviewer_yH6G", "summary": "This paper presents the controlled Monte Carlo diffusion sampler (CMCD). CMCD builds upon previous work in the literature by adapting both the forward and the backward dynamics. The authors also posit a regularised iterative proportional fitting for schrodinger bridges that improve upon standard IPF. Finally, the authors link CMCD to Jarzinsky identity and demonstrate its performance through numerical experiments.", "questions": "Is eq. (10) the same as eq (1)? I failed to find the difference.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 1, "contribution": 2, "strengths": "The technical contributions of this paper is unquestionable. The authors have managed to link key and fundamental ideas from various fields. I especially liked the connection between regularised IPF and EM algorithm, which is illuminating. I think the paper has good contribution and would be helpful to researchers in the field if presented correctly.", "weaknesses": "Unfortunately, I think the presentation of the paper needs to be revised extensively. The paper is confusing, and posits ideas without forewarning or motivation. The introduction is especially difficult to read, and not because of its technicality. In general, I believe the authors' math is quite readable. It is the writing and motivation that needs improvements. As an example, nowhere in the introduction the authors explicitly state the problem they are tackling (CMCD). It takes until page 7, eq. (22) for the authors to state the explicit problem that they address. \n\nAs another example, the authors never explicitly state the necessity of various connections they draw. For example, the text above Proposition 3.1 explains why IPF fails to perform well. However, there is no explanation afterwards to say why going through the lens of EM solves this issue. Rather, the text proceeds to describe optimality conditions in what seemed to me to be like a tangent. In summary, although I could follow the math, it was extremely hard for me to follow the context of the paper.\n\nIt pains me to say that I cannot suggest the authors any single thing that could drastically improve the paper. I appreciate the technical contribution. However, given how difficult it is to follow its context, I don't believe that the paper would be helpful to the wider community in its current state.\n\nAt the very least, I think the first two pages of the introduction should be thoroughly revised to focus more on the problem addressed in the paper, rather than overloading on related technical terms and connections. I also think the propositions which are not key contributions of the paper (for instance, 2.1) could be better placed in the supplements. \n\nI also list some minor typos:\n\n1. A bracket has not been closed in page 2, paragraph 3, after \"(Section 2.2\"\n2. Page 7, last paragraph, \"Proposition2.2\" is missing a white-space.\n\nPost Rebuttal: I thank the authors for the extensive revision. I think it addresses most of my comments and the authors truly went above and beyond. I have revised my scores to reflect this. I think that in its current form the paper can definitely be useful and has interesting insights which are presented clearly. I believe, if accepted this will be an important contribution to the community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iFVGoJzxBK", "reviewer": "Reviewer_LiYc", "summary": "This paper is divided in two parts. In the first part, the authors provide a unifying framework for VAEs and diffusion models using stochastic differential equations. In the second part connect the iterative proportional fitting procedure with the expectation maximization algorithm and provide a new objective for solving IPF that avoids the mode forgetting phenomenon that happens when one implements the diffusion schrodinger bridge. The added regularization term is null iff the HJB equation holds. Finally, the authors introduce the MCMD algorithm which makes uses of forward and backward SDEs that admit as marginals a prescribed sequence of distributions $(\\pi_t)_t$.", "questions": "i have no further questions", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- I enjoyed reading this paper; it is highly pedagogical and sheds light on interesting connections between variational inference, diffusion models, stochastic control as well as optimal transport. The connection with the expectation maximization algorithm is particularly interesting. \n- The proposed algorithm CMCD generalizes existing methods such as Monte Carlo VAE aswell as MCD. I believe that the adopted point of view here is more rigourous and elucidates what the backward process should \"look like\", which was not clear in the Monte Carlo VAE paper for example. \n- Extensive experiments show that the CMCD method outperforms existing methods in relatively high dimensional examples.", "weaknesses": "- It is quite unfortunate that there are no experiments for the loss proposed in Proposition 3.2 to back the fact that the proposed loss does not suffer from mode forgetting. Furthermore, it is not clear how this is implemented in practice since a Laplacian term is involved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PrsY5UsWIr", "reviewer": "Reviewer_2i1s", "summary": "The authors propose to revisit sampling through the lens of generative modeling. In particular, they show that the EM algorithm can be formulated as an Iterative Proportional Fitting when there is no restriction on the kernels underlying the generative models. Based on their presentation and observations, they propose a score-based annealing (Controlled Monte Carlo Diffusion) that betters the state of the art on two datasets: funnel and GMM", "questions": "I would be very happy to have more information on the relation between the correspondence between EM and IPF as discussed in limitations.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I find the paper very clear and the presentation very interesting. I believe the main novelty to be in the score-based annealing CMCD. I am unsure how novel the restricted correspondence between EM and IPF is; however, presenting such a relation is very interesting.", "weaknesses": "I am unsure how novel the restricted correspondence between EM and IPF is; in particular, it seems to me that the correspondence only holds when there is no restriction on $p^\\theta(x\\vert z)$; in the general setting, EM does not allow for the forward generative model $\\pi^{2n+1}(x,z)$ to have as a marginal $\\pi_x(x) = \\mu(x)$ and the backward to have as a marginal $\\pi_z(z) = \\nu(z)$. If I am not mistaken, I think it would be very helpful to have some comments on this point in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q8UHyQlsnM", "reviewer": "Reviewer_nqRv", "summary": "The authors propose a novel framework for generative modelling and Bayesian computation that bridges variational inference and optimal transport using the theory of diffusion models. Starting from the Schrödinger problem and iterative proportional fitting, the authors draw a connection to the EM-algorithm, and finally present a novel regularized diffusion objective with a unique minimizer. They validate their method on several experimental benchmarks where it outperforms competing state-of-the-art methods.", "questions": "- In my opinion the paper could benefit from a shorter introduction and motivation and instead an extended experimental section with method limitations, experimental details, computational trade-offs, etc. for more applied audiences. Since this is likely not possible due to page limits, an extended treatment in the appendix would be very welcome.\n- What are the training times of the method, for instance, in comparison to running SMC?\n- How were the neural networks architectures chosen and how is the hyperparameter selection and other details impacting the experimental results (e.g., see Table F1).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper is as far as I (the reviewer) can judge highly original and will be greatly valued by the community. \n- The paper is well written and generally clear. The contextualization of their work with respect to other frameworks is well done. It unifies and generalizes several previously introduced frameworks, such as Monte Carlo diffusions or unadjusted Langevin annealing. The motivation and lead-up to their final results is in my opinion very convincing.\n- The proposed method outperforms other state-of-the-art-methods on a variety of experimental evaluations. The experimental section is sufficiently exhaustive wrt compared data sets and competing methods.", "weaknesses": "I don't have major comments regarding weaknesses other than that the source code for the experimental validation needs to be provided in the supplement and not only via some possibly non-permanent web link.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xb89MEUPaf", "reviewer": "Reviewer_ndxa", "summary": "This paper present a framework to perform variational inference with a score-based algorithm. Specifically, it introduces HJB-regulariser to the diffusion normalizing flow process and provide some justification for the framework.", "questions": "In Variational inference literature, the hidden dimension $z$ is a low dimensional manifold that can be interpretable factors? The latent of VAE uses smaller latent dimension than the data. However, in diffusion model/transport-based algorithms. The space of $x$ and $z$ are in the same dimension, which makes $z$ encode all the information of $x$ (and loses the ability to perform inference). How can you justify this compensation? What are the benefits of the (almost) lossless VI?\nMoreover, it is not clear that how good is this “VI” compared to other VAEs, such as image generation, latent variable inference. The current setting is relatively toy.\n\n\nFrom the transport side, although it is understandable that Figure 2 is more well behaved than non-regularized transport, we still found that Figure 3 is good enough in real practice. Can you find an example that the non-regularized way fail?\n\nIn general, I do believe that the HJB-regulariser is useful in some cases, but given the current justifications and evidence, the significance is not clear at this moment.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This connection between score-based algorithm and variational inference is an interesting topic. It remains open to make these score-based model really works in VI literature.", "weaknesses": "The idea is relatively direct and the whole motivation looks incremental.\n\nThe empirical results are not strong enough from both VI/transport side.\n\nThe ablation study of HJB-regulariser is weak.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)