Transfer learning based on atomic feature extraction for the prediction of experimental $^{13}$C chemical shifts

Žarko Ivković, Jesús Jover, Jeremy Harvey

Submission Track: Full Paper

Submission Category: AI-Guided Design + Automated Chemical Synthesis

Keywords: Machine Learning, Atomic Representation, Transfer Learning, Graph Neural Networks, NMR, Chemical Shifts, Feature Extraction, Low-data, Atomic Embeddings

TL;DR: Atomic embeddings extracted from forcefields improve model's performance in low-data regimes

Abstract:

This study indicates that atomic features derived from a message passing neural network (MPNN) forcefield are robust descriptors for atomic properties. A dense network utilizing these descriptors to predict $^{13}$C shifts achieves a mean absolute error (MAE) of 1.68 ppm. When these features are used as node labels in a simple graph neural network (GNN), the model attains a better MAE of 1.34 ppm. On the other hand, embeddings from a self-supervised pre-trained 3D aware transformer are not sufficiently descriptive for a feedforward model but show reasonable accuracy within the GNN framework, achieving an MAE of 1.51 ppm. Under low-data conditions, all transfer-learned models show a significant improvement in predictive accuracy compared to existing literature models, regardless of the sampling strategy used to select from the pool of unlabeled examples.
We demonstrated that extracting atomic features from models trained on large and diverse datasets is an effective transfer learning strategy for predicting NMR chemical shifts, achieving results on par with existing literature models. This method provides several benefits, such as reduced training times, simpler models with fewer trainable parameters, and strong performance in low-data scenarios, without the need for costly ab initio data of the target property.

Submission Number: 5