Augmenting Molecular Graphs with Geometries via Machine Learning Interatomic Potentials

TL;DR

This paper develops machine learning interatomic potential models trained on a large dataset to predict molecular geometries and properties, offering a faster alternative to expensive quantum chemistry methods with promising accuracy and transferability.

Contribution

It introduces a large-scale dataset and pre-trained MLIP models that can generate approximate geometries and enhance molecular property predictions, advancing the integration of geometry prediction in ML workflows.

Findings

Pre-trained models can generate low-energy geometries improving property prediction.

Geometry fine-tuning enhances downstream molecular property accuracy.

Models trained on relaxation data transfer well to various molecular tasks.

Abstract

Accurate molecular property predictions require 3D geometries, which are typically obtained using expensive methods such as density functional theory (DFT). Here, we attempt to obtain molecular geometries by relying solely on machine learning interatomic potential (MLIP) models. To this end, we first curate a large-scale molecular relaxation dataset comprising 3.5 million molecules and 300 million snapshots. Then MLIP pre-trained models are trained with supervised learning to predict energy and forces given 3D molecular structures. Once trained, we show that the pre-trained models can be used in different ways to obtain geometries either explicitly or implicitly. First, it can be used to obtain approximate low-energy 3D geometries via geometry optimization. While these geometries do not consistently reach DFT-level chemical accuracy or convergence, they can still improve downstream performance compared to non-relaxed structures. To mitigate potential biases and enhance downstream predictions, we introduce geometry fine-tuning based on the relaxed 3D geometries. Second, the pre-trained models can be directly fine-tuned for property prediction when ground truth 3D geometries are available. Our results demonstrate that MLIP pre-trained models trained on relaxation data can learn transferable molecular representations to improve downstream molecular property prediction and can provide practically valuable but approximate molecular geometries that benefit property predictions. Our code is publicly available at: https://github.com/divelab/AIRS/