[Paper Review] Pre-training via Denoising for Molecular Property Prediction

Sheheryar Zaidi, Michael Schaarschmidt+7 more|arXiv (Cornell University)|May 31, 2022

Machine Learning in Materials Science30 citations

TL;DR

The authors propose a self-supervised pre-training method that denoises 3D molecular structures to learn representations, achieving state-of-the-art results on QM9 and other benchmarks.

ABSTRACT

Many important problems involving molecular property prediction from 3D structures have limited data, posing a generalization challenge for neural networks. In this paper, we describe a pre-training technique based on denoising that achieves a new state-of-the-art in molecular property prediction by utilizing large datasets of 3D molecular structures at equilibrium to learn meaningful representations for downstream tasks. Relying on the well-known link between denoising autoencoders and score-matching, we show that the denoising objective corresponds to learning a molecular force field -- arising from approximating the Boltzmann distribution with a mixture of Gaussians -- directly from equilibrium structures. Our experiments demonstrate that using this pre-training objective significantly improves performance on multiple benchmarks, achieving a new state-of-the-art on the majority of targets in the widely used QM9 dataset. Our analysis then provides practical insights into the effects of different factors -- dataset sizes, model size and architecture, and the choice of upstream and downstream datasets -- on pre-training.

Motivation & Objective

• Motivate and develop a pre-training approach for molecular property prediction from 3D structures where labeled data are scarce.
• Leverage denoising/score-matching to learn a meaningful force-field-like representation from equilibrium structures.
• Demonstrate that denoising-based pre-training improves downstream task performance across diverse datasets and architectures.
• Provide practical insights into how upstream dataset size, model size, architecture, and dataset similarity affect transfer.

Proposed method

• Formulate a denoising pre-training objective on 3D molecular structures by perturbing atomic coordinates with Gaussian noise and training a GNN to predict the noise.
• Interpret the denoising objective as learning a force field corresponding to a Gaussian-mixture approximation of the Boltzmann distribution around equilibrium structures.
• Apply this pre-training to Graph Net Simulator (GNS) and a Tailored Activation Transform (TAT) enhanced variant (GNS-TAT).
• Demonstrate architecture-agnosticity by also applying denoising pre-training to TorchMD-NET, a transformer-based model with invariant/equivariant components.
• Use Noisy Nodes as an auxiliary loss to complement pre-training and analyze its impact.
• Analyze dataset size, model size, and upstream/downstream dataset relationships on transfer performance.

Experimental results

Research questions

• RQ1 Does pre-training via denoising improve downstream molecular property prediction compared to random initialization?
• RQ2 How does the benefit of denoising pre-training depend on the relationship between upstream (pre-training) and downstream datasets?
• RQ3 Can denoising pre-training be effective across different architectures (GNNs and transformers) and tasks (QM9, OC20, DES15K)?

Key findings

• Pre-training via denoising significantly improves downstream performance across multiple datasets and targets, achieving new state-of-the-art on the QM9 majority of targets.
• GNS-TAT with Noisy Nodes after denoising pre-training yields state-of-the-art results on 10 of 12 QM9 targets; some targets may be sensitive to the noise scale.
• Pre-training on PCQM4Mv2 (3.4M 3D structures) generally transfers well to QM9 and DES15K, with varying transfer depending on element overlap and task similarity.
• Pre-training on OC20 did not improve OC20 IS2RE performance when pre-trained on PCQM4Mv2, highlighting transfer limitations due to dataset dissimilarity, but did speed convergence when pre-trained on OC20 itself.
• Increasing upstream data generally improves downstream representations; larger models benefit from pre-training and can surpass models trained from scratch.
• Pre-training also improves force prediction (e.g., MD17 aspirin) and accelerates training when fine-tuning decoder components only.