Title: Stochastic Optimal Control for Diffusion Bridges in Function Spaces
Venue: NeurIPS 2024
Accepted: yes
Keywords: stochastic optimal control

bridge matching

diffusion model
Abstract: Recent advancements in diffusion models and diffusion bridges primarily focus on finite-dimensional spaces, yet many real-world problems necessitate operations in infinite-dimensional function spaces for more natural and interpretable formulations. 
In this paper, we present a theory of stochastic optimal control (SOC) tailored to infinite-dimensional spaces, aiming to extend diffusion-based algorithms to function spaces. 
Specifically, we demonstrate how Doob’s $h$-transform, the fundamental tool for constructing diffusion bridges, can be derived from the SOC perspective and expanded to infinite dimensions. 
This expansion presents a challenge, as infinite-dimensional spaces typically lack closed-form densities. 
Leveraging our theory, we establish that solving the optimal control problem with a specific objective function choice is equivalent to learning diffusion-based generative models. 
We propose two applications: 1) learning bridges between two infinite-dimensional distributions and 2) generative models for sampling from an infinite-dimensional distribution. 
Our approach proves effective for diverse problems involving continuous function space representations, such as resolution-free images, time-series data, and probability density functions.

Stochastic Optimal Control for Diffusion Bridges in Function Spaces

Abstract

Recent advancements in diffusion models and diffusion bridges primarily focus on finite-dimensional spaces, yet many real-world problems necessitate operations in infinite-dimensional function spaces for more natural and interpretable formulations. In this paper, we present a theory of stochastic optimal control (SOC) tailored to infinite-dimensional spaces, aiming to extend diffusion-based algorithms to function spaces. Specifically, we demonstrate how Doob's h -transform, the fundamental tool for constructing diffusion bridges, can be derived from the SOC perspective and expanded to infinite dimensions. This expansion presents a challenge, as infinite-dimensional spaces typically lack closed-form densities. Leveraging our theory, we establish that solving the optimal control problem with a specific objective function choice is equivalent to learning diffusion-based generative models. We propose two applications: 1) learning bridges between two infinite-dimensional distributions and 2) generative models for sampling from an infinite-dimensional distribution. Our approach proves effective for diverse problems involving continuous function space representations, such as resolution-free images, time-series data, and probability density functions. Code is available at https://github.com/bw-park/DBFS .

1 Introduction

Stochastic Optimal Control (SOC) is designed to steer a noisy system toward a desired state by minimizing a specific cost function. This methodology finds extensive applications across various fields in science and engineering, including rate event simulation [33, 35], stochastic filtering and data assimilation [47, 55, 72], non-convex optimization [9], modeling population dynamics [8, 43]. SOC is also related to diffusion-based sampling methods that are predominant in machine learning literature. Specifically, if we choose the terminal cost of a control problem as the log density ratio between a target distribution and a simple prior distribution, solving the optimal control reduces to learning a diffusion-based generative models [56, 73, 77] built upon the Schr¨ odinger bridge problem [12, 54].

While SOC associated diffusion-based generative models have been well-established for finitedimensional spaces [5, 10, 11, 44], their theoretical foundations and practical algorithms for infinitedimensional spaces remain underexplored. There is a growing interest in developing generative models in function spaces. Examples include learning neural operators for partial differential equations (PDEs) [37, 39, 40], interpreting images as discretized functions through implicit neural representations (INRs) [21, 63], and operating in function spaces for Bayesian neural networks (BNNs) [65, 74]. Models that operate in function spaces are more parameter-efficient as they avoid resolution-specific parameterizations. In response to this demand, several extensions of

diffusion-based generative models for infinite-dimensional function spaces have been proposed [3, 25, 32, 41, 42, 53]. However, SOC theory for building diffusion bridges in function spaces is still demanding in the community of generative modeling.

To address these challenges, this work introduces an extension of SOC for diffusion bridges in infinite-dimensional Hilbert spaces, particularly focusing on its applications in sampling problems. Specifically, we demonstrate the idea of Doob's h -transform [13, 59] can be derived from SOC theory and extend its formulation into Hilbert spaces. Due to the absence of a density with respect to the Lebesgue measure in infinite-dimensional spaces, building a diffusion bridge between function spaces is a nontrivial task separated from the finite-dimensional cases [51, 61]. To this end, we propose a Radon-Nikodym derivative relative to a specified Gaussian reference measure in Hilbert space. Leveraging the infinite-dimensional Doob's h -transform and SOC, we then formulate diffusion bridgebased sampling algorithms in function spaces. While the infinite-dimensional Doob's h -transform has already been derived in [2, 26, 31], our main goal is not merely to derive it. Instead, we aim to generalize various finite-dimensional sampling problems [51, 61, 73, 77] into the infinite-dimensional space by exploiting the theory of infinite-dimensional SOC.

To demonstrate the applicability of our theory, we present learning algorithms for two representative problems. First, we introduce an infinite-dimensional bridge-matching algorithm as an extension of previous methods [51, 61] into Hilbert spaces, which learns a generative model to bridge two distributions defined in function spaces. As an example, we show that our framework can learn smooth transitions between two image distributions in a resolution-free manner. Second, we propose a simulation-based Bayesian inference algorithm [73, 77] that operates in function space. Instead of directly approximating the target Bayesian posterior, our algorithm learns a stochastic transition from the prior to the posterior within the function space. We demonstrate the utility of this approach by inferring Bayesian posteriors of stochastic processes, such as Gaussian processes. We summarize our contributions as follows:

Based on the SOC theory, we derive the Doob's h -transform in Hilbert spaces. We propose a h function as a Randon-Nikodym derivative with respect to a Gaussian measure in infinitedimensional space.

Based on the infinite-dimensional extension of the Doob's h -transform, we present the diffusion bridge and simulation-based Bayesian inference algorithm in function spaces.

Wedemonstrate our method for various real-world problems involving function spaces, including resolution-free image translation and posterior sampling for stochastic processes.

Notation. Consider a real and separable Hilbert space H with the norm and inner product denoted by ∥·∥ H and ⟨· , ·⟩ H . Throughout the paper, we consider a path measure denoted by P ( · ) on the space of all continuous mappings Ω = C ([0 , T ] , H ) . The stochastic processes associated with this path measure P ( · ) are denoted by X ( · ) , and their time-marginal distribution at time t ∈ [0 , T ] as push-forward measure µ ( · ) t := ( X ( · ) t ) # P ( · ) . Furthermore, for a function V : [0 , T ] ×H → R , we define D x V , D xx V as the first and second order Fr´ echet derivatives with respect to the variable x ∈ H , respectively, and ∂ t V as the derivative with respect to the time variable t ∈ [0 , T ] .

2 A Foundation on Stochastic Optimal Control for Diffusion Bridges

In this section, we first present a brief introduction to the theory of stochastic optimal control (SOC) in infinite dimensions and the Verification Theorem (Lemma 2.1), which is the key to understanding the theoretical connection between SOC and the diffusion bridges. Then we propose Doob's h -transform in Hilbert spaces based on the SOC theory (Theorems 2.2 and 2.3).

2.1 Preliminaries

Gaussian Measure and Cameron-Martin Space. Let (Ω , F , Q ) and ( H , B ( H )) be two measurable spaces and consider H -valued random variable X : Ω → H such that the push-forward measure µ := X # Q induced by X is Gaussian, i.e., a real-valued random variable ⟨ u, X ⟩ H follows Gaussian distribution on R for any u ∈ H . Then, there exist a unique mean m X ∈ H given by ⟨ m X , u ⟩ H = E µ [ ⟨ X , u ⟩ H ] and a nonnegative, symmetric, and self-adjoint covariance operator Q : H → H defined by ⟨ u, Qv ⟩ H = E µ [ ⟨ X -m X , u ⟩ H , ⟨ X -m X , v ⟩ H ] for any u, v ∈ H .

We write µ = N ( m X , Q ) and say X is centered if m X = 0 . For a covariance operator Q , we assume H -valued centered X follows the distribution N (0 , Q ) supported on H so that Q is guaranteed to be compact. Hence, there exists an eigen-system { ( λ ( k ) , ϕ ( k ) ) ∈ R ×H : k ∈ N } such that Q ( ϕ ( k ) ) = λ ( k ) ϕ ( k ) holds and Tr ( Q ) = ∑ ∞ k =1 λ ( k ) < ∞ . We define the Cameron-Martin space by H 0 := Q 1 / 2 ( H ) , H 0 ⊆ H is a separable Hilbert space endowed with the inner product ⟨ u, v ⟩ H 0 := ⟨ Q -1 / 2 u, Q -1 / 2 v ⟩ H .

Stochastic Differential Equations in Hilbert Spaces. The standard R d -valued Wiener process has independent increments w t +∆ t -w t ∼ N (0 , ∆ t I d ) . In the case of infinite-dimensional Hilbert space H 0 , however, such identity covariance may not be a trace class. We consider a larger space H 0 ⊂ H 1 such that H 0 is embedded into H 1 with a Hilbert-Schmidt embedding J : H 0 →H 1 . Let Q := JJ ∗ and we define H 1 -valued Q -Wiener process [17, Proposition 4.7] as W Q t = ∑ ∞ k =1 Q 1 / 2 ϕ ( k ) w ( k ) t . We focus on the Cameron-Martin space H 0 where the Wiener process has increments N (0 , ∆ t I H ) , and a larger space H 1 = H , where the Wiener process has increments N (0 , ∆ t Q ) . Then we define a path measure P and associated stochastic differential equation (SDE) in H as follows

where A : H → H is a linear operator, a constant σ > 0 and a Q -Wiener process W Q defined on a probability space (Ω , F , ( F t ) t ≥ 0 , P ) . For a more comprehensive understanding, see [17, 29].

2.2 Stochastic Optimal Control in Hilbert Spaces

From the uncontrolled SDE introduced in equation (1), various sampling problems in R d including density sampling [5, 73, 77] and generative modeling [10, 11] can be solved by adjusting the SDE with proper drift (control) function α ∈ R d . Motivated by these approaches, introducing stochastic optimal control (SOC) to solve real-world sampling problems, we aim to introduce SOC to the infinite-dimensional Hilbert space H . We consider that a controlled path measures P α is induced by following infinite-dimensional SDE defined as follows:

where α ( · ) : [0 , T ] ×H → H is infinite-dimensional Markov control (see Section A.1.2 for more details). We refer to the SDE in equation (2) as controlled SDE. The controlled SDE can be exploited to the various problems. In general, it can be done by finding the optimal control function that minimizes the objective functional with suitably chosen cost functionals depending on the problem:

where R : H → R are running cost and G : H → R is terminal cost. The measurable function J ( t, x , α ) , representing the total cost incurred by the control α over the interval [ t, T ] , given that the control strategy from the interval [0 , t ] has resulted in X α t = x . The objective is to minimize the objective functional in (3) over all admissible control policies α ∈ U , where U is the Hilbert space of all square-integrable H -valued processes adapted to W Q defined on [0 , T ] . Then we define the value function V ( t, x ) = inf α ∈U J ( t, x , α ) , the optimal costs conditioned on ( t, x ) ∈ [0 , T ] ×H . By using the dynamic programming [23], we can solve the Hamilton-Jacobi-Bellman (HJB) equation,

where LV t := ⟨ X α t , A D x V t ⟩ H + 1 2 Tr [ σ 2 QD xx V t ] . We demonstrate that with specific choices of cost functionals R and G in (3), the optimal control α ⋆ of the minimization problem aligns with the proper drift function for sampling problems we will discuss later. To do so, we start with how the HJB equation can characterize the optimal controls.

Lemma 2.1 (Verification Theorem) . Let V be a solution of HJB equation (4) with R ( α ) := 1 2 ∥ α ∥ 2 H satisfying the assumptions in A.1. Then, we have V ( t, x ) ≤ J ( t, x, α ) for every α ∈ U and ( t, x ) ∈ [0 , T ] ×H . Let ( α ∗ , X α ∗ ) be an admissible pair such that

for almost every s ∈ [ t, T ] and P -almost surely. Then ( α ∗ , X α ∗ ) satisfying V ( t, x ) = J ( t, x , α ∗ ) .

Lemma 2.1 demonstrate that with specific choices of running costs, the solution to the HJB equation in (4) is the optimal cost of the minimization problem in (3) with the closed-form optimal control α ∗ t = -σQ 1 / 2 D V t . In the subsequent subsection, we reveal the connection between the optimal controlled process, characterized by the controlled SDE with optimal control α ⋆ , and the conditioned SDE in H . Additionally, we show various problems depending on the choice of the terminal cost G .

2.3 Doob's h -transform for Diffusion Bridges in Hilbert Spaces

The conditioned SDE is a stochastic process that is guaranteed to satisfy a set constraint defined over the interval [0 , T ] . For instance, a Diffusion Bridge is a stochastic process that satisfies both terminal constraints X 0 = x 0 , X T = x T for any x 0 , x T ∈ H . We will show that with a proper choice of the terminal cost G , the conditioned SDE is a special case of controlled SDE. For this, let us define the function h : [0 , T ] ×H→ R :

where N e t A x ,Q t is a Gaussian measure with mean function e t A x and covariance operator Q t = ∫ t 0 e ( t -s ) A Qe ( t -s ) A ds and ˜ G := e -G for the function G in (3). The function h evaluates the future states X T using ˜ G which propagated by (1) for a given initial x at time t ∈ [0 , T ] . It can be shown that the function h satisfies the Kolmogorov-backward equation [17] with terminal condition,

Now, we employ the Hopf-Cole transformation [24] to establish an inherent connection between two classes of PDEs, the linear PDE in (7) and the HJB equation in (4), which provide us a key insight to deriving the Doob's h -transform in function spaces utilizing the SOC theory.

Theorem 2.2 (Hopf-Cole Transform) . Let V t = -log h t . Then V t satisfies the HJB equation:

According to Theorem 2.2, the solution of linear PDE in equation (7) is negative exponential to the solution of the HJB equation in (8). Given that it already verified that the optimal control α ⋆ results the value function V t , which has explicit form as described in (8), we find that the relationship of two PDEs through V t = -log h t leads to a distinct form of optimal control α ⋆ = -σQ 1 / 2 D x V = σQ 1 / 2 D x log h , where D x log h := D x h/h . Consequently, it yields another class of SDE as follows:

where ˆ W Q t is a Q -Wiener process on P h . This representation is consistent with infinite-dimensional conditional SDE [2, 26] which induces an expansion of Doob's h -transform [59] in Hilbert space H .

The Doob's h -transform in finite-dimensional spaces is well-established to construct the diffusion bridge process under the assumption that h ( t, x ) = P ( X T ∈ d x T | X t = x ) has an explicit RadonNikodym density function [46, 60], enable to simulate the bridge SDEs. In contrast, although the choice of ˜ G ( x ) = 1 d x T ( x T ) in (6) yields same representation of h ( t, x ) = P ( X T ∈ d x T | X t = x ) , we cannot easily define the Radon-Nikodym density in H due to the absence of an equivalent form of Lesbesgue measure which hinder the computation of D x log h ( t, X h t ) in (9) explicitly. Hence, to define the diffusion bridge process in H , we need to identify an explicit density form of h ( t, x ) = P ( X T ∈ d x T | X t = x ) . The following theorem reveals the explicit form of h function and becomes a key ingredient in deriving infinite dimensional diffusion bridge processes.

Theorem 2.3 (Explicit Representation of h ) . For any t > 0 and any x ∈ H , the measure N e t A x ,Q t and N 0 ,Q ∞ are equivalent, where N 0 ,Q ∞ is an invariant measure of P in (1) as t → ∞ where Q ∞ = -1 2 Q A -1 . Moreover, for any x , y ∈ H , the Radon-Nikodym density d N e t A x ,Q t d N 0 ,Q ∞ ( · ) = q t ( x , · ) is given by

where Θ t = Q 1 / 2 ∞ ( Q -1 / 2 t e t A ) ∗ ( Q -1 / 2 ∞ Q 1 / 2 t ) ∗ ( Q 1 / 2 ∞ ( Q -1 / 2 t e t A ) ∗ ( Q -1 / 2 ∞ Q 1 / 2 t ) ∗ ) ∗ , t ≥ 0 .

Theorem 2.3 states that the marginal distribution of certain classes of SDEs described in (1) has an explicit Radon-Nikodym density with respect to their invariant measure. Therefore, by the time-homogeneity of the process in (1), it allows us to define the h function explicitly:

This framework enables the construction of an infinite-dimensional bridge process, by selecting ˜ G in (12) properly. Below, we demonstrate the infinite-dimensional diffusion bridge.

Example 2.4 (Diffusion Bridge in H ) . Let { ( λ ( k ) , ϕ ( k ) ) ∈ R ×H : k ∈ N } be an eigen-system of H . Then for each k ∈ N , the SDE system in equation (1) can be represented as:

where A ϕ ( k ) = -a k ϕ ( k ) , Qϕ ( k ) = λ ( k ) ϕ ( k ) , X ( k ) t = ⟨ X t , ϕ ( k ) ⟩ H and W ( k ) t = ⟨ W t , ϕ ( k ) ⟩ H . Then, for any x T ∈ H , the conditional law of x ( k ) T given X ( k ) t is a Gaussian N ( m ( k ) T | t X ( k ) t , Σ ( k ) T | t ) with

Now, by setting the terminal condition in (6) as ˜ G ( x ) := 1 x T ( x ) (i . e ., Dirac delta of x T ) then h ( t, x ) = q T -t ( x , x T ) . Thus, for each coordinate k , we get following representation:

with two end points conditions X ( k ) 0 = x ( k ) 0 and X ( k ) T = x ( k ) T .

2.4 Approximating path measures

Since the function h is intractable for a general terminal cost ˜ G in (12), simulating the conditioned SDEs in (9) requires some approximation techniques. As observed in Theorem 2.2, finding the function h is equal to learning the control function α such that P α is equal to P ⋆ := P α ⋆ = P h . Therefore, with a parametrization α := α ( · , θ ) the approximation can be done by neural network parameterization i . e ., α ⋆ ≈ α θ ⋆ with local minimum θ ⋆ = arg min θ D ( P α || P ⋆ ) , where D ( P α || P ⋆ ) is a divergence between P α and P ⋆ . For example, the cost functional described in equation (3) can be represented as relative-entropy loss D rel ( P α || P ⋆ ) = E P α [ log d P α d P ⋆ ] 2 . Therefore, the training loss for θ can be estimated by first simulating the parameterized control path and then calculating (3) for a specified cost functional R,G . Moreover, if we can access to the P ⋆ , one can define the variational optimization [69] where the loss is defined as cross-entropy loss D cross ( P α || P ⋆ ) = E P ⋆ [ log d P ⋆ d P α ] . See [20, 48] for more details about the approximation technique and other loss functions.

3 Simulating Diffusion Bridges in Infinite Dimensional Spaces

Leveraging the SOC theory within H , we show how our approach generalizes existing diffusionbased sampling methods. Specifically, incorporating the relation between controlled SDEs (2) and conditioned SDEs (9), we introduce two learning algorithms that allow us to simulate various diffusion bridge-based sampling algorithms.

3.1 Infinite Dimensional Bridge Matching

In this section, our objective is to learn a control α that yields P α such that { X α t } t ∈ [0 ,T ] satisfies µ α t ≈ µ ⋆ t for all pre-specified µ ⋆ t over the interval t ∈ [0 , T ] . Specifically, we assume that the end-point marginals µ ⋆ 0 and µ ⋆ T follow the laws of two data distributions π 0 and π T , respectively, and the intermediate marginals { µ ⋆ t } t ∈ (0 ,T ) are defined as a mixture of diffusion bridge paths. This

learning problem is referred to as the Bridge Matching (BM) algorithm [45, 50, 61] and can be expressed as a solution of the SOC problem structured as

In (16), various divergences can be chosen for the same learning problem, as discussed in Section 2.4. Here, we will choose the cross-entropy because the relative entropy requires the appropriate selection of the terminal cost G in (3), which is intractable since we do not have access to the distributional form of π 0 , π T [44]. Furthermore, keeping the entire computational graph of P α with parameterized α can become resource-intensive, especially for higher-dimensional datasets like images [11].

Now, we specify the optimal path measure P ⋆ for a problem in (16). Let P | 0 ,T be a path measure induced by (15) and µ t | 0 ,T be a marginal distribution of P | 0 ,T . Moreover, let P ⋆ = P | 0 ,T Π 0 ,T for an independent coupling Π 0 ,T = π 0 ⊗ π T . Then the optimal path measure P ⋆ is defined as Mixture of bridge . Under regular assumptions, the optimal control α ⋆ that induces the optimal path measure P ⋆ can be constructed as a mixture of functions h in (9) by choosing G ( x ) = 1 x T ( x ) .

Theorem 3.1 (Mixture of Bridges in H ) . Let us consider a marginal distribution of P ⋆ at t ∈ [0 , T ] , µ ⋆ t ( d x t ) = ∫ µ t | 0 ,T ( d x t )Π 0 ,T ( d x 0 , d x T ) has density p ⋆ t with respect to some Gaussian reference measure µ ref i . e ., µ ⋆ t ( d x t ) /µ ref ( d x t ) = p ⋆ t ( x t ) . Then the optimal path measure P ⋆ associated with:

Objective Functional for Bridge Matching. With the structure of P ⋆ specified in (17) we can estimate the divergence between the optimal target path measure P ⋆ and a path measure P α . To accomplish this, we first define

Then, by applying the Girsanov theorem 3 which provides us the Radon-Nikodym derivative between P ⋆ and P α , we can derive the cross-entropy loss in equation (16):

Then, under the neural network parameterization of control function α θ , we can reformulate the SOC problem in (16) as a learning problem with the training loss function represented by:

The L BM ( θ ) in (20) yields the infinite-dimensional BM summarized in Alg 1.

3.2 Bayesian Learning in Function Space

In the previous section, we observed that by appropriately defining a terminal cost functional G in (3), the SOC problem aligns with the sampling problem, where optimal control effectively steers

the distribution from π 0 to the target distribution π T , where we can access samples from π 0 and π T . However, accessing samples from unknown target distribution π T is generally not feasible. For instance, for a π T , a posterior distribution over function. In this case, although direct samples from π T are unattainable, its distributional representation is given as [3, 66]:

where U is a energy function. Here, our primary objective is to sample from a distribution over function π T := µ ⋆ T by simulating the controlled diffusion process { X α t } t ∈ [0 ,T ] over finite horizon [0 , T ] with T < ∞ . It can be represented as a solution of the following SOC problem:

The following theorem implies that with a suitable terminal cost functional G in (3), it is possible to achieve X α ⋆ T ∼ π T as an expansion of [54, 71] for infinite dimensional space H .

Theorem 3.2 (Exact sampling in H ) . Consider that the initial distribution µ 0 is given as the Dirac measure δ x 0 for some x 0 ∈ H and the following objective functional

where µ T = N ( e T A x 0 , Q T ) as a marginal distribution of X T in (1) with a well-defined terminal cost dπ T dµ T by Theorem 2.3. Then, X α ⋆ T ∼ π T .

Objective Functional for Bayesian Learning. Unlike problem in (16) where we can access to the target path measure P ⋆ directly, it is not feasible here because h ( t, x ) = E P [ -dπ T dµ T ( X T ) | X t = x ] does not have an explicit solution. Therefore, we will use the relative-entropy loss as our training loss function:

The key difference from previous algorithms [73, 77] is that the Radon-Nikodym derivative dπ T dµ T ( X α T ) may not be well-defined on H due to the absence of the Lesbesgue measure. However, the Theorem 2.3 suggests that by choosing certain classes of Gaussian measure i . e ., µ prior := N ( e t A x , Q t ) , the Radon-Nikodym density of dπ T dµ prior and dµ T dµ prior has explicit form since µ prior and N (0 , Q ∞ ) are equivalent. Thus, using the chain rule, the terminal cost for any x ∈ H can be computed as follows:

With L Bayes ( θ ) in (24), the infinite-dimensional bayesian learning algorithm is summarized in Alg 2.

4 Related Work

Most diffusion models operate within the framework of time-reversal [1], where the generation process is learned from its corresponding time-reversed SDEs [64]. In contrast, diffusion models based on conditioned SDEs, such as diffusion bridges, built upon the theory of Doob's h -transform, offer a conceptually simpler approach as they solely rely on a forward process. [50] proposes generative models with this concept, showing that the mixture of forward diffusion bridge processes effectively transports between couplings of two distributions. [76] introduces a family of first hitting diffusion models that generate data with a forward diffusion process at a random first hitting time based on Doob's h -transform. Combining time-reversal with the h -transform, [46] proposes a diffusion bridge process on constrained domains. Moreover, [51, 61] presented that the Schr¨ odinger bridge problem can be solved by an iterative algorithm, which is improved by [18] to enhance efficiency. Furthermore, [44] generalizes the Schr¨ odinger bridge matching algorithm by introducing an approximation scheme with a non-trivial running cost. Compared to prior works, which primarily focus on finite-dimensional spaces, our work extends the formulation of Doob's h -transform into Hilbert space, enabling the development of various sampling algorithms in function spaces.

5 Experiments

This section details the experimental setup and the application of the proposed Diffusion Bridges in Function Spaces (DBFS) for generating functional data. We interpret the data from a functional perspective, known as field representation [75, 79], where data are seen as a finite collection of function evaluations { Y [ p i ] , p i } N i . Here, a function Y maps points p i from a coordinate space X to a signal space Y , i.e., Y : X → Y . Additional experimental details are provided in Appendix A.8.

5.1 Bridge Matching

First, we present empirical results for the infinite-dimensional BM algorithm discussed in Sec 3.1, applied to 1D and 2D data. For 1D data, we consider X = R and Y = R . For 2D data, we assume X = R 2 and Y = R for probability density or grayscale images, and Y = R 3 for RGB images.

Bridging Field. We begin by validating our bridge matching Algorithm in Alg 1 on bridging probability density function within H . Specifically, we set π 0 := δ p 0 with a ring-shaped density function p 0 and π T := δ p T characterized by a Gaussian mixture density function p T . The functions map each grid points p i to the probability in Y = R . Therefore, both density functions can be represented as their field representations { p 0 [ p i ] , p i } N i , { p T [ p i ] , p i } N i , respectively. Figure 1 illustrates the progressive propagation of the target optimal bridge process P ⋆ from p 0 to p T . Despite the α ⋆ is trained on

X

0

=

p

0

X

t

∼

µ

t

|

0

,T

X

T

=

p

T

o

the functions generated from P ⋆ which are evaluated on a coarse grid { p i } 32 2 i , P α ⋆ is capable of producing accurate functional evaluations on a finer grid { p i } 256 2 i . This resolution-invariance property indicates that our method is adept at learning continuous functional representations, rather than merely memorizing the discrete evaluations.

10

5

0

5

10

100

50

0

50

100

Data

10

5

0

5

10

100

50

0

50

100

Samples

(a) Quadratic

10

5

0

5

10

0

2

4

Data

10

5

0

5

10

0

2

4

Samples

(b) Melbourne

10

5

0

5

10

2

1

0

1

2

Data

10

5

0

5

10

2

1

0

1

2

Samples

(c) Gridwatch

1D function generation. We conducted an experiment on a 1D function generation task, comparing our baseline methods [22, 52] on three datasets: Quadratic, Melbourne, and Gridwatch, following the setup from [52]. For gen-

erative modeling, we set the initial distribution π 0 as N (0 , Q ) with RBF kernel for the covariance operator Q and the terminal distribution as data distribution π T , respectively. We employing the bridge matching algorithm in Alg 1. For quantitative evaluation, we used the power of a kernel two-sample hypothesis test to distinguish between generated and ground-truth samples. Table 1 shows that our method performs comparably to baseline infinite-dimensional methods. Additionaly, The generated samples compared to the ground-truth for each dataset are provided in Figure 2.

Unpaired Image Transfer. We compare our proposed model with a finite (fixed)-dimensional baseline through an experiment on unpaired image transfer between the MNIST and EMNIST datasets at 32 2 resolution, as well as wild and cat images from the AFHQ dataset [14], downsampled to 64 2 resolution (AFHQ-64). Specifically, we evaluate the performance of [51, 61] and our DBFS model. For a fair comparison, we follow the iterative training scheme of [51] based on

(a)

32

2

(observed resolution)

(b)

64

2

(unseen resolution)

(c)

128

2

(unseen resolution)

(d)

64

2

(observed resolution)

(e)

128

2

(unseen resolution)

A

a

a

B

B

B

b

a

h

a

B

A

D

d

七

b

4

3

3

6

H

7

3

9

C

P

0

7

C

b

d

D

a

B

C

q

7

1

3

6

2

廿

A

a

b

4

b

d

d

0

a

A

A

d

4

7

3

u

4

C

B

B

B

C

A

A

a

P

b

2

3

b

8

7

4

4

4

7

7

9

S

9

d

b

一

7

b

B

4

3

B

C

p

2

6

の

7

。

B

#4

B

44

3

a

A

B

4

3

the public repository 4 , where two forward and backward control networks are trained alternately. For quantitative evaluation, we estimate the FID score between the generated samples and real datasets. We set σ = 1 for both [51] and our method, while FID scores for [61] are taken from [18]. Table 2: Test FID on unpaired image transfer task. (A) EMNIST → MNIST, (B) AFHQ-64 Wild → Cat. Method (A) (B) IDBM [51] 8.2 -DSBM † [61] 6.0 25.4 DBFS (Ours) 9.1 44.4 † result from [18]. Table 2 shows that our method performs comparably to the finitedimensional method. Additionally, we provide generated samples at various unseen resolutions in Figure 3 to demonstrate the resolution-invariant property of our infinite-dimensional models. We note that our method may have slightly lower FID scores compared to finite-dimensional baselines, which may align with the observation in [79] that resolution-agnostic methods tend to have lower FID scores compared to resolution-specific ones. This could be because resolution-specific methods can incorporate domain-specific design features in their score networks. Samples generated from the reverse direction can be found in Figure A.2.

5.2 Bayesian Learning

We validate our Bayesian learning algorithm for modeling functional data. Specifically, we will consider the temporal data as a function. We denote Y [ O ] = { Y [ p i ] } | O | i =1 as a collection of a function evaluation on a set of 1-dimensional observation grid O = { p i } | O | i where 0 ≤ p 0 < · · · < p | O | ≤ I . We assume that each observed time series approximates a corresponding underlying continuous function X : R → R d as the number of observations increases i . e ., { Y [ p i ] } | O |→∞ i =1 ≈ X . For given observations Y [ O ] , our goal is to infer the posterior distribution on some set of unobserved grid T = [0 , I ] -O i . e ., P ( Y [ T ] | Y [ O ]) and therefore modeling distribution over X on [0 , I ] . Please refer to Section A.8.2 for further details.

Functional Regression To verify the effectiveness of the proposed DBFS in generating functions in the 1D domain, we conducted regression experiments using synthetic data generated from the Gaussian Process (GP) by following the experimental settings in [38]. Figure 4 shows the sampled trajectories of a controlled dynamics X α t for t ∈ [0 , T 2 , T ] trained on data generated from GP with RBF covariance kernel. The stochastic process begins from the deterministic function X α 0 = x 0 at t = 0 and propagates towards the conditional posterior distribution X α T ∼ P ( Y [ T ] | Y [ O ]) at t = T .

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

2.0

3.0

2.5

2.0

1.5

1.0

0.5

0.0

0.5

1.0

Missing

Observed

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

2.0

3.0

2.5

2.0

1.5

1.0

0.5

0.0

0.5

1.0

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

2.0

3.0

2.5

2.0

1.5

1.0

0.5

0.0

0.5

1.0

0

10

20

30

40

0.15

0.10

0.05

0.00

0.05

0.10

0.15

Missing

Observed

0

10

20

30

40

0.15

0.10

0.05

0.00

0.05

0.10

0.15

0

10

20

30

40

0.15

0.10

0.05

0.00

0.05

0.10

0.15

Imputation. We evaluate our method against recent diffusion-based imputation method where the goal is to infer the conditional distribution p ( Y [ T ] | Y [ O ]) of unobserved grid Y [ T ] give observations Y [ O ] . CSDI [68] utilizes DDPM[34] to learn the reverse process by treating the temporal data Y [ O ] as a R | O |× d dimensional feature. Extending this, DSDP-GP [6]

† results from [6].

enhances CSDI by incorporating noise derived from a stochastic process, instead of simple Gaussian noise. We maintained the same training setup as these models, including random seeds and the model architecture for control α θ in (2). Consistent with their methodology, we employed the Physionet dataset [30], which comprises medical time-series data collected on an hourly rate. Since the dataset inherently contains missing values, we selected certain degrees of observed values to create an imputation test set for evaluation. We then reported the results on this test set, varying the degrees of missingness. Table 3 shows that we outperform the previous methods even though it solely relies on forward propagation of controlled SDEs in (9) without denoising procedure.

6 Conclusion and Limitation

In this work, we shed light on the application of the infinite-dimensional Doob's h -transform, exploiting SOC theory in infinite-dimensional spaces. By developing an explicit Radon-Nikodym density, we address the challenge posed by the absence of an equivalent to the Lebesgue measure. With specified cost functions for control objectives, it enables us to extend previous algorithm based on the finite-dimensional Doob's h -transform into infinite-dimensional function spaces, such as resolution-free unpaired image transfer and functional Bayesian posterior sampling.

Compared to the recent infinite-dimensional score-based diffusion model [42], our work restricts the coefficients for the stochastic dynamics to be time-independent. This limitation prevents us from defining a noise schedule for the diffusion model [78], which may hinder performance improvements. Additionally, in Bayesian learning, computing the gradient of the proposed training loss function (24) can be computationally demanding. Thus, developing a more scalable algorithm would be an interesting direction for future work. Furthermore, as our model can be applied to any functional domain, we have limited our experiments to regular 1D and 2D domains, leaving the extension to more general domains for future work.

Broader Societal Impact

Similar to other works in the literature, our proposed method holds the potential for both beneficial outcomes, such as automated data synthesis, and adverse implications, such as the deep fakes, depending on how it is used. We adhere to ethical standards for using our model in generative AI.

References

Brian D.O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

Elizabeth Louise Baker, Gefan Yang, Michael L Severinsen, Christy Anna Hipsley, and Stefan Sommer. Conditioning non-linear and infinite-dimensional diffusion processes. arXiv preprint arXiv:2402.01434 , 2024.

Lorenzo Baldassari, Ali Siahkoohi, Josselin Garnier, Knut Solna, and Maarten V de Hoop. Conditional score-based diffusion models for bayesian inference in infinite dimensions. Advances in Neural Information Processing Systems , 36, 2024.

Y.I. Belopolskaya and Y.L. Dalecky. Stochastic Equations and Differential Geometry . Mathematics and its Applications. Springer Netherlands, 2012.

Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusionbased generative modeling. Transactions on Machine Learning Research , 2024.

Marin Biloˇ s, Kashif Rasul, Anderson Schneider, Yuriy Nevmyvaka, and Stephan G¨ unnemann. Modeling temporal data as continuous functions with stochastic process diffusion. In International Conference on Machine Learning , pages 2452-2470. PMLR, 2023.

Vladimir Bogachev, Giuseppe Da Prato, and Michael R¨ ockner. Uniqueness for solutions of fokker-planck equations on infinite dimensional spaces. Communications in Partial Differential Equations , 36(6):925-939, 2011.

R. Carmona and F. Delarue. Probabilistic Theory of Mean Field Games with Applications I: Mean Field FBSDEs, Control, and Games . Probability Theory and Stochastic Modelling. Springer International Publishing, 2018.

Pratik Chaudhari, Adam Oberman, Stanley Osher, Stefano Soatto, and Guillaume Carlier. Deep relaxation: partial differential equations for optimizing deep neural networks. Research in the Mathematical Sciences , 5:1-30, 2018.

Tianrong Chen, Jiatao Gu, Laurent Dinh, Evangelos Theodorou, Joshua M. Susskind, and Shuangfei Zhai. Generative modeling with phase stochastic bridge. In The Twelfth International Conference on Learning Representations , 2024.

Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood training of schr¨ odinger bridge using forward-backward SDEs theory. In International Conference on Learning Representations , 2022.

Yongxin Chen, Tryphon T. Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schr¨ odinger bridge. SIAM Review , 63(2):249-313, 2021.

Rapha¨ el Chetrite and Hugo Touchette. Nonequilibrium markov processes conditioned on large deviations. In Annales Henri Poincar´ e , volume 16, pages 2005-2057. Springer, 2015.

Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2020.

G. Da Prato and J. Zabczyk. Differentiability of the feynman-kac semigroup and a control application. Atti della Accademia Nazionale dei Lincei. Classe di Scienze Fisiche, Matematiche e Naturali. Rendiconti Lincei. Matematica e Applicazioni , 8(3):183-188, 10 1997.

G. Da Prato and J. Zabczyk. Second Order Partial Differential Equations in Hilbert Spaces . London Mathematical Society Lecture Note Series. Cambridge University Press, 2002.

G. Da Prato and J. Zabczyk. Stochastic Equations in Infinite Dimensions . Encyclopedia of Mathematics and its Applications. Cambridge University Press, 2014.

Valentin De Bortoli, Iryna Korshunova, Andriy Mnih, and Arnaud Doucet. Schr \ ' odinger bridge flow for unpaired data translation. arXiv preprint arXiv:2409.09347 , 2024.

Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schr¨ odinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709, 2021.

Carles Domingo-Enrich, Jiequn Han, Brandon Amos, Joan Bruna, and Ricky TQ Chen. Stochastic optimal control matching. arXiv preprint arXiv:2312.02027 , 2023.

Emilien Dupont, Adam Goli´ nski, Milad Alizadeh, Yee Whye Teh, and Arnaud Doucet. Coin: Compression with implicit neural representations. arXiv preprint arXiv:2103.03123 , 2021.

Vincent Dutordoir, Alan Saul, Zoubin Ghahramani, and Fergus Simpson. Neural diffusion processes. In International Conference on Machine Learning , pages 8990-9012. PMLR, 2023.

Giorgio Fabbri, Fausto Gozzi, and Andrzej Swiech. Stochastic optimal control in infinite dimension. Probability and Stochastic Modelling. Springer , 2017.

Wendell H Fleming and Halil Mete Soner. Controlled Markov processes and viscosity solutions , volume 25. Springer Science & Business Media, 2006.

Giulio Franzese, Giulio Corallo, Simone Rossi, Markus Heinonen, Maurizio Filippone, and Pietro Michiardi. Continuous-time functional diffusion processes. Advances in Neural Information Processing Systems , 36, 2024.

Marco Fuhrman. A class of stochastic optimal control problems in hilbert spaces: Bsdes and optimal control laws, state constraints, conditioned processes. Stochastic processes and their applications , 108(2):263-298, 2003.

Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and S. M. Ali Eslami. Conditional neural processes. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 1704-1713. PMLR, 10-15 Jul 2018.

Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J. Rezende, S. M. Ali Eslami, and Yee Whye Teh. Neural processes, 2018.

L. Gawarecki and V. Mandrekar. Stochastic Differential Equations in Infinite Dimensions: with Applications to Stochastic Partial Differential Equations . Probability and Its Applications. Springer Berlin Heidelberg, 2010.

A. L. Goldberger, L. A. N. Amaral, L. Glass, J. M. Hausdorff, P. Ch. Ivanov, R. G. Mark, J. E. Mietus, G. B. Moody, C.-K. Peng, and H. E. Stanley. PhysioBank, PhysioToolkit, and PhysioNet: Components of a new research resource for complex physiologic signals. Circulation , 101(23):e215-e220, 2000 (June 13). Circulation Electronic Pages: http://circ.ahajournals.org/content/101/23/e215.full PMID:1085218; doi: 10.1161/01.CIR.101.23.e215.

Ben Goldys and Bohdan Maslowski. The ornstein-uhlenbeck bridge and applications to markov semigroups. Stochastic processes and their applications , 118(10):1738-1767, 2008.

Paul Hagemann, Sophie Mildenberger, Lars Ruthotto, Gabriele Steidl, and Nicole Tianjiao Yang. Multilevel diffusion: Infinite dimensional score-based diffusion models for image generation. arXiv preprint arXiv:2303.04772 , 2023.

Carsten Hartmann, Omar Kebiri, Lara Neureither, and Lorenz Richter. Variational approach to rare event simulation using least-squares regression. Chaos: An Interdisciplinary Journal of Nonlinear Science , 29(6), 2019.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Lars Holdijk, Yuanqi Du, Priyank Jaini, Ferry Hooft, Bernd Ensing, and Max Welling. Path integral stochastic optimal control for sampling transition paths. In ICML 2022 2nd AI for Science Workshop , 2022.

Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier J Henaff, Matthew Botvinick, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver IO: A general architecture for structured inputs & outputs. In International Conference on Learning Representations , 2022.

Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Research , 24(89):1-97, 2023.

Juho Lee, Yoonho Lee, Jungtaek Kim, Eunho Yang, Sung Ju Hwang, and Yee Whye Teh. Bootstrapping neural processes. Advances in neural information processing systems , 33:66066615, 2020.

Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. Transactions on Machine Learning Research , 2023.

Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations , 2021.

Jae Hyun Lim, Nikola B Kovachki, Ricardo Baptista, Christopher Beckham, Kamyar Azizzadenesheli, Jean Kossaifi, Vikram Voleti, Jiaming Song, Karsten Kreis, Jan Kautz, et al. Score-based diffusion models in function space. arXiv preprint arXiv:2302.07400 , 2023.

Sungbin Lim, Eunbi Yoon, Taehyun Byun, Taewon Kang, Seungwoo Kim, Kyungjae Lee, and Sungjoon Choi. Score-based generative modeling through stochastic evolution equations in hilbert spaces. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou. Deep generalized schr¨ odinger bridge. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky T. Q. Chen. Generalized schr¨ odinger bridge matching. In The Twelfth International Conference on Learning Representations , 2024.

Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. I 2 sb: Image-to-image schr¨ odinger bridge. arXiv preprint arXiv:2302.05872 , 2023.

Xingchao Liu, Lemeng Wu, Mao Ye, and qiang liu. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations , 2023.

S.K. Mitter. Filtering and stochastic control: a historical perspective. IEEE Control Systems Magazine , 16(3):67-76, 1996.

Nikolas N¨ usken and Lorenz Richter. Solving high-dimensional hamilton-jacobi-bellman pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications , 2:1-48, 2021.

William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 4195-4205, 2023.

Stefano Peluchetti. Non-denoising forward-time diffusions, 2022.

Stefano Peluchetti. Diffusion bridge mixture transports, schr¨ odinger bridge problems and generative modeling. Journal of Machine Learning Research , 24(374):1-51, 2023.

Angus Phillips, Thomas Seror, Michael Hutchinson, Valentin De Bortoli, Arnaud Doucet, and Emile Mathieu. Spectral diffusion processes. arXiv preprint arXiv:2209.14125 , 2022.

Jakiw Pidstrigach, Youssef Marzouk, Sebastian Reich, and Sven Wang. Infinite-dimensional diffusion models for function spaces. arXiv e-prints , pages arXiv-2302, 2023.

Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied Mathematics and Optimization , 23:313-329, 1991.

Sebastian Reich. Data assimilation: the schr¨ odinger perspective. Acta Numerica , 28:635-711, 2019.

Lorenz Richter. Solving high-dimensional PDEs, approximation of path space measures and importance sampling of diffusions . PhD thesis, BTU Cottbus-Senftenberg, 2021.

Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations , 2024.

Severi Rissanen, Markus Heinonen, and Arno Solin. Generative modelling with inverse heat dissipation. In The Eleventh International Conference on Learning Representations , 2023.

L Chris G Rogers and David Williams. Diffusions, Markov processes and martingales: Volume 2, Itˆ o calculus , volume 2. Cambridge university press, 2000.

S. S¨ arkk¨ a and A. Solin. Applied Stochastic Differential Equations . Institute of Mathematical Statistics Textbooks. Cambridge University Press, 2019.

Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schr¨ odinger bridge matching. Advances in Neural Information Processing Systems , 36, 2024.

Isabel Sim˜ ao. Regular transition densities for infinite dimensional diffusions. Stochastic Analysis and Applications , 11(3):309-336, 1993.

Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Implicit neural representations with periodic activation functions. Advances in neural information processing systems , 33:7462-7473, 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2021.

Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. FUNCTIONAL VARIATIONAL BAYESIAN NEURAL NETWORKS. In International Conference on Learning Representations , 2019.

Taiji Suzuki. Generalization bound of globally optimal non-convex neural network training: Transportation map estimation by infinite dimensional langevin dynamics. Advances in Neural Information Processing Systems , 33:19224-19237, 2020.

Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 7537-7547. Curran Associates, Inc., 2020.

Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. Csdi: Conditional score-based diffusion models for probabilistic time series imputation. In Advances in Neural Information Processing Systems , 2021.

Evangelos A. Theodorou, George I. Boutselis, and Kaivalya Bakshi. Linearly solvable stochastic optimal control for infinite-dimensional systems. In 2018 IEEE Conference on Decision and Control (CDC) , pages 4110-4116, 2018.

Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research , 2024. Expert Certification.

Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In COLT , 2019.

Ramon Van Handel. Stochastic calculus, filtering, and stochastic control. Course notes., URL http://www. princeton. edu/rvan/acm217/ACM217. pdf , 14, 2007.

Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil D Lawrence, and Nikolas N¨ usken. Bayesian learning via neural schr¨ odinger-f¨ ollmer flows. Statistics and Computing , 33(1):3, 2023.

Ziyu Wang, Tongzheng Ren, Jun Zhu, and Bo Zhang. Function space particle optimization for Bayesian neural networks. In International Conference on Learning Representations , 2019.

Yiheng Xie, Towaki Takikawa, Shunsuke Saito, Or Litany, Shiqin Yan, Numair Khan, Federico Tombari, James Tompkin, Vincent Sitzmann, and Srinath Sridhar. Neural fields in visual computing and beyond. Computer Graphics Forum , 41(2):641-676, 2022.

Mao Ye, Lemeng Wu, and qiang liu. First hitting diffusion models for generating manifold, graph and categorical data. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations , 2022.

Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. Denoising diffusion bridge models. In The Twelfth International Conference on Learning Representations , 2024.

Peiye Zhuang, Samira Abnar, Jiatao Gu, Alex Schwing, Joshua M. Susskind, and Miguel ´ Angel Bautista. Diffusion probabilistic fields. In The Eleventh International Conference on Learning Representations , 2023.

A Appendix

A.1 Verification Theorem and Markov control

For the derivation, we will use the following assumption

Assumption A.1. The function V : [0 , T ] × H → R and its derivatives D x V , D xx V , ∂ t V are uniformly continuous on bounded subsets of [0 , T ] ×H and (0 , T ) ×H , respectively. Moreover, for all ( t, x ) ∈ (0 , T ) ×H , there exists C 1 , C 2 > 0 such that

The HJB equation (4) can be derived by the following theorem.

Theorem A.2. Let assumptions A.1 hold and let the function V with V ( T, x ) = G ( x ) satisfying the dynamic programming principle for every 0 < t < t ′ < T , x ∈ H

Then V is a solution of the following equation:

Proof. The proof can be found in [23, Theorem 2.34].

A.1.1 Proof of Lemma 2.1

Proof. To begin the proof, we can formally compute the minimum of F ( D x V ) . Since [15]

Therefore, F ( σQ 1 / 2 D x V ) takes the infimum at α ∗ = -σQ 1 / 2 D x V . Next, applying Itˆ o's formula [23, Proposition 1.165] to V and taking expectation on both sides, we get

where we denote E t, x P α [ · ] = E P α [ ·| X t = x ] By incorporating the fact that V satisfies the equation in (A.3), we can derive the following by adding E t, x P α [ ∫ T t 1 2 ∥ α s ∥ 2 ds ] to both terms. The LHS of the equation (A.5) becomes:

And for the RHS of the equation (A.5):

where the last equation can be derived by adding and subtracting E t, x P α [ ∫ T t F ( D x V ( s, X α s ) ds ] and incorporating the fact that V satisfies the equation in (A.3) again. Hence, we get the following equation:

Since, by definition

Therefore, by taking the infimum over α ∈ U in the RHS of (A.8),

Moreover, since we already verified that the F ( σQ 1 / 2 D x V ) has infimum at α ∗ = -σQ 1 / 2 D x V . Therefore, by choosing u = α ∗ , we have

Thus, we get

Therefore, together with (A.8), this implies that ( α ∗ , X α ∗ ) is optimal at ( t, x ) ∈ [0 , T ] ×H This concludes the proof.

A.1.2 Markov Control Formulation

Now, we introduce the following corollary that states the Markov control formulation.

Corollary A.3 (Markov Control [23]) . Let us consider the measurable function ϕ t : ( t, T ) ×H→U which admit a mild solution X ϕ t of the following closed-loop equation:

Then the pair ( α ϕ t , X ϕ t ) , where the control α ϕ t is defined by the Markov feedback law α ϕ t s = ϕ ( s, X ϕ t s ) is admissible and it is optimal at ( t, x ) for all s ∈ [ t, T ] .

Therefore, in the context of the initial value problem, such as in our case, we consider the form of the Markov control α s := α ϕ 0 s = ϕ ( s, X ϕ 0 s ) for s ∈ [0 , T ] . The proof and details can be found in [23, Chap 2.5.1].

A.2 Proof of Theorem 2.2

Proof. Let us consider the function V ( t, v ) = -log h ( t, v ) .

Recall that h satisfy the KBE in equation (7):

Since h = e -V , hence ∂ t h = -L h = ∂ t e -V = -∂ t V h . Then,

We can simplify the last equation as

Following [69], the second term of RHS can be derived as follows

Therefore, combining the above results, we have

Since (A.28) coincides with (4) with ψ ( · ) := 1 ∥·∥ 2 , this concludes the proof.

A.3 Proof of Theorem 2.3

Proof. A proof of Theorem 2.3 is based on [16, Chap. 10.3]. Since Q ∞ = -1 2 Q A -1 is a trace class, we define a trace class operator Θ t as follows:

for all t ≥ 0 . Since Q t = Q ∞ -e t A Q ∞ e t A ∗ , we can rewrite Q t in terms of Θ t ,

Thus, we have (1 -Θ t ) x = Q -1 / 2 ∞ Q t Q -1 / 2 ∞ x for all x ∈ H 0 . It implies that ⟨ (1 -Θ t ) x , x ⟩ H 0 ≥ 0 , the non-negativity of (1 -Θ t ) . Moreover it also implies that (1 -Θ t ) -1 is invertible:

Consequently, it yields the following formula [16, Proposition. 1.3.11]

Now, for the general case, by using the chain rule, we have:

2 H

and utilizing Cameron-Martin theorem [16, Theorem. 1.3.6], we get:

where ( i ) follows from (A.33), (1 -Θ t ) -1 Q -1 / 2 ∞ = ( Q -1 / 2 t Q 1 / 2 ∞ ) ∗ Q -1 / 2 t = Q 1 / 2 ∞ Q -1 t . Thus, by substituting (A.34) and (A.39) into (A.35), we obtain the following result:

It concludes the proof.

A.4 Derivation of Example 2.4

For a diffusion bridge process, let us define the h function as:

If we choose ˜ G ( x , y ) = 1 d y ( x ) . Then, for any y ∈ H and t ∈ [0 , T ] , Theorem 2.3 implies that

Moreover, with an eigen-system of H , { ( λ ( k ) , ϕ ( k ) ) ∈ R × H : k ∈ N } , q t can be represented as [62]:

where for each coordinated k , q ( k ) t ( x ( k ) , y ( k ) ) has following representation:

Therefore, since D x log h ( t, T, x , x T ) = D x log q T -t ( x t , x T ) , by projecting D x log q T -t ( x t , x T ) to each coordinate ϕ ( k ) , we obtain the following results:

A.5 Deriving Divergence Between Path Measures

Here we present an infinite-dimensional generalization of Girsanov's theorem [17, Theorem 10.14], which plays a crucial role in estimating the divergence between two path measures discussed in Sec 2.4. The theorem is formulated as follows:

Theorem A.4 (Girsanov's Theorem in H ) . Let γ be a H 0 -valued F t -predictable process such that

Then the process ˜ W Q t = W Q t -∫ T 0 γ s ds is a Q -Wiener process with respect to {F t } t ≥ 0 on the probability space (Ω , F , Q ) where

Or we can derive alternative formulation, by substituting W Q t = ˜ W Q t + ∫ T 0 γ s ds to (A.52) ,

Now, we will apply Girsanov's theorem to path measures related to (1) and (9). Let Q := P α in (A.53). Then γ s := Q 1 / 2 α s and we get

Proof. Proof can be founded in [17, Theorem 10.14]

Since our goal is find an optimal control α ⋆ such that X α ⋆ T satisfying terminal constraints which is represented by function ˜ G in (6), we may define our target path measure P ⋆ as d P ⋆ d P = 1 Z ˜ G ( · ) [57], where Z = E P [ ˜ G ( X T ) ] Then, we can compute the logarithm of Radon Nikodym derivative as

Since ˜ W Q t is Q -Wiener process on P α , we can compute the relative entropy loss in Sec 2.4:

where we denote ∥ ( · ) ∥ 2 H 0 = ∥ ∥ Q -1 / 2 ( · ) ∥ ∥ 2 H This representation matches with (3) when R ( · ) := 1 2 ∥·∥ 2 H . Similarly, the cross entropy loss in (19) can be derived by set Q := P ⋆ and P := P α in (A.53). Then γ s can be defined as (18) and ˆ W Q t = ˜ W Q t + ∫ T 0 γ s ( θ ) ds is a Q -Wiener process on P ⋆ where P α satisfies the Radon-Nikodym derivative:

Hence, the cross-entropy loss can be computed as in (19).

A.6 Proof of Theorem 3.1

Proof. Let us consider that the marginal distributions { µ ⋆ t } t ∈ 0 ,T satisfying the following relation in a weak sense [4]:

where L ⋆ t is adjoint operator of L t . Now, let us denote µ t | 0 ,T ( x t ) = µ ⋆ t ( x t | x 0 , x T ) and consider factorizable marginal distribution µ ⋆ t ( d x t ) = ∫ Π µ ⋆ t ( x t | x 0 , x T )Π( d x 0 , d x T ) , where µ ⋆ t ( x t | x 0 , x T ) satisfying the following relation:

where we denote h t | 0 ,T ( x t ) := σ 2 QD x log h ( t, x t ) . Now, we can obtain the Kolmogorov operator associated with the diffusion process associated with the marginal distributions µ ⋆ t ( d x t ) as follows

where we have defined ∫ H h ⋆ t ( x t ) µ ⋆ t ( d x t ) := ∫ Π ∫ H h t | 0 ,T µ ⋆ t ( d x t | x 0 , x T )Π( d x 0 , d x T ) . It implies that the diffusion dynamics d X ⋆ t associated with the Kolmogorov operator L := ⟨A x t , D x f ( x t ) ⟩ H + ⟨ h ⋆ t ( x t ) , D x f ( x t ) ⟩ H + 1 2 Tr [ σ 2 QD xx f ( x t ) ] also associated with Fokker-Planck equation [7] L ∗ t µ ⋆ t ( d x t ) = 0 , meaning X ⋆ t ∼ µ ⋆ t for all t ∈ [0 , T ] . Now, for some reference measure µ ref where the Radon-Nikodym derivatives dµ ⋆ t dµ ref ( x t ) = p t ( x t ) , dµ t | 0 ,T dµ ref ( x t ) = p t | 0 ,T ( x t ) exist. Then, under µ ref , h ⋆ t can be defined as follows:

Therefore, the diffusion process associated with marginal distributions µ ⋆ t ( x t ) has following representation:

This concludes the proof.

A.7 Proof of Theorem 3.2

Proof. Let us assume the initial condition is fixed to deterministic point x 0 ∈ H and define a reference measure µ ref = N (0 , Q ∞ ) . Since our goal is X α T satisfying the terminal condition G , define P ⋆ as d P ⋆ d P = 1 Z ˜ G ( · ) , where Z = E P [ ˜ G ( X T )] . Therefore, we have the following relation for marginal distributions dµ ⋆ T = 1 Z ˜ G ( X T ) dµ T . Moreover, since we have defined G ( X T ) = -log dπ T dµ T ( X T ) in (23), then it result h ( t, x ) = E P [ dπ T dµ T ( X T ) | X t = x ] by following Theorem 2.2, we get ˜ G ( X T ) = h ( T, X T ) . Hence, we have for any Borel set B ∈ B ( H ) ,

Now, given that we have confirmed that the conditioned SDE in (9) correspond to the controlled process (9) with optimal control, we can establish the result P ( X α ⋆ T ∈ B ) = π ( B ) . This concludes the proof.

Encoder

𝑸

× L

𝑲

𝑽

Cross

Attention

Latent array

𝑸

Self

Attention

𝑲

𝑽

Decoder

Cross

Attention

𝑸

𝑲

𝑽

Query encoding

Projection

GFFT

…

…

…

𝑸

𝑲

𝑽

Self

Attention

× L

A.8 Experimental Details

A.8.1 Experiment on 2D-Domain

Synthetic Experiment. In a synthetic experiment, we computed the log-probability of p 0 and p T across a uniformly sampled grid of 64 2 points, with each point p i ranging within [ -7 , 7] 2 . For p 0 , the log-probability was generated using an 8-Gaussian mixture model as specified in [70]. For p T , we employed the following log-density function:

Training was conducted using the Adam optimizer with a learning rate of 1 e -3 . The network was trained with a batch size of 24 for a total of 1000 iterations. We set σ = 0 . 2 in (1) for this experiment and set 100 discretization steps. We use a single A 6000 GPU for this experiment.

The control function was parameterized using a 4 -layer FNO-2D [40], with the cutoff number of Fourier modes set at 8 and each convolution layer having a width of 32 .

Simulation of DBFS. We follow the simulation scheme introduced in [42, 58]. For x = ( x 1 , x 2 ) ∈ R 2 , we use the discrete cosine transformation ( DCT ) for projection. Specifically, the eigenvector ϕ ( k ) ( x ) and eigenvalue λ ( k ) of the negative Laplacian operator -∆ , which is positive definite and Hermitian, and satisfies the zero Neumann boundary condition, are given by:

It implies that the orthonormality of ϕ ( k ) ( · ) with respect to the associated inner product, thereby enables computations in (A.47-A.50). Now, considering a rectangular domain with Cartesian coordinates, where pixels in the image are sampled from an underlying regular domain, the eigenbasis is given as a separable cosine basis:

Then, the negative Laplacian -∆ can then by represented by an eigen decomposition -∆ = EDE T , where E T is the projection matrix for the DCT i . e ., ˜ X t = E T X t = DCT ( X t ) , and D is a digonal

Algorithm 3 DBFS sampling for Bridge Matching

Input: Linear operator A , trace-class operator Q , learned control α , initial distribution π 0 , discrete and inverse discrete cosine transforms DCT , iDCT , target resolution grid T 2 , trained resolution grid O 2 .

̸

matrix containing the eigenvalues λ ( n,m ) . Hence the controlled SDEs (2) with Q = I can be rewritten as

Sampling Algorithm. The sampling algorithm for DBFS in bridge matching, discussed in Section 5.1, is provided in detail in Algorithm 3.

Unpaired dataset Transfer Experiment. For the experiment involving transfer between dataset, we followed the setup described in [51].

(A) For the EMNIST and MNIST datasets, the initial distribution, π 0 , was set as the MNIST dataset, while for the terminal distribution, π T , we used the EMNIST dataset with the first five lowercase and uppercase characters, as outlined by [19]. The iterative training scheme proposed by [51]. was adopted, which involved two neural networks, α t ( t, x , θ ) and α t ( t, x , ψ ) , each with around 20.7 million parameters. These networks approximate mixtures of bridges for the forward ( π 0 → π T ) and reverse ( π T → π 0 ) directions, respectively. The SDE was discretized into 30 steps without a noise schedule. The DBFS model was trained for 60 iterations, with each iteration comprising 5,000 gradient updates. Additionally, 2,560 cached images were used for training each network, updated every 250 steps. We used the Adam optimizer with a learning rate of 1e-4 and a batch size of 128, with the EMA rate set to 0.999. The complete DBFS training for the MNIST experiment took approximately 15 hours on a single A6000 GPU.

(B) For the AFHQ dataset [14] 5 , we evaluated DBFS between the wild and cat classes on a 64 2 grid, with each class containing approximately 5,000 samples. The control networks each contained about 120.3 million parameters. We discretized the SDE into 100 steps without using a noise schedule. The Adam optimizer was used with a learning rate of 1e-4, and the EMA rate was set to 0.999. We used a batch size of 64 and trained for 20 iterations, with a total of 400,000 gradient steps. Additionally, 2,560 cached images were used for training each network, updated every 1,000 steps. The training took approximately 8 days, using 8 A6000 GPUs for the experiment.

For each control network α , we used a transformer network architecture inspired by PerceiverIO [36] from public repository 6 to model functional representation. This transformer architecture was chosen for its efficiency in evaluating field representations over a large number of grid points and its

ability to map arbitrary input arrays to arbitrary output arrays in a domain-agnostic way. Additionally, we adopted the attention mechanism proposed in [49]. The decoding cross-attention mechanism was also modified, inspired by [39], with the output query set to the target grid points, which were transformed as Gaussian Fourier features [67].

In practice, we start by passing the evaluations X [ p ] through a single MLP layer and combine it with the Fourier feature embeddings of the grid points p . This combined representation is then input into the encoder blocks as keys and values, where QKV attention is first applied with the latent array as the query, followed by self-attention for each block. We implement the time-modulated attention block proposed by [49], embedding the time t into latent space. Next, the target grid points, represented as Fourier feature embeddings of the grid points p , are fed into the decoder blocks as queries. Here, QKV attention is applied with the encoded latent array serving as keys and values, followed by self-attention in each decoder block. Finally, the decoded array is mapped to grey scale or RBF channels using a single linear layer. Conceptual illustrations of the proposed network are presented in Figure A.1, and detailed network hyperparameters are listed in Table A.1.

(a)

32

2

(observed resolution)

(b)

64

2

(unseen resolution)

(c)

128

2

(unseen resolution)

2

2

十

S

】

6

4

1

8

B

6

8

2

6

a

6

e

D

C

b

3

a

c

C

A

D

2

A

C

7

5

e

3

1

5

se

e

a

A

C

b

7

9

b

c

O

5

6

G

8

9

C

B

e

C

6

8

3

9

2

b

C

A

D

3

9

B

e

A

7

3

2

o

6

7

B

B

b

5

0

3

6

3

8

9

2

E

C

a

B

e

C

e

C

D

9

7

S

(d)

64

(observed resolution)

(e)

128

(unseen resolution)

2

2

64

(observed resolution)

(b)

128

(unseen resolution)

(a)

A.8.2 Experiment on 1D-Domain

For the experiments on 1D-domain, we consistently set A := -1 2 and Q := exp( -∥ p -p ′ ∥ 2 /γ ) and set γ = 0 . 2 and γ = 0 . 02 for GP regression and imputaion, respectively. The choice of γ is hyper-parameter, we search over the set [0 , 01 , 0 . 1 , 0 . 2 , 0 . 5 , 1 . 0] and find optimal value for GP regression. For imputation, we set γ = 0 . 02 by following [6].

Terminal Cost Computation For all experiments conducted in the 1D domain, we implemented a parameterized initial condition which takes as input the observed sequences X θ 0 = x θ ( Y [ O ]) . We employed the energy functional U as the Gaussian negative log-likelihood (NLL). For each evaluation point on T , U can be computed as follows:

where σ 2 θ is set as an output from the neural network in accordance with [38] for GP regression, and fixed as σ 2 θ = 0 . 5 for imputation, to establish a loss function analogous to [6]. Additionally, we specified a learnable prior distribution µ prior = N ( e -T 2 x θ , Q T ) . Consequently, the terminal cost retains only the NLL term, simplifying the computation.

GP regression For the GP regression, we borrow the experiment setting from [38]. The model trained with curves generated from GP with RBF kernel and tested in various settings such as data generated from GP with other type of kernel (Mat´ ern 5/2, Periodic). We generated p uniformly on the interval [ -2 , 2] and generated Y [ p ] from using RBF kernel κ ( p i , p j ) = l 2 1 exp( -∥ p i -p j ∥ 2 /l 2 2 ) with l 1 ∼ Unif (0 . 1 , 1 . 0) and l 2 ∼ Unif (0 . 1 , 0 . 6) and the white noise ξ ∼ N (0 , 1 e -2) is added. We set | O | randomly from Unif (3 , 37) and | T | from Unif (3 , 50 -| O | ) . For the other test data, we define κ ( p i , p j ) = l 2 1 (1+ √ 5 d/l 2 +5 d 2 / (3 l 2 2 )) exp ( - √ 5 d/l 2 ) with d = ( ∥ p i -p j ∥ ) , l 1 ∼ Unif (0 . 1 , 1 . 0) and l 2 ∼ Unif (0 . 1 , 0 . 6) for Mat´ ern kernel and κ ( p i , p j ) = l 2 1 exp( -2 sin 2 ( π ∥ p i -p j ∥ 2 /p ) /l 2 ) with l 1 ∼ Unif (0 . 1 , 1 . 0) , l 2 ∼ Unif (0 . 1 , 0 . 6) and p ∼ Unif (0 . 1 , 0 . 5) for periodic kernel.

We set batch size of 100 and trained for 100 , 000 iterations. The Adam optimizer is used, the initial learning rate 5e-4 decayed with cosine annealing scheme. For testing, we evaluated the trained models using 3,000 batches, each consisting of 16 samples. We report the mean and standard deviation for five runs. We a single A 6000 GPU for this experiment.

The architectures for NP [28] and CNP [27], we use the same setting as described in [38] 7 . In our approach, we adapted the CNP architecture to incorporate a parameterized initial condition x θ (add one linear layer to output x θ ). The total number of parameters is similar across all three models.

Physionet Imputation The Physionet [30] contains 4000 clinical time series with 35 variables for 48 hours from intensive care unit. Following [68], we preprocess this datasets to hourly timeseries which have 48 time steps. Since the dataset already contains around 80% of missingness, we randomly choose 10 / 50 / 90% of observed values as test data.

In the imputation experiments, we employed the same experimental setup as [6] 8 which is slight modification of original CSDI model. In this setup, the Gaussian noise of the DDPM model was replaced with GP noise employing an RBF kernel. Additionally, we adjusted the measurement approach to record the actual elapsed time rather than rounding to the nearest hour, better capturing the inherent timing characteristics of the Physionet dataset. We use a single A 6000 GPU for this experiment.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Yes, the main claims made in the abstract and introduction are consistent with the scope of the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations are discussed in the conclusion section.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Yes, we believe our proof and assumptions are both sufficient and correct.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Yes, we believe we have provided all the necessary information to reproduce our results in Appendix A.8.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have provided the codebase used for our experiments, particularly for unpaired image transfer.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Yes, we have provided the experimental details.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Yes, when available, we conducted the experiments five times and reported the mean and standard deviations.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Yes, we provided the required resources in the experimental details section.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We support the NeurIPS Code of Ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The broader impacts are discussed in the conclusion section.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [Yes]

Justification: We adhere to ethical standards for using our model in generative AI.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Yes, the license and terms of use are noted.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We do not involve crowdsourcing or research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: We do not involve crowdsourcing or research with human subjects.

Peer reviews:
comments: SUMMARY:
The paper uses a stochastic optimal control to derive Doob's h-transform in infinite dimensions, and it shows the relation between solving the optimal control problem and learning diffusion generative models. The approach applies both to bridge sampling and for generative modelling. The approach is demonstrated on infinite dimensional problems, including bridges between images and bridges between probability distributions.

STRENGTHS:
Strengths:
- well-written and interesting paper
- the stochastic optimal control approach to deriving Doob's h-transform is well-founded and interesting
- the authors derive a Bayesian inference algorithm using a reference measure
- the method is tested on simple examples

WEAKNESSES:
Weaknesses:
- Doob's h-transform in the infinite dimensional setting has been derived using other methods in previous papers (both in the linear and non-linear cases, e.g. ref [2,47], https://arxiv.org/abs/math/0610386). I believe the list of contributions and the introduction does not clearly show that the current paper is not the first to do this, e.g. [2] is first mentioned much later in the paper. I am not sure the introduction and list of contributions adequately reflects this, something that should be addressed before acceptance

QUESTIONS:
no questions

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors investigate the notion of h-transform in infinite dimensional state spaces and provide a novel representation (Theorem 2.3) based on connections to stochastic optimal control.\n\nThe authors introduce two approaches to using this h transform derivation - firstly in something resembling bridge matching whereby both marginals are known and secondly by simulating the process with network parameterized h transform and taking gradients through the simulation.\n\nThe authors then apply this to image super resolution and Bayesian inference tasks in function space.\n\nSTRENGTHS:\n- Derivations appear correct\n- Although the h-transform has been described for infinite dimension through Hilbert spaces in the context of diffusion models in Baker et al 2024 (https://arxiv.org/pdf/2402.01434); as far as I am aware this connection to optimal control is novel.\n- Experiments are reasonable compared to other infinite dimensional methods (see below) but still not on the same level of fixed dimension methods for e.g. superresolution.\n  - Spectral diffusion processes, Phillips et al 2023: https://arxiv.org/abs/2209.14125\n  -  Neural Diffusion Processes, Dutordoir et al 2022, https://arxiv.org/abs/2206.03992\n  - Baker et al 2024 (https://arxiv.org/pdf/2402.01434)\n\nWEAKNESSES:\n- Motivation for infinitedimensional diffusion bridge is not very strong and experiments are not very convincing. \nI am not so familiar with the Bayesian inference experiments and what is SOTA. There are a few baselines missing as noted below. For the superresolution task there are stronger and simpler methods which have not been discussed. I think some stronger use-case in scientific applications would be needed for a higher score.\n\n- More discussion with Baker et al 2024 (https://arxiv.org/pdf/2402.01434) would be appreciated\n- As the authors note, the second training method for Bayesian learning problems (Alg 2) requires taking gradients through the simulated diffusion which can be slow, unstable and memory intensive. This goes against much of the diffusion model philosophy of splitting the generative problem into smaller problems through time and solving each jointly. I fear this will not be very scalable beyond 2D.\n\n**Experiments**\n- FID score or other quantitative metrics are note provided for superresolution tasks.\n- There are no baseline or comparisions to other methods. There are many superresolution, and infinite dimensional diffusion methods.\n- The authors compare to neural processes but there are more recent and comparable baselines for similar infinite dimensional / functional/ Bayesian experiments  which do not rely on gradients through the simulated process, such as:\n  - Spectral diffusion processes, Phillips et al 2023: https://arxiv.org/abs/2209.14125\n  -  Neural Diffusion Processes, Dutordoir et al 2022, https://arxiv.org/abs/2206.03992\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}

comments: SUMMARY:
This article proposes a perspective on diffusion-based generative models based on stochastic optimal control, with objective functions based on the log density ratio between objectives.

STRENGTHS:
As far as I could evaluate, the mathematics are correct, and this particular mathematical perspective is new (to the best of my knowledge).

WEAKNESSES:
I found this submission to have a weak presentation. It reads more like a stochastic calculus journal article than a machine learning conference submission. 

This perspective is not clearly motivated: what is gained by considering an infinite dimensional perspective compared to the wide literature already approaching diffusion-based approaches through the length of stochastic optimal control? Since there are several elements that are infinite-dimensional in nature in this problem (distribution of random variables, score matching functions, etc), some early-on explanation and clarification of the approach considered here would be helpful.

Further, while a lot of the writing is centered around an infinite-dimensional perspective, this is then converted to a parametric model, with finitely many parameters. How much of the infinite-dimensional perspective is then lost? Is this important?

QUESTIONS:
Thanks to the authors for addressing my comments during the rebuttal.

LIMITATIONS:
Presentation and motivation - adressed by authors during rebuttal.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-18

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents stochastic control in function spaces with applications in diffusion bridges and Bayesian learning. Since the Lebesgue measure does not exist in infinite dimensional space, the authors derive Doob-h function with the Radon-Nikodym density with respect to a   suitable Gaussian measure and conduct bridge matching experiments under this setup.\n\nSTRENGTHS:\nOverall, the paper is well-motivated and well-written. The paper reviews the stochastic control in function space and the connection of Doob's h-transform with stochastic control and bridge matching in Section 2. It transits smoothly to Section 3, where it proposes an algorithm for diffusion bridges in function space, and an extension for Bayesian learning.\n\nWEAKNESSES:\nAs the theory exists for stochastic control in function space, and there is a recent work on h-transform [1] and generative model in infinite dimensional space, the novelty mainly lies in the application to bridge matching and Bayesian learning. These applications are interesting and important, however, the weaknesses are in the discussion on Bayesian learning, and the experiments on bridge matching. In particular, there should be a comparison with finite-dimensional bridge matching; several arguments in section 3.2 about Bayesian learning need more clarification (see questions for the details of this point).\n\n[1] Baker, Elizabeth Louise, et al. \"Conditioning non-linear and infinite-dimensional diffusion processes.\" arXiv preprint arXiv:2402.01434 (2024).\n\nQUESTIONS:\nComments and major questions: \n\n1. As the paper introduced in section 2.1, one can instead consider the cylindrical Wiener processes on the Cameron-Martin space. What will break down in the current results? Will the cylindrical Wiener processes set-up bring convenience to the experiments as it is implemented in finite dimensions?\n2. How do you arrive at equation(21)? What assumptions are required, and what are the regularity requirements for the energy function? It would also be good to remind the readers what mu_T is here.\n3. In equation(25), how are energy function U and covariance operator Q determined in general? What are the particular choices used in the presented experiments? \n4. What are the challenges of using time-dependent diffusion processes in the current method? \n\nMinor points: \n\n1. The paper repeatedly refers to Lemma 2.2, but the authors seem to be referring to Theorem 2.2.\n2. The paper should include a proof of Theorem 2.3 for its completeness and rigor.\n3. Why is it H_0 instead of H in Theorem 3.2?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-24"}