Because the capabilities of generative AI models have grown, you have probably seen how they’ll transform easy text prompts into hyperrealistic images and even prolonged video clips.
More recently, generative AI has shown potential in helping chemists and biologists explore static molecules, like proteins and DNA. Models like AlphaFold can predict molecular structures to speed up drug discovery, and the MIT-assisted “RFdiffusion,” for instance, may also help design recent proteins. One challenge, though, is that molecules are consistently moving and jiggling, which is significant to model when constructing recent proteins and medicines. Simulating these motions on a pc using physics — a method generally known as molecular dynamics — will be very expensive, requiring billions of time steps on supercomputers.
As a step toward simulating these behaviors more efficiently, MIT Computer Science and Artificial Intelligence Laboratory (CSAIL) and Department of Mathematics researchers have developed a generative model that learns from prior data. The team’s system, called MDGen, can take a frame of a 3D molecule and simulate what’s going to occur next like a video, connect separate stills, and even fill in missing frames. By hitting the “play button” on molecules, the tool could potentially help chemists design recent molecules and closely study how well their drug prototypes for cancer and other diseases would interact with the molecular structure it intends to affect.
Co-lead creator Bowen Jing SM ’22 says that MDGen is an early proof of concept, but it surely suggests the start of an exciting recent research direction. “Early on, generative AI models produced somewhat easy videos, like an individual blinking or a dog wagging its tail,” says Jing, a PhD student at CSAIL. “Fast forward a number of years, and now we’ve got amazing models like Sora or Veo that will be useful in all varieties of interesting ways. We hope to instill an identical vision for the molecular world, where dynamics trajectories are the videos. For instance, you’ll be able to give the model the primary and tenth frame, and it’ll animate what’s in between, or it will probably remove noise from a molecular video and guess what was hidden.”
The researchers say that MDGen represents a paradigm shift from previous comparable works with generative AI in a way that permits much broader use cases. Previous approaches were “autoregressive,” meaning they relied on the previous still frame to construct the subsequent, ranging from the very first frame to create a video sequence. In contrast, MDGen generates the frames in parallel with diffusion. This implies MDGen will be used to, for instance, connect frames on the endpoints, or “upsample” a low frame-rate trajectory along with pressing play on the initial frame.
This work was presented in a paper shown on the Conference on Neural Information Processing Systems (NeurIPS) this past December. Last summer, it was awarded for its potential industrial impact on the International Conference on Machine Learning’s ML4LMS Workshop.
Some small steps forward for molecular dynamics
In experiments, Jing and his colleagues found that MDGen’s simulations were just like running the physical simulations directly, while producing trajectories 10 to 100 times faster.
The team first tested their model’s ability to absorb a 3D frame of a molecule and generate the subsequent 100 nanoseconds. Their system pieced together successive 10-nanosecond blocks for these generations to achieve that duration. The team found that MDGen was capable of compete with the accuracy of a baseline model, while completing the video generation process in roughly a minute — a mere fraction of the three hours that it took the baseline model to simulate the identical dynamic.
When given the primary and last frame of a one-nanosecond sequence, MDGen also modeled the steps in between. The researchers’ system demonstrated a level of realism in over 100,000 different predictions: It simulated more likely molecular trajectories than its baselines on clips shorter than 100 nanoseconds. In these tests, MDGen also indicated a capability to generalize on peptides it hadn’t seen before.
MDGen’s capabilities also include simulating frames inside frames, “upsampling” the steps between each nanosecond to capture faster molecular phenomena more adequately. It will probably even “inpaint” structures of molecules, restoring details about them that was removed. These features could eventually be utilized by researchers to design proteins based on a specification of how different parts of the molecule should move.
Toying around with protein dynamics
Jing and co-lead creator Hannes Stärk say that MDGen is an early sign of progress toward generating molecular dynamics more efficiently. Still, they lack the info to make these models immediately impactful in designing drugs or molecules that induce the movements chemists will need to see in a goal structure.
The researchers aim to scale MDGen from modeling molecules to predicting how proteins will change over time. “Currently, we’re using toy systems,” says Stärk, also a PhD student at CSAIL. “To reinforce MDGen’s predictive capabilities to model proteins, we’ll need to construct on the present architecture and data available. We don’t have a YouTube-scale repository for those varieties of simulations yet, so we’re hoping to develop a separate machine-learning method that may speed up the info collection process for our model.”
For now, MDGen presents an encouraging path forward in modeling molecular changes invisible to the naked eye. Chemists could also use these simulations to delve deeper into the behavior of medication prototypes for diseases like cancer or tuberculosis.
“Machine learning methods that learn from physical simulation represent a burgeoning recent frontier in AI for science,” says Bonnie Berger, MIT Simons Professor of Mathematics, CSAIL principal investigator, and senior creator on the paper. “MDGen is a flexible, multipurpose modeling framework that connects these two domains, and we’re very excited to share our early models on this direction.”
“Sampling realistic transition paths between molecular states is a serious challenge,” says fellow senior creator Tommi Jaakkola, who’s the MIT Thomas Siebel Professor of electrical engineering and computer science and the Institute for Data, Systems, and Society, and a CSAIL principal investigator. “This early work shows how we would begin to handle such challenges by shifting generative modeling to full simulation runs.”
Researchers across the sector of bioinformatics have heralded this technique for its ability to simulate molecular transformations. “MDGen models molecular dynamics simulations as a joint distribution of structural embeddings, capturing molecular movements between discrete time steps,” says Chalmers University of Technology associate professor Simon Olsson, who wasn’t involved within the research. “Leveraging a masked learning objective, MDGen enables modern use cases similar to transition path sampling, drawing analogies to inpainting trajectories connecting metastable phases.”
The researchers’ work on MDGen was supported, partly, by the National Institute of General Medical Sciences, the U.S. Department of Energy, the National Science Foundation, the Machine Learning for Pharmaceutical Discovery and Synthesis Consortium, the Abdul Latif Jameel Clinic for Machine Learning in Health, the Defense Threat Reduction Agency, and the Defense Advanced Research Projects Agency.
