Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling
Authors: Yuanqi Du, Michael Plainer, Rob Brekelmans, Chenru Duan, Frank Noe, Carla P. Gomes, Alan Aspuru-Guzik, Kirill Neklyudov
NeurIPS 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We investigate the capabilities of our approach across a variety of different settings. We first illustrate features of our method on toy potentials before continuing to real-world molecular systems, including a commonly-used benchmark system, alanine dipeptide, and a small protein, Chignolin.Our method requires fewer energy evaluations (1 million vs. 1 billion) while finding paths with similar energy and likelihood. |
| Researcher Affiliation | Collaboration | 1Cornell University 2Zuse School ELIZA 3Technische Universität Berlin 4Freie Universität Berlin 5Berlin Institute for the Foundations of Learning and Data 6Vector Institute 7Massachusetts Institute of Technology 8Deep Principle, Inc. 9Rice University 10Microsoft Research AI4Science 11University of Toronto 12Université de Montréal 13Mila Quebec AI Institute |
| Pseudocode | Yes | Algorithm 1: Training (Single Gaussian) Input: Reference drift bt, diffusion matrix Gt Conditioning endpoints while not converged do Sample t U(0, T) Sample xt q(θ) t|0,T using (15) Calculate u(q,θ) t|0,T (xt) using (13) Calculate v(q,θ) t|0,T (xt) using u(q,θ) t|0,T (xt), bt(xt), (14) Calculate L = v(q,θ) t|0,T (xt), Gt v(q,θ) t|0,T (xt) (Thm. 1) Update θ optimizer(θ, θL) end while return θ. Algorithm 2: Sampling Trajectories def get_drift(xt, t): Evaluate µ(θ) t|0,T , Σ(θ) t|0,T at t return drift u(q,θ) t|0,T (xt) using (13) Sample initial state x0 N(A, σ2 min) return SDESolve(x0, get_drift, T)Algorithms for training with a single Gaussian path (Alg. 1) and sampling or generating transition paths at test time (Alg. 2). Note that we sample from the marginals qt|0,T during training, but generate paths by simulating the SDE (10). |
| Open Source Code | Yes | The code behind our method is available at https://github.com/plainerman/variational-doob. |
| Open Datasets | Yes | We investigate the capabilities of our approach across a variety of different settings. We first illustrate features of our method on toy potentials before continuing to real-world molecular systems, including a commonly-used benchmark system, alanine dipeptide, and a small protein, Chignolin.The underlying Müller-Brown potential that has been used for our experiments can be written as U(x, y) = 200 exp (x 1)2 10y2 100 exp x2 10 (y 0.5)2 170 exp 6.5 (0.5 + x)2 + 11 (x + 0.5) (y 1.5) 6.5 (y 1.5)2 + 15 exp 0.7 (1 + x)2 + 0.6 (x + 1) (y 1) + 0.7 (y 1)2. |
| Dataset Splits | No | The paper describes the training procedure and data generation, but does not specify explicit training/validation/test splits for a pre-existing dataset in the conventional sense. Training involves sampling from distributions defined by the model, rather than splitting a fixed dataset. |
| Hardware Specification | Yes | All our experiments involving training were conducted on a single NVIDIA A100 80GB. The baselines themselves were computed on a M3 Pro 12-core CPU. |
| Software Dependencies | No | The paper mentions software like AMBER14, Open MM, DMFF, and JAX, but does not provide specific version numbers for these dependencies. |
| Experiment Setup | Yes | We used a first-order Euler integration scheme to simulate transition paths with 275 steps and a dt of 10 4s. ξ was chosen to be 5 and 1,000 transition paths were simulated. We have used an MLP with four layers and a hidden dimension of 128 each, with swish activations. It has been trained for 2,500 steps with a batch size of 512 and a single Gaussian.For the concrete simulations, we ran them with the timestep dt = 1fs, with T = 1ps, γ = 1ps, and Temp = 300K.Neural Network Parameterization. We parameterize our model with neural networks, a 5-layer MLP with Re LU activation function and 256/512 hidden units for alanine dipeptide and Chignolin, respectively. The neural networks are trained using an Adam optimizer with learning rate 10 4. |