JAMUN: Transferable Molecular Conformational Ensemble Generation with Walk-Jump Sampling

Abstract: Conformational ensembles of protein structures are immensely important both to understanding protein function, and for drug discovery in novel modalities such as cryptic pockets. Current techniques for sampling ensembles are computationally inefficient, or do not transfer to systems outside their training data. We present walk-Jump Accelerated Molecular ensembles with Universal Noise (JAMUN), a step towards the goal of efficiently sampling the Boltzmann distribution of arbitrary proteins. By extending Walk-Jump Sampling to point clouds, JAMUN enables ensemble generation at orders of magnitude faster rates than traditional molecular dynamics or state-of-the-art ML methods. Further, JAMUN is able to predict the stable basins of small peptides that were not seen during training.

• The paper introduces JAMUN, leveraging walk-jump sampling to generate transferable molecular conformational ensembles with superior performance over standard MD.
• It employs point-cloud representations and SE(3)-equivariant neural networks to capture spatial symmetries, enhancing sampling accuracy.
• Experimental results demonstrate faster convergence and robust predictions of stable conformational basins, even for peptides outside the training set.

Overview of JAMUN: Transferable Molecular Conformational Ensemble Generation with Walk-Jump Sampling

The study presented in this paper explores the development and implementation of JAMUN, a model designed to generate conformational ensembles of molecular structures, especially proteins. This task is of significant importance in both understanding protein dynamics and facilitating drug discovery. Traditional approaches, such as molecular dynamics (MD), suffer from inefficiencies and limitations in terms of system transferability, motivating the exploration of advanced ML models like JAMUN.

Main Contributions

The paper introduces JAMUN (walk-Jump Accelerated Molecular ensembles with Universal Noise), leveraging an extension of Walk-Jump Sampling (WJS) to address these challenges. Built on point-cloud representations of molecular structures, JAMUN enables faster and more efficient generation of conformational ensembles compared to standard MD and contemporary ML techniques. The distinctive contribution of JAMUN lies in its demonstrated ability to predict stable conformational basins for small peptides not included in its training set, highlighting its transferability.

Walk-Jump Sampling Paradigm

WJS operates by applying noise to the input data and training a denoising network to recover the original structure. This decouples the walk (sampling) and jump (projection back to the original distribution) steps, performed using Langevin dynamics. Unlike many ML models which transform the data into a simple Gaussian distribution, JAMUN uses a denoiser to ensure the generative process remains efficient and physically interpretable.

Transferability and Results

The experiments conducted demonstrate remarkable transferability of JAMUN on datasets of two amino acid peptides. The model achieves convergence in generating conformational ensembles at speeds dramatically faster than conventional MD. The introduction of noise levels ($\sigma = 0.4 \si{\angstrom}$ for capped peptides and $\sigma = 0.6 \si{\angstrom}$ for uncapped peptides) is a pivotal innovation, smoothing the manifold and enabling more rapid sampling.

Comparisons against Transferable Boltzmann Generators (TBG) further validate JAMUN’s superior performance, showing enhanced efficiency and coverage in sampling important conformational states across molecular structures.

Methodological Insights

The paper outlines a rigorous methodological approach encompassing the assignment of physical priors informed by MD data. The application of $SE(3)$-equivariant neural networks within the denoising architecture stands out as a means to appropriately capture the spatial symmetries of molecular structures. This ensures better accuracy and efficiency in generating molecular conformations.

Implications and Future Directions

The implications of JAMUN extend into both practical and theoretical domains. Practically, the model provides a more efficient tool for simulating molecular dynamics, with potential applications in drug design, such as identifying cryptic pockets or optimizing antibody conformations. Theoretically, JAMUN serves as a basis for future exploration into larger biological systems, possibly integrating with enhanced sampling techniques to address the complexities of protein dynamics.

Future developments could focus on extending the denoising architecture or exploring alternative noise handling techniques to further refine sampling accuracy. Moreover, the paper hints at the possibility of incorporating classical enhanced sampling methods to improve the traversal of the noisy space.

In summary, JAMUN represents a significant advancement in the efficient simulation of molecular conformational ensembles. Its adaptability and speed offer a promising avenue for research and practical application in molecular dynamics and related fields, fostering further innovation in the study of protein dynamics and drug discovery efforts.