Stochastic Optimal Control for Diffusion Bridges in Function Spaces

We would like to express our gratitude to the reviewer for their thorough evaluation of our work. We appreciate your recognition of the merits of our research.

1.Clarification on prior work

• We agree that clarifying the relationship between previous papers on Doob's h-transform and our work is essential to avoid any unintended confusion. Our contribution involves leveraging stochastic optimal control to derive Doob's h-transform and extend finite-dimensional SOC problems based on conditional diffusion into infinite-dimensional spaces. In the revised manuscript, we will explicitly outline how our contribution compares to previous works, particularly [1].

[1] Baker et al., "Conditioning non-linear and infinite-dimensional diffusion processes."

We appreciate the recognition of our paper's strengths and extend our thanks to the reviewer for their comprehensive review and insightful comments. Below, we provide detailed responses to address each valuable comment.

1.Motivation, Comparison with baselines

• We agree with the reviewer’s concerns regarding motivation and experiments. To address these, we have included a PDF file in the general response with additional experiments comparing unpaired image transfer methods with finite-dimensional baselines [1] and 1D functional generation with infinite-dimensional baselines [2, 3]. Additionally, we have further clarified our motivation. Please kindly refer to that section for more details.

2.Comparison with [4]

• [4] primarily focused on developing the conditional diffusion process in function space. To achieve this, they defined a Doob’s h transform in infinite-dimensional space using Itô’s lemma and Girsanov’s theorem. This approach has its merits, such as enabling the conditioning of non-linear SDEs (while our work has considered the linear SDEs). However, simulating conditioned non-linear SDE often faces challenges because the conditional distribution of such SDE is generally intractable. Therefore, they require the approximation algorithm presented in [5].

• While our approach shares the idea of developing an infinite-dimensional Doob’s h transform, as the reviewer already mentioned, our primary goal is not merely to derive Doob’s h transform but to generalize various finite-dimensional sampling problems [6, 7] into the infinite-dimensional space by exploiting the theory of infinite-dimensional stochastic optimal control.

• In practice, the choice of linear SDEs to develop the relevant theory might appear to be a strict limitation for modeling complex distributions. However, similar to most recent diffusion-based models, the linear form may be sufficient for modeling. Moreover, this choice can be beneficial as it allows for more scalable algorithms due to the closed-form solution of the conditional distribution. In this light, by leveraging the theory of stochastic optimal control along with the choice of linear dynamical systems, our contribution also includes proposing tractable learning algorithms for real-world sampling problems.

3.Computational concerns

• As the reviewer pointed out and as we stated in our paper, Algorithm 2 may induce computational difficulties. While it is possible to consider a more computationally favorable approach, such as implementing the adjoint solver [8] for memory efficiency or using the variance reduction technique proposed in [9], in the current work we focus on the theoretical property that the optimal control still yields Bayesian posterior sampling despite being defined on function space. Proposing a more scalable algorithm will be an interesting direction for future work.

[1] Peluchetti, “Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling.”
[2] Phillips et al., “Spectral Diffusion Processes”
[3] Durordoir et al., “Neural Diffusion Processes”
[4] Baker et al., “Conditioning non-linear and infinite-dimensional diffusion processes”
[5] Heng et al., “Simulating Diffusion Bridges with Score Matching”
[6] Zhang et al., “Path Integral Sampler: A Stochastic Control Approach For Sampling”
[7] Shi et al., “Diffusion schrodinger bridge matching”
[8] Li et al., “Scalable Gradients for Stochastic Differential Equations”
[9] Xu et al., “Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations”

We gratefully thank the reviewers for their valuable feedback and suggestions. Here, we address the concerns raised by the reviewer.

1.Early-on explanation and clarification

• Following the reviewers’ suggestions, we have further clarified our motivation in the general response. Please kindly refer to that section for more details. We hope this provides the clarity you need.

2.Finite-dimensional Approximation

• We would like to point out that the model having finite number of parameters does not mean that it should be modeling finite-dimensional models; for instance, Gaussian process regression, one of the most popular infinite-dimensional stochastic process models, effectively requires finite number of parameters to be estimated from the observed data. Usually, the function is evaluated only at a countable set of sampling points that are assumed to be generated from an infinite-dimensional stochastic process. In this case, we can approximate the infinite-dimensional function by fitting the models to the finite sampling points with finite-dimensional parameters.

• The “finite dimensional approximation” happening in our model would be in the part where we are approximating the covariance operator Q. Specifically, we approximate Q via the truncation, that is, choosing the finite number of eigenfunctions among infinitely many ones. This approximation may incur approximation error, but it does not alter the nature of our model as an infinite-dimensional model; To see this, note that our model can model image data in a resolution-agnostic way. This is possible because ours deals with the infinite-dimensional stochastic processes, so it can model a varying number of finite sampling points.

We sincerely appreciate your interest in our research and acknowledgment of its significant contributions. We are also grateful for the insightful questions raised by the reviewer, to which we have provided detailed responses in the subsequent text.

1.Comparison with finite-dimensional bridge matching

• Thank you for your valuable suggestion to improve our work. In general response, we have included a PDF file containing additional experiments and a comparison with finite-dimensional method [1]. We will incorporate these results into the paper.

2.Cylindrical Wiener process

• In theory, selecting a cylindrical Wiener process for infinite-dimensional SDEs can also be a viable option. Indeed, when prior knowledge of the data domain isn't available, opting for a cylindrical Wiener process, as noted by the reviewer, can facilitate the construction of a bridge model. However, the choice of Q determines the geometric structure of the Hilbert space where functional data resides. Therefore, If we model Q so that our Hilbert space contains characteristics of the target data, such as smoothness or curvature information, it can be beneficial.
• In practice, where our objective is to model complex data as a function, the choice of the covariance operator Q can significantly impact the ability of learned neural network operators (control in our case). For instance, in time-series imputation, our baselines [2, 3] differ only in their selection of Q (Technically “kernel” since they are finite-dimensional models) —[2] using a standard Wiener process and [3] employing an RBF kernel to construct the Wiener process— leading to performance improvements. Furthermore, in generative modeling, studies [4] empirically demonstrate that selecting an appropriate operator Q enhances the correctness of generation. Specifically, [4] highlights that opting for a cylindrical Wiener process may lead to issues such as mode collapse.

3.Regularity condition for U and choice of U and Q

• Equation (21) is coming from a Bayes formula [5, Section 2] $\frac{d\pi_T}{d\mu_{prior}} \propto \exp(-\mathcal{U}(\mathbf{X}))$, where $\mu_{prior} = \mathcal{N}(m_{prior}, Q_{prior})$ and the potential $\mathcal{U}$ is a measurable mapping $\mathcal{H} \to \mathbb{R}$ for a given $\mathbf{X} \sim \mu_{prior}$. Hence $\mu_{prior}$ in equation (21), equation (25) and lines (245-246, 751) should be changed into. We apologize for the typo and for any confusion it may have caused in understanding the paper. The potential $\mathcal{U}$ typically chosen as a negative log-likelihood function (also referred to as the potential energy); in our case, we set it as a negative Gaussian log likelihood function. For Bayesian learning, we choose $\mathcal{A}=-\frac{1}{2}$ and covariance operator $Q$ as an RBF kernel. We have detailed the setting in Appendix A.9.2.

4.Time-dependent SDEs

• The main challenges of using a time-dependent diffusion process in our method are proving the existence and uniqueness of the invariant measure. For an explicit form of the h-function, as stated in Theorem 2.3, we need to define a certain class of Gaussian measures, where the collection of time-dependent Gaussian measures is equivalent to its invariant measure over a long-term period. This class of Gaussian measures should be defined by a linear SDE.

5.Minor comments

• We appreciate the reviewer pointing out the typo related to Lemma 2.2 and $\mathcal{H}_0$​ in Theorem 3.2. It should be expressed with respect to the norm in $\mathcal{H}$. We will make the necessary corrections in the revised manuscript. Moreover, as reviewer suggested, we will include the proof of Theorem 2.3.

[1] Peluchetti, “Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling.”
[2] Tashiro et al., “CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation”
[3] Bilos et al., “Modeling Temporal Data as Continuous Functions with Stochastic Process Diffusion”
[4] Hagemann et al., “Multilevel Diffusion: Infinite Dimensional Score-Based Diffusion Models for Image Generation“
[5] Hairer et al., "Signal Processing Problems on Function Space: Bayesian Formulation, Stochastic PDEs and Effective MCMC Methods"