Reducing Down(stream)time: Pretraining Molecular GNNs using Heterogeneous AI Accelerators

TL;DR

This paper demonstrates that using Graphcore IPUs accelerates pretraining of molecular GNNs on large datasets, significantly reducing training time and enabling efficient transfer learning for molecular tasks.

Contribution

It introduces the use of heterogeneous AI accelerators for pretraining molecular GNNs, achieving substantial speedups and enabling practical transfer learning in chemistry.

Findings

Training time reduced from 2.7 days to 1.2 hours on large datasets.

Finetuning for downstream tasks completed in under 30 minutes.

Pretrained models effectively transfer to different molecular tasks.

Abstract

The demonstrated success of transfer learning has popularized approaches that involve pretraining models from massive data sources and subsequent finetuning towards a specific task. While such approaches have become the norm in fields such as natural language processing, implementation and evaluation of transfer learning approaches for chemistry are in the early stages. In this work, we demonstrate finetuning for downstream tasks on a graph neural network (GNN) trained over a molecular database containing 2.7 million water clusters. The use of Graphcore IPUs as an AI accelerator for training molecular GNNs reduces training time from a reported 2.7 days on 0.5M clusters to 1.2 hours on 2.7M clusters. Finetuning the pretrained model for downstream tasks of molecular dynamics and transfer to a different potential energy surface took only 8.3 hours and 28 minutes, respectively, on a single GPU.