Rare event sampling in dynamical systems is a fundamental problem arising in the natural sciences, which poses significant computational challenges due to an exponentially large space of trajectories. For settings where the dynamical system of interest follows a Brownian motion with known drift, the question of conditioning the process to reach a given endpoint or desired rare event is definitively answered by Doob's h-transform. However, the naive estimation of this transform is infeasible, as it requires simulating sufficiently many forward trajectories to estimate rare event probabilities. In this work, we propose a variational formulation of Doob's h-transform as an optimization problem over trajectories between a given initial point and the desired ending point. To solve this optimization, we propose a simulation-free training objective with a model parameterization that imposes the desired boundary conditions by design. Our approach significantly reduces the search space over trajectories and avoids expensive trajectory simulation and inefficient importance sampling estimators which are required in existing methods. We demonstrate the ability of our method to find feasible transition paths on real-world molecular simulation and protein folding tasks.
Paper
Similar papers
Peer review
Summary
The paper is concerned with sampling trajectories with a terminal condition. For stochastic processes governed by a Brownian Motion, Doob's h-transform gives an posterior SDE that leads to samples with the final condition. However, estimating the h-function that is needed for the posterior SDE usually involves simulating trajectories, which is inefficient if the terminal condition is rarely reached. The authors propose a simulation-free variational optimization method to estimate the h-function based on a least action principle and Gaussian approximations to the marginal densities.
Strengths
1. The authors provide a solid and clear background on Doob's h transform that gets supported by the provided proofs in the appendix. 2. The paper clearly highlights the challenges of the optimization problem and proposes an efficient solution adressing these challenges. 3. The related work section gives a good overview and nicely connects to related topics.
Weaknesses
1. I found the path histograms in Figure 2 to be too cluttered. I would propose adding less samples. 2. The paper misses a learning curve. It would in general be interesting to have more training details. Minor Weaknesses: In Chapter 3 and in the appendix, the authors change from trajectory length T to the unit interval. There should be a sentence that explains this change.
Questions
I could not totally follow the derivation in the appendix. Can you explain how to get to Equation 25 from the previous equation in line 548? I do not see why second term in equation 25 gets subtracted, while in line 548 only the last term has a minus.
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
2
Limitations
The authors adequately adressed the limitations of their work.
Summary
This paper proposes a variational formulation of Doob's h-transform, which characterizes the distribution over paths with a given endpoint. Instead of relying on potentially wasteful sampling approaches, the authors propose directly optimizing a tractable variational distribution over transition paths which satisfy the initial and terminal conditions by design. This approach reduces the search space over trajectories and avoids trajectory simulation. Experiments on real-world molecular simulation and protein folding tasks demonstrate the applicability of the approach.
Strengths
The paper clearly motivates the problem it tackles, describes the challenges well and nicely introduces the idea behind their method in an illustrative fashion. Overall the paper is very well-written and structured. In particular, it builds up the method piece-by-piece explaining the choices along the way. As far as I can judge the related work seems to be exhaustive. The experiments are done on interesting problems as far as I can tell and I also appreciated that the authors made their code publicly available.
Weaknesses
My main concern comes from trying to interpret the experimental results and in particular judging the performance compared to MCMC. It seems quite clear from Tables 1 and 2 that the variational approach requires fewer calls to the potential energy function than MCMC, however their performance differences are harder to judge in my opinion. More specifically, in Table 1 the standard deviations are so large that there is basically no meaningful statistical difference between the shown results especially for the Max Energy, but also the Log-likelihood. Now this might mean that MCMC and the presented method both perform well, but its surprising to me that there is that much inherent variation. Similarly in Table 2, the variance for the Max Energy of MCMC in the first line is huge. I'm also a bit confused by the Max Energy increases when using a mixture in Table 2. To better understand how the performance of the variational approach improves during training, it would be nice to see a plot that shows the Max Energy as a function of the training epochs. Finally, I would have liked to see a discussion of the limitations of the approach. Section 6 has Limitations in the title, but does not actually discuss them in any way beyond extensions of the proposed method.
Questions
- How do you choose the number of mixture components in practice? - Why does the MaxEnergy increase when using a mixture distribution as a variational approximation in Table 2? - How do you explain the huge variance of the Max Energy of MCMC (variable length) in the first row of Table 2? - What are the main limitations of the current approach?
Rating
7
Confidence
2
Soundness
3
Presentation
4
Contribution
3
Limitations
Section 6 has "Limitations" in the title, but limitations are not discussed, only future work.
Summary
The Authors of this paper tackle the problem of sampling conditioned SDEs with a specific interest in "transition path sampling", i.e. sampling a Langevin-type SDE undergoing a transition between an initial state (or set of states) $A$ and a final target set $B$. Sampling transition paths efficiently can provide a tremendous boost to research in catalysis or drug design. In this paper, the Authors model the transition path distribution with either (i) a parametrized Gaussian process or (ii) a parametrized mixture of Gaussians processes. These families of priors are then optimized by leveraging a variational formulation of the problem developed by the Authors starting from Doob's $h$-transform.
Strengths
The Authors tackle a challenging problem with a novel variational formulation, which is apt to be optimized by leveraging techniques developed by the generative modeling community within ML. A clear strength of the paper is the sound theoretical analysis justifying the proposed method. Interestingly, the choice of Gaussian process (or a mixture of Gaussian processes) variational priors allows the Authors to simplify the algorithm using analytical results and sidestep lengthy calculations at the deployment phase.
Weaknesses
I find the paper overall clearly written, but I had to go through section 3.2 (Computational Approach) several times to grasp how the method can be deployed in practice. Specifically, I think that it can be improved the explanation of the fact that, because of the Gaussian prior, you only need to model the transition path probability and not the $h$-transform itself, as well as the explanation of the actual optimization step (Reparametrization of Gradients). The subscript $0,T$ is introduced for the first time in Eq. (6) without explanation and used throughout the manuscript to label variables related to the conditioned process. It might be useful to clarify this notation explicitly. The experimental evaluation is somewhat limited, especially concerning the Chignolin protein. Over the years, many transition path sampling strategies have been developed, as well as very much related enhanced sampling techniques. It would be nice to have a comparison also to different baselines, as well as a discussion on the computational complexity and actual running times of the baselines.
Questions
Is the $h$-transform related to the committor function in some way? If so, how does your Thm 1 relates to the variational formulation of the committor (see e.g. Eq. 20 of "Transition Path Theory and Path-Finding Algorithms for the Study of Rare Events" by Weinan E and Eric Vanden-Eijnden)
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
Limited experimental evaluation
Summary
The submitted manuscript presents a variational formulation of Doob’s h-transform, leading to a novel (simulation-free) computational approach for rare event sampling in transition paths. The task of interest involves conditioning a dynamical system driven by Brownian motion with a known drift term to reach a given endpoint. In theory, this terminal condition can be addressed using Doob’s h-transform, resulting in an associated SDE for the conditional dynamic. However, practical implementation requires knowledge of the h-transform. The authors propose a variational problem formulation that provides the necessary information about the h-transform as its solution. Solving this variational problem results in a computational approach for simulating the conditional dynamic. Compared to existing methods, it is claimed that the proposed approach avoids importance sampling estimators and expensive trajectory simulations. The method is tested on both synthetic and real datasets.
Strengths
The paper introduces a novel variational objective to describe Doob’s h-transform, leading to a promising and innovative computational approach for simulating the conditional dynamics of transition paths. Although I appreciate the general approach presented in Section 3.1, I am not convinced by the computational approach proposed in Section 3.2 due to several reasons pointed out below. If these issues can be addressed, I believe that the approach can be further utilized to construct highly efficient sampling strategies.
Weaknesses
- The proposed approach in Section 3.2.1 introduces certain issues. The connection to Doob’s h-transform established in Theorem 1 is only guaranteed if an exact solution to the optimization problem (9) or (11) is found. However, by introducing the Gaussian parametrization of $q_{t|0,T}$, this guarantee is lost. There is a lack of discussion regarding the implications and potential effects of this parametrization on the overall accuracy and validity of the method. - Moreover, the formulation of the task of solving equation (12) given $q_{t|0,T}$ is misleading. In reality, you are solving an inverse problem here. Given the solution of the PDE in (12), your goal is to reconstruct $u_{t|0,T}$, which is generally an ill-posed problem. For instance, introducing uncertainties into the model description via the drift $b_t$ or diffusion $\Xi_t$ can lead to significant challenges. This scenario is likely when applying the approach in practice under model misspecification. A critical question to address is the robustness of the proposed framework against model misspecification, such as small perturbations in $b_t$. - The claim that the approach is simulation-free is not entirely clear. In line 173, it is stated that the Gaussian parametrization allows for the generation of arbitrary samples. However, this assumes that the states $x_{t|0,T}$ are independent for all $t$, which is not true when $x_{t|0,T}$ is a solution of the SDE in (10). The dependence between the states needs to be accounted for in the sampling process, and this aspect seems to be overlooked. - I am wondering what the measured quantities in the numerical experiments actually reveal about the correctness of the proposed approach. Evaluating the maximum energy and the likelihood might not provide sufficient information about the accuracy of the estimated distribution. If the sampled paths only follow high likelihood regions, this does not necessarily indicate that the correct distribution has been captured. The objective should be to estimate quantities of interest related to the conditioned SDE in (6). However, it is not clear whether the simulated paths correspond to accurate simulations of (6). For instance, how does the method perform when estimating rare event probabilities or other Monte Carlo estimators with respect to (6)? This aspect needs to be thoroughly addressed to validate the effectiveness of the proposed approach.
Questions
- To enhance the practical applicability of the proposed framework, it is crucial to study the impact of inexactly solving the variational problems (9) or (11). Specifically, how do errors propagate to the conditioned dynamical system driven by the SDE (6) when an approximate solution of (9) or (11) is used? Understanding this error propagation is essential for assessing the robustness and reliability of the proposed method in practical scenarios. - Is the Optimization problem in (9) well-posed? This means, is the optimal solution unique? - In several instances, there are missing commas in the notation for the inner product (e.g., equations (9b) and (12)). - Many important mathematical details and assumptions are missing. Most importantly, what are the assumptions on the drift vector field $b_t$ to ensure the well-posedness of the proposed scheme? - Regarding Section 3.2.2: Instead of introducing $\xi_{\min}$ I suppose that one could also directly work with a pseudoinverse when defining $G_t^{-1}$. - In Proposition 4, should it be $u_{t|0,T}^{(k)}$ on the right hand side of the equation for $u$? - When referring to probability density functions, it's important to use proper notation. Instead of $\rho(x_t=x)$, it would be clearer to use $\rho_t(x)$ depending on the context. - In the abstract, you claim that no "inefficient" importance sampling estimators are required. However, it would be beneficial to compare this approach with such estimators to demonstrate its advantages more clearly.
Rating
6
Confidence
3
Soundness
2
Presentation
3
Contribution
3
Limitations
- There is no discussion about error propagation arising due to the Gaussian parametrization. - The approach is limited to conditioning the sample path on point sets $x_T = B$. - Limited statistical evidence is presented in the numerical experiments.
References
[1] Blei, Kucukelbir, McAuliffe. “Variational Inference: A Review for Statisticians”, 2016. [2] Jónsson, H., Mills, G. and Jacobsen, K.W., 1998. Nudged elastic band method for finding minimum energy paths of transitions. In Classical and quantum dynamics in condensed phase simulations (pp. 385-404). [3] Weinan, E., Ren, W. and Vanden-Eijnden, E., 2004. Minimum action method for the study of rare events. Communications on pure and applied mathematics, 57(5), pp.637-656. [4] C. Dellago, P. G. Bolhuis, and P. L. Geissler, 2006, “Transition Path Sampling Methods,” Computer Simulations in Condensed Matter Systems: From Materials to Chemical Biology Volume 1. Springer Berlin Heidelberg, pp. 349–391. [5] Brekelmans, Rob, and Kirill Neklyudov, 2023, "On Schrödinger Bridge Matching and Expectation Maximization." In NeurIPS 2023 Workshop Optimal Transport and Machine Learning. [6] B. Jamison, 1975, „The Markov processes of Schrödinger“, Zeitschrift für Wahrscheinlichkeitstheorie und Verwandte Gebiete, Bd. 32, Nr. 4. Springer Science and Business Media LLC, S. 323–331. [7] Holdijk, L., Du, Y., Hooft, F., Jaini, P., Ensing, B. and Welling, M., 2024. Stochastic optimal control for collective variable free sampling of molecular transition paths. Advances in Neural Information Processing Systems, 36.
Summary
This paper proposes a variational formulation of the Doob's h-transform to change the problem from the expensive simulations of trajectories to an optimization problem over possible trajectories from given initial point and end point. The model parameterization introduced imposes the desired boundary conditions by design and uses a mixture of Gaussian as the variational family. The variational formulation offers an alternative to expensive simulations and MCMC which is hard to scale. The proposed framework is applied on a simulated case study and real world problem in material science/computational chemistry of transition path sampling which is a study of molecule transitions betweeb local energy minima and metastable states under random fluctuations. The results are along expected lines where it is possible to scale the solution on computationally intractable problems with MCMC and other sampling based methods.
Strengths
1. The paper is well written, with sound math, a lot of references on recent work and older literature on the subject. I found the introduction to be a great summary of the work. 2. The idea is sound and through the use of variational framework to solve a complex problem of estimating rare event probabilties by sampling forward trajectories brings out the desired qualities of sample efficiency, reducing search space and matching a complex target distribution with an easier to sample variational distribution. As also observed in Table1 and 2, MCMC cannot be scaled to certain large scale experiments where Variational inference can give good results. 3. The work finds application for solving transition path sampling problem which is important in material science and chemistry from what it seems(I am not so well acquanited with those domains.) The case study on protein is well documented and explained. 4. The figures and illustrations are clean and support the narrative.
Weaknesses
Minor things - Increase font sixe for Table1, bold the important results. - I would have liked to see a bit more dicussion on the dimensionality, what are the typical values of D ..
Questions
1. I am curious, similar to Figure 3, the authors can show an experiment with familiy of dtistributions other than Gaussian. How does the expressivity then come to play ? 2. as for any SDE problem solution, what is the effect of discretization values and parameters in practice as given in line 184-185. 3. How big is the constraint of fixed length(T) on transition paths in practice as solved in this method ? 4. What is the behaviour of the problem when there are multiple basins of attraction ?
Rating
7
Confidence
3
Soundness
3
Presentation
4
Contribution
3
Limitations
The authors have addressed the challenges quite well and what the future research directions could be taken.
I thank the Authors for their thorough rebuttal and clarifications. I will keep my evaluation to "Accept".
Thanks for the response
Thank you very much for the detailed response. I might be misunderstanding the entire simulation-free approach, but don’t you sample independently at each time $t$ from the marginal distributions to generate paths of the conditioned SDEs? (Please correct me when I am misunderstanding something) While this sampling scheme might be consistent under certain assumptions on the drift and diffusion coefficients, it certainly isn’t universally applicable. Generally, how do you ensure the continuity of the sampled paths? This concern remains relevant regardless of whether you employ Gaussian or Gaussian mixture parameterizations. I would be less critical of my concern if there were clear empirical experiments to support the correct sampling distribution. In particular, the synthetic data example presents an opportunity to conduct a statistically robust case study. Additionally, I remain concerned about the lack of explicit mathematical assumptions and the challenges associated with conditioning on point sets. Doesn’t it impose technical challenges when applying Doob‘s h-transform for conditioning on null sets? For example, the application of Jamison (1975) requires a positive function h (as also noted in line 514). While I find the proposed framework very promising, I am maintaining my score due to the concerns I've outlined.
Clarifying Misunderstanding of Test vs. Train Sampling
**Testing vs. Training** > I might be misunderstanding the entire simulation-free approach, but don’t you sample independently at each time 𝑡 from the marginal distributions to generate paths of the conditioned SDEs? Simulation-free refers to our *training* method, where our objective in Thm 1 only requires samples from the time-marginals $q_{t|0,1}$ of the conditioned SDE. This justifies neglecting full trajectory information and sampling directly from $q_{t|0,1}$ (without SDE simulation) during training. See Algorithm 1 of the general-response PDF or cell 20 of the [anonymized code](https://anonymous.4open.science/r/TPS-Doob-843E/notebooks/tps_gaussian.ipynb). > … don’t you sample independently at each time $𝑡$ from the marginal distributions to generate paths of the conditioned SDEs? Generally, how do you ensure the continuity of the sampled paths? We **do not** sample each marginal independently to generate paths of the conditioned SDE. Instead, our Algorithm 2 in the general-response PDF outlines our sampling approach. For a given set of time marginals $q_{t|0,1}$, we can calculate the appropriate drift $u_{t|0,1}$ for an SDE with given diffusion coefficients. **At generation (test) time, we simulate this SDE**: $dx_{t|0,1} = u_{t|0,1}(x_{t|0,1}) dt + \Xi_t dW_t$ starting from $x_0 \sim \mathcal{N}(A, \sigma_{\text{min}} \mathbb{I})$. Thus, our sampled paths are continuous (up to the SDE solver error). Note that $u_{t|0,1}$ implies a control drift term $v_{t|0,1}$ (in Eq. 10) as in Eq. 14, but we do not need to evaluate the costly drift term $b_t$ at generation time. See cells 24-25 of [anonymized code](https://anonymous.4open.science/r/TPS-Doob-843E/notebooks/tps_gaussian.ipynb). > (Please correct me when I am misunderstanding something) Thank you for pinpointing this misunderstanding! We hope that the above points help to address these concerns. Based on your feedback, we agree that the following changes will greatly improve the final manuscript: - Discussion of the sampling procedure (and Algorithm box) will be included in the main text - Distinctions between the time marginals $q_{t|0,1}$ of the conditioned process and the full path measure or SDE will be carefully emphasized throughout. - We will visualize individual trajectories in Figure 2 to show that the paths are continuous and similar to MCMC trajectories. - We will modify line 47 to “our training method is simulation-free”. It appears that all other usage of “simulation-free” explicitly refers to training or optimization, rather than generation. **Evaluating Sampling Trajectories** > I would be less critical of my concern if there were clear empirical experiments to support the correct sampling distribution. Thank you for this suggestion. For the Müller-Brown experiment in Figure 2, we report Wasserstein-1 distances between (i) samples from expensive MCMC simulation (which we treat as ground-truth) and (ii) sample trajectories generated using Algorithm 2 with our learned model. We report the mean and std of the W1 distance across discrete $t$: | Wasserstein W1| Value | |-----------------------|-------| | Mean | 0.1251 | | Std | 0.0392 | | Median | 0.1130 | | Min | 0.0393 | | Max | 0.2115 | We will report further numbers and plots in the revised version of the manuscript. **Conditioning on Dirac Deltas** > …challenges associated with conditioning on point sets. Doesn’t it impose technical challenges when applying Doob‘s h-transform for conditioning on null sets? In the case of the conditioning on a point mass $\delta(x-B)$, the h-function becomes a density, i.e. $h(y,t) := \rho(x_T = B | x_t = y)$ is the density of the transition probability $\mathbb{P}(x_T \in dx| x_t = y)$. Conditioning on point sets is commonly used [1 Thm 7.11,2,3], and indeed, all the derivations in Appendix A hold if we take $h(y,t):=\rho(x_T = B | x_t = y)$. We apologize for the confusion, we will clarify it in the next version. [1] Särkkä, Simo, and Arno Solin. Applied stochastic differential equations. Vol. 10. Cambridge University Press, 2019. [2] Heng, J., De Bortoli, V., Doucet, A. and Thornton, J., 2021. Simulating diffusion bridges with score matching. arXiv preprint arXiv:2111.07243. [3] Liu, Xingchao, Lemeng Wu, Mao Ye, and Qiang Liu. "Let us build bridges: Understanding and extending diffusion generative models." arXiv preprint arXiv:2208.14699 (2022).
Raising the score
Thank you for your prompt response and thorough clarification of my concerns. I am happy to raise my score to 6. I did not give a higher score because I believe further effort is needed to enhance the presentation of the proposed framework by incorporating the points discussed in the rebuttal.
Decision
Accept (spotlight)