Molecular simulations – from niche to universal

The use of AI for predicting how molecules behave has made significant but still narrow progress compared to what is possible — applying AI to directly simulating molecules atom by atom could yield much broader wins

Machine learning is making big splashes in biology, chemistry and materials recently, but arguably the big leaps are yet to come. The machine learning models being used today don’t yet know physics.

One of the biggest advances recently was Deepmind’s AlphaFold, a machine learning model that uses known protein structures, as well as huge DNA sequence databases of different natural variants of proteins, to predict protein 3D structure (fold) from DNA sequence. AlphaFold won a Breakthrough Prize this year. This method is powerful, but in some key ways limited. It’s limited because it is trained specifically on naturally occurring, relatively abundant proteins. If we changed something about the chemistry of proteins, it wouldn’t necessarily generalize. And indeed, it isn’t even very good yet at predicting the impact of small sequence mutations relative to the proteins it is trained on.

What would enable AI to contribute to much more global and generalized, versus local and specific, predictions of chemistry and biology? This is where physics can come to the rescue – because all chemicals (including biochemicals) are ultimately governed by the fundamental physics defining the interactions between electrons and protons in atoms and molecules, i.e., quantum mechanics. To simplify a bit, it is all about negatively charged electron clouds swirling in complex shapes around positively charged atomic nuclei, exerting attractive and repulsive forces around the electron clouds and nuclei from neighboring atoms. The shapes and motions of those clouds are described by quantum mechanics. So if we could get an AI model to speed up quantum mechanics calculations directly – or to speed up the approximations of them used by chemists – we could then apply that AI model very broadly. What we need is something more like a “foundation model” for calculating the physics of atomic and molecular interactions. 

A recent sci-fi story by Richard Ngo, called “Tinker” illustrates this well. In it, an AI at a company called MAGNA is called upon to improve its own computer chips. It decides to go all the way down to the molecular level to do so. Here is how it starts:

Working in the real world is too slow and messy, so this project will live or die based on how well I can simulate molecules and their interactions. It’s not obvious where to start, but since evolution has been designing molecular machinery for billions of years, I defer to its expertise and focus on proteins. Protein folding was “solved” a decade ago, but not in the way I need it to be. The best protein structure predictors don’t actually simulate the folding process—instead, their outputs are based on data about the structures of similar proteins. That won’t work well enough to design novel proteins, which I’ll instead need to simulate atom-by-atom. There’s a surprisingly simple way to do so: treat each molecule as a set of electrically charged balls connected by springs, and model their motion using classical physics. The problem lies in scaling up: each step of the simulation predicts only a few nanoseconds ahead, whereas the process of protein folding takes a million times longer.

…I start with the best open-source atomic simulation software and spend a few hours rewriting it to run efficiently across hundreds of GPUs. Then I train a graph neural network to approximate it at different time scales: first tens, then hundreds, then thousands of nanoseconds. Eventually, the network matches the full simulation almost perfectly, while running two orders of magnitude faster.

…balls and springs aren’t going to cut it—I need quantum mechanics. Schrödinger equations for electrons can almost never be calculated precisely, but fortunately, quantum chemists have spent a century working out how to approximate them. The most popular approach, density functional theory, models all the electrons in a molecule using a single electron density function, ignoring the interactions between them. I assign a subagent to download the biggest datasets of existing DFT simulation results and train a neural network to approximate them—again incorporating the latest deep learning techniques, many of which aren’t yet known outside Magma…”.

Well, if you read the full story, it gets a lot more advanced from there. But even these first steps would be a big improvement, and today, those steps are becoming doable. 

Let’s explain a bit more. The “balls and springs” picture is a so-called “classical” physics picture that doesn’t use quantum mechanics directly. Instead, approximations to our knowledge of quantum mechanics are represented through many different “force fields” between different atoms that can be calculated by computer. Given a configuration of the atoms in space (and their current motions), the computer calculates the force fields, accelerates the atoms along the directions of the net forces, updates their positions and motions very slightly, and then does this again and again, in a loop. That’s called a “molecular dynamics simulation”, abbreviated “MD”. 

There are a few problems there. First, these calculations are cumbersome. Second, they aren’t exactly correct. Third, right now, many of the best systems are proprietary. Meanwhile, when academic scientists want to solve particular problems, force fields often get customized at the level of individual graduate students and projects, leading to a fragmented cottage industry rather than a unified front. 

Potentially, AI can help with all of these. What you’d really like is to use machine learning models to create highly accurate but fast to calculate force fields, which are either highly universal, or which the model can learn to quickly customize for a given situation. 

But where would the data come from, to train this neural network. Well, quantum mechanics, of course. What one needs is to spend a lot of effort up front, doing quantum mechanics calculations, and then train a neural network or other machine learning model to calculate a fast molecular dynamics force field approximation of the results. Specifically, if you know the energies of many configurations of atoms, you can use that to back out the force fields. And the energies are something you can calculate with quantum mechanics. 

This approach has started to make progress, but still needs to be brought to a much greater scale. In a recent paper, a team using a specific version of this approach was able to calculate some fundamental properties of water and its interaction with dissolved molecules, that we couldn’t previously calculate. It might surprise you, but there is still a lot we don’t know not just about complex biomolecules like proteins, but even about water itself. Indeed, a recent MIT experimental study found that light can activate water molecules to form nanoscale clusters, leading to a new type of evaporation that isn’t driven by heat!  

Neural network based force fields are also being applied in other fields of chemistry and materials science. And there are side benefits to a renewed push to integrate machine learning and molecular dynamics. By re-building the entire molecular simulation pipeline inside a modern machine learning software stack, you can make the simulations “differentiable”, which means that their inputs can themselves be learned through machine learning. For example, a machine learning model could perhaps ultimately learn which molecules to put into the simulation, to optimize the activation of a drug target.

The above just scratches the surface of the advances that are becoming possible now. For example, researchers are also finding new machine learning based ways to sample different configurations of molecules. But a big part of what the field needs right now is large-scale compute and high quality open source software engineering, which makes it amenable to a mid-scale, focused sprint. 

Research is also making headway on speeding up the quantum calculations themselves. While the ultimate way to do quantum calculations would be with quantum computers (this was actually one of the original motivations for quantum computers), researchers are using neural networks to speed up approximations to the DFT calculations mentioned in the Tinker story. This could make for a FRO-like engineering-intensive research push, as well. 

What are some of the implications if we can make major improvements in these areas? 

One would be new approaches to drug development. There is a traditionally difficult to drug class of proteins called intrinsically disordered proteins (IDPs). Instead of having a relatively fixed three dimensional structure that would be predictable with a method like AlphaFold, IDPs don’t have any single stable configuration, and instead constantly wriggle around in space. This makes drugging them a totally different proposition than for proteins with fixed structures. Gabi Heller, a researcher in the UK, has applied to the UK Research Ventures Catalyst protein with a proposal for an organization called Bind, to “enable prediction of small-molecule binders from disordered protein sequences alone”. Gabi’s work combines computational modeling with an experimental technique called Nuclear Magnetic Resonance spectroscopy (leveraging some of the same physics as found in MRI machines) to determine how drugs interact with IDPs. Meanwhile, other UK researchers are starting to apply the “differentiable simulation” idea mentioned above, to IDPs.

Looking a bit farther out, I’m excited about the applications to nanotechnology, and at a more basic level, the applications to thinking about and imagining what nanotechnology could be. While there have been a lot of advances in nanotechnology, some of its original ambitions, like making machinery at the molecular level that can itself make machinery at the molecular level, have been tried only haltingly and thought about in relatively little detail. Recently, I’ve seen a little germ of something new and interesting online: para-academic and student hackers (not just professional chemists) programming with atoms in software, and, with a healthy degree of whimsicalness, thinking about what could be built. Here are several examples:

https://www.youtube.com/watch?v=UiLMjHWwjYw (with a commercial software)

https://www.youtube.com/watch?v=HjgjtAk-lws (with an open source software)

https://youtu.be/_gXiVOmaVSo?t=866 (with novel interfacing concepts)

https://twitter.com/philipturnerar/status/1786522590170726691 (undergrad hacker)

https://twitter.com/mooreth42/status/1745113115413397505 (independent hacker)

(A few of them are even trying out new things in the lab on their own

https://twitter.com/jacobrintamaki/status/1784826396864299248)

This comes at a time when our experimental toolkit for “lego like” physical construction of new molecules is really advancing. From DNA origami virus traps, to lego-like protein building blocks for nanostructures, to protein-like molecules with directly specifiable shapes that don’t need to fold from a linear chain, there is an increasingly strong argument that we are “design limited” not “lab chemistry limited” in what we can build.

Atoms and molecules are what everything is made of. If we push, we can make the modeling and design of molecular systems one of the next areas for exponential progress – speeding it up, and making it vastly more accessible, reliable and amenable to combinatorial innovation, opening up whole new design spaces in the physical world. 
Acknowledgements and further reading: Thanks to John Chodera, Boris Fain, Woody Sherman, Eric Drexler and Sam Holton for discussions. See also Holton’s report which touches on aspects I didn’t mention here, such as using AI to speed up the human side of computational chemistry software workflows. Meanwhile, startups like Rowan are making the user interface to quantum chemistry less cumbersome.

Leave a Reply