MMGNN: A Molecular Merged Graph Neural Network for Explainable Solvation Free Energy Prediction

1 Introduction

Understanding solute-solvent interactions in specific solvents is pivotal for areas in physical chemistry, including chemical reactions and electrochemistry [Varghese and Mushrif, 2019; D'Souza et al., 2011]. A comprehensive grasp of these interactions is vital not only for explanatory experimental outcomes but also for guiding the design and control of reactions and properties [Chung et al., 2022; Fang et al., 2024b;

Xia et al., 2023b]. In this context, the solvation Gibbs free energy ∆Gsolv emerges as a critical physicochemical property, dictating a molecule's behavior in solution [Low et al., 2022]. This property is intricately linked to the solute's partition coefficient between gas and solvent phases [Chung et al., 2022]. However, empirically testing solute-solvent free energies for all combinations is impractical due to high costs and extensive time requirements. This challenge necessitates an increased reliance on deep learning models to predict these energies more efficiently [Varghese and Mushrif, 2019].

The landscape of solute-solvent interaction modeling is currently dominated by two primary approaches. The first, termed embedding concatenation, employs two separate Graph Neural Networks (GNNs) to represent solute and solvent molecules. The individual embeddings from these GNNs are then concatenated for subsequent prediction tasks, as depicted in Figure 1 (a). While effective, this method may not adequately capture the intricate coupling between solute and solvent. In contrast, the embedding merging approach, shown in Figure 1 (b), seeks to address this limitation. It also utilizes two Graph Encoders for initial representation but differs in the subsequent processing of these embeddings. Here, the embeddings are merged through advanced interaction strategies, such as Transformer-based techniques or interactive pruning algorithms, to better reflect the complex interactions before making predictions.

Although the embedding merging approach marks an advancement, it falls short in accurately simulating the real chemical processes involved in dissolution, as depicted in Figure 1 (c). The dissolution process is characterized by the substitution of intramolecular forces within solute and solvent with intermolecular forces between them. This is exemplified by the interaction of acetonitrile and ethanol, where individual acetonitrile molecules are gradually surrounded by ethanol molecules. The formation of hydrogen bonds between atoms (nodes) of these molecules leads to a homogeneous solution. This natural dissolution process suggests a need for an algorithm that explicitly models these intermolecular node interactions, rather than relying solely on embedding-level interactions. Such a model, by more closely mirroring actual chemical processes, is crucial for both scientific accuracy and the effectiveness of the predictions.

In response to the limitations of existing models, we propose the Molecular Modeling Graph Neural Network (MMGNN), a novel framework designed to closely align with the actual chemical processes in solute-solvent interactions. This framework significantly improves the prediction accuracy of Gibbs free energy. Illustrated in Figure 1 (d), MMGNN begins by indiscriminately connecting intermolecular atoms to enhance interactions, such as hydrogen bonds, between molecules. After that, MMGNN assigns variable weights to these connections, reflecting the different constraints of various chemical bonds. An attention-based aggregation method is then employed, enabling the framework to adaptively learn from diverse solute-solvent combinations. The result is a merged graph representation for each pair, highlighting the most significant atomic interactions.

Meanwhile, this framework also addresses potential challenges: (1) the presence of non-existent or repulsive atomic interactions, and (2) increased complexity and convergence difficulties in merged complete graphs for large molecular pairs. Inspired by graph explainability algorithms, MMGNN incorporates interpreters for both nodes and solute-solvent edges in the merged graph. This approach effectively reduces redundancy, focusing only on relevant interactions. The explanatory subgraph is then encoded and utilized in regression models, such as Fully Connected Neural Networks (FCNN), to predict Gibbs free energy with greater precision.

In conclusion, the main contributions of this paper can be summarized as follows:

• Introduction of MMGNN: We present the Molecular Modeling Graph Neural Network (MMGNN), a pioneering framework designed to explicitly align with actual chemical processes, enabling more accurate modeling of solute-

solvent interactions.

• Advancement in Explainability: Our approach integrates explainability into the model, allowing for the extraction of key subgraphs. This feature not only enhances the understanding of MMGNN's predictions but also aids in further research into solute-solvent dynamics.
• Validation through Extensive Experiments: We conducted comprehensive experiments to validate MMGNN's performance. The results demonstrate both the accuracy and explainability of our model's predictions, underscoring its effectiveness in the field of physical chemistry.

2 RelatedWork

2.1 Molecular Relational Learning

Generally, message passing in GNNs or graph convolutional networks (GCNs) refers to the utilization of trainable interaction layers to facilitate the exchange of information among atoms within a local neighborhood [Xia et al., 2023a; Wang et al., 2024; Wang et al., 2023]. Behler and Parrinello are pioneers in modelling interatomic properties [Behler, 2015; Behler, 2016] by summing per-atom contributions in neural network predictions. Tensormol [Yao et al., 2018] employs Behler and Parrinello's approximation to accurately compute the dipole moment in water dimers as one water molecule rotates about the O-H bond, resulting in reasonable outcomes. This method is subsequently refined by integrating the atomic-pairwise into the neural network [Glick et al., 2020] which led to a reduction in the binding energy error by a factor of five. [Lee et al., 2023] provide a molecular relational learning framework that predicts the interaction behavior based on graph information bottleneck theory. In brief, the original method is summing overall the atoms, or fusing information after message passing process but a more rigorous way should consider the atoms or fragments and introducing atom-pair symmetry message passing functions.

2.2 Gibbs Free Energies PredictionMethods

Solvation free energies have been of great interest for many years and have spurred the development of numerous predictive methods [Hildebrand, 1981]. These methods encompass a range of techniques, including molecular dynamics and quantum chemistry methods, as well as empirical or datadriven approaches. While quantum chemistry methods such as the SMx and COSMO(-RS) models [Klamt and Eckert, 2000; Marenich et al., 2007] are based on first-principles calculations of all promising relevant conformers of the solute and solvent molecules, they are often computationally expensive and labor-intensive. In contrast, empirical or data-driven methods could provide a faster alternative for predicting solvation properties. Recent studies have aimed at improving these models for predicting ∆Gsolv. [Vermeire and Green, 2021] utilized a transfer learning method to achieve a mean absolute error of 0.21 kcal/mol when predicting experimental results from a standard quantum computing dataset. similarly, [Zhang et al., 2022] calculated water solvation Gibbs free energies of more than 100,000 organic compounds, and utilized a graph neural network (GNN) to predict these values achieving an error of 0.4 kcal/mol.

Step1: Molecule Merging Step3: Merged Molecular Explainability Injection

3 Method

3.1 Overview of the Framework

In this study, we introduce a novel merged molecular graph model, wherein the weighted connection edges adhere to the principles of similarity and compatibility in chemistry. To process an input solute-solvent pair, we initially convert it into a machine-readable topological graph and incorporate relevant weighted edges to enhance molecular information interaction, constructing a merged graph as detailed in Section 3.2 and illustrated in Figure 2. Within this merged graph, we facilitate both inter- and intra-molecular message passing processes. In the initial layer of message passing, each atomic feature of the solute molecule is updated independently and subsequently acquires all corresponding atomic information of the solvent molecule, and vice versa. In the second stage of message passing, each atom of the solute molecule receives information beyond the atomic details in the solvent, including bond and node information with one hop. The final realization of the merged molecular representation is provided in Section 3.3. Additionally, we consider three distinct molecular explainability injection methods, as outlined in Section 3.4, to highlight the most significant atomic interactions. This augmentation aims to enhance the model's generalization capability and align it more closely with biochemical realities.

3.2 Molecule Merging

The molecular structure can be naturally depicted as a topology graph, wherein atoms serve as the nodes and bonds act as the edges [Wen et al., 2021; Zhou et al., 2023b]. The molecular graph can be formally denoted by:

where E is the set of edges (bonds), and V is the set of nodes (atoms). U is the global feature vector which is extracted from each molecule. The detailed features used in molecular representation are recorded in the appendix.

Specifically, all of bonds are enumerated in a molecule. The composition of E comprises ek, pk and qk, where ek represents the edge feature, and pk and qk, is the two atoms connected by ek.

Similarly, V enumerates all of atoms in a molecule. vi is atomic features vector for atom i (e.g., atomic species, valence, etc.) and N v is the total number of atoms in the molecule. So far, the single molecular graph has been fully defined. Next, a merged graph ˜ G could be generated by establishing a weighted edge between two molecules (solutesolvent) which connects every atom of each molecule to all atoms in the other molecule.

The set of relation edges are R :

where i ∈ {1, 2, 3, . . . ,N a }, j ∈ {1, 2, 3, . . . ,N b }. N a and N b is the total number of atoms in each of the two molecules, use rk to represent this new type of edge.

3.3 MergedMolecular Representation

Generally, in this merged molecular graph ˜ G , the message passing process is firstly run in single molecule (i.e. solute and solvent molecules). During the message passing phase, ek is updated to the new vector e′ij by aggregating the initially bond features, as well as the features of the two atoms, vi and vj, and the global features u. In addition, the atomic feature vector, vi, for each atom i is updated into v′i :

where v′i and e′ij are the updated node vectors and edge vectors. τ is the LeakyReLU activation function [Maas et al., 2013]. FC is fully connected layer. ⊙ denotes the Hadamard product. Then, the intramolecular atomic features could be updated. σ(·) is the sigmoid activation function, and ϵ is a fixed constant (0.0001) added for numerical stability. The next atomic update step would achieve intermolecular information passing process. In a molecule a,the atomic feature vector vaifor each atom i is further updated as follows:

where β are the hyperparameter that control the update rate of information, which makes information differentiated to avoid information decay or overwrite the original information. δ is an additional function that measures the feature similarity between two atoms, which is defined as follows:

Finally, the global feature vector u and relation bond vector r′i is updated:

where N v and N e are the number of atoms and bonds in the molecule, respectively. vai and vbi are two atoms connected by the relation, u is the global feature.

3.4 MergedMolecular Explainability Injection

In this part, we elaborate the process of injecting explainability into the model. The input for the explanation module is the representation of the merged graph. Specifically, for a node, it receives the embedding of the node to output an importance score, and for an edge, it processes the concatenated embeddings of the two nodes connected by the edge to output the importance of the edge. Specifically, we get inspiration from the significant success of the post-hoc explanation methods of the graph neural networks [Ying et al., 2019; Fang et al., 2023b; Fang et al., 2024a] and design our explainers from three perspectives as follows. Note that Figure 2 illustrates the workflow of the explainer designed from the second approach, namely the global mask-based perspective.

• The first one is the local mask-based perspective. Specifically, local mask-based methods endeavour to multiply the features v′ai and v′bi with the corresponding masks (all are initialized to 1) pai and pbi to get v′′ ai = v′ai pai and v′′ bi = v′bi pbi. Then, v′′ ai and v′′ bi is sent into the model to obtain the updated output ∆Gsolv. Then they attempt to find the optimal score pi by minimizing the difference between this processed output ∆Gsolv and the label Y . To limit the size of explanatory subgraphs, they apply l1 regularization to the value of mask. In this case, the loss function is:

where D denotes the distance function; λ is the trade-off parameter.

• The second one is the global mask-based perspective, and its only difference from the first approach lies in the method of mask generation. Here, during the process of the adding the mask, the trainable mask pai and pbi in local mask-based methods is replaced with a trainable MLPψ (i.e., pai = MLPψ(v′ai ), pbi = MLPψ(v′bi )). Meanwhile, following the theory of graph information bottleneck (GIB) [Wu et al., 2020; Yu et al., 2021], these methods instantiates the information constraint (ℓp) proposed by [Miao et al., 2022], where ℓp is defined as:

where k is the pre-defined hyperparameter. After applying ℓp regularization to the value of mask, the expression of the final loss function is:

where D denotes the distance function; λ is the trade-off parameter.

• The third one is the gradient-based perspective. For the v′ai and v′bi , these methods first calculate the absolute values of the features in the derivative of ∆Gsolv w.r.t v′ai and v′bi . After that, their importance score pai and pbi are defined as the normalized sum of these values. Then, at the end of each training phase, we mask some nodes and edges according to the above derivatives, and input the remain part to the next training phase. More formally:

After that, we conduct another merged molecular message passing process as introduced in Section 3.3 to conduct the final interactive prediction. Specifically, the readout phase occurs after the desired number of message-passing layers has taken place (as determined via hyperparameter optimization). To aggregate the set of final feature vectors of v′solute and v′solvent into a single vector suitable for input into a set-toset (set2set) module, where the solute and solvent molecular vectors are input separately rather than incorporating them as a whole vector that could result in a substantial loss of information during the pooling process. The global feature vector, being a single feature vector, is appended after the set2set steps. This produces a final feature vector which is invariant to permutation of atom, relation and bond indices:

where F is a set2set layer and ∥ denotes vector concatenation. ginteract map is then input into a FCNN to obtain the result.

4 Experimental Results

In this section, we conduct extensive experiments to answer the following questions:

• RQ1: Can MMGNN improve the prediction accuracy of ∆Gsolv?
• RQ2: How does the MMGNN perform when explaining the importance of intramolecular and intermolecular features simultaneously?
• RQ3: How effective is MMGNN in terms of generalization capabilities?

4.1 Experimental Setup

Datasets. The data set utilized in this study is the CombiSolvExp database, which is compiled by Vermeire and Green [Vermeire and Green, 2021], and combined the experimental data from multiple sources. And the FreeSolv database is published by Mobley and Guthrie [Mobley and Guthrie, 2014], the CompSol database fromMoine et al. [Moine et al., 2017], and remaining dataset is Abraham collected by the Abraham group [Grubbs et al., 2010].

Baselines. We compared MMGNN with some advanced methods, such as D-MPNN [Vermeire and Green, 2021], SolvBert [Yu et al., 2023], SMD [Meng et al., 2023], Explainable GNN [Low et al., 2022], GAT [Velickovic et al., 2017], GROVER [Rong et al., 2020], Uni-Mol [Zhou et al., 2023a], Gem [Fang et al., 2022], these baselines solely concatenate the molecular representations of solute and solvent, but ignore the information interaction process between them. CIGIN [Pathak et al., 2020] and CGIB [Lee et al., 2023] implicitly considered the inter-molecular interactions.

Experimental settings. MMGNN is trained using the Adam optimizer [Kingma and Ba, 2014] (1 × 10-4 → 0.5) via batch size 50. We employ mean squared error (MSE) as our training loss. The training process was halted if the validation error failed to reduce in 150 epochs, or if the maximum training limit of 1000 epochs was reached. MMGNN is implemented in PyTorch framework via Tesla A100 40GB. Experiment results are presented in following metrics: the mean absolute error (MAE) and root mean squared error (RMSE). The mean and standard deviation are recorded in the table.

4.2 The Prediction Performance (RQ1)

Similar to previous studies, the random split is a widely used method for evaluating model performance. In each iteration, all methods were trained, validated, and tested on identical datasets in an 8:1:1 ratio. For the three explainability methods: local mask-based, global mask-based, and gradientbased, we conducted separate tests on various datasets and documented the optimal results by MMGNN (this can be found in the appendix). The experimental results are recorded in Table 1. The sensitivity experiment analysis is presented in Table 2. Based on these outcomes, we can delineate three key observations:

Obs.1: MMGNN exhibits the optimal predictive performance. The performance of MMGNN surpasses that of other baselines across three test datasets. In terms of the MAE indicator, MMGNN exhibited an average reduction of 0.025, 0.16, and 0.57 on the CompSol, Abraham, and CompSolExp datasets, respectively. Although on the FreeSolv dataset, MMGNN ranks second, slightly behind the CGIB model, this still underscores its notable capability in accurately capturing the characteristics of solute-solvent molecular combinations. Overall, MMGNN consistently demonstrates leading scores across the majority of evaluated datasets, highlighting its effectiveness in precisely characterizing intricate interactions within solute-solvent molecular combination.

Obs.2: β = 0.2 yielded the best results to control the extent of information exchange. Our experiments revealed that setting β = 0.2 achieved the best results in controlling the extent of information exchange in Table 2. Various β values were tested to assess the impact on the degree of information exchange. The β value of 0 implies exclusive consideration of intramolecular aspects, resulting in the highest prediction error, as indicated by both MAE and RMSE metrics. Incrementally increasing the interaction level (β ranging from 0 to 0.2) significantly enhanced predictive performance, leading to an 11.5% reduction in MAE. However, further escalation of β resulted in increased MAE, potentially due to information redundancy and repetition causing interference with the model's ability to capture crucial details. Our findings underscore the importance of maintaining a balance between intrinsic and interactive information.

Obs.3: MMGNN is a versatile framework that can be employed with various GNN backbones. We conducted additional tests using GAT, GCN, and GIN to extract intermolecular information. In this scenario, the intramolecular information interaction method utilized the optimal β value of 0.2 from the aforementioned experiments. It is evident that the GCN framework outperformed GAT and GIN, resulting in a reduction of approximately 2.2% and 14.7 % in MAE indicators, respectively. Therefore, we executed the experiments with GCN framework for intramolecular message passing process and hyperparameter β = 0.2 for intermolecular message passing process.

4.3 Qualitative Evaluation (RQ2)

Due to the current lack of consensus on the key functional groups involved in the energy release process of solution solvents, we have collected some potentially important groups from relevant literature to serve as ground truth for the validation of algorithm explainability [Fang et al., 2023c; Fang et al., 2023a]. The intensity of the color corresponds to the strength of the enhancing (green) or decreasing (pink) effect on ∆Gsolv prediction.

Obs.4: The significant substructures identified by MMGNN closely align with the ground-truth of their graphs, likely stemming from the class-specific patterns captured by the explanatory subnetwork. However, the baseline explainers could select certain edges that do not belong to the groundtruth. This validates MMGNN's ability to capture substructure interactions aligned with chemistry principles, showcasing its capacity to extract implicit scientific knowledge in solute-solvent molecular interactions.

Obs.5: In contrast to the substructures chosen by baseline explainers, the subgraphs generated by MMGNN exhibit enhanced connectivity. We attribute these distinctions to the alternating and iterative selection processes, a capability not present in baseline explainers.

Obs.6: For specific nodes in molecules (e.g., N, O), certain connected nodes can introduce interference during explanation generation. MMGNN is designed to circumvent these challenges, ensuring explanations are free from interference. In contrast, subgraphs generated by baseline methods may include these disruptive nodes, highlighting the robustness and reliability of MMGNN.

4.4 Generalization Test (RQ3)

To evaluate the model's generalizability on unseen solvents or solute category, a solvent holdout and solute scaffold split cross-validation test are conducted, and the result are present in Figure 4 and Table 3. Here, the four datasets are merged into a single dataset.

Obs.7: MMGNN has superior results on a variety of excluded solvents such as benzene (0.17 kcal mol-1 ), ethanol (0.23 kcal mol-1 ), hexane (0.115 kcal mol-1 ), acetone (0.23 kcal mol-1 ) (in Figure 4 (a)), etc. Nevertheless, it is noteworthy that the water solvent exhibits a higher error compared to the other solvents. This difference is intricately connected to factors such as the molecule's molar mass and element type. A comprehensive breakdown of results and in-depth analysis can be found in the appendix.

Obs.8: MMGNN model demonstrates effective generalization performance across various types of molecular scaffolds as shown in Figure 4 (b). Heterocyclic and polycyclic molecules exhibit better predictive performance than monocyclic molecules. For instance, phenanthrene and morpholine have predictivity values of 0.57 and 0.56, respectively, while cyclopentene has a predictivity result of 1.31. We hypothesize that monocyclic molecules tend to display nonpolar characteristics, whereas polycyclic and heterocyclic molecules often exhibit polar properties, which makes them easier to differentiate and learn for model. 1-yne and nitrate is significate polar molecular groups which has the relatively low prediction error.

Obs.9: MMGNN outperforms other baselines in generalization test. We selected three representative baseline models for the same generalization testing experiments (in Table 3). MMGNN showed significant superiority over other baseline models. This is mainly attributed to (1) its more fitting modeling of molecular combinations through the use of molecular merge graphs, which aligns better with the underlying physicochemical processes. The explicit enhancement of molecular interaction processes is achieved. (2) The molecular explainability module injection retains core substructures, removes redundant information, and further enhances generalization performance. Constrained by space limitations, additional solute test results are presented in the appendix.

5 Conclusion

In this study, we present MMGNN, a molecular graph neural network designed for predicting solvation free energy. MMGNN exploits the complete connectivity among atoms in solute and solvent molecules to facilitate information transfer and improve model explainability. Demonstrating superior performance on various datasets compared to existing methods, MMGNN highlights the significance of balancing intrinsic and interactive information through our experiments. The model excels in generalization tasks, including solvent holdout and solute scaffold split, superior than other baselines. Its transparent explanations align with well-established chemical insights, which instills confidence among experimental chemists. Furthermore, understanding the interactions between molecules could provide valuable perspectives for understanding drug compound stability and guiding the development of new synthetic routes and catalytic processes.

Ethical Statement

There are no ethical issues.

Acknowledgments

This paper is partially supported by the Project of Stable Support for Youth Team in Basic Research Field, CAS(YSBR005), Natural Science Foundation of China National Major Research Instrument Development Project (No.12227901), the Innovation Program for Quantum Science and Technology (2021ZD0303303), the National Natural Science Foundation of China (22025304, 22033007, 62072427).

Contribution Statement

Wenjie Du and Shuai Zhang are co-first authors; Yang Wang and Junfeng Fang are corresponding authors.