Peering inside the black box: Learning the relevance of many-body functions in Neural Network potentials

Abstract

Machine learned potentials are becoming a popular tool to define an effective energy model for complex systems, either incorporating electronic structure effects at the atomistic resolution, or effectively renormalizing part of the atomistic degrees of freedom at a coarse-grained resolution. One of the main criticisms to machine learned potentials is that the energy inferred by the network is not as interpretable as in more traditional approaches where a simpler functional form is used. Here we address this problem by extending tools recently proposed in the nascent field of Explainable Artificial Intelligence (XAI) to coarse-grained potentials based on graph neural networks (GNN). We demonstrate the approach on three different coarse-grained systems including two fluids (methane and water) and the protein NTL9. On these examples, we show that the neural network potentials can be in practice decomposed in relevance contributions to different orders, that can be directly interpreted and provide physical insights on the systems of interest.

1 Introduction

Molecular simulations have emerged in the last 75 years as a valuable tool to recover or even predict interesting physical phenomena at the microscopic scale and provide a detailed mechanism for grasping the underlying molecular processes [1]. In principle, the most accurate description of a molecular system is given by the solution of the associated Schrödinger’s equation. However, it is common practice to invoke the separation of scales between electrons and nuclei (Born-Oppenheimer approximation) and define an effective energy function for the nuclei that should take into account the electronic effects [2]. Historically, this has been done empirically in the definition of classical atomistic force-fields, which have been designed, refined, and used for the study of molecular systems [3, 4]. Classical force fields assume that the energy of a molecular system can be described as a function of “bonded” terms (e.g. bonds, angles, dihedrals) and “non-bonded” pairwise potentials (e.g., Van der Waals, Coulomb) [2, 1]. All the potential energy terms are defined by fixed functional forms, with parameters tuned to reproduce experimental data and/or first principle calculations on small test systems [3, 1].

Recent advances in machine learning have triggered a step-change in the development of data-driven force fields. In particular, Artificial Neural Networks (ANNs) have been proposed to more accurately capture the electronic effects in the potential energy functions for the nuclei [5, 6]. While classical, non-bonded potential terms are generally limited to 2-body interactions, the use of ANNs, and more specifically Graph Neural Networks (GNNs) [7] defining connections between neighboring atoms, significantly increases the expressivity of the energy function and allows a flexible parameterization of many-body interactions [8, 9, 10, 11, 12, 13].

While leaps in the development of GNN-based models have shown great promise in studying complex macromolecular systems [14] and predicting material properties [15], the results and the models themselves are often seen as black boxes. In a machine learned force-field, a molecular conformation is given as input to the neural network and only the total energy and its derivatives are obtained as output. The increased model accuracy comes at the cost of insight into the nature and strength of molecular interactions: In a classical force field each term in the energy function can be dissected, but deciphering which terms in the potential energy are important for stabilizing certain physical states or interpreting a prediction is significantly more difficult in a GNN-based model.

This black box problem is not unique to molecular systems, but rather ubiquitous in the application of machine learning. As a response, the new area of “Explainable Artificial Intelligence (XAI)” has emerged to start providing tools to tackle the interpretation of deep neural networks [16]. Different approaches have been proposed to interpret the results obtained with ANNs, ranging from self-explainable architectures [17, 11, 18, 19] to post-hoc explanations  [20, 21, 22, 23]. Some of those approaches are starting to find use also in physical and chemical applications, e.g., for explaining predictions regarding toxicity or mutagenicity [24, 25, 26], predictions of electronic-structure properties [27, 28], guiding strategies in drug discovery [29, 30], analyzing protein-ligand binding [31, 32, 33], or uncertainty attribution [34, 35]. XAI approaches have also been utilized to provide a better understanding of the error introduced in coarse-grained molecular models [36].

In principle, an “interpretable model” should allow a researcher to extract scientific knowledge in a successful application and to identify the sources of deficiencies/anomalies when the model fails. In this work, we use ideas from XAI and implement them in the context of machine learned molecular models for molecular dynamics simulations. In particular, we focus our attention on the understanding of coarse-grained (CG) models. In parallel to the development of atomistic force-fields, GNNs have been successfully employed in the definition of models at reduced resolutions [37, 38, 39, 40, 41], where some of the atomistic degrees of freedom are renormalized into a reduced number of effective “beads” to speed up simulation time. The difficulty in the definition of CG models lies in the fact that many-body terms play an important role, as a reduction in the number of degrees of freedom is associated with increased complexity in the effective CG energy function. It has been shown that to reproduce either experimentally measured free energy differences [42] or the thermodynamics of a finer-grained model [43, 44], many-body terms need to be included. The need for many-body terms and the difficulty in capturing them make CG models an ideal test ground for GNN interpretation.

In the present study, we train a GNN energy function at CG resolution from atomistic simulation data of several systems of different complexity and then interpret the model to provide a deeper understanding, rather than a mere energy prediction. By using explainable AI in the context of CG models, we provide evidence that machine learned force-fields are indeed learning physically significant interactions.

To interpret the energy prediction of a CG molecular model, we use the method of “Layerwise-relevance propagation” (LRP) for GNNs, which has been recently proposed to explain a model prediction by decomposing the output (in our case, the energy of the system) in terms of the contribution (or “relevance”) of groups of nodes (i.e. CG beads) [28]. In a sense, we can see the GNN-LRP as a multi-body expansion of the total energy of the system provided by the GNN.

As a first example, we compare different classes of GNN architectures to obtain CG models of bulk fluids and show that an interpretation of accurate CG models provide a meaningful physical information. Interestingly, two GNN architectures, even if different from each other, convey the same physical interpretation: at least in terms of 2-body and 3-body terms, the two networks offer different functional representations of the same underlying energy landscape.

As a second example, we examine a machine learned CG model for the protein NTL9 and show that its interpretation allows us to pinpoint the stabilizing and destabilizing interactions in the various metastable states, and even interpret the effects of mutations.

Results

1.1 Methane and Water

We start by analyzing and comparing CG models for bulk methane (CH₄) and water (H₂O). Methane has been previously studied with various coarse-graining methods since its non-polar and weak Van-der-Waals interactions make it a simple test system [37]. On the other hand, water is capable of forming complex hydrogen bonding structures and much research has been devoted to its modelling, both at the atomistic scale [45] and on the CG level [46, 47, 48, 49, 38].

For both systems, we define CG models by integrating out the hydrogen atoms and positioning an effective CG bead in place of the central carbon or oxygen atom. For each system, we train two CG models with different choices of GNN architectures, PaiNN [50] and SO3Net [51], for the definition of the CG effective energy. Both architectures are trained using the force-matching variational principle for coarse-graining, to create a thermodynamically consistent CG model from the atomistic data [52, 53, 47, 54, 39]. A brief discussion of the specific features of these GNNs is provided in the Methods section.

To interpret the network predictions, we use the method of “Layerwise-relevance propagation” (LRP) for GNN, called GNN-LRP [28]. The idea of GNN-LRP is schematically described in Fig. 1 and

Figure 1: Concept of GNN-LRP illustrated for a system of four particles (i.e. CG beads, in the present context). a) In GNNs, the input graph is defined by a cutoff radius that determines the direct neighbors for each input node. By stacking several message aggregations in multiple layers, information can be exchanged between more distant nodes (outside of the cutoff region). The model output is obtained by passing the learned feature representations through a multilayer perceptron and final pooling. Obtaining the relevance involves propagating the output back through the network, by considering the connections between each node in one layer and the nodes in the previous layer. This procedure defines “walks” across the network layers. b) The walks involving the same subset of $n$ nodes are aggregated to obtain a decomposition of the output into $n$-body contributions.

a more detailed description is provided in the Methods section as well as in the Supplementary Section S5. Essentially, GNN-LRP provides us with a “relevance score” for each subset of $n$ CG beads (with $n=1,\dots,N_{l}+1$, where $N_{l}$ is the number of layers in the GNN) in each configuration of the system, indicating the contribution of their interaction to the total energy.

Figure 2: Comparison of radial distribution functions resulting from simulations with an atomistic or CG model and corresponding 2-body relevance. Panels a) and c) correspond to water and b) and d) to methane models. Panels a) and b) show the results for PaiNN-based, and panels c) and d) for SO3Net-based models. The relevance, shown in red, is normalized by the absolute total relevance over the number of walks of the respective model and rescaled for each model type. A negative value implies a stabilizing interaction as the model output is the energy.

The ability of the two CG models to reproduce structural features of the two systems is shown in Figs. 2-3. In particular, the radial distribution function (RDF) as obtained in the CG models is shown in Fig. 2 against the atomistic reference model for methane (right column) and water (left column). Note that the SO3Net model exhibits irreducible representations up to an angular momentum of $l_{\mathrm{max}}=2$, while PaiNN utilizes a maximum angular momentum of $l_{\mathrm{max}}=1$. As a consequence, in comparison to SO3Net, PaiNN requires a larger cutoff to accurately reproduce the RDF of water. For a comparison between PaiNN models with different cutoff radii, please refer to Supplementary Fig. S2.

For methane, the smooth oscillatory behavior of the RDF is similar to a Lennard-Jones (LJ) fluid  [55], suggesting that many-body interactions may not be very relevant in a CG model of this molecule. In contrast, for water, the height of the first solvation shell is more sharply peaked and decays more rapidly than in the case of methane. For both systems, both architectures reproduce the corresponding RDF.

In Fig. 2, the average relevance score for the 2-body contributions is plotted (red curves) as a function of their distance, alongside the RDF for each model. Since the relevance score corresponds essentially to a decomposition of the output energy, a positive (negative) relevance score implies an increase (decrease) in the energy, thus a destabilizing (stabilizing) effect of the associated interaction.

For both methane and water, both PaiNN and SO3Net show a 2-body relevance score that diverges as the distance between two beads goes to zero. This observation matches our intuition that at distances below a certain “effective radius”, the network should learn a repulsive excluded-volume interaction to avoid the overlapping of the CG beads. For all models, the relevance score decays to zero as the distance between two atoms approaches the cutoff value considering two atoms connected in the corresponding GNN. This corresponds to the intuition that the interactions between two atoms become weaker as the atoms move further apart, and is enforced by the cosine shape of the cutoff function.

Differences between the two systems emerge at intermediate distances. Fig. 2b) and d) shows that, for methane, the relevance score is mostly positive. The PaiNN model displays a slight stabilizing well around the first peak of the RDF, that is washed out in the SO3Net model due to the proximity between the network cutoff and the location of the first solvation shell. On the other hand, the relevance score for both water models (Fig. 2 a) and c)) dips significantly below 0 at $\sim$3\text{\,}\mathrm{\SIUnitSymbolAngstrom}$$, indicating a stabilizing interaction between molecules at a distance corresponding to the first solvation shell.

Figure 3: 3-body relevance for water and methane with similar layout to Fig. 2 but examining the angular distribution between triplets of neighboring atoms. The relevance is separated into walks containing a self-loop (dashed red line) and walks without self-loops (full red line). Angle distributions are computed with a cutoff after the first solvation shell, at 3.5$\mathrm{\AA}$ for water and 5.6$\mathrm{\AA}$ for methane. Only relevance values for triplets within this cutoff are plotted, more details can be found in Supplementary Section S3.

Fig. 3 examines the angle distributions between triplets of interacting atoms within the first solvation shell, and reports the average 3-body relevance score as a function of this angle (as red curves). As illustrated in Fig. 1, the number of steps in the “walks” of connections between the nodes (CG beads) from one layer to the previous one in the GNN is determined by the number of message passing interaction blocks of the ML model (see Methods section for details). Here, we use models with three interaction blocks in both PaiNN and SO3Net models, thus we obtain walks of four steps in the input graph (with corresponding relevance attribution $R_{ijkl}$). As the architecture of both models include skip connections in their interaction blocks, a given graph node in a network layer may be connected to the same node in the previous layer. That is, a node may appear more than once in a single network walk. We indicate the case where a node is connected to itself in two subsequent layers as a “self loop” (see Fig. 1). Since features involved in self loops are expected to dominate the respective walk relevance, here, we distinguish between 3-body walks containing a self-loop (eg. connecting nodes [i, i, j, k] in Fig. 1) and walks containing no self-loops (eg. connecting nodes [i, j, k, i] in Fig. 1).

For methane, both PaiNN and SO3Net can reproduce the local atomistic angle distribution and they both have an associated relevance score very close to zero that does not particularly favor any angular configuration, indicating that the 3-body terms are not very important for the CG methane. For water, both models produce a similar behavior for the averaged 3-body relevance (shown in red). The walks without self-loop are stabilizing and favor angles around  $50$-$60$ degrees, corresponding to the population of water molecules sitting interstitially inside the tetrahedral arrangement [56, 57]. The walks containing self-loops are mostly destabilizing and likely correct for an overstructuring of the 2-body interactions. Indeed, in Supplementary Fig. S5, one can see that the relevance from the walks with self-loops varies much more as a function of the distance than the angle, whereas the relevance of walks without self-loops depends strongly on the angles and is not merely dominated by pairwise distances. Having 3-body interactions correct for 2-body interactions is a known effect when parametrizing explicit $n$-body functions for constructing CG models [48, 49, 58]. In both models presented here, 3-body terms are crucial to recover structural properties of CG water.

To further support this point, a 2-body-only model is shown in the Supplementary Section S3, where the Inverse Monte Carlo (IMC) method is used for parametrizing a pair potential on the system’s RDF. Supplementary Fig. S1 shows that, for methane, the IMC model is capable of reproducing the correct distributions of the relevant features, whereas for water, this 2-body-only model fails at reproducing the angular distribution. Additionally, we also observe that a non-equivariant GNN such as SchNet performs well on methane, but fails at fully reproducing the atomistic distributions for water (see Supplementary Section S3).

It’s worth noting that while the average relevance is shown in Fig. 2 and Fig. 3 as a function of distances and angles, the individual relevance scores cover a broad range. Supplementary Figs. S3 and S4 show the entire distribution of 2-body relevance values over all the configurations of the water models. High (destabilizing) relevance values correspond to short distances between beads and low (stabilizing) relevance values to distances corresponding to the first peak of the RDF. A relevance score of almost zero corresponds to distances approaching the network cutoff.

1.2 NTL9

Finally, we show the GNN-LRP interpretation in a coarse-grained protein model, specifically that of 39 residue NTL9 (PDB ID: 2HBA, residues 1-39) which has been a system of interest in a recent study by Majewski et. al. [40]. We use the PaiNN architecture introduced in the previous section to learn a CG model of NTL9 from the same atomistic reference data as used in the previous study [40], following the procedure introduced by Husic et. al. [39]. More details on the training procedure can be found in the Supplementary Section S2. In this study, only $C_{\alpha}$ atoms are kept in the CG resolution and each amino-acid type is represented by a unique bead embedding. We selected this protein due to its well-characterized folded/unfolded state and non-trivial folding pathways. Comparison of the Free Energy surface (FES) projected onto the first two TICA components (collective variables capturing the slow motions of the system) [59] is shown in Fig 4.

Figure 4: Comparison of the FES from the all-atom (left) to the CG (right) shown as a function of the first two TICA components [59]. Four regions of interest in TICA space are labeled by F (folded), U (unfolded), P1 (folding pathway 1), and P2 (folding pathway 2). The structures used for the interpretation of the respective states are shown on the left.

The folded state of NTL9 contains 3 $\beta$ sheets that are formed by the residues along the C- and N-terminal regions as well as a central $\alpha$ helix (shown on the bottom left in Fig. 4). The stability of this short fragment of the N-terminal domain of the Ribosomal Protein L9 is likely due to the strong hydrophobic core between the $\beta$-sheets and the $\alpha$-helix [60, 61]. To test whether the model is learning these interactions, we compute the relevance score for 2- and 3-body interactions in the trained model, for structures taken from different metastable states. The contact map in Fig. 5a) shows the 2-body relevance scores for each pair of amino-acids in the folded (upper right) and unfolded (lower left) states. Interestingly, the 2-body interactions stabilizing the folded state correspond to contacts associated with the main secondary structure elements, while in the unfolded state both stabilizing and destabilizing interactions are found also outside of the secondary structure. The strongest 2- and 3-body interactions inside the folded state are shown in Fig. 6. The strongest 2-body contributions are found in the $\beta_{13}$ sheet, most of which are stabilizing interactions. Notably, the VAL3-GLU38 interaction is destabilizing, which indicates that the CG model learns the side-chain specific interaction between the charged Glutamate and the hydrophobic Valine. The strongest 3-body interactions in the folded state, Fig. 6 panel (b), stabilize the helix and the tertiary structure of the protein.

Figure 5: Relevance contact maps for wild-type NTL9. a) mean relevance of amino-acid pairs in the folded state (upper right) and the unfolded state (lower left); b) mean relevance difference between the folded state and the P1 (upper right) and P2 (lower left) states, respectively, i.e. $R_{P1/P2}-R_{F}$. The regions bordered in black correspond to the contacts associated with the main secondary structure elements.

Figure 6: Snapshots of folded NTL9 with the most relevant 2- and 3-body interactions learned by the CG model. Panels a) and b) show the five most important 2- and 3-body interactions respectively. Blue lines indicate stabilizing interactions (negative relevance) and red lines destabilizing interactions (positive relevance). Darker color shades indicate stronger interaction strength (darker blue/red lines corresponds to stronger repulsive/attractive interactions). Visualizations generated with PyMOL [62].

NTL9 can fold by two different pathways, which appear as two distinct ”branches” in the free energy landscapes in Fig. 4. We examine the differences in relevance patterns between the two pathways connecting the folded to unfolded states. In panel b) of Fig. 5 we show the difference in relevance for both intermediate states relative to the folded state. Here, a positive difference means that the interaction has a lower relevance in the folded state than in the intermediate state and thus that this interaction is less stable in the intermediate state. For the state indicated as P1 in Fig. 4, many destabilizing interactions are located inside the $\beta_{12}$ sheet, whereas the difference to the folded state is essentially null in the $\beta_{13}$ sheet and in the $\alpha$-helix. This indicates that P1 corresponds to an intermediate state where the $\beta_{13}$ sheet and the $\alpha$-helix are native-like but the $\beta_{12}$ is less stable than in the native state. Indeed, analysis of the structures in P1 reveal that $\beta_{12}$ is register-shifted. The P2 state shows the opposite behavior with destabilizing interactions in the $\beta_{13}$ sheet and in the $\alpha$-helix, indicating that in the P2 state, only the $\beta_{12}$ sheet is correctly formed. These two folding pathways with the same characteristics are also found in previous computational studies of NTL9 [63, 64, 65]. Note that if the relevance attribution in the folded state shows that the network captures the interaction decay with the distance between residues, the relevance attribution in a given state provides more information than contained merely in the contact maps for these states, as can be seen by comparing the relevance attribution of both intermediate states to their distance contact maps shown in Supplementary Fig. S6.

The interpretation of the learned interactions can be pushed a step further by considering the effects of mutations on the relevance analysis. We consider mutations of residues deemed stabilizing in the folded state. In the machine learned C_α CG model of the protein employed here, one can straightforwardly perform a mutation by changing the aminoacid identity, that is by changing the embedding of the corresponding C_α bead. We select two mutations to illustrate the ability of the CG model to learn specific interactions such as hydrophobic/hydrophilic interactions between side-chains and side-chain specific packing. In particular, the mutation ILE4ASN is chosen to disrupt the hydrophobic interaction of the $\beta$-sheets, and the mutation LEU30PHE is chosen to disrupt the central $\alpha$-helix as well as the tight packing between the $\alpha$-helix and the $\beta$-sheets. These residues are flagged as important in the analysis above and have been shown in mutation experiments to play a role in the stabilization of the folded state [61, 66].

Figure 7: Effect of mutations on the model prediction. Panel a) shows the mean relevance difference for each amino-acid pair to the wild-type prediction, i.e. $R_{mut}-R_{WT}$. The upper right half corresponds to the LEU30PHE and the lower left to the ILE4ASN mutation. A red (positive) interaction corresponds to a higher relevance in the mutated state than in the wild-type, thus a destabilized interaction. Panels b) and c) show an example structure with the mutated residue highlighted in orange as well as strongest interactions flagged by the network interpretation in the same fashion as in Fig. 6.

In Fig. 7a) we show the 2-body relevance difference between the mutated states (LEU30PHE on the upper right and ILE4ASN on the lower left half) and the wild-type folded state. Replacing the identities of hydrophobic Isoleucine by polar Asparagine of about the same size at position 4 introduces a strong destabilization of all contacts with hydrophobic residues in the neighboring sheets, as it is also visualized in panel c). In the crystal structure of NTL9, LEU30 is tightly packed between the $\alpha$-helix and the $\beta$-sheets and mutation studies suggest that even a small change in side-chain size has a destabilizing effect [61]. Indeed, replacing the identity of Leucine 30 by the bigger Phenylalanine in our CG model induces a destabilization of the entire $\alpha$-helix (see Fig. 7, panels a) and b)). Interestingly, the disruption also has an effect on contacts inside the $\beta_{13}$ sheet that do not directly involve the mutated residue, indicating that the model has indeed learned non-trivial many-body interactions. These findings further corroborate the ability of the CG model to learn amino-acid specific interactions in a $C_{\alpha}$-only representation.

Discussion

In this work, we propose an extension of GNN-LRP to interpret the effective energy of machine-learned CG protein models in terms of a multi-body decomposition. We have shown on the application to CG fluids that the learned interactions are physically meaningful and consistent even if different ML architectures are used. The explanations provided by GNN-LRP indicate when multi-body interactions are required to recover the thermodynamics of the fine-grained system, showing the suitability of this higher-order explanation method to ML CG force-fields. Moreover, the multi-body relevance contributions show that the different ML models have learned similar physically relevant interatomic interactions, indicating that these models effectively learn the same underlying potential energy surface. The application of this idea on ML CG protein models allows to disentangle the strength of the different interactions between residues. It can also be used to evaluate the effect of mutations in the model and shows that bottom-up $C_{\alpha}$-based ML CG models capture amino-acid specific interactions without explicit representation of the side-chains.

These results provide reassurance that bottom-up machine learned CG force-fields indeed learn physically relevant terms by approximating the many-body potential of mean-force associated with the integration of part the degrees of freedom [47]. We note that the methods introduced are model agnostic and can be used in general to interpret machine learned potentials of different systems at different resolutions.

Future work is still needed to refine these concepts to provide greater insight into ML models and allow researchers to be more systematic in their choice of architecture and functionalization. It is our hope that this work helps lay the groundwork for researchers to better understand the outputs of their models as well as give the coarse-graining community a way to probe these learned many-body effects more explicitly.

Methods

Machine learned CG Potentials

Graph Neural Networks (GNNs) have been proposed as a promising method to learn interatomic potentials [10], and many different model architectures have been developed in recent years [67]. GNNs represent the underlying atomic details using structured graph data where nodes represent atoms and use the idea that locality dominates the energy landscape to draw edges between nodes if two nodes are within a pre-defined cutoff distance. In GNNs there are generally three steps that go into producing the network output based on the input positions and node identities: (i) a message-passing step, where neighboring nodes (connected by edges) exchange information about their respective feature values, (ii) an update step, where the node features are modified based on the received messages, and (iii) a final readout step, where the features of each node are used to predict the target property [10]. Once the node features are fed through the readout layer the network can make use of backpropagation to extract the force by taking derivatives of the energy based on molecule positions that can then be used to train during force matching or to propagate the dynamics.

In this manuscript, we examine two different GNN architectures for the effective CG energy: PaiNN [50] and SO3Net [51], two equivariant message-passing architectures that mainly differ on the order of their SO3-equivariant features ($l_{\mathrm{max}}=1$ for PaiNN and $l_{\mathrm{max}}=2$ for SO3Net). Both preserve the basic euclidean symmetries of the system, notably those of translation, rotation and reflection. Both architectures are parametrized to reproduce the CG potential of mean force using the force-matching approach [52, 53, 47, 54, 39]. As a comparison to more classical methods, Inverse Monte Carlo (IMC) [68, 69] is also performed using the votca library [70] with results shown in the Supplementary Section S3.

Layerwise-relevance propagation

Layerwise-relevance propagation (LRP) has emerged as a method to explain model predictions in a post-hoc and model agnostic manner [71, 21, 72, 73]. Originally, LRP has been used to obtain first-order explanations in the form of relevance attributions (also referred to as relevance scores) in the input domain. E.g., for image classification tasks, the relevance attributions would indicate to what extent a respective pixel is responsible for the network decision [21, 74]. The relevance attributions can be visualized in the input domain in form of a heat map, where large relevance attributions highlight features of the classified object that predominantly lead to the respective decision of the neural network [21]. Recently, efforts have been made to adapt LRP and other explanation methods to regression tasks [27, 75, 76], such as, e.g., the prediction of atomization energies [27, 28].

Pixel-wise relevance attributions have contributed enormously to a better understanding of the inner workings of ANNs. However, in some cases, restricting explanations to first-order (input features) may result in oversimplified explanations. Especially for problems where the interaction between several input nodes is considerably strong, the relevance information of higher-order features becomes increasingly important. This is the case for the coarse-grained systems considered in this study, where multi-body interactions are essential [77, 43, 44, 42]. To this end, a variety of higher-order explanation frameworks have been introduced [78, 79, 28, 80, 81, 82, 83, 84, 85]. One of those methods is GNN-LRP, which extends LRP to higher-order explanations for GNNs [28, 80]. In the following, we will give a brief summary of the methodology of GNN-LRP, first describing first-order relevance propagation and then extending it to higher-order. For an in-depth introduction, please refer to [71, 28].

In general, relevance attributions are obtained by propagating relevance from the model output back to the input features, with the relevance being a conserved quantity much the same as the mass flow through a pipe or electrical current through a junction. This means that the sum relevance attribution of a layer in the network remains the same as it is propagated backwards through the network. The relevance propagation from the neurons $\{j\}$ to a lower-layer neuron $i$ reads

$$R_{i}^{l}=\sum_{j}\frac{q_{ij}}{\sum_{i}q_{ij}}\cdot R_{j}^{l+1}~{},$$ (1)

where $q_{ij}$ quantifies the contribution of neuron $i$ to the activation of neuron $j$. This propagation approach can be applied subsequently from the output neuron until all relevance is attributed to the input neurons. There are multiple ways to define $q_{ij}$ associated with different propagation rules (see also [71]). According to the deep Taylor decomposition [72], LRP is equivalent to a decomposition of the neural network into several first-order Taylor expansions. In this picture, the different LRP rules implicitly help find suitable root points for the respective Taylor expansions [72].

In contrast to the first-order node-wise relevance attributions, GNN-LRP explains the network prediction based on so-called relevance “walks”. Each walk is a collection of connected edges and nodes of the input graph. The length of the walks is dictated by the number of interaction layers (message passes) in the model. Hence, e.g., a model with two interaction layers would allow for up to 3-body walks. The concept of GNN-LRP is illustrated in Fig. 1a) for a model with two interaction layers applied to an exemplary system composed of four particles (e.g., CG beads). The top part of the figure shows the forward pass of a common GNN starting from the feature embedding on the input graph until the model output, while the bottom figure illustrates how the walks with their associated relevance attributions are constructed. Note that the connections for message passing in the GNN are defined by the graph structure. Depending on the number of interaction layers and the size of the node features, each graph node will have multiple corresponding neurons in the GNN. We distinguish between graph nodes ($i,j,k,...$) and neurons ($a,b,c,...$) by using subscript and superscript, respectively. Similar to Eq. 1, the relevance of a walk $\mathcal{W}=[j,k,l]$ is obtained by propagating back the relevance from upper-layer neurons

$$R_{jkl}^{b}=\sum_{c}\frac{\lambda_{jk}Q_{jk}^{bc}}{\sum_{b,j}\lambda_{jk}Q_{jk}^{bc}}R_{kl}^{c}$$ (2)

where $\lambda_{ij}$ denotes the edge feature between graph nodes $i$ and $j$, and $Q_{jk}^{bc}$ is the contribution of node $j$ with its respective neuron $b$ to neuron $c$ with the associated node $k$. Depending on the propagation rule, again, $Q_{jk}^{bc}$ takes different forms. For more information on how $Q_{jk}^{bc}$ is obtained considering the propagation rules used in this work, please refer to the Supplementary Section S5. Note that Eq. 2 yields the walk relevance for a specific neuron $b$. The total relevance of the respective walk is given by summing over all corresponding neurons:

$$R_{jkl}=\sum_{b}R_{jkl}^{b}~{}.$$ (3)

The relevance of each walk can be seen as a contribution to the model output resulting from interactions between the subset of particles in the graph corresponding to the respective walk. More precisely, the relevance of a walk is associated with a sequence of interactions between pairs of particles or a single bead with itself. However, such a sequence of interactions is rather abstract and cannot be associated with a physical quantity in a meaningful way. In order to obtain a quantity that we can interpret as a multi-body decomposition of the output, we aggregate a collection of walks to so-called $n$-body contributions, as depicted in Fig. 1b). As suggested by Schnake et. al. [28], this can be achieved by summing up the relevance of all walks $\left\{\mathcal{W}\right\}$ that are part of a certain graph substructure $\mathcal{S}$. Hence, the relevance of a substructure $\mathcal{S}$ is given as

$$R_{\mathcal{S}}=\sum_{\mathcal{W}\in\mathcal{S}}R_{\mathcal{W}}~{}.$$ (4)

In this way, for the example of Fig. 1, we obtain 2-body and 3-body relevance attributions by summing over the relevance attributions associated to the walks which are part of the respective 2-body or 3-body substructure. In the case of energy predictions, this yields the energy contribution of the respective $n$-body substructures.

Data Availability

Simulation data and the code to reproduce the analysis and the plots shown in the manuscript are accessible at https://box.fu-berlin.de/apps/files/?dir=/peering_inside_the_black_box_supplementary_data_and_code&fileid=555049372