Neural Guided Diffusion Bridges

Claims and Evidence

The guided bridge relies on training the neural network. Once trained, independent samples are obtained. The quality of the guided proposal depends much on the nonlinearity of the diffusion and number of pCN-steps required in the MCMC-algorithm. This is case dependent. For the performance evaluation, please refer to the table in our reply to Reviewer 8sR1.

Methods and Evaluation Criteria Inclusion of the guiding term from the guided proposal is necessary to ensure the guided neural process will be a bridge. Without it, the neural net has to learn the unbounded term which behaves near the endpoint like $(v-x)/(T-t)$ (in case of a full observation and a uniformly elliptic diffusion).

Theoretical Claims

1. A proof that $E_t$ in Eq. (5) is a local martingale is given in Palmowski & Rolski: it follows from partial integration and involves Dynkin's martingale. We believe this to be out of scope for this paper as it heavily depends on stochastic calculus.
2. For Eq. (10): the ``weight'' is obtained by integrating $\frac{({\mathcal{A}}-\tilde{\mathcal{A}})\tilde{h}}{\tilde{h}}=\sum_i(b_i-\tilde{b}_{i})\frac{\partial_i\tilde{h}}{\tilde{h}}+\frac{1}{2}\sum_{i,j} (a_{ij}-\tilde{a}_{ij})\frac{\partial^2_{ij}\tilde{h}}{\tilde{h}}$ over $[0, t]$.
3. Eq. (11) and (12) are taken from Section 2 of Mider et al.

Due to the space limitation, we can't include more details about 2 and 3 in this reply, but we are happy to proceed them in the conversation.

Experimental Designas or Analyses

• "Cell dynamics": Yes, this is true. Empirical results show multimodal marginal distributions emerge after $t=2.0s$ (Fig. 6), with neural-guided bridges evolving primarily unconditionally beforehand. Peak splitting coincides with intensified drift forcing toward $v$, matching guided proposal dynamics. Neither approach fully replicates the three unconditional sampling modes, though all three become distinct during $t=2.0–3.0s$ under weakening constraints – despite potential intermediate convergence failures toward $v$.

Supplementary Material

We acknowledge the importance of additional experimental results, and will prioritize space allocation for key figures and tables while maintaining methodological rigor during the revision of the manuscript.

Essential References Not Discussed

We agree that Proposition is known in the literature. However, we did not find another reference where this is spelled-out for our purposes in detail. We simply want to have one clean statement that explains the change-of-measure to $\mathbb{P}^\star$ for this particular $h$-function, in the setting of SDEs.

Thanks for pointing out the ``non-essential recent work'':

1. Denker et al.'s diffusion framework (Proposition 2.2) shares our Proposition 3.2's $h$-function concept (termed "generalized $h$-transform"), though methodologically distinct through denoising score matching. We will add reference to contextualize prior appearances of this construct.
2. While our focus remains on KL divergence's mode-seeking properties, we acknowledge log-variance divergence's potential benefits for mode collapse mitigation. This comparison will be expanded in future research.
3. The adjoint matching method shows promising scalability for high-dimensional applications (images/point clouds). We plan to investigate adapting our objective function into this framework.

Other comments or suggestions

1. (a) We agree that it is useful to the reader to remind them of the definition of $b^\circ$ in Eq. (6). There is no $\tilde{b}^\circ$ in our paper.

   (b) We agree with your suggestion and will add a remark at the end of Section 3.2 to explain this.

   Indeed, if $\sigma$ is invertible, then for $\tilde{X}_t = \sigma W_t$ transition densities $\tilde{p}$ are Gaussian and $\nabla_x \log \tilde{p}(t,x; T,v) = (\sigma\sigma^T)^{-1} \frac{v-x}{T-t}.$ Therefore, the bridge has drift $(v-x)/(T-t)$. We added this to the text.

2. We will clarify this in the revision of the manuscript. The following hierarchical model is meant: $v\mid y\sim q(v\mid y), y\sim p(T,y \mid 0,x_0).$ Here, $y$ is the parameter that gets the prior $p( T,y \mid 0,x_0)$ assigned; $v$ is the observation. In this model, the likelihood can be written as $L(y; v)=q(v\mid y)$.

3. We are unsure what is meant by the remark on $b(t,v)$. Regarding the reason it is useful for $d'<d$: we may only observe certain components of the diffusion, for example, if we have a 2 dimensional example and observe $v \sim N(x_1, \sigma^2)$, then the dimension of the observation is lower than that of the diffusion. The Fitzhugh-Nagumo example is an illustration of this.

4. Yes, we will stress this in the revision and clarify the explanation given in lines 347-355. Indeed, for the guided proposal we always use MCMC with pCN-steps, contrary to the other methods.

5. We think we have addressed this comment in the previous item.

Theoretical Claims

The formulation was imprecise and we propose to reformulate it as follows:

If $\theta_{\rm opt}$ is a local minimizer of $L$ and $L(\theta_{\rm opt})=-\log \frac{\tilde{h}(0,x_0)}{h(0,x_0)}$, then $\theta_{\mathrm{opt}}$ is a global minimizer. This implies from which we obtain D_{KL}(\mathbb{P}^{\bullet}_{\theta_{\mathrm{opt}} | \mathbb{P}^{\star})=0, from which we obtain $\mathbb{P}^{\bullet}_{\theta_{\mathrm{opt}}} = \mathbb{P}^{\star}$.

Experimental designs or analyses

We cannot explicitly choose the target mode to sample from. Theoretically, the variational class $\{\mathbb{P}^{\bullet}_{\theta};\theta\in\Theta\}$ is large enough to contain the target $\mathbb{P}^{\star}$, and the global optimal $\mathbb{P}^{\bullet}_{\theta_{\mathrm{opt}}} = \mathbb{P}^{\star}$ can be guaranteed.

However practically, the inherent non-convexity of the objective function induced by $G(s, x)$ predisposes $\vartheta_{\theta}$ to converge to local minima, as empirically evidenced by sample concentration within individual modes. However, directing the optimization trajectory within this complex landscape remains nontrivial, and we currently do not have a reliable method steer the process towards preferred local minima.

Questions for Authors

We don't fully understand what you mean by "It would be better if you could give proof or reasoning about the experiments' result.". We kindly ask you to clarify so we can answer your question accurately.

Methods and Evaluation Criteria

1. All $\vartheta_{\theta}$ implementations use fully-connected networks with $(1+d)$-dimensional inputs (time $t$ and state $x$). Time integration varies by dimensionality: direct concatenation for low-dimensional systems (Brownian/OU/cell/FHN) versus sinusoidal time embeddings with feature-wise modulation (Perez et al.) in high-dimensional landmark processes. Architectures employ tanh activations to ensure Lipschitz continuity (critical for Assumption 4.1 via gradient clipping). Layer counts/hidden dimensions were determined via loss-minimizing parameter sweeps.
2. To our knowledge, no established quantitative metrics exist for high-dimensional performance evaluation. Only visual assessment by brute force unconditional sampling (cell and FHN)—prohibitively inefficient for rare events, are available. In contrast, real-world applications (e.g., image/video) permit rigorous benchmarking through domain-specific metrics like FID/IS scores used in translation tasks.

Experimental designs or analyses

1. We took two representative examples to compare the computational cost of the methods considered in more detail in the following table.

2. Our proposed loss eliminates matrix inversion of $a=\sigma\sigma^T$. Precomputable terms $L(t)$, $M(t)$, and $\mu(t)$ are solved once before training, while $\tilde{r}(t,x)$ (Eq. 11b) and $G(t,x)$ (Eq. 10b) are computed during integration - all without requiring $a^{-1}$. This contrasts with the canonical score matching loss (Heng et al.) which contains $\Sigma(t,x)=a(t,x)$. Such inversions prove numerically unstable for near-singular $a$, and critically fail in our FitzHugh-Nagumo example where $a$ is singular. Our method remains applicable where score matching becomes undefined. To show more insight why the (neural) guided bridge works for hypo-elliptic diffusions, consider the SDE $dY_t = \{b(t, Y_t) + \sigma(Y_t)f(t, Y_t)\}dt + \sigma(Y_t)dW_t,$ and denote the law of $Y$ to be $\mathbb{Q}$. $\mathbb{P}^{\star}$ is absolutely continuous with respect to $\mathbb{Q}$ if there exists a bounded solution $\eta$ to the equation $b^\star(t,x)- b(t,x)-\sigma(x) f(t,x) = \sigma(x) \eta(t,x).$ Recall that $b^\star(t,x)=b(t,x) + \sigma(x) \sigma(x)^T \nabla_x \log h(t,x)$. Then the preceding display can be rewritten to $\sigma(x) \left( \sigma(x)^T\nabla_x \log h(t,x) -f(t,x)\right) = \sigma(x) \eta(t,x).$ Hence, one can easily see such an $\eta$ exists by the specific form of the additional drift. In this way, we circumvent inversion of $\sigma$. This also explains why the additional drift in the neural guided bridge contains the $\sigma$ premultiplication.

Questions for Authors

1. We cannot ensure this, similarly to other applications of Variational Inference. But some heuristic methods can be helpful in our case. For example, we can initialize the neural bridge differently, and test whether it learns the same drift term.
2. This "matching condition" appeared first in Theorem 1 of Schauer et al. This paper deals with the fully observed, uniformly elliptic case. It says that $\tilde\sigma$ should be $\tilde{a}(T)=a(T,v)$, see also our discussion in Section 3.3 of our paper. The claim that all examples in Section 5 satisfy the matching conditions needs adjustment. It is true, except for the Fitzhugh-Nagumo (FHN) model. Thank you for drawing our attention to this: we adjusted the formulation. The matching condition specifically refers to Assumption 2.4 in Bierkens et al. It consists of verifying 4 inequalities. As the diffusivity is constant, the 4th of the these is trivially satisfied. The first and third assumptions can be verified similar to Example 3.2 in Bierkens et al. (with $\Delta(t)=(T-t)^{-1}$). The second assumption in there is on the difference $b(t,x) -\tilde{b}(t,x)$. Inspecting the proof, it suffices that the second assumption only needs to hold true for $t=T$ and those $x$ for which $Lx=v$. With our choice of $\tilde b(t, x)$, it reads $b(T,x) -\tilde{b}(T,x)=0$, therefore the second assumption also suffices.
3. Yes, please refer to our reply to Review ysG7's Question 5.
4. We now simply choose $B$ and $\beta$ in a rather simple way to ensure the neural bridge will end up in the right point. We don't see any direct complication from parametrising these functions by a neural net, apart from computational resources. Presently, it is unclear if the additional training time is worth the effort.
5. Guided proposals in the case of multiple partial observations are discussed in Mider et al. We would start from this approach, and simply add a drift term on each of the segments between observations times, as we propose in the present paper for only one segment.

Claims and Evidence / Methods and Evaluation Criteria

• While neural bridges and MCMC-guided proposals differ methodologically—complicating direct cost comparisons—their forward simulation costs are comparable: for example, in the landmark process ($d=100$), the forward simulation time is 9.81ms (neural) vs 6.85ms (guided). Training complexities diverge as $O(Nd^3)$ (score-matching/adjoint) vs $O(N^2d^2)$ (neural), where $N$ is the number of time steps. Due to the space limitation, we can not include benchmark results here, but please refer to our reply to Reviewer 8sR1, and we also have added a more comprehensive summary into the manuscript per suggestion.

Other Comments or Suggestions

We have fixed the equation, thank you!

Questions for authors

1. We are unsure what you mean by ``discrete densities''. Proposition 3.2 requires $\nabla_x \log h(s,x)$ to be well defined. The simplest way to ensure this is to assume existence of smooth transition densities. Note however that we can write $h(t,x)= \int q(v\mid y) P(T, d y\mid t,x)$, where $P(T, d y\mid t,x)$ is the Markov kernel. This expression also makes sense if the kernel does not admit (smooth) densities with respect to some dominating measure. Informally, the case of conditioning on a state without noise corresponds to taking $q(v\mid y)$ like a Dirac mass in $v$ and then the above display should be interpreted as $h(t,x) = p(T,v\mid t,x).$ Clearly, existence of $\nabla_x \log h(s,x)$ requries that $h$ is strictly positive and that its gradient exists. Throughout the paper we have assumed that the distribution of $v$ conditional on $y$ is Gaussian. It is however possible to relax this assumption.
2. The non-noisy case is most challenging. Intuitively, the larger the noise, the less the process needs to be guided in a certain direction. So we tested the approaches on the more difficult case. We formulated Proposition 3.2 deliberately for the non-noisy case; in case we condition on the full state without noise, we can simply take $h(t,x) = p(T,v\mid t,x)$. However, in case of a partial observation, such as with the FitzHugh-Nagumo model, the form of $h$ reads more pleasantly when noise is assumed. We refer to Section 1.3.2. of Bierkens et al. for the somewhat more involved formulas that appear when no noise is assumed. So the motivation is merely to present a ``clean'' statement. In the examples, we assume $q(v\mid y) = \psi(v; L y, \epsilon^2 I)$, where $\epsilon$ is very small. For example, in the FHN-example $L=[1, 0]$, which corresponds to only observing the first component. Taking $\epsilon$ nonzero makes guided proposals also numerically better behaved near the time of conditioning.
3. The drift of the bridge behaves in a very specific way near the point we condition on. This is most easily seen in the uniformly elliptic case and when we condition on the full state $v$ at time $T$. Essentially, the drift of the true bridge behaves for $t\approx T$ as $(v-x)/(T-t)$. Failure to replicate this behaviour breaks absolute continuity between true and proposal bridge laws. Our guided proposal network explicitly replicates this asymptotic behavior. The neural bridge's auxiliary drift (confined to $\sigma$'s range and bounded) enhances proposals on $[0, T-\eta)$ for small $\eta$, preserving absolute continuity when added. Empirical validation appears as Fig. 6 in Bierkens et al.
4. In general yes. The additional learned drift in our method is the difference between the true score $\nabla_x\log h(s, x)$ and the proposed score $\nabla_x\log \tilde{h}(s, x)$ (with $\sigma$ as a scaling). The objective function is based on a closed-form expression for the likelihood ratio between the target measure $\mathbb{P}^{\star}$ and the proposal measure $\mathbb{P}^{\circ}$ (As Eq. (10) in the paper). The key challenges in applying our method are to select $\mathbb{P}^{\circ}$ that satisfies two criteria: the likelihood ratio between the target $\mathbb{P}^{\star}$ and $\mathbb{P}^{\circ}$ is tractable, and sampling paths under $\mathbb{P}^{\circ}$ can be done efficiently.
5. Yes, mode collapse occurs in both cell diffusion and FitzHugh-Nagumo (FHN) experiments, showing multimodal marginal distributions. This stems from the forward KL divergence's mode-seeking objective. We've also considered reverse KL alternatives, but it will introduce problematic stochastic integrals and unstable optimization. Additionally, we see three possible mode collapse explanations: (i) local optimum convergence; (ii) misspecified endpoint guidance from $\tilde X$; (iii) Oversimplified $\vartheta$. We conjecture (i) to be most likely. We implemented strategies: (i) optimizer/hyperparameter sweeps (Adam/RMSprop/SGD); (ii) dropout regularization. Neither approach yielded significant improvements. We hypothesize that the presence of $G(t,x)$ in the objective function introduces significant non-convexity, rendering the optimization landscape inherently challenging.