Can Graph Neural Networks Understand Chemical Elements?

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 pKa 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.