Conditioning Diffusions Using Malliavin Calculus

Thanks for your review! We're glad you think the method is interesting. We hope we will be able to address your concerns on the scalability and efficiency of our proposed method, which we discuss below!

Section q3Kt.1: Larger Experiments and Flow-Matching

Thank you for the great suggestions. Initially, we mainly targeted the case of conditioning diffusion on a single endpoint, since this can be seen as the most singular, and hence in principle most challenging, reward case (a Dirac delta). However, due to the fact that our method works in these extremal cases, it is actually applicable to a wide range of settings. We followed your proposal and applied our methodology to condition a diffusion model trained on images (could be similarly done for flow matching)!

We first trained a diffusion model on Fashion-MNIST and then used our methodology to condition on the upper left corner of an image. You can see the results here:

• https://postimg.cc/bsbNgMY2

The upper left image of the right grid is the ground truth. We extracted the upper left quarter of that image as conditioning input (seen as the upper left image in the left grid). We then generated samples from the posterior distribution after conditioning on that upper left image. Training the conditioning drift for this particular problem took us about 30-60 minutes on a high-end GPU. Note that after this training time one can condition on any upper left corner.

In general, we can condition on any observation of the form $Y = G(X).$ Here, $Y$ can be discrete (classification) or continuous (upscaling, inpainting, ...), and $G$ can be smooth, but also non-differentiable and even non-continuous, since for diffusion bridges we have the extreme case of $G$ being a Dirac delta. We support both deterministic (noise-free) and stochastic (noisy) versions of G; adding artificial noise is not required but is possible to account for uncertain observations depending on the application. Of course, when $G$ has nice properties, like gradients, it might make sense to include these, but due to our lack of assumptions on the reward, our algorithm is quite generally applicable.

We understand that our paper wasn't to clear on this, and will put more emphasis on this point in the updated version. However, nearly nothing has to change, one only needs to replace $x_T$ by a general condition $y$ in the neural network input during training and evaluation, since we never make use of the form of the observation $x_T = Id(x_T)$ in any of our steps.

Section q3Kt.2: Schroedinger Bridges

Indeed, bridges can be used in the calculation of Schroedinger Bridges. Note that if you can solve the static Schroedinger Bridge problem (Section 3.1 of the paper you linked), and one has access to the Bridges of the reference process, one can interpolate the static coupling of the marginals with the bridges to get the dynamic Schroedinger bridge solution. Therefore, if one has already solved for the bridge process, the problem of solving Schroedinger bridges on path space reduces to just finding the optimal coupling of the marginals!

I hope we were able to dispel your concerns about the applicability of our method to different settings and high dimensional problems. Thank you for motivating us to apply the method to a diffusion model, we are very happy with the results and think this showcases the generality of our method well.

Section q3Kt.3: Lack of Analysis on Efficiency and Complexity

We agree that it maybe was not fully clear how costly the evaluation of the score process is. Therefore we updated our paper to be more detailed on the calculation of the score process. We explain the complexity of the algorithm in the answer to reviewer 4Bsi, you can find it by searching for "Section 4Bsi.3" on this page.

We hope this answers all questions. Furthermore, we hope our experiments on problems with very high energy barries (double well) as well as high dimensions (images), showcase the applicability and efficiency of our method for different kinds of problems.

Summary

Thanks again for the review. We believe we have managed to resolve the issues you raised, but please let us know if this is not the case! We hope you agree that the changes to the paper of including a high-dimensional image experiment alongside more details on computing the loss and its complexity will improve it!

We really appreciate your positive review and that you think our paper is well written!

One weakness is that the experiments are fairly toy

We based our experiments on setups from the most closely connected literature and are glad you found them more substantial than in similar papers!

We think the included experiments are good tests for our method since

1. In the double well experiment, there is a high energy barrier between transitioning between wells, meaning it happens very rarely. Therefore, conditioning to go from one well to another is challenging.
2. In the shape models, points close to one another are very correlated making transitioning to another given shape challenging.

However, we agree these are fairly toy. In order to resolve this we have now conducted an experiment on images, which we hope will convince you of the applicability of our method to numerous situations, including high dimensions!

We use the following setup:

1. As the base SDE we take a pretrained diffusion model SDE trained to go from the noise distribution to the data distribution (for this we have used Fashion MNIST).

2. We condition the SDE on the upper left corner of the image to take a given specific value. Practically, this corresponds to the task of completing partial images.

3. We link our results here: https://postimg.cc/bsbNgMY2
   The true "left corner" is from the first image in the top left corner. We see all the samples match well, generating only images with the same left upper corner.

Strictly speaking, this experiment goes slightly beyond the considered set-up, as conditioning on a part of the image translates into "pinning down" only some of the coordinates of the endpoint. Indeed the proposed method is very flexible, and applies to conditioning on any observation $Y= G(X_T)$, with no assumptions on the function $G$ (e.g. $G$ need not be differentiable and $Y$ could be discrete or continuous). For more details, please see our answer to Reviewer q3Kt (see Section q3Kt.1). We realise this is not so clear in the paper right now and will include this discussion!

We hope we have managed to address your concern.

The order of the numbers of some Figures/Tables is strange. Thanks for pointing this out! We've updated this now so that the tables and figures are named in order of appearance.

What is the computation time like? Is this method practical?

Training the conditioning drift for the image task took us about 30-60 minutes on a high-end GPU, so yes, this method can potentially be applied in large-scale settings. Note that after this training time one can condition on any upper left corner (for example jackets, shoes, etc). Please do let us know if you want anymore information on this!

For a more detailed complexity analysis of the algorithm (including the number of forward and backward passes) please refer to our answer to Reviewer 4Bsi (Section 4Bsi.3).

Thanks again for your review. We hope we have managed to resolve your concerns by including the extra experiment and an analysis on the complexity of the method.

Thanks a lot for your positive review! We're very happy to hear that you really enjoyed this paper! And we were especially happy that you found the general quality to be on par with some of the seminal papers in this field. We appreciate your detailed comments and questions, and answer them below.

Section 4Bsi.1

"Can the claim "recall the connection of diffusion bridges to Doob’s -transform in Section 2.1, in particular the relevance of conditional score functions" be a contribution?"

Indeed, we agree it is well known. Our intention was to also to use this section to provide an outline to the paper, but we see how this could be confusing. To make it clearer, we will write "recall the (well-known) connection of diffusion bridges" and cite the relevant literature. We hope you agree this makes it clearer.

Section 4Bsi.2: "Footnote 1 is incomplete."

You're right -- the footnote currently continues on the next page. We will fix this! Thank you.

Section 4Bsi.3: Computation of the Score Process (Q2)

Thanks for the excellent questions. Indeed, the paper currently lacks some detail on how we compute the score estimator, which may have caused confusion. As you noted, vector-Jacobian products are cheaper via reverse-mode autodiff, and adjoint methods avoid computing the full Jacobian—this is exactly what we do. We simulate paths of the reference SDE (1), then compute the integral (6) in reverse, akin to adjoint ODE methods, but adapted for our setting, picking up $dB_t$ terms along the way instead of only backpropagating a final gradient. We present the full algorithm below for $\sigma = 1$ and Euler-Maruyama discretization, and will include a detailed explanation in the paper. Thank you for the helpful references—we will cite them in the revised version.

Given:

• $(X_0, X_{\delta t}, \ldots, X_T)$
• $(Z_0, Z_{\delta t}, \ldots, Z_T)$, which is the noise which was used to produce $X_{t + \delta t} = X_t + \delta t b(X_t) + \sqrt{\delta t}~Z_t$

Output:

• $(S_0, S_{\delta t}, \ldots S_T)$, the score process at all times $t$

Method: Initialize: $S_T = 0$ Propagate: $S_t = S_{t+\delta t}^T (I + \delta t \nabla b(X_t))$ At the end: Divide $S_t$ by $T - t$ to normalize it.

We hope it's now clear that since we left-multiply with the Jacobian, we never need the full Jacobian—just one backprop pass per timestep. This yields $N = T/\delta t$ regression targets $S_t$ for gradient updates.

Note that the Jacobian in (6) is transposed before right-multiplying with $dB_t$, which is equivalent to left-multiplication. We agree this wasn’t clear and will revise it.

Thanks again for the great question—highlighting this has improved the paper. To show scalability to higher dimensions, we’ve added image experiments (see “Section q3Kt.1” in our reply to Reviewer q3Kt).

Section 4Bsi.4: Choice of $\alpha$ (Q6 and Q7)

We found the choices "first", "optimal" and "average" all to do quite well on most problems (however, average and first are easier to implement than optimal). Which one worked the best was problem dependent. Note that "optimal" is only theoretically optimal in a specific setting (a forward process which is a Brownian motion). There are two interesting ways to study optimality:

• Getting a deeper theoretical understanding of optimal $\alpha$ choices for different classes of SDEs.
• Optimizing $\alpha$ numerically for a given problem.

However, these were out of the scope of the current work, but they are certainly very interesting for future research.

Note that in Section 4.2 we do not compare different $\alpha$ schemes. We just picked the best-performing one for shape spaces (which was "average"), and compare it to other algorithms from other publications. We observe that we outperform them. We did not observe this problem to behave very differently with regard to the choice of $\alpha$ than the other problems.

Section Other Questions

-- 1. -- Indeed, it's the derivative with respect to $t$. We will add this to the theorem statement to make it clear!

-- 3. -- Thanks for bringing this for our attention. We have updated this to be "for $i=1$ to $N$ ...".

-- 4. -- Yes, thanks - Equation (17) should read $\alpha'_t$. We've corrected this!

-- 5. -- That's a very interesting comment. In our setting (diffusion bridges), there should be no oscillation between the wells as the conditioning "pins down" the endpoint; transitioning back and forth is highly unlikely (even more unlikely than transitioning once!) under the original dynamics. Oscillating back and forth would be expected if the dynamics were set up to equilibrate and target a (spatial) distibution of interest, for example in MCMC type sampling algorithms. Potentially our method could be modified to enhance sampling, but since this is quite different, we believe it is better reserved for future work.

Thanks again for the feedback! We hope we managed to answer your questions sufficiently and address your concerns.

Thanks for your review! We're glad you see our method as unifying and extending diffusion bridge approaches. Below, we address your questions and concerns.

For theorem 2.1 as $\sigma(X_t)$ becomes singular, wouldn't you have stability issues with its inverse? & The method assumes that the matrix A and diffusion coefficient $\sigma(X_t)$ are always invertible...

Thanks for the question!

For $A$: Since $\alpha$ is user-chosen, it can always be selected to make $A$ invertible. All natural choices for $\alpha$ do so.

For $\sigma$: It is helpful to discuss two slightly different cases:

1. If $\sigma$ is small but positive definite, stability can be partly managed via small step sizes or adjusting $\alpha$. However, small noise makes the bridge problem harder in general (regardless of the method), as the dynamics become rigid; for $\sigma = 0$, the problem is ill-defined.

2. The diffusion coefficient is sometimes genuinely rank deficient, for example in the underdamped Langevin dynamics, $\mathrm{d}q_t = p_t \, \mathrm{d}t,$ $\mathrm{d} p_t = - \nabla V(q_t) \, \mathrm{d}t - \gamma p_t \, \mathrm{d}t + \sqrt{2\gamma} \mathrm{d}W_t,$ of great interest in molecular dynamics (here the diffusion matrix would be \begin{pmatrix} 0 & 0 \\ 0 & \sqrt{2 \gamma} \end{pmatrix} with noise only acting on the $p$-variable). In this case (and in similar other cases), it is possible to extend Theorem 2.1 using pseudo-inverses, and the proposed methodology generalises straightforwardly. On a technical level, this extension requires "hypoellipticity" which guarantees smooth transition probabilities of the underlying dynamics. Since underdamped Langevin dynamics is important for applications, we suggest adding a short explanation in the appendix.

There seems to be a implicit circular dependency...

We are a bit unsure exactly which proof you mean, but we assume you mean in Theorem 2.1. To clarify, the proof steps are as follows:

1. By Doob’s $h$-transform (Proposition 2.2) for eq (8) to coincide with the conditional law, we need to show that $u^*(x; x_T) = \nabla \log p_{T\mid t}(X_T = x_T \mid X_t = x)$.
2. By Proposition 2.3, we know $\nabla \log p_{T\mid t}(X_T = x_T \mid X_t = x) = \mathbb{E}[\mathcal{S}_t \mid X_t=x, X_T=x_T]$.
3. In the proof of Theorem 2.1 we show that the minimiser of $\mathcal{L}(u)$ is given by $\mathbb{E}[\mathcal{S}_t \mid X_t=x, X_T=x_T]$. Since this is equal to the term arising from Doob’s $h$-transform, eq (8) coincides with the SDE under the conditional law.

In particular, in equations (14) and (15), $\mathcal{S}_t$ is already fixed and does not depend on $u$, hence there is no circularity.

If this is what you mean, we’d be glad to include the above steps in proof of 2.1 making it more explicit!

The quality of the writing still needs to be improved...

Thanks for pointing this out! We will fix this immediately.

Lack of real world experiments conducted leaving applicability beyond theory uncertain. We based our experiment setup on papers in the literature that also concentrate on similar problems [1, 2, 3, 4, 5]. Indeed, the double well experiment is a challenging problem in conditioning due to the high energy barrier between wells and therefore the transitioning events are very rare. However, we agree it is interesting to consider image generation, and have since experimented with this!

We condition a pretrained model to only generate images matching with a given top left corner. Please see our answer to Reviewer q3Kt (Section q3Kt.1) for details, and see the samples linked here: https://postimg.cc/bsbNgMY2.

We hope we've conveyed the method's applicability and flexibility! We'll include this experiment to show it scales to high-dimensional SDEs and has practical relevance beyond the theory.

[1] https://arxiv.org/abs/2111.07243\

[2] https://arxiv.org/abs/2407.15455\

[3] https://link.springer.com/article/10.1007/s42985-021-00102-x\

[4] https://arxiv.org/pdf/2312.02027v4 \

[5] https://projecteuclid.org/journals/bernoulli/volume-23/issue-4A/Guided-proposals-for-simulating-multi-dimensional-diffusion-bridges/10.3150/16-BEJ833.full.

Figure 1 hardly convey the message...start by introducing a motivating example...

We agree we should include more information in this first image, and update it as follows: https://postimg.cc/NK7gq2JB. We hope you agree that this better illustrates the motivation and challenge of the paper, and it also gives us the opportunity of referring to it in the introduction when explaining the reward (and its potential singularity).

...do you think your method will be a good fit for high-dimensional task like image generation?

Yes, as we have seen from our experiment on images it does scale to high-dimensional tasks! The Jacobian calculation can be done very efficiently using vector-matrix products; for more details on this and on computing $\mathcal{S}$ please see our answer to Reviewer 4Bsi (Section 4Bsi.3).