Transfer learning using attentions across atomic systems with graph neural networks (TAAG)

Abstract

Recent advances in Graph Neural Networks (GNNs) have transformed the space of molecular and catalyst discovery. Despite the fact that the underlying physics across these domains remain the same, most prior work has focused on building domain-specific models either in small molecules or in materials. However, building large datasets across all domains is computationally expensive; therefore, the use of transfer learning (TL) to generalize to different domains is a promising but under-explored approach to this problem. To evaluate this hypothesis, we use a model that is pretrained on the Open Catalyst Dataset (OC20), and we study the model’s behavior when fine-tuned for a set of different datasets and tasks. This includes MD17, the *CO adsorbate dataset, and OC20 across different tasks. Through extensive TL experiments, we demonstrate that the initial layers of GNNs learn a more basic representation that is consistent across domains, whereas the final layers learn more task-specific features. Moreover, these well-known strategies show significant improvement over the non-pretrained models for in-domain tasks with improvements of 53% and 17% for the *CO dataset and across the Open Catalyst Project (OCP) task, respectively. TL approaches result in up to 4× speedup in model training depending on the target data and task. However, these do not perform well for the MD17 dataset, resulting in worse performance than the non-pretrained model for few molecules. Based on these observations, we propose transfer learning using attentions across atomic systems with graph Neural Networks (TAAG), an attention-based approach that adapts to prioritize and transfer important features from the interaction layers of GNNs. The proposed method outperforms the best TL approach for out-of-domain datasets, such as MD17, and gives a mean improvement of 6% over a model trained from scratch.