Stochastic Optimal Control for Diffusion Bridges in Function Spaces

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

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

yes

Summary

The 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. The 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. The authors then apply this to image super resolution and Bayesian inference tasks in function space.

Strengths

- Derivations appear correct - 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. - 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. - Spectral diffusion processes, Phillips et al 2023: https://arxiv.org/abs/2209.14125 - Neural Diffusion Processes, Dutordoir et al 2022, https://arxiv.org/abs/2206.03992 - Baker et al 2024 (https://arxiv.org/pdf/2402.01434)

Weaknesses

- Motivation for infinitedimensional diffusion bridge is not very strong and experiments are not very convincing. I 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. - More discussion with Baker et al 2024 (https://arxiv.org/pdf/2402.01434) would be appreciated - 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. **Experiments** - FID score or other quantitative metrics are note provided for superresolution tasks. - There are no baseline or comparisions to other methods. There are many superresolution, and infinite dimensional diffusion methods. - 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: - Spectral diffusion processes, Phillips et al 2023: https://arxiv.org/abs/2209.14125 - Neural Diffusion Processes, Dutordoir et al 2022, https://arxiv.org/abs/2206.03992

Questions

See weaknesses.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses.

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.

Rating

5

Confidence

3

Soundness

3

Presentation

1

Contribution

2

Limitations

Presentation and motivation - adressed by authors during rebuttal.

Summary

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

Strengths

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

Weaknesses

As 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). [1] Baker, Elizabeth Louise, et al. "Conditioning non-linear and infinite-dimensional diffusion processes." arXiv preprint arXiv:2402.01434 (2024).

Questions

Comments and major questions: 1. 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? 2. 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. 3. 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? 4. What are the challenges of using time-dependent diffusion processes in the current method? Minor points: 1. The paper repeatedly refers to Lemma 2.2, but the authors seem to be referring to Theorem 2.2. 2. The paper should include a proof of Theorem 2.3 for its completeness and rigor. 3. Why is it H_0 instead of H in Theorem 3.2?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations.

Discussion

Please participate to the discussion