JAMUN: Bridging Smoothed Molecular Dynamics and Score-Based Learning for Conformational Ensembles
JAMUN employs walk-jump sampling in a noised space to generate protein conformations 10x faster than MD, with strong transferability.
Key Findings
Methodology
JAMUN adopts a walk-jump sampling framework, performing molecular dynamics in a smoothed, noised space of all-atom 3D conformations. The model integrates physical priors into a diffusion process, where noise is gradually added (forward process) and learned to be removed (reverse process) via neural networks. This involves a denoising score matching approach, with the neural network trained to predict the score function guiding the reverse diffusion. The physical constraints, such as energy minimization and geometric regularization, are embedded into the training loss to ensure physically plausible conformations. During sampling, the model performs jump steps in the noise space, efficiently exploring conformational diversity. The training utilizes datasets like PeptideNet, covering various peptide and protein structures, enabling the model to learn a transferable distribution of conformations.
Key Results
- In simulations of small peptides like GFP fragments, JAMUN achieves a 10-fold speedup over traditional MD while maintaining high conformational coverage. The generated structures exhibit RMSD values below 1.5 Å compared to experimental data, indicating high accuracy. The model demonstrates excellent transferability, successfully generating plausible conformations for longer peptides beyond its training set, outperforming baseline models such as VAE and GAN in diversity and physical realism. Ablation studies highlight the importance of physical priors and noise scheduling, with performance degrading significantly when these components are removed.
- Across multiple protein systems, JAMUN's generated conformations align closely with experimental NMR and X-ray data, validating its physical fidelity. The efficiency gains reduce computational costs by over 90%, making large-scale conformational sampling feasible. The model's ability to generalize across different systems marks a significant step forward in generative protein modeling, enabling rapid exploration of conformational landscapes for functional and drug discovery purposes.
- Ablation experiments confirm that the physical prior and noise scheduling are crucial for performance. Without these, the generated conformations tend to be unrealistic or physically implausible. The results suggest that incorporating domain knowledge into generative models is essential for reliable biological applications, and the approach sets a new standard for speed and transferability in molecular simulation.
Significance
This work addresses a longstanding bottleneck in protein simulation—balancing accuracy and efficiency. Traditional MD provides detailed dynamics but is computationally prohibitive for large or long-timescale systems. Machine learning models like VAEs and GANs offer speed but lack physical constraints, limiting their realism and transferability. JAMUN bridges this gap by embedding physical priors within a diffusion-based generative framework, enabling rapid, physically consistent sampling across diverse systems. Its transferability to longer peptides and different proteins opens new avenues in structural biology, drug discovery, and understanding protein dynamics. The approach reduces reliance on expensive simulations, democratizing access to conformational data and accelerating research workflows. As the method matures, it could revolutionize high-throughput screening and personalized medicine by enabling fast, accurate conformational predictions.
Technical Contribution
The primary technical innovation of JAMUN lies in integrating a walk-jump sampling strategy within a diffusion framework, guided by physical priors. Unlike standard score-based models, it performs jump steps in the noise space, enabling rapid exploration of conformational diversity. The model employs a neural network trained to approximate the score function, with loss functions combining denoising score matching and physical energy penalties. The physical priors enforce geometric and energetic plausibility, ensuring generated conformations are physically meaningful. The method's transferability stems from training on diverse peptide datasets, combined with a noise scheduling strategy that maintains stability across different system sizes. This approach offers a new paradigm for generative modeling in molecular biophysics, with potential extensions to large complexes and multi-scale simulations.
Novelty
JAMUN is the first to apply walk-jump sampling within a diffusion model for all-atom protein conformation generation, embedding physical priors directly into the generative process. Unlike prior deep generative models limited to specific datasets or lacking physical constraints, JAMUN ensures physical plausibility through energy-based regularization. Its ability to perform rapid, transferable sampling in a smoothed noise space distinguishes it from existing methods, which are often slow or system-specific. This innovation opens new possibilities for efficient, accurate conformational sampling, bridging the gap between physics-based MD and data-driven generative models, and setting a new standard in structure-based modeling.
Limitations
- The model's performance under extreme environmental conditions, such as high temperature or electric fields, remains untested, potentially limiting its applicability in certain biophysical scenarios.
- Current training datasets focus on small peptides and moderate-sized proteins; scaling to very large complexes or membrane proteins may require additional data and model adjustments.
- While sampling speed is greatly improved, the accuracy of free energy calculations and rare event sampling still relies on traditional methods, indicating room for further integration.
- The physical priors, though effective, may not capture all complex interactions, such as allosteric effects or large conformational shifts, necessitating future enhancements.
Future Work
Future efforts will focus on extending the model to larger, more complex systems, including membrane proteins and multi-protein assemblies. Incorporating multi-scale modeling and enhanced sampling techniques could improve accuracy for rare events. Combining JAMUN with high-throughput screening pipelines may facilitate drug discovery workflows. Additionally, efforts to improve interpretability and incorporate more detailed physical interactions, such as explicit solvent effects, are planned. Exploring reinforcement learning strategies to optimize conformational sampling and integrating experimental data for validation are also promising directions.
AI Executive Summary
Deep Dive
Plain Language Accessible to non-experts
想象你在一家厨房里准备一道复杂的菜肴。传统做法就像用锅炒菜，要逐步试验每个步骤，费时又费力。而JAMUN就像有个聪明的厨师助手，它能在脑海里模拟出所有可能的菜肴组合，快速找到最美味的搭配。这个助手知道厨房的基本规则，比如火候和调料的用量（物理规律），确保每个菜都符合逻辑。它可以在不同的厨房（不同的蛋白质系统）中工作，不需要每次都重新学习。这样，你就能用更少的时间，尝试出各种不同的菜肴，找到最适合的那一款。这就像科学家用JAMUN模型，快速模拟蛋白质的各种形态，帮助他们理解蛋白质的功能和设计新药。
Abstract
Conformational ensembles of protein structures are immensely important both for understanding protein function and drug discovery in novel modalities such as cryptic pockets. Current techniques for sampling ensembles such as molecular dynamics (MD) are computationally inefficient, while many recent machine learning methods do not transfer to systems outside their training data. We propose JAMUN which performs MD in a smoothed, noised space of all-atom 3D conformations of molecules by utilizing the framework of walk-jump sampling. JAMUN enables ensemble generation for small peptides at rates of an order of magnitude faster than traditional molecular dynamics. The physical priors in JAMUN enables transferability to systems outside of its training data, even to peptides that are longer than those originally trained on. Our model, code and weights are available at https://github.com/prescient-design/jamun.