As the capabilities of generative ai models have grown, you've probably seen how they can transform simple text messages into hyper-realistic images and even extended video clips.
More recently, generative ai has shown potential to help chemists and biologists explore static molecules, such as proteins and DNA. Models like AlphaFold can predict molecular structures to accelerate drug discovery, and the “MIT-assisted”radio frequency transmission”, for example, can help design new proteins. However, one challenge is that the molecules are constantly moving and shaking, which is important to model when building new proteins and drugs. Simulating these movements on a computer using physics (a technique known as molecular dynamics) can be very expensive, requiring billions of time steps on supercomputers.
As a step toward simulating these behaviors more efficiently, MIT's Computer Science and artificial intelligence Laboratory (CSAIL) and researchers in the Department of Mathematics have developed a generative model that learns from past data. The team's system, called MDGen, can take a frame of a 3D molecule and simulate what will happen next as a video, connecting separate still images and even filling in missing frames. By pressing the “play button” on molecules, the tool could help chemists design new molecules and closely study how well their prototype drugs for cancer and other diseases would interact with the molecular structure they aim to impact.
Co-senior author Bowen Jing SM '22 says MDGen is an early proof of concept, but it suggests the beginning of an interesting new research direction. “At first, generative ai models produced somewhat simple videos, like a person blinking or a dog wagging its tail,” says Jing, a PhD student at CSAIL. “Fast forward a few years, we now have amazing models like Sora or Veo that can be useful in all kinds of interesting ways. We hope to instill a similar vision for the molecular world, where dynamic trajectories are the videos. For example, you can give the model the first and tenth frames, and it will animate what's in between, or it can remove noise from a molecular video and guess what was hidden.”
The researchers say MDGen represents a paradigm shift from previous comparable work with generative ai in a way that allows for much broader use cases. Previous approaches were “autoregressive”, meaning they relied on the previous still frame to build the next one, starting from the first frame to create a video sequence. In contrast, MDGen generates the frames in parallel with the diffusion. This means that MDGen can be used to, for example, connect frames at endpoints, or “sample” a low frame rate trajectory in addition to pressing play on the initial frame.
This work was presented in a paper shown at the Neural Information Processing Systems Conference (NeurIPS) last December. Last summer, it was awarded for its potential commercial impact at the ML4LMS Workshop at the International Conference on Machine Learning.
Some small steps forward for molecular dynamics
In experiments, Jing and his colleagues found that MDGen simulations were similar to running physics simulations directly, while producing trajectories 10 to 100 times faster.
The team first tested their model's ability to capture a 3D frame of a molecule and generate the next 100 nanoseconds. Their system put together successive blocks of 10 nanoseconds so that these generations reached that duration. The team found that MDGen could compete with the accuracy of a reference model, while completing the video generation process in about a minute, a mere fraction of the three hours it took for the reference model to simulate the same dynamics.
When given the first and last frames of a nanosecond sequence, MDGen also modeled the intermediate steps. The researchers' system demonstrated a degree of realism in more than 100,000 different predictions: it simulated molecular trajectories more likely than their baselines in clips of less than 100 nanoseconds. In these tests, MDGen also indicated an ability to generalize about peptides it had not seen before.
MDGen's capabilities also include frame-within-frame simulation, “up-sampling” the steps between each nanosecond to more adequately capture faster molecular phenomena. It can even “paint” structures of molecules, restoring information about them that was removed. Over time, researchers could use these features to design proteins based on a specification of how different parts of the molecule should move.
Playing with protein dynamics
Jing and co-senior author Hannes Stärk say MDGen is an early sign of progress toward generating molecular dynamics more efficiently. Still, they lack data for these models to have an immediate impact on the design of drugs or molecules that induce the movements that chemists will want to see in a target structure.
Researchers aim to scale MDGen from modeling molecules to predicting how proteins will change over time. “We currently use toy systems,” says Stärk, also a PhD student at CSAIL. “To improve MDGen's predictive capabilities for protein modeling, we will need to leverage the current architecture and available data. “We don’t have a YouTube-scale repository for those types of simulations yet, so we hope to develop an independent machine learning method that can speed up the data collection process for our model.”
For now, MDGen presents an encouraging path toward modeling molecular changes invisible to the naked eye. Chemists could also use these simulations to delve deeper into the behavior of prototype drugs for diseases such as cancer or tuberculosis.
“Machine learning methods that learn from physical simulation represent a burgeoning new frontier in ai for science,” says Bonnie Berger, Simons Professor of Mathematics at MIT, CSAIL principal investigator, and senior author of the paper. “MDGen is a versatile, multipurpose modeling framework that connects these two domains, and we are very excited to share our first models in this direction.”
“Sampling realistic transition pathways between molecular states is a significant challenge,” says senior author Tommi Jaakkola, the Thomas Siebel Professor of Electrical and Computer Engineering at MIT and the Institute for Data, Systems and Society, and a CSAIL principal investigator. . “This initial work shows how we could begin to address these challenges by shifting generative modeling to full simulation runs.”
Researchers in the field of bioinformatics have heralded this system for its ability to simulate molecular transformations. “MDGen models molecular dynamics simulations as a joint distribution of structural embeddings, capturing molecular motions between discrete time steps,” says Chalmers University of technology associate professor Simon Olsson, who was not involved in the research. “Leveraging a masked learning objective, MDGen enables innovative use cases, such as sampling transition paths, establishing analogies with painting trajectories that connect metastable phases.”
Researchers' work on MDGen was supported, in part, by the National Institute of General Medical Sciences, the U.S. Department of Energy, the National Science Foundation, the Machine Learning Consortium for Pharmaceutical Discovery and Synthesis, the Abdul Latif Jameel Clinic for Machine Learning in Health, the Defense Threat Reduction Agency and the Defense Advanced Research Projects Agency.
