The world’s leading publication for data science, AI, and ML professionals.

Microsoft’s New AI Method to Predict How Molecules Move and Function

Termed the "Distributional Graphormer", it could kick-start the next step in the revolution that AI is bringing to the chemical and…

"Can we predict not only static protein structures but also their structural diversity?"—I asked around a year ago when all the excitement around AlphaFold 2 was settling:

Can we predict not only static protein structures but also their structural diversity?

Well, apparently Microsoft’s teams applying AI to science (together with a renowned professor, see notes at the very end) might have a first "Yes" as an answer to my question. They have just presented their novel "Distributional Graphormer" which can predict not just single molecular structures (of proteins, or other molecules, or materials) but also actually the multiple alternative structures (or "conformations") that a molecule or piece of material can adopt in 3D. That is, the alternative possible arrangements that their atoms can adopt in space. Moreover, the new AI model also "understands" that different structures will have different energies and thus be populated to different extents; thus, the new model can be trained to predict the underlying thermodynamics that govern how the different conformations exchange as the molecule moves.

As implemented in its first version, the Distributional Graphormer can parse any kind of molecule. In fact, the preprint presenting the model describes example applications to the tasks of predicting how proteins move (what I called "structural diversity" in my question above), how proteins bind small molecules (which covers from enzyme substrates to medicaments, metabolites, etc.), also how molecules adsorb on the surface of catalysts, and how carbon structures can be tuned in their electron-conducting properties.

Read on to understand all this in detail, know more about the model, and grasp the impact it may have on modern chemistry and biology.

An evolution of the Graphormer

Microsoft’s new Distributional Graphormer evolves from the Graphormer they introduced in 2021 with the aim to accelerate research and application in AI for molecule sciences from pharma to materials. The Graphormer itself is a Deep Learning algorithm and package that facilitates researchers and developers to train custom models for molecule modeling tasks.

As an algorithm, the Graphormer is a transformer adapted to better work with graph representations, which are especially good to describe molecules in neural networks, by encoding their structures into the models. See just for an example how such descriptors help neural networks to handle molecules here or here. To know more, check this preprint where Microsoft teams introduces the method proposing several simple yet effective structural encoding methods and characterizing the expressive power of the Graphormer, including concrete examples where it is used to encode structures:

Do Transformers Really Perform Bad for Graph Representation?

As a package, the Graphormer provides support for various molecular simulation tasks such as molecular dynamics simulations and property prediction, both elements used in the new work: simulations to probe how molecules move to change the relative positions of their atoms, thus changing shape as required for their function as in the case of protein conformation exploration; and property prediction for example to relate molecular structure to properties as in the case of the tuning of electron conductivity by carbon structures. You can know about this package, including information about multiple tests that researchers and developers have put it through, here:

Graphormer – Microsoft Research

After predicting static structures, AI goes for predictions of molecular motions

If AI methods applied to science were so far focused mostly in predicting "the" (static) structures of molecules, as AlphaFold 2 was so successful in doing (see articles here), the holy grail is now in an even more difficult goal: understanding how molecular structure relates to function.

This requires understanding how properties emerge from structures, i.e. from the 3D arrangements of atoms in space, and how these arrangements change over time. The latter is what’s known as "internal molecular dynamics", which refers to the displacement of atoms relative to each other as opposed to simple diffusion where all atoms move and rotate together thus overall not changing the shape of the molecule. See here two examples of molecular systems experiencing such "internal dynamics" as simulated through classical physics -first a piece of carbon nanotube vibrating, and second a protein getting folded:

Internal molecular dynamics, or just "molecular dynamics" or "structural conformations" for short, emerge from the coupling of the very fast and almost random thermal atom motions into the slower but larger-scale and collective motion of multiple atoms in a coordinated fashion. See, two atoms connected to each other may vibrate in a timescale of femtoseconds to picoseconds without being much "aware" of their neighborhoods, but a bunch of interconnected atoms that move due to thermal fluctuations will slowly evolve collective motions in say the nanosecond or microsecond timescale, that might end up moving a part of the molecule substantially such that its overall shape, and hence its properties, change.

For example, in proteins, a region that may be inaccessible in the average structure might become temporarily exposed due to its internal dynamics, allowing for example access of substrate molecules. This is widely studied in proteins because it’s at the core of how most if not all proteins work—a topic I’ve touched on from the viewpoint of "classical" molecular dynamics calculations as applied to proteins here:

How computer modeling, simulations, and artificial intelligence impact protein engineering in…

Simulations of the type described in that article, which essentially model physics and integrate Newton’s equations of motion, are very powerful but also exceedingly demanding in terms of computer power, hence very slow to run, let alone to converge. Therefore, it would be a huge advance if one could quickly predict the alternative conformations that a molecule can adopt, and how it will move, without actually running any molecular simulation. That’s what many large AI labs are trying to pursue now, and where Microsoft’s new model could make a difference.

Microsoft’s Distributional Graphormer provides a first approach to this, in the form of a deep learning framework that can handle not one but multiple structures of a molecule or piece of material, and also thermodynamic landscapes, which measure how likely it is to observe each structure. When we talk about a protein and its structure, we usually refer to one very stable conformation that was identified experimentally or modelled. But a real solution of a protein has trillions of copies of it, each adopting a structure drawn from a pool of possible states that have different probabilities of being observed. Within this "landscape" of options, the most stable ones (in terms of free energy, i.e. a combination of energetics and entropy) are the most populated ones, which means that more copies of the molecule will adopt the corresponding structure in a given moment. More importantly, due to its internal dynamics, a protein can change its structure thus effectively "jumping" from one populated state to another populated state, very temporarily transitioning through an "unlikely" state unfavored by free energy. Normally, due to thermal fluctuations a protein will be continuously changing structure, in average staying more time in the most stable ones. The collection of all possible structures and their populations is the "thermodynamic conformational landscape" that all biologists would like to predict without having to do experiments -and that even through experimentation are very hard to obtain.

How the Distributional Transformer plays in

While Machine Learning has very much improved structure prediction for molecules in the last 3–4 years since AlphaFold broke in with its specific application to proteins, many real-world applications require understanding macroscopic observations derived from the distribution of structures rather than individual structures -i.e. the conformational landscape I just explained above. Regular methods employed by scientists to sample these distributions, usually rooted in molecular dynamics simulations using some kind of physics-only or data-based potentials, are computationally expensive. Very expensive.

Microsoft’s Distributional Graphormer is first of its kind to address this challenge, and it does so by employing deep neural networks to transform a simple starting distribution such as a Gaussian into its prediction of the equilibrium distribution through transformations conditioned on descriptors of the molecular system. The expansion from a simple distribution to a rough landscape happens through a diffusion process similar to a simulated annealing simulation rooted in thermodynamics.

To guide learning of the distribution of states, the architecture can be trained using different types of data or information during the transformation, such as energy functions built ad hoc, data produced through molecular dynamics simulations, etc. During training, the distributional graphormer seeks to minimize the discrepancy between the probabilities predicted for each state and the probabilities provided as training.

Notably, this approach allows leverage prior knowledge of the system, which for cases of widely studied molecules may be quite substantial. But apparently, the Distributional Graphormer also shows good generalizing abilities, likely because it inherits the advantages of advanced deep-learning architectures and applies them to the new and challenging task of distribution prediction.

Versatility proven with varied use cases

The Distributional Graphormer efficiently generates diverse conformations and provides estimations of state densities, thus approaching a description of the conformational landscape of the molecule in question. The preprint and blog post presenting the Distributional Graphormer already show applications of the framework to several molecular tasks, including sampling of protein conformations, of ligand binding to proteins structure, of molecules adsorbing on catalysts, and the generation of structures guided by tuning of material properties.

In protein conformation sampling, where I can dig in more deeply as you know if you follow me, the Distributional Graphormer is shown to successfully generate structures and conformational landscapes for some well-studied proteins, displaying good coverage of the conformational landscape compared to molecular dynamics simulations—but much faster to generate. Some examples also demonstrate the model’s ability to generate functional conformation states for proteins with multiple metastable structures that are known experimentally. Furthermore, the examples show how in some cases the conformational landscapes outputted by the Distributional Graphormer predict plausible pathways for the interconversion of different conformations, enabling a deeper understanding of protein dynamics.

In ligand binding structure generation, the Distributional Graphormer is shown to predict the binding poses of ligands within pockets of a number of example proteins. A similar task but dealing with completely different types of atoms and attesting to the flexibility of the model’s framework, the preprint and blog show how it can efficiently sample adsorbate configurations on catalytic surfaces, crucial for understanding catalytic processes.

Finally, the preprint and blog show how the model allows for property-guided structure generation, exemplified through an electronic band gap predictor. By gradually refining a simple distribution, the model generates carbon structures with target band gap values, including some known carbon crystals such as diamond and graphite.

The above examples show that at least the core ideas implemented into the Distributional Graphormer work, and that with sufficient training data it can become a very powerful tool. Certainly, one that could foster further advancement of the field of deep learning applied to molecular systems with direct implications to pharma, biotechnology, materials science, and beyond. Being the model foundational, it can also open up new research opportunities in molecular science and pave new ways to approach statistical understanding of complex molecular systems by building on data.

References

The preprint at arXiv, led by Researchers from Microsoft’s AI4Science and Quantum divisions and including scientist Frank Noe. The later, a Microsoft partner research manager from the Free University of Berlin, is one my most admired modern-ray scientists. He is among the younger world leaders in the already decades-old field of molecular simulations, bringing together serious theoretical approaches with pragmatic applications. These days he explores of course AI methods quite a lot, having put forward the first elements for deep learning systems to assist molecular simulation, such as the Boltzmann Generator. In the past he made all kinds of contributions, among the most interesting ones he pioneered the use of Markov state models to describe molecular mechanics.

Towards Predicting Equilibrium Distributions for Molecular Systems with Deep Learning

Two blog-like entries by Microsoft, explaining the work in a more accessible format:

Towards Predicting Equilibrium Distributions for Molecular Systems with Deep Learning

Distributional Graphormer: Toward Equilibrium Distribution Prediction for Molecular Systems


www.lucianoabriata.com I write and photoshoot about everything that lies in my broad sphere of interests: nature, science, technology, programming, etc.

Tip me here or become a Medium member to access all its stories (I get a small revenue without cost to you). Subscribe to get my new stories by email. Consult about small jobs on my services page here. You can contact me here.


Related Articles