Imagine you’re baking a cake, but instead of following a recipe, you're trying to figure out which ingredients work best by guessing and experimenting. It might take a few tries (and some pretty bad cakes), but eventually, you start to recognize patterns: more sugar makes it sweeter, and too much flour makes it dry. Now, imagine if you could spot these patterns instantly, without any trial and error—that's the kind of insight AI brings to chemistry. The real advantage? It doesn’t stop at finding patterns; it enables us to predict behavior and even design molecules tailored to our specific needs.

Chemistry, like baking, is all about structure. Atoms are the ingredients, and the way they’re arranged in a molecule determines the final product. For centuries, chemists have used these molecular structures to predict how substances will behave in reactions, relying on cause-and-effect models. When you know how atoms are connected, the type of bonds they make, you can figure out what happens next. This works because chemistry, at its core, follows rules we’ve spent years mastering

But here’s the twist: what if, instead of relying only on these rules, we could find hidden patterns in chemistry’s data? Patterns no human has noticed before. That’s where probabilistic models and neural networks come into play. These AI models don't just rely on the well-established theories—they learn from vast amounts of chemical data to make predictions that can surprise even the experts.

Neural networks, a type of AI inspired by the human brain, can take on complex problems without needing step-by-step instructions. When we feed these models molecular structures—the detailed maps of how atoms are arranged—something powerful happens. They learn, adapt, and eventually make predictions about chemistry at a level we once thought was only possible with years of expertise

Unfortunately, neural networks are designed in a way that puts this trust into serious question. These models mimic the brain by using neurons—tiny problem solvers that pass information along and work together to solve problems. But how many of us know how our own brains operate? I can’t even remember if I... turned off the oven this morning! But hey, our neurons keep working anyway, solving life's problems—just like neural networks, though we have no clue how they're doing it!"

But here’s the big problem: these AFVs are cryptic numerical data—essentially lists of information—that represent each atom in the molecule. Think of them as the GNN’s "thoughts" about each atom: how it behaves, how it interacts with other atoms, and how it contributes to the overall molecule. However, this all happens in a framework that the neural network understands—not one that we do. It’s like the inner workings of a master chef: their brain firing the right signals, fine-tuning each decision, and performing the right actions. And much like a chef, even they can’t fully explain how they achieve their mastery.

And here’s another huge problem: these AFVs are packed with details—128 parameters worth for each atom! That’s a lot of information and trying to analyze it directly is impossible—it’s just too overwhelming. It’s like analyzing the human brain for decisions. So, we need to simplify things, and we do that with a tool called t-distributed stochastic neighborhood embedding (t-SNE). This technique reduces the complexity of the data from a 128-D space to a 2-D, while losing as little information as possible. This helps us to visualize the data in a way that’s easy to digest.

t-SNE does this using distances; by measuring the distances between various AFVs and then mapping those distances to a 2-D space, we can organize the high-dimensional space onto a 2-D shadow projection, which allows us to glean at how the core knowledge of the GNN is organized. Note, however, this is just a shadow of the original information stored, and so we always have to use caution when analyzing the results of these techniques.

Once t-SNE works its magic, something pretty cool happens (see figure below): the AFVs we pulled out for each atom per molecule in the dataset organize themselves based on molecular substructures or functional groups—kind of like how you’d group together ingredients that serve similar purposes in a recipe (flour, sugar, and butter for cakes; spices and herbs for savory dishes). The network has learned to recognize these chemical building blocks and group them, accordingly, giving us a glimpse into how it models chemistry and predicts molecular behavior. For example, in one place, all the alcohols—primary, secondary, and tertiary—are grouped together, in another place, all the carboxylic acids, carbonates, carbamates, and so on. We show this for SchNet here, but in our article [1], we show it for AIMNet as well.

Fig 4 - GNN diagram. Highlighting (in red) the nodal (atom) feature vectors (AFVs) for every atom per molecule. We extract these AFVs as part of our approach to explainability of GNNs.

Here’s how it works: in a GNN, each atom is a node, and the interactions between them are the edges. Both nodes and edges are fine-tuned to figure out how these atoms combine to determine a molecule’s properties. After all, it’s the structure of a molecule that gives it its unique flavor—just like how the right balance of ingredients makes the cake perfect. Graph neural networks are based on this graph structure (nodes & edges), but the exact algorithm and models implemented that use this graph structure modelling can vary widely.

Take SchNet [2-5] and AIMNet [6] for example—both are graph neural networks, but how they fine-tune the nodes (atoms) and edges (bonds) is strikingly different. SchNet operates like a message-passing neural network. It focuses on learning the relationships (edges) between atoms, passing information between them to understand the molecule's behavior. There’s a lot of math involved, but the key idea is that the atoms (nodes) need to interact, constantly communicating to build an understanding of the whole molecule. (I’ll leave the detailed citations [2-5] for those interested in diving deeper!)

AIMNet [6], on the other hand, uses priors that establish a graph structure for the molecule before feeding it into the algorithm. The network then learns to fine-tune nodes and edges and focus on the most critical parts of this graph structure.

If we take this to the kitchen, SchNet is like a baker carefully coordinating between ingredients: the flour informs the butter how to blend, and the sugar adjusts based on the flour’s texture. AIMNet, however, behaves more like a chef with laser focus, zoning in only on the key elements of the recipe.

However, we’re still faced with the challenge of explainability. Just like a chef’s creative genius is hidden beneath layers of experience, we understand what these models do with the data (nodes & edges), but not always why they make certain decisions. What is their underlying framework? What drives their choices? Can we somehow "peek inside" their decision-making process and uncover the logic behind their predictions?

Of course, I’m talking about GNNs here—not chefs—and trying to do it ethically! 😉 So, we've trained both graph neural networks, SchNet and AIMNet, using a vast amount of organic chemistry data from the QM9 dataset [8]. And our job is to explain every decision these algorithms make.

So, to understand how a neural network operates is a complex task on par with understanding neural signals in the brain, albeit less complex as they are man-made. Naturally, just as in the human brain, one way to understand what is going on is to break down the neural signals firing, where are they coming from? where they are going? what they are activating? This way we can take a peek at what the AI "thinks" before it solves a problem. Is it considering the right chemical interactions? Can we connect its thought process to deeper scientific theories? This transparency is what allows us to trust AI’s decisions, much like how you might trust a baker who explains every step of the recipe.

But here’s the tricky part: the ingredients these neural networks use vary widely, just like how chemistry itself can be understood in many ways. They might use compound names, chemical properties, or even textual descriptions. However, the secret ingredient that has worked best so far is the actual molecular structure—the arrangement of atoms and bonds.

Why? Because the structure of a molecule is the key to understanding its behavior, just like how the right ratio of ingredients can make or break your cake. Many of our chemical theories are grounded in molecular structure, and by feeding neural networks this information, they can make predictions that are on par with the most advanced theories out there.

This result reassures us that the AFVs aren’t just random numbers—they are well-structured representations that reflect the underlying chemistry. The model’s high accuracy in identifying functional groups shows that the clusters in the 128-D space are resolutely associated with functional groups, without overlap or mixing. This means our PCA dimension reduction of Fig 6 is misleading, as it had to compress the AFV data in such a way that it over clustered the original 128 dimensions, due to attempting to look at all these dimensions through a 2-D window. In the original 128-D AFV space everything is clearly separated based on functional groups, much like t-SNE projection shows in Fig 5. Only through the LDA results in Fig 8 were we able to confirm that the 128-D space is highly resolute and faithfully represents functional group distinctions.