Can Graph Neural Networks Understand Chemical Elements?

Abstract

Classical machine learning (ML) has proven itself successful in a vast range of chemical applications. These successes often only come after substantial effort in feature engineering to obtain an effective data representation for the model input. Graph neural networks (GNNs) limit this feature engineering by learning their own representation directly from molecular graphs. However, the lack of an expert-generated representation may put larger requirements on training data set size compared to classical ML approaches. GNNs often still require vertex and edge features to provide information that cannot be learned from graph connectivity alone, such as chemical elements and bond orders. In this work, we investigate if chemically inspired representations of these features (specifically chemical elements) can combat GNN data requirements and/or improve performance or generalizability of the models. To this end, nine representations of chemical elements were created that encode trends in elemental reactivity and behavior. We tested these representations across GNN model size, data set size and additional degrees of featurization, utilizing the prediction of molecular pK _a values as a benchmark. Our results suggest that the models largely classify elements, converting all representations into a largely orthogonal representation used internally. This mechanism explains our findings of minimal differences in model performance and generalizability between elemental representation. However, slight performance benefits are observed with more orthogonal elemental representations, probably because they better match the internal representation of the models. These findings led us to investigate the use of representations with increased orthogonality in the form of (i) atom-typing (effectively including information on atomic neighbors of the graph vertices), which yielded modest performance gains, and (ii) similarity typing as a novel atomic representation (encoding for similarity between atom types), which provided further performance gains but with substantial increase in computational cost for data set preparation.

Introduction

In recent years, machine learning (ML) has been increasingly used within chemical research, with successful application in areas as diverse as de novo structure generation, resolving and predicting molecular spectra, (retro)­synthesis, , and property prediction. ,, However, ML models typically require a fixed-sized numerical vector as model input. This input format is not always directly suited for the variable-sized and interconnected structure of molecules, thus necessitating substantial effort in feature selection and engineering of the input, e.g., using molecular descriptors and fingerprints. , In contrast to such fixed-sized vector based input, Graph Neural Networks (GNNs) − are able to learn directly on graph data. Graphs – with vertices representing atoms and edges as bonds – offer an intuitive representation of molecules and have been used as molecular representations for more than a century. The ability of GNNs to learn on graph data has proven successful in predicting biological and physiochemical properties including binding affinity, solubility ,, and pK _a values. ,−

Because GNNs construct their own internal representation of molecules by learning on their molecular graphs, they generally require less feature engineering than traditional, fixed-sized vector based ML methods. However, by being forced to learn an internal representation, GNN-based learning may require substantial training time and can put high demands on training data in terms of quantity and diversity. ,,

Furthermore, for most applications, GNNs still require featurization of their input. In the case of molecular data, this is generally in the form of vertex features such as the element or charge of an atom and edge features such as the bond order. Categorical or discrete features are often represented as a one-hot encoding of possible values. ,,,,, The resulting bit-vector provides a presegmented encoding of the data, which could accelerate learning by not requiring the model to separate feature-values during training. However, one-hot encoding is completely orthogonal and removes any underlying structure in feature space, potentially limiting the transferability and generalizability of GNN models. The inability of one-hot encoding to encode structures present in feature space may exacerbate the already high data requirements of GNNs for both the size and the diversity of training data.

In the current work, we study whether GNNs, when presented with molecular data, can benefit from a feature representation that is designed to encode known structures and relations in feature space. An ideal case study for this is the representation of chemical elements, which is typically done by using a fully orthogonal one-hot encoding. ,,,,, However, reactivity and behavior of the elements are known to have clear periodic trends that are lost in orthogonal encoding. Here, we compare the fully orthogonal representation of elements, i.e., one-hot encoded atomic numbering, to other elemental representations that incorporate information on the known periodic trends. These include discrete representations that directly encode information from the periodic table (in terms of period and group of the elements) or on the quantum numbers of the outermost electrons of the elements. These different discrete feature sets are used here in a numerical or one-hot-encoded fashion. We will additionally use representations with continuous-valued features in the form of elemental properties such as electronegativity and polarizability. The performance of these models will also be compared to a model without any elemental representation.

In the first part of our study, the GNN models with different elemental representations are benchmarked on the task of pK _a prediction of Brønsted acids. pK _a prediction is a graph-level regression task (predicting one value per molecule) and it was chosen as our test case because of its relevance across many fields of research. Furthermore, pK _a values characterize one of the simplest chemical reactions, with a well-understood molecular mechanism and individual atomic contributions that aid in the interpretation of data and predictions. In addition, local structure (i.e., the deprotonatable group) and hence, atomic contributions govern pK _a values more than general global factors do (such as molecular shape). Also for this reason pK _a prediction is likely suited to highlight effects in elemental representation on model performance. Finally, previous work has demonstrated that pK _a models with good predictive performance can be obtained, both via ML-based ,− or structure–property relationship approaches, as well as by using quantum chemical calculations. ,

The second part of our study analyzes the effect of the tested elemental representations on the generalizability of our GNN models toward chemical elements that are rare or even absent in the training set. Encoding trends in element behavior may allow a model to infer the reactivity of elements that are underrepresented or absent in the training set from those that are abundantly present in the training compounds. For example, the possible effects of an iodine on pK _a values could be inferred due to its similarity to other halogens such as chlorine, and possibly refined based on trends in electronegativity and polarizability in the periodic table. In addition, we investigated whether vertex featurization – especially elemental representation – can enhance the task of learning and predicting the site-of-deprotonation (SoDP) in our training and test sets of acidic compounds. Because SoDP prediction is a vertex-level classification task, it may be affected differently by vertex featurization than the graph-level regression task of pK _a prediction.

The results of our investigations indicate that the GNN models perform an internal classification of elements, essentially converting any elemental representation adopted into a fully orthogonal one. Utilizing this knowledge, we constructed in the last part of this work a highly orthogonal atom-typed representation, which yielded modest performance benefits compared to the elemental representations introduced above, especially for smaller models in terms of layer and channel count. In addition, we introduced similarity typing, which attempts to provide generalizability to the orthogonal atom-typed representations by including information on atomic neighbors in the vertex featurization, yielding further performance benefits.

Methodology

Data Preparation

A curated data set of 2374 molecules and experimental pK _a values was assembled from the literature. ,− Tautomeric and protomeric states of the data set moleculesincluding the conversion of bases to their corresponding conjugate acidswere resolved using semiempirical tight-binding calculations in standard workflows provided by xTB. , These calculations were performed with the GFN2-xTB functional with a generalized born surface area (GBSA) solvation term. The data set was further filtered toward compounds with pK _a values in the well-represented range of 1–13, and 8 arsenic compounds were excluded from the data set due to issues in atomic (Gasteiger) charge assignment. The final data set contains 2374 compounds (only in their [conjugate] acidic form) and is available at https://github.com/KylleV/pka-dataset. The data set provides a broad coverage of pK _a values in the range from 1 to 13. The distribution of values in the data set is shown in Figure S1 of the Supporting Information [SI]. A variety of deprotonatable moieties (Figure S2) is covered, including oxygen-, nitrogen-, sulfur-, and phosphorus-containing acids (see Figure S3 for deprotonatable groups in the data set as defined in section 1.1 of the SI, together with their pK _a values). For model training 5-fold cross-validation was adopted using an 80/20 training/test split with nonoverlapping test-sets.

Graph Featurization

Our GNN models for pK _a prediction are trained on hydrogen-deprived graphs, using one of eight different methods to represent the elements present in our data set [C, N, O, F, P, S, Cl, Br, I, and (implicitly) H]. Each elemental representation is described in detail in the following subsection (Elemental Representations). The construction and featurization of graph input from molecular data was performed using an early version of an in-house RDKit-Based python library (Kenochroma).

In addition to evaluating the effect of using different elemental representations, three different levels of featurization are tested (Cf. Figure ): (i) the highest degree of featurization used is referred to in the following as maximally featurized graphs. For this degree of featurization, edges are featurized with a one-hot encoded bond-order, and a Boolean indicating whether or not the bond is in a conjugated system. Vertices are featurized with one of the elemental representations detailed below, together with a min-max scaled hydrogen count and valence, a standard scaled Gasteiger charge a one-hot encoded hybridization, and a Boolean indicator for aromaticity (see Figure S4 for distributions of feature values in our data set). Note that Gasteiger charges were used here because we found that this improved pK _a predictions on the order of 0.2–0.4 pK _a units over the use of formal atomic charges. Encoding of features and scaling of values was performed based on training set values only. (ii) Data sets were also created using a minimal featurization in which no edge features were included, and one of the elemental representations below was used to featurize vertices, with neighboring-hydrogen count as only additional vertex feature. (iii) In addition, we tested the effect of using elemental representation as only feature (i.e., without any of the other vertex or edge features listed in Figure ). Note that use of further featurization with quantum chemical descriptors such as Fukui reactivities, i.e., descriptors for chemical reactivity, was investigated as well, but this did not increase model performance sufficiently to justify the computational cost for calculation of quantum chemical descriptors (as indicated by the relatively small improvement in test-set mean absolute error [MAE] seen in Figure S5).

1.

Schematic representation of the levels of graph featurization deployed in this work for a hypothetic molecule. The scheme indicates which vertex and edge features are included in either our minimal featurized models (comprising one of the vertex elemental representations and number of hydrogens only) and fully/maximal featurized models (including all vertex and edge features listed above). See the subsection Elemental Representations for the description of the Elemental Representation feature, and see the text below for explanation of the different other feature types deployed.

Elemental Representations

Three discrete elemental representations were deployed as vertex features, using either (i) atomic number, (ii) period and group of the elements, and (iii) quantum numbers of an element’s outermost electron. These three discrete representations were encoded either numerically (i.e., as a set of min-max scaled values), or using one-hot encoding of the values present in the training set (with an extra bit set true for novel values in the test set when absent in the training set). Thus, in total 2 ·3 = 6 different discrete elemental representations were considered. The elemental representations are illustrated for carbon in Table .

1. Elemental Representations and Their Features as Used in this Work (First Two Columns) .

representation (Discrete)	feature	value	min-max	one-hot
atomic number	Z	6	0.02	[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
period/group	period	2	0	[1, 0, 0, 0]
	group	14	0.25	[0, 1, 0, 0, 0]
quantum number	N	2	0	[1, 0, 0, 0]
	m j	0	1	[0, 1]
	m s	$\frac{1}{2}$	1	[0, 1]

representation (Continuous)	feature	value	min-max	standard scaled
elemental properties	r w	1.7	–	0.41
	χ	2.55		–0.48
	1st IE	11.26		–0.45
	α D	11.3		0.29
	E ea	1.26		–0.043
t-SNE	t-SNE 1	–11.42	0.80	–
	t-SNE 2	10.18	0.27

^aCarbon is used in the last three columns to exemplify feature values and encodings, using numerical values (min-max) or one-hot encoding for discrete representations, and min-max scaled or standard scaler values for continuous representations.

^bAtomic number (Z).

^cPrinciple quantum number (N).

^dmagnetic quantum number (m_j ).

^espin quantum number (m_s ).

^fvan der Waals radius (in Å) (r_w ).

^gPauling electronegativity, in Pauling units (χ).

^hFirst ionization level in eV (1st IE).

^IDipole polarizability in a.u. (α_D).

^jElectron affinity in eV (E _ea).

^kThe two principle axis of the t-SNE based dimensionality reduction of elemental properties (see Figure ).

^lNote that one-hot encoding is performed on the training set, but values presented here are based on the full (training + test) data set.

^mFor quantum numbers the Azimuthal quantum number is ignored in this work because the data set only contains p-block elements (and hydrogen).

In addition to these 6 discrete representations, we also deployed 2 continuous-valued representations (Table ) using (i) a vector of given elemental properties or (ii) a two-dimensional alternative made via t-SNE based dimensionality reduction. Both the property-vector and t-SNE representation contain continuous values. The property-vector representation incorporates 5 intrinsic and electronic elemental properties: Pauling electronegativity, , electron affinity, polarizability, van der Waals radius, and the first ionization level, each scaled using a standard scaler, which removes the mean and scales value to the unit variance of the data for values in the training set. The t-SNE representation contains the two components resulting from utilizing t-SNE dimensionality reduction of the property-vector values to give a 2-dimensional coordinate, positioning elements fairly in accordance with existing trends in the periodic Table (Figure ). To preserve the relative distances in t-SNE feature space, a min-max scaling was used in this representation.

2.

t-SNE coordinates of the elements 1–92 of the periodic table, color coded per group (except for f-block elements) and with annotation for different segments of the periodic table.

Figure S6 visualizes the distributions of and variations in feature values for our 6 discrete and 2 continuous-value elemental representations, and in Figure S7 we present a quantification of how (dis)­similar the representations are.

Furthermore, we trained and evaluated models without any elemental representation (referred to as “No Representation” in figures presenting model performances). Thus, in total 9 different elemental representations were deployed.

Graph Neural Networks

The GNNs were created and trained with PyTorch Geometric, using 2, 4, or 8 identical GATv2Conv layers and with 2, 8, or 16 channels, each followed by one-dimensional batch normalization. On the final layers of the networks a global mean pool is applied and a linear layer produces the final output. Each model is trained with the Adam optimizer using an L1 loss function (i.e., the mean absolute error - MAE). 5-fold cross-validation was performed for each of the 3 layer sizes, 3 channel counts, 9 elemental representations and 3 featurization levels, yielding in total 5 ·3 ·3 ·9 ·3 = 1215 models trained.

Results

Model Performance

Models were trained on the maximally featurized data set (Figure ), with increasing GNN model size in terms of layer and channel count. Our largest models, consisting of 8 GATv2Conv layers each with 16 channels, show a test-set MAE of 0.7–0.9 pK _a units (Figure ). These models outperform our models that are smaller in terms of layers or channel count or of their level of featurization (individual model performance on test and train set in Figures S8 and S9, respectively) and they approach performance of previously reported pK _a predictors. ,, For a more direct comparison with literature, we additionally trained maximally featurized GNN models (consisting of 8 layers and 16 models) on the pK _a data set from ref (which we prepared for our model training as described in section 5.2.1 of the SI). Performance of these models was evaluated via 5-fold cross-validation and using two external test sets employed in ref , with results shown in Figure S10. The best MAEs reported in ref for test-set cross-validation and for the two external data sets were 0.68, 1.15, and 0.53 respectively, while our models using the quantum-number elemental representation achieved values of 0.86, 1.49, and 0.91. The worse performance of GNNs compared to classical ML models, such as the random forest used in ref , is discussed in the Data Requirements section.

3.

Test-set performance in terms of MAE values for predicted pK _a values over 5-fold cross-validation for 8-layer 16-channel models trained on maximally featurized data with different elemental representations (listed on the y-axis). Data for model performance are presented as box plots across 5-fold cross validation, and as points for the individual models.

Performance Differences Between Representations

Figure shows only minor differences between models trained with different elemental representations, with all models having MAE values within 0.2 pK _a units of each other. However, one-hot encoding seems to consistently give a slight improvement in performance over using the numerical version of the same representation, by around 0.1 pK _a unit (Figure , top six boxplots). The finding that all models provide similarly good performance, including those without elemental representation, suggests that the effects of using any elemental representations may be masked by the extensive degree of additional featurization used for the maximally featurized graphs (Figure ). To isolate the effect of the elemental representations employed, further models were trained on data sets with minimally featurized graphs. Test-set performances of the resulting models are presented in Figure (left panel) and highlight two trends: (i) the benefit of one-hot encoding (compared to the equivalent numerically encoded counterpart) that was observed in Figure , is also present for the minimally featurized models in Figure (left), and (ii) models trained with no elemental representation or with numerical atomic-number representation saw a substantial decrease in model performance (with MAEs increasing by 0.60 and 0.37 pK _a units, respectively) upon decreasing the featurization level of the data. This is in contrast with the other elemental representations, for which reducing the level of featurization only led to an MAE increase of on average 0.15 pK _a units. We also note that for the minimal-featurized models our best performing elemental representations (i.e., all models in Figure (left) excluding atomic number (numerical) and no representation) show a similarly small spread in test-set MAE. This trend is similar to what was observed for the maximal-featurized analogues (Figure ). It should further be noted that for some small compounds (with two heavy atoms only, especially the hydrazonium ions), pK _a values were poorly predicted by minimal-featurized models or when only using elemental representations, as discussed in further detail in section 3.1.1 of the SI.

4.

Performance in terms of MAE values for predicted pK _a values over 5-fold cross-validation of 8-layer 16-channel models trained on minimally featurized data with different elemental representations (listed on the y-axis), for (left panel) the entire test set, or (right panel) subsets of test compounds for which predictions are defined to be either (gray) trivial or (red) nontrivial (see text below for an explanation of the corresponding definition). Data for model performance are presented as box plots across 5-fold cross validation, and as points for the individual models.

In addition to a masking effect of extensive featurization, differences between elemental representation on model performance can also be masked by data points (here: pK _a values) that are relatively straightforward (trivial) to predict. For example, our data set contains 556 carboxylic acids that have pK _a values in a relatively narrow range of 4.1 ± 0.9. Hence, the straightforward predictions of pK _a values for carboxylic acids will have a major contribution to the MAE of a model. If values for large parts of a data set are trivial to predict, all models are likely to perform well on this subset, possibly overshadowing the contribution from more difficult (nontrivial) predictions for which differences in performance between elemental representations may be more pronounced. This effect is further exaggerated by the fact that the more samples of a given type of data (in this case acid type) there are in a data set, the more data the model has to learn on, and thus the model becomes better at predicting this given type of data. At the same time, the more samples there are of a given, probably well-predicted data type, the more likely it is to occur in a randomly selected test set and thus contribute more to the MAE or another evaluation metric of the model. A similar observation was made on the task of link prediction in graphs.

To investigate the effect of triviality on our GNN model predictions, we split up our sets of test compounds into compounds for which pK _a prediction can be considered trivial or nontrivial. We defined predictions to be trivial for acid types where the protonatable group is both abundant in the data set (more than 25 occurrences) and generally has similar pK _a values (with a standard deviation of less than 2 pK _a units), as detailed in section 4.1 of the SI. The 15 acid types thus defined as trivial are shown in Figure S11, and the distribution of their pK _a values is given in Figure S12. The results for models trained on the minimally featurized data (Figure , right panel) show that all models, except those without elemental representation or using atomic number (numerical) as representation, perform similarly on trivial predictions, whereas for nontrivial predictions we again observe the previously noted advantage of one-hot encoding with respect to model predictive performance.

We hypothesize that the poor performance of using no elemental representation or numerically encoded atomic numbers (Figure ) is due to an inadequate understanding or identification of deprotonatable functional groups, which in turn is caused by a poor understanding of chemical elements. Models without representation and minimal featurization rely purely on connectivity and hydrogen count to infer elements, which may well lead to cases for which there are multiple elements that could be assigned to a single vertex (such as halogens). The poor understanding of elements by the numerically encoded atomic number representation may originate from largely unused sections of feature space, by squeezing elements within the same period together – with 4 elements in the atomic number range 6 to 10 – and only two elements in the range 17–53 (as shown in the visualization of the representation in Figure S6a), making them difficult to distinguish.

We also investigated effects of elemental representation, and especially of different encoding methods, on the performance of models trained on two external data sets. For this purpose we used the pK _a data set from ref and, for a non-pK _a related task, the solubility data set from ref . Data preparation for model training was performed using the RDKit implementation of MolVS, see sections 5.1.1 and 5.2.1 of the SI. Performance of the resulting models are shown in Figures S13 and S14 for the two sets, respectively, and confirm the observed poor performance of atomic number (numerical) and no-representation, as well as the performance benefit of one-hot encoding over numerical. It should be noted, however, that for one of the two additional test sets used in ref , we found the period/group numerical representation to perform slightly better than its one-hot encoded counterpart, Figure S14d.

Model Ranking

To further inspect differences in performance between representations, we compared the relative performances across all (1215) models. Because of differences in model hyperparameters, MAEs cannot be directly compared between models of different size or featurization level. However, when considering the relative rankings for models with the same hyperparameters and featurization level but with different elemental representation, these rankings give further insight into possible effects of elemental representations on predictive performance. Interestingly, when analyzing this relative ranking across all 1215 models (Figure , individual rankings in Figure S15b) we found a trend of increasing performance with feature count in the elemental representation for numerical representations, but a reversed trend for one-hot encoded representations. Note, however, that for maximally featurized models alone, the one-hot encoded period-group representation slightly outperforms the use of one-hot encoded atomic numbers (Figure S15a).

5.

Relative rankings of representations (ranked based on test-set MAEs) across models with same sizes and featurization levels but with different elemental representation (left y-axis), sorted by the number of features per elemental representation (right y-axis). Data are presented as box plots with point clouds of data points for individual models (colored by featurization level), which are positioned around each rank for each representation.

Random Elemental Representations

The observed correlations between performance and feature count could be caused directly by the feature count (i.e., by the architecture of the representation), or by the contents of the representation. To study which factor is on the basis of the observed (reversed) correlations for numerical and one-hot encoding representations, we created a set of representations with different numbers of features and feature space dimensions, that mimic the structure of the discrete elemental representations, either numerically or one-hot encoded (see section 7.1.1 of the SI for methodological details). By randomly assigning feature values to the elements in our molecular data set, we ensured that no information on elemental behavior was encoded in this set of representations. Figure S16 shows that for these randomized representations we find a slight decrease in performance at higher feature count for one-hot encoding. However, a clear increase in the performance of the numerical representation with increased feature count can be seen. Figure S16 also shows that at 3–4 features, the performance of the numerical representation and the one-hot encoding converges, which is also where their architecture converges. Together with the superior performance of the one-hot over the numerical encoding of the same representation at lower feature counts (Figures and S16), this result suggests that the architecture and structure of the representation is more important than the content of what is represented.

For our 4-layer 8-channel minimally feauturized models trained on the external pK _a and solubility data sets, we also observed an increasing performance with number of features (going from single featured atomic number to the property-vector representations with five features, Figures S13 and S14). No clear correlation was observed for one-hot encoded representations with feature count, similarly to our findings for 4-layer models using random representations displayed in Figure S16. However, unlike the results in Figure , the one-hot quantum number representations often outperformed the one-hot period/group models when trained on the external data.

Shuffled Elemental Representations

The randomized representations contain no information on the elements and can hence only be used to investigate effects of the architecture of the representation on model performance. To also investigate the relative importance of the contents of representations, a set of shuffled representations was generated in which information on the properties of the elements within the data set was shuffled over the elements according to the predefined scheme outlined in Table S2 and section 7.2.1 of the SI. If the contents of the representation had importance, then these shuffled representations should lead to underperformance compared to the models presented in Figure . However, as seen in Figure (and for GNN models with other hyperparameters in Figure S17), no substantial decrease in performance can be seen when shuffling the information on elemental properties. The exception to this is the numerically encoded atom-number representation which has strongly increased performance upon shuffling. This may be explained by our previous hypothesis on the subpar performance of this representation being caused by gaps in feature space. In the baseline (unshuffled) representation the abundant and important period 2 elements are pushed together in feature space, but in the shuffled representations these are better distributed (for example, carbon and iodine swapped atom number, causing the abundant carbon to be easily isolated in feature space). This suggests that the contents of an elemental representation does not contribute measurably to model performance in pK _a prediction.

6.

Test-set MAE of a minimally featurized 4-layer 16-channel model compares the model performances for the shuffled elemental representations (red boxes) to the ones presented in Figure (i.e., with unmodified elemental representation, green boxes). Data for model performance are presented as box plots across 5-fold cross validation, and as points for the individual models. Note that the x-axis scale is reduced here for clarity, leaving out one badly performing outlier of the atomic number (numerical) model.

Model Generalizability

The results of our model rankings (Figures and ) indicate that the architecture of an elemental representation (and not the information content) is important for model performance. If content is not important, little benefit would be gained from understanding and learning trends and patterns, thus little or no model generalizability should be attainable. To verify this, models were trained on data sets in which molecules with a specific element were completely (Figure S18) or partly removed from the training set (Figure S19). The performance of the models on the compounds with the removed elements could then be utilized as a measure for generalizability of a model. Performance for completely and partially removed elements, relative to baseline models, are shown in Figures S18 and S19, respectively. Alternatively, an element was left as the only member of its group, and generalizability was measured as the performance increase when other group members were added, with results comparing performance for compounds with selected elements presented in Figure S20.

Elements in the same group behave and react in a similar manner, and thus occur in similar functional groups (e.g., alcohols and thiols). Hence, it is likely that the addition of more, similar substructures to the data set – rather than true elemental generalizability – will be on the basis of increasing model performance. This is an issue in all three experiments. To account for this, atomic number (one-hot) can be utilized as a baseline, since it is fully orthogonal and thus has no elemental generalizability. For any representation to show signs of elemental generalizability, the performance increase from addition of other elements must be above that of atomic number (one-hot) which was not consistently the case for any representation in any of the three generalizability experiments (Figures S18, S19 and S20).

Data Requirements

Next, we investigated the effect of using our elemental representations on the data requirements of the GNNs. To establish the data-requirements of GNNs, we compared GNN performance to that of three vector-input ML methods (random forest (RF), multiplayer perceptron (MLP), and support vector machines (SVM)) as implemented in scikit-learn across different training-set sizes. The results are shown in Figure S21 and highlight that GNN performance is more sensitive to decreases in training-set size than the vector-input models. Upon decreasing the training set from 1898 to 250 compounds, the GNN MAE increased by 0.7 (97%) which is to be compared to an MAE increase for the RF model of 0.56 (77%), with more extreme differences when decreasing the training set to only 100 compounds (Figure S21). The results in Figure S21 also show that when classical vector-based ML methods are provided with appropriately engineered input (consisting of both fingerprints , and descriptors , ), these models consistently outperform 8-layer, 16-channel maximally featurized GNNs, at a fraction of the computational cost (see Figure S22 for a comparison of computational costs). This is consistent with previous findings, suggesting that GNNs are better suited for more complex tasks such as multiclass classification than for regression. In addition, the high data-requirements of GNNs were not affected by different elemental representations, as shown in Figure S23. This result is in agreement with the lack of generalizability observed, as the ability of models to better generalize data would be on the basis of potentially decreasing their data requirements.

Internal Representation

In conclusion, the results presented here suggest that the GNN models learn to associate specific feature values with specific elements, without understanding the meaning of those features or of structures and trends present in feature space. This is consistent with the models performing an internal classification step of the elements, using a completely orthogonal internal representation (i.e., by converting any elemental representation into one-hot encoded atomic numbering). This explains why one-hot encoded representations outperform their numerical counterparts, considering that the former are closer to the internal representation of the model. It also explains the (reversed) patterns observed in performance ranking between the differently encoded representations in Figure . As one-hot encoding comes presegmented, it does not require the model to segment feature space itself. Increasing the number of features for numerical representations also present more axes to segment feature space (Cf. Table S1), making segmentation easier and possibly leading to the increase in performance with increasing feature count for numerical representations. In the case of one-hot encoding, the decrease in performance with feature count can be explained by the same reasoning, where any deviation from the fully orthogonal atomic number (one-hot) encoding requires more efforts to segment the feature space. If the model has a completely orthogonal internal representation upon learning, this would also explain the lack of generalizability and why changing elemental representation has no effect on data requirements of GNNs. Furthermore, this could explain the finding that the contents of representation have no significant effect on model performance: the contents are only used to classify the elements, as indicated by the lack of performance decrease when shuffling the contents (Figure ).

To investigate the internal representation of our models, the node embeddings of 8-layer 16-channel minimally featurized models were visualized using UMAP dimensionality reduction. The resulting plots are presented in Figure S24 and generally show that elements are separated into different sections of embedded space in the first 7 layers of the models, with different behavior of the final layer, possibly because its representation may be optimized for the final linear layer.

For a more quantitative evaluation of node embeddings, atoms were grouped by their element and similarity matrices were computed between the node embeddings of these groups. The results are presented in Figure S25 and show differences between the representations at the early layers, but this changes to increasing similarity in later layers (excluding the final eighth layer). For example, with increasing layer size a decreasing trend in the internal similarity of oxygen and nitrogen is observed, likely because these frequently occurring deprotonatable atomic sites are separated into multiple subcategories based on the functional group they are part of. Another trend is a similarity between some halogens in later layers, see Figure S25. However, this similarity is also found when using no elemental representation, indicating that this similarity is likely learned from the unique connectivity of halogens (being the only monovalent elements when implicit hydrogens are used) and not from trends inferred from elemental representation.

Furthermore, the models were generally capable of inferring aromaticity of an atom from element and hydrogen count alone, as indicated by aromatic and aliphatic carbons generally occupying separate areas of embedded space, as shown in the UMAP plots in Figure S24.

Our results indicate that a classification step is applied to the elements, or that the model essentially overfits to element space. This is likely exacerbated by the small number of elements used here (as is generally the case for ML on molecular data), and the discrete nature of chemical elements. Therefore, we recommend using fully orthogonal (i.e., one-hot encoded atomic numbering) as elemental representation when training molecular GNN models. This representation is already the most prevalent representation used in literature ,,,,, and shows small benefits in performance without costing generalizability. Our observation is that there is no generalizability to be gained by preinforming the models with domain knowledge in the form of a more chemically inspired elemental representation. This recommendation to use orthogonal representations may be relevant in general for representing features consisting of small sets of discrete values.

Site-of-Deprotonation Prediction

After studying the effects of different elemental representations on the graph-level regression task of pK _a prediction, we also verified their effects on a vertex-classification task. The chosen task here is predicting which atom is (most) prone to deprotonation, referred here to as predicting the site-of-deprotonation (SoDP; see section 11.1 of the SI for methodological details). Because vertex classification focuses on the nodes, the featurization of the individual node (i.e. atom) may well be of higher importance for this task than for the more holistic graph-level task of pK _a prediction. Nevertheless, the differences in recall of the models with different elemental representation (shown in red in Figure ) show similar trends as observed for the differences in pK _a prediction shown in Figure : models with atomic number (numerical) or no elemental representation perform poorly, and one-hot encoded models outperform their numerically encoded counterparts. Note that these trends are also observed in Figure S26 where we used other metrics to evaluate model performance, such as the node-rank decent metric which was designed here specifically for this type of classification task (as described in section 11.3 of the SI).

7.

Top-1 score (indicating how often the correct SoDP [Site-of-Deprotonation] is ranked highest by a model) for minimally featurized 4-layer 16-channel models using either graphs with implicit (green boxplot) or explicit hydrogen atoms (red boxplot), with results for explicit-hydrogen models converted to the same task of SoDP prediction as for implicit-hydrogen models (as described in section 11.1 of the SI). Data for model performance are presented as box plots across 5-fold cross validation, and as points for the individual models.

Hydrogens can be represented as explicit atoms or implicitly as a numerical feature in hydrogen-deprived graph. It is generally expected that inclusion of explicit hydrogens adds no additional information to the graph as hydrogens on the same atom are equivalent. In line with this we observed for the graph-level regression task of pK _a prediction no significant differences in model performance between graphs with hydrogens either implicitly or explicitly included, Figure S27. Interestingly, for the vertex-classification task of SoDP prediction, we observed a small performance benefit by using explicit instead of implicit hydrogens, as seen in Figure . This difference is observed when for explicit hydrogens, the task is changed to identifying which hydrogen is lost instead of which atom loses it. The two models are compared in Figure on their top-1 score, which indicates how often a true SoDP of a molecule is also the highest ranked one by the model. The performance is compared by converting the result of the explicit hydrogen model to the task of the implicit hydrogen model (as described in section 11.1 of the SI). The results in Figure show a slight performance benefit of explicitly representing hydrogens, with a top-1 metric increase of ≈ 0.1–0.2 over using implicit hydrogens. In contrast, the No representation models perform better with implicit hydrogens, probably because its encoding provides hydrogen count as a feature for the (heavy) atoms, whereas explicit hydrogen models cannot distinguish hydrogens from heavy atoms due to the lack of any elemental information in the No representation models. The improved performance of using explicit hydrogens for all other models is also present across other metrics designed for the task presented (see Figure S26 for performance evaluation using these metrics). Since hydrogens on the same atoms are topologically equivalent, no information gain is provided in the explicit-hydrogen models. This suggests that the performance benefit is due to the different representation of the same data or because of increased computational resources allocated around the areas of interest. The performance increase of using explicit hydrogens may also be transferable to other tasks when the prediction task can be moved to hydrogen vertices.

Atom and Similarity Typing

Our findings on GNN-based learning on pK _a data indicate that increasing orthogonality (and thus decreasing generalizability) of elemental representations results in slightly improved predictive performance, as exemplified by results for the fully orthogonal one-hot encoded atomic number representation (Figure ). If we extrapolate this trend further, performance gains could be attainable by using even more orthogonal (or specific) representations. Such an increase in orthogonality could be obtained by not only distinguishing atoms by their element, but also by their surroundings, separating for example an alcohol oxygen from a carbonyl oxygen. This form of differentiation is known as atom-typing, which is a classification system commonly used in force fields. , Based on the findings described above this may also prove useful for GNNs.

In force fields, atom typing is generally established by using a static predefined set of possible types. Creating a set of presumed relevant atom types provides implicit information to the model regarding which atoms are important. This could be beneficial for model performance, but can also bias comparisons to the elemental representations employed above. Therefore, we instead generated atom types here in an unsupervised manner, using one of the following two ways. The first method uses a radial definition, giving an atom a type based on itself and on atoms within two bonds. This method is referred to as Radius-2 atom typing in Figure . The second method assigns atoms to functional groups according to the atom-typing strategy of Ertl. See section 12.1.1 of the SI for further details on the two employed approaches for atom typing. In both cases, atom types were only generated for the molecules present in the training set. This atom-typing provides similar information on each atom as the maximally featurized graph, but in a highly orthogonal manner (for example by encoding an aromatic and aliphatic version of the same atom as two different types).

8.

MAE values for trivial (green) and nontrivial pK _a predictions (red) for GNN models using atom- and similarity-typing (upper panel), compared to use of elemental representations for otherwise maximally featurized graphs (lower panel). In all cases, 4-layer 8-channel models are used. Data for model performance are presented as box plots across 5-fold cross validation, and as points for the individual models.

While little generalizability is attained with any of the elemental representations, the use of atom typing may lead to a further decrease in generalizability of the trained models, as it would not be able to generalize a known element to a new environment (such as generalizing oxygen behavior as from acyclic to cyclic ethers). To overcome this, we introduce similarity typing, which replaces the zeros in the atom-typing one-hot encoding with the average similarity of that atom type to the ones of the atom in question; see section 12.1.2 of the SI for methodological details. This ensures that a novel atom type will not be an all-zero vector, but it will be described in terms of its similarity to known atom types, possibly aiding generalizability.

Our test set predictions show that radial segmentation outperforms Ertl’s functional-group based atom typing (Figure ). This may be because in radial segmentation an atom can be represented in multiple atom types, and/or it is due to a larger set of atom types generated (on average 2154 ± 37 for radial segmentation vs 1307 ± 8 for Ertl’s approach, where variations in these averages are due to differences in train/test set splits among the models trained per atom typing approach). Figure also shows that Radius-2 atom typing performs similarly to our elemental representations in terms of MAE values obtained, and that similarity-typing leads to a modest but clear performance benefit. The latter especially holds true when nontrivial predictions are considered (red box plots in Figure ). Furthermore, both atom- and similarity-typing are more robust to decreases in model sizes, outperforming elemental representations at fewer channels (Figure S28).

The atom and similarity typing process is similar to a single convolution step in a GNN layer, where the structure of the neighborhood surrounding a node is combined into a single vector on that node. Hence, the introduction of atom and similarity typing can be considered similar to adding an additional GNN layer to a model with elemental representation only. As an additional test, we therefore compared model performance of our atom- and similarity-typing models to elemental representation models with an additional layer. After this addition of an additional layer to the elemental representation models they are still outperformed by atom typing, and especially by similarity typing as shown in Figure S29.

We also tested a combination of elemental representations (in the form of atomic number-based one-hot encoding) and atom or similarity typing, which seemed to provide an additional modest increase in performance that is most clearly seen in test-set MAE values for minimally featurized models as seen in Figure S30. This effect was more pronounced for atom-typing than for similarity-typing, and more pronounced for the functional-group based segmentation compared to radial segmentation. Interestingly, for minimally featurized graphs the performance of radial segmentation-based models was negatively affected by the addition of elemental representations, but not generally so for fully featurized graphs (Figure S30).

The generalizability of models with atom- or similarity-typed representations was tested by training the model on a data set where compounds with a specific atom type (representing a common SoDP) were removed, and then evaluating them for compounds with this removed atom type. This was performed for radial segmentation only, as it overlaps with our acid-type definition. This test was performed for multiple atom-types as listed in Table S3. The results are presented in Figure S31 and show some difference for specific atom types, such as better performance (lower MAE) for elemental representations on arylamonium ions (Figure S31h,i), while similarity typing outperforms atom typing for phenolic alcohols (Figure S31d–f), but atom-typing performs better on sulfonamides (Figure S31j–l). However, no consistent trend is present as to which of the approaches is better across the (four) atom types tested. The results indicate that some encodings can be better for certain substructures, likely based on how many similar substructures exist in the data set, and how similar their behavior and reactivity are. Using different similarity metrics for similarity typing may thus be a direction of optimization of the models.

From the above, the atom- and especially the similarity-typing based models can provide modest performance benefits. However, this comes at a substantial increase in computational costs for graph featurization, from matching each new substructure to all existing atom-types, which causes the computational costs of atom-typing to increase with the square of the number of atom types. Utilizing smaller, premade sets of atom types such as those used in force fields or by constructing one based on domain knowledge or statistical analysis may prove in future research to be a more practical approach for the use of similarity or atom typing in GNNs.

Conclusion

In the current work, we investigated the use of different modes of elemental representations in GNNs on the performance of models in learning on molecular pK _a data. The evidence provided here suggests that, with the exception of the poorly performing atomic-number based numerical elemental representation, little difference in model performance exists between the elemental representations deployed. Furthermore, no gain in generalizability of elemental behavior and reactivity was attained with any of the investigated elemental representations. Shuffling feature values among elements did not affect model performance, which suggests that the content of the elemental representations is not measurably relevant to the learning task at hand. Our results indicate that GNN models learn to associate a specific set of feature values to specific elements, without understanding which information they represent or any of the trends and structure present in feature space. The evidence provided here suggests that the models instead perform a classification of elements, converting any elemental representation into an orthogonal one. Orthogonal representations do not encode any trends, which explains the lack of generalizability observed. Furthermore, if the models -as suggested here- internally convert all representations into a fully orthogonal one, this can also explain why only minor differences in performance are observed when using different elemental representations. These minor differences align with how easily a representation can be converted into an orthogonal one during training, explaining the slightly better performance of the more orthogonal representations employed. As a result, our recommendation for elemental representations is to use a fully orthogonal atomic-number based one-hot encoding, which is already the most commonly used in the literature for GNNs learning on molecular data. This representation is very similar to the presumed internal representation. The use of a fully orthogonal representation may also be optimal to encode other discrete features with relatively few values.

Utilizing this obtained knowledge of the internal representation of the trained models, we investigated atom typing and introduced similarity typing as alternatives to elemental representations. We found that our similarity-typing approach outperformed atom typing (in terms of pK _a prediction), which in turn outperformed our models using elemental representations (especially at smaller model sizes). The performance of atom or similarity typing was improved further by also including elemental representation. It should be noted, however, that the clear but modest performance improvement of atom and similarity typing over the elemental representations comes at the cost of a significant increase in processing time for data preparation in our current implementations. The results presented here highlight how in this case, predictive performance increases by optimizing featurization that is not based on domain knowledge, but on knowledge of the models and their architecture instead.

Our results suggest that the commonly used one-hot orthogonal representation of elements is most appropriate for learning molecular properties or reactivity. These results may not necessarily extend to tasks in which large numbers of elements are included, such as inorganic or organometallic cases, where further research would be needed. While atom and similarity typing shows promising results, their full potential will require substantial optimization, both in terms of computational efforts needed for featurization, but also in terms of which atom typing algorithm yields the most accurate results. However, we believe that optimization of similarity- and atom-typing can lead to further performance increases, and utilization of premade, task-specific atom-types may improve performance and reduce or even eliminate the preprocessing costs of defining atom and similarity types.

The full molecular data set (with hydrogen atoms explicitly included, and including quantum chemical features) is freely available at https://github.com/KylleV/pka-dataset. The script for computing node-classification metrics is available at https://github.com/KylleV/noclame

The Supporting Information is available free of charge at https://pubs.acs.org/doi/10.1021/acs.jcim.6c00384.

• Overview of the data set structure and contents, as well as detailed analyses of each investigation performed; and a methodological description of the deployed atom-/similarity-typing and node-classification metrics (PDF)

∥.

Department of Computer Science, Åbo Akademi University, European Laboratory for Learning and Intelligent Systems (ELLIS), 20500 Turku, Finland

V.K. performed model training, data analysis, method development, and manuscript drafting. In addition, D.A.P. III was responsible for the data set construction, and D.P.G., H.M., M.C., and D.A. contributed to the method selection and interpretation of results. All authors were involved in the design and direction of the study and in manuscript revision. All authors approved the final manuscript.

This research was financially supported by an Innovation PhD grant from the Department of Computer Science at Vrije Universiteit, and by the Dutch Research Council (NWO) as the Virtual Human Platform for Safety Assessment (VHP4Safety), NWA 1292.19.272. Dimitrios Alivanistos and Michael Cochez were partially funded by the Elsevier Discovery Lab. Michael Cochez was partially funded by the Graph-Massivizer project, funded by the Horizon Europe program of the European Union (grant 101093202), and supported by a gift from Accenture LLP. His work on this publication is in part based upon work from COST Action CA23147 GOBLIN - Global Network on Large-Scale, Cross-domain and Multilingual Open Knowledge Graphs, supported by COST (European Cooperation in Science and Technology, https://www.cost.eu).

The authors declare no competing financial interest.